[論文レビュー] When Vision Transformers Outperform ResNets without Pre-training or Strong Data Augmentations

Xiangning Chen, Cho‐Jui Hsieh+1 more|arXiv (Cornell University)|Jun 3, 2021

Advanced Neural Network Applications参考文献 60被引用数 103

ひとこと要約

本論文は、Vision Transformers および MLP-Mixers をゼロから訓練すると、同規模の ResNets を、sharpness-aware optimizer (SAM) と組み合わせることで上回ることができると示す。大規模な事前学習や強いデータ拡張を用いず、損失地形を分析してその理由を説明している。

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}.

研究の動機と目的

• ResNets と比較したときの最適化と一般化のギャップを理解するため、ViTs および MLP-Mixers をゼロから訓練して調査する。
• 畳み込みを用いないアーキテクチャにおける鋭い極小値を特徴付けるため、損失地形と Hessian を分析する。
• さまざまなタスクにおける精度、頑健性、学習ダイナミクスに対する sharpness-aware optimizer (SAM) の影響を評価する。
• 大規模な事前学習や強力なデータ拡張なしに、SAM が競争力のあるまたは優れた性能を可能にすることを示す。

提案手法

• ImageNet で収束時の ViTs および MLP-Mixers の損失地形と Hessian スペクトルを研究する。
• ViTs および MLP-Mixers に sharpness-aware minimizer (SAM) を適用し、損失幾何学における平坦な領域を促進する。
• 同じ基本的な前処理を用いて、SAM 訓練済みの ViTs / Mixers をゼロから訓練された ResNets と比較し、精度と頑健性を評価する。
• SAM を用いた内部的なモデル変化を分析する。 Hessian ブロック、重みノルム、ニューロンの活性化のスパーシティを含む。
• SAM による解釈可能性の改善を評価するため、アテンションマップを可視化して、SAM による解釈可能性の改善を評価する。

実験結果

リサーチクエスチョン

• RQ1 損失幾何学を意識した最適化手法を用いて、ViTs および MLP-Mixers は、事前学習なし・過度なデータ拡張なしで、ゼロから訓練された ResNets を上回ることができるのか？
• RQ2 SAM は、畳み込みを用いないアーキテクチャの損失地形、Hessian の曲率、そして学習可能性にどのように影響するのか？
• RQ3 SAM に伴う内在的なアーキテクチャの変化（例：スパーシティ、重みノルム）は何で、それらが一般化と頑健性とどう関連するのか？
• RQ4 ImageNet-R および ImageNet-C において、SAM で訓練された ViTs および Mixers は、ResNets と比較して精度と頑健性の観点でどのように性能を示すか？

主な発見

• SAM は、ViTs および MLP-Mixers をゼロから訓練した際の精度と頑健性を著しく向上させる（例：ImageNet の ViT-B/16 および Mixer-B/16 の top-1 大幅向上）。
• SAM なしでは ViTs および MLP-Mixers ははるかに鋭い局所極小値に収束し、大きな Hessian 固有値を持つが、SAM は損失地形を滑らかにし、平均および最悪ケースの曲率を低減する。
• SAM は初期層で活性化ニューロンをよりスパースにし、特に MLP-Mixers でそれが、滑らかな損失幾何学とより良い一般化と関連している。
• ImageNet でゼロから訓練された ViTs は、SAM を用いると同程度またはより大きなサイズの ResNets を、同等のスループットで上回すことができる。大規模な事前学習や強い拡張を必要としない。
• SAM は、クリーンな精度を改善するよりも、腐敗（ImageNet-C）や分布シフト（ImageNet-R）に対する頑健性を向上させる（ベースラインモデルと比較して）。
• SAM は、ViTs を SGD（モーメンタム）で訓練して Adam/SAM の組み合わせに近い性能を達成させ、最適化手法間のギャップを狭める。