[논문 리뷰] Early Convolutions Help Transformers See Better

Tete Xiao, Mannat Singh외 4명|arXiv (Cornell University)|2021. 06. 28.

Advanced Neural Network Applications참고 문헌 56인용 수 352

한 줄 요약

ViT의 patchify stem을 가벼운 합성곱 stem으로 교체한 ViT C는 최적화 안정성을 크게 향상시키고, 정확도 손실 없이 SGD 또는 AdamW를 가능하게 하며, 모델 크기 및 데이터 규모에 걸쳐 ImageNet에서 top-1이 1–2% 상승합니다.

ABSTRACT

Vision transformer (ViT) models exhibit substandard optimizability. In particular, they are sensitive to the choice of optimizer (AdamW vs. SGD), optimizer hyperparameters, and training schedule length. In comparison, modern convolutional neural networks are easier to optimize. Why is this the case? In this work, we conjecture that the issue lies with the patchify stem of ViT models, which is implemented by a stride-p p*p convolution (p=16 by default) applied to the input image. This large-kernel plus large-stride convolution runs counter to typical design choices of convolutional layers in neural networks. To test whether this atypical design choice causes an issue, we analyze the optimization behavior of ViT models with their original patchify stem versus a simple counterpart where we replace the ViT stem by a small number of stacked stride-two 3*3 convolutions. While the vast majority of computation in the two ViT designs is identical, we find that this small change in early visual processing results in markedly different training behavior in terms of the sensitivity to optimization settings as well as the final model accuracy. Using a convolutional stem in ViT dramatically increases optimization stability and also improves peak performance (by ~1-2% top-1 accuracy on ImageNet-1k), while maintaining flops and runtime. The improvement can be observed across the wide spectrum of model complexities (from 1G to 36G flops) and dataset scales (from ImageNet-1k to ImageNet-21k). These findings lead us to recommend using a standard, lightweight convolutional stem for ViT models in this regime as a more robust architectural choice compared to the original ViT model design.

연구 동기 및 목표

• ViT의 최적화 가능성이 왜 CNN보다 뒤처지는지 동기 부여하고 이해한다.
• 초기 시각 처리(패치화 stem)가 ViT의 최적화 문제를 야기하는지 조사한다.
• 최소한의 합성곱 stem을 드롭인 교체로 평가하고 원래 ViT stem과의 학습 동역학을 비교한다.
• 다양한 모델 크기와 데이터 규모에 걸쳐 수렴 속도, 옵티마이저 호환성, 하이퍼파라미터 안정성, 그리고 최대 성능을 평가한다.

제안 방법

• ViT patchify stem (stride-p, p×p)을 약 5층 규모의 가벼운 합성곱 stem으로 대체한다.
• FLOPs 및 런타임의 동등성을 유지하기 위해 트랜스포머 블록 하나를 제거한다.
• 다양한 크기(1G–36G FLOPs)와 ImageNet 스케일(1k–21k)에서 ViT P(원래 patchify stem)와 ViT C(합성곱 stem)를 학습시킨다.
• 에러 분포 함수(EDF)를 통해 학습 길이 안정성, 옵티마이저 안정성(AdamW 대 SGD), 하이퍼파라미터 안정성(lr, wd)을 평가한다.
• EMA 가중치와 표준 증강 방법으로 제어된 학습 설정에서 피크 성능을 비교한다.

실험 결과

연구 질문

• RQ1 패치형 stem을 합성곱 stem으로 대체하면 ViT의 최적화 안정성과 수렴 속도에 영향이 있는가?
• RQ2 합성곱 stem을 가진 ViT 모델이 significant accuracy 손실 없이 SGD와 함께 효과적으로 학습될 수 있는가?
• RQ3 합성곱 stem이 패치형 stem과 비교해 학습률과 가중치 감소 안정성에 어떤 영향을 미치는가?
• RQ4 다양한 모델 복잡도와 데이터 세트 규모에서 피크 ImageNet top-1 성능에 어떤 영향이 있는가?
• RQ5 합성곱 stem이 표준 학습 방식에서 ViT를 CNN과 경쟁적으로 만들 만큼 충분한가?

주요 결과

• ViT C는 모델 복잡도 전반에서 ViT P보다 더 빠르게 수렴한다.
• ViT C는 AdamW 또는 SGD 중 어느 것을 사용하더라도 큰 정확도 손실 없이 학습이 가능하게 한다, 반면 ViT P는 SGD에서 고전한다.
• ViT C는 학습률과 가중치 감소에 대한 안정성이 더 커 하이퍼파라미터 민감도를 CNN과 비슷한 수준으로 감소시킨다.
• 1G–36G FLOPs 및 ImageNet 1k–21k에 걸쳐 ViT C는 ImageNet-1k에서 ViT P보다 약 1–2 포인트 낮은 top-1 오차를 나타낸다.
• 공정한 피크 성능 비교에서 중간 복잡도 구간에서 ViT C가 ViT P 및 경쟁 CNN보다 우수하며, ImageNet-21k 사전훈련에서 ViT C가 모든 것을 능가한다.