[논문 리뷰] When Vision Transformers Outperform ResNets without Pre-training or Strong Data Augmentations

Xiangning Chen, Cho‐Jui Hsieh외 1명|arXiv (Cornell University)|2021. 06. 03.

Advanced Neural Network Applications참고 문헌 60인용 수 103

한 줄 요약

이 논문은 Vision Transformers와 MLP-Mixers가 처음부터 학습된 상태에서 SAM(sharpness-aware optimizer)과 함께 사용할 때 비슷한 규모의 ResNets를 능가할 수 있음을 보여주고, 대규모 사전 학습이나 강한 증강 없이도 가능하며, 손실 지형(loss landscapes)을 분석해 그 이유를 설명한다.

ABSTRACT

Vision Transformers (ViTs) and MLPs signal further efforts on replacing hand-wired features or inductive biases with general-purpose neural architectures. Existing works empower the models by massive data, such as large-scale pre-training and/or repeated strong data augmentations, and still report optimization-related problems (e.g., sensitivity to initialization and learning rates). Hence, this paper investigates ViTs and MLP-Mixers from the lens of loss geometry, intending to improve the models' data efficiency at training and generalization at inference. Visualization and Hessian reveal extremely sharp local minima of converged models. By promoting smoothness with a recently proposed sharpness-aware optimizer, we substantially improve the accuracy and robustness of ViTs and MLP-Mixers on various tasks spanning supervised, adversarial, contrastive, and transfer learning (e.g., +5.3\% and +11.0\% top-1 accuracy on ImageNet for ViT-B/16 and Mixer-B/16, respectively, with the simple Inception-style preprocessing). We show that the improved smoothness attributes to sparser active neurons in the first few layers. The resultant ViTs outperform ResNets of similar size and throughput when trained from scratch on ImageNet without large-scale pre-training or strong data augmentations. Model checkpoints are available at \url{https://github.com/google-research/vision_transformer}.

연구 동기 및 목표

• ViTs와 MLP-Mixers가 ResNets와 비교했을 때 최적화 및 일반화의 차이를 이해하기 위해 처음부터 학습된 상태로 조사한다.
• 합성곱 없이 구성된 아키텍처에서 샤프 미니마를 특성화하기 위해 손실 지형과 해시안 행렬을 분석한다.
• 샤프니스 인식 최적화기(SAM)가 정확도, 강건성, 그리고 작업 전반의 학습 역학에 미치는 영향을 평가한다.
• SAM이 대규모 사전 학습이나 강한 증강 없이도 경쟁력 있거나 우수한 성능을 가능하게 함을 보여준다.

제안 방법

• ImageNet에서 수렴 시점의 ViTs와 MLP-Mixers의 손실 지형과 Hessian 스펙트럼을 연구한다.
• ViTs와 MLP-Mixers에 샤프니스-어웨어 미니마이저(SAM)를 적용해 손실 기하에서 평탄한 영역을 유도한다.
• 동일한 기본 전처리(preprocessing)를 사용해 처음부터 학습된 ResNets와 SAM으로 학습된 ViTs/ Mixers를 비교하고 정확도와 강건성을 평가한다.
• SAM에 따른 모델의 고유 변화(해시안 블록, 가중치 노름, 뉴런 활성화 희소성 등)를 분석한다.
• SAM로 인한 해석 가능성 향상을 평가하기 위해 주의(attention) 맵을 시각화한다.

실험 결과

연구 질문

• RQ1 손실 지형-의식형 옵티마이저를 사용하여 ViTs와 MLP-Mixers가 pre-training이나 강한 증강 없이 처음부터 학습된 ResNets를 능가할 수 있는가?
• RQ2 SAM이 합성곱 없는 아키텍처의 손실 지형, 해시안 곡률, 학습 가능성에 어떤 영향을 미치는가?
• RQ3 SAM이 수반하는 고유한 아키텍처 변화(예: 희소성, 가중치 노름)와 이것이 일반화 및 강건성과 어떻게 연결되는가?
• RQ4 ImageNet-R와 ImageNet-C 하에서 SAM으로 학습된 ViTs와 Mixers가 ResNets에 비해 정확도와 강건성 측면에서 어떤 성능을 보이는가?

주요 결과

• SAM은 처음부터 학습된 ViTs와 MLP-Mixers의 정확도와 강건성을 크게 향상시킨다(예: ViT-B/16 및 Mixer-B/16에서 ImageNet의 Top-1 증가가 크게 나타남).
• SAM 없이도 ViTs와 MLP-Mixers는 훨씬 날카로운 국소 최소점으로 수렴하며 큰 Hessian 고유값을 가지지만, SAM은 손실 기하를 매끄럽게 하고 평균 및 최악의 곡률을 모두 감소시킨다.
• SAM은 특히 MLP-Mixers에서 초기 계층의 활성 뉴런을 더 희소하게 만들어, 더 매끄러운 손실 기하와 더 나은 일반화와 상관된다.
• ImageNet에서 SAM으로 학습된 ViTs는 동등하거나 더 큰 크기의 ResNets보다 유사 처리량에서 더 좋은 성능을 낼 수 있으며, 대규모 사전 학습이나 강한 증강 없이도 가능하다.
• SAM은 깨짐(ImageNet-C) 및 분포 변화(ImageNet-R)에 대한 강건성을 깨끗한 정확도 향상보다 더 크게 개선한다, 기준 모델과 비교했을 때.
• SAM은 SGD(모멘텀)로 ViTs를 훈련시켜 Adam/SAM 조합에 가까운 성능을 달성하게 하여 옵티마이저 간 격차를 줄인다.