[논문 리뷰] DeepViT: Towards Deeper Vision Transformer

Daquan Zhou, Bingyi Kang외 6명|arXiv (Cornell University)|2021. 03. 22.

Advanced Neural Network Applications참고 문헌 46인용 수 348

한 줄 요약

이 논문은 Vision Transformers를 깊이 확장했을 때 성능이 왜 포화되는지 분석하고, Re-attention을 도입하여 다양한 주의 맵을 재생성함으로써 매우 깊은 ViT(예: 32 블록)의 안정적인 학습을 가능하게 하며 추가 데이터 없이 ImageNet 정확도를 향상시킨다.

ABSTRACT

Vision transformers (ViTs) have been successfully applied in image classification tasks recently. In this paper, we show that, unlike convolution neural networks (CNNs)that can be improved by stacking more convolutional layers, the performance of ViTs saturate fast when scaled to be deeper. More specifically, we empirically observe that such scaling difficulty is caused by the attention collapse issue: as the transformer goes deeper, the attention maps gradually become similar and even much the same after certain layers. In other words, the feature maps tend to be identical in the top layers of deep ViT models. This fact demonstrates that in deeper layers of ViTs, the self-attention mechanism fails to learn effective concepts for representation learning and hinders the model from getting expected performance gain. Based on above observation, we propose a simple yet effective method, named Re-attention, to re-generate the attention maps to increase their diversity at different layers with negligible computation and memory cost. The pro-posed method makes it feasible to train deeper ViT models with consistent performance improvements via minor modification to existing ViT models. Notably, when training a deep ViT model with 32 transformer blocks, the Top-1 classification accuracy can be improved by 1.6% on ImageNet. Code is publicly available at https://github.com/zhoudaquan/dvit_repo.

연구 동기 및 목표

• 비전 트랜스포머를 깊이 확장할 때 성능이 포화되는 이유를 조사한다.
• 깊은 ViT에서 주의 붕괴의 원인을 식별한다.
• 레이어 간 주의를 다양화하기 위한 경량화된 메커니즘(Re-attention)을 제안한다.
• 처음부터 학습된 더 깊은 ViT가 ImageNet-1k에서 정확도를 향상시킬 수 있음을 보여준다.

제안 방법

• 레이어 간 주의 맵의 유사성을 관찰하기 위해 ImageNet에서 ViT 깊이 확장의 실증적 연구를 수행한다.
• 크로스-레이어 주의 유사성을 통해 주의 붕괴를 정의하고 정량화한다.
• 주 의 헤드 간 정보를 교환하는 학습 가능한 헤드 대 헤드 변환으로 Re-attention을 도입한다.
• ViT 블록에서 MHSA를 Re-attention으로 대체하여 DeepViT 아키텍처를 구성한다.
• 추가 데이터나 증강 없이 ImageNet-1k에서 DeepViT를 SOTA CNN 및 ViT와 비교한다.
• 임베딩 차원 대 깊이 및 대안 주의 샤프닝 기준선에 대한 어블레이션을 제공한다.

실험 결과

연구 질문

• RQ1 ViT가 CNN처럼 더 깊은 아키텍처에서 의미 있게 이익을 얻을 수 있는가, 아니면 성능이 포화되는가?
• RQ2 깊은 ViT에서 주의 맵이 층 간에 비슷해지게 하는 원인은 무엇인가?
• RQ3 가벼운 메커니즘이 주의 헤드 간 정보를 재활용하여 다양성을 회복하고 더 깊은 ViT를 가능하게 할 수 있는가?
• RQ4 ImageNet-1k에서 처음부터 학습된 DeepViT 모델이 비슷한 컴퓨트로 기존의 SOTA 모델을 능가하는가?

주요 결과

• 표준 MHSA를 사용할 때 ViT 깊이를 직접 증가시키면 ImageNet에서 정확도가 포화되거나 감소한다.
• 더 깊은 계층의 주의 맵이 매우 유사해져(주의 붕괴), 특징 진화의 정체와 상관된다.
• Re-attention은 주의 맵을 헤드 간으로 선형 혼합하는 학습 가능한 매트릭스로 다양성을 보존하고 층 간 주의 붕괴를 제거한다.
• DeepViT 모델(32 블록)은 일관된 정확도 향상을 달성하고 ImageNet-1k에서 기초 ViT 및 여러 CNN/ViT SOTA 모델을 추가 데이터나 학습 트릭 없이 능가한다.
• MHSA를 Re-attention으로 대체하면 0개의 유사-block 주의 맵이 형성되고 32-block 모델의 Top-1 정확도가 최대 1.6 퍼센트 포인트 증가한다.
• DeepViT-S와 DeepViT-L은 여러 동시대 모델보다 더 적은 파라미터로 경쟁력 있거나 우수한 정확도를 달성한다(예: DeepViT-L은 변형에 따라 82.2–83.1% 도달).