[論文レビュー] How to train your ViT? Data, Augmentation, and Regularization in Vision Transformers

Andreas Steiner, А. И. Колесников+4 more|arXiv (Cornell University)|Jun 18, 2021

Advanced Neural Network Applications被引用数 228

ひとこと要約

本論文は、データ量、 augmentation、正則化、モデルサイズ、計算予算が Vision Transformers にどのように影響し合うかを系統的に検討し、AugReg が大規模データの代替となり得ること、上流データセットを大きくすることからの転移学習がしばしばよりコスト効果が高いことを示している。

ABSTRACT

Vision Transformers (ViT) have been shown to attain highly competitive performance for a wide range of vision applications, such as image classification, object detection and semantic image segmentation. In comparison to convolutional neural networks, the Vision Transformer's weaker inductive bias is generally found to cause an increased reliance on model regularization or data augmentation ("AugReg" for short) when training on smaller training datasets. We conduct a systematic empirical study in order to better understand the interplay between the amount of training data, AugReg, model size and compute budget. As one result of this study we find that the combination of increased compute and AugReg can yield models with the same performance as models trained on an order of magnitude more training data: we train ViT models of various sizes on the public ImageNet-21k dataset which either match or outperform their counterparts trained on the larger, but not publicly available JFT-300M dataset.

研究の動機と目的

• トレーニングデータサイズ、augmentation、正則化が ViTs のモデルサイズ別でどのように相互作用するかを理解する。
• ViTs の事前学習データ量と下流タスクへの転移のための計算資源とデータのトレードオフを定量化する。
• AugReg を用いた下流データセットの転移学習とゼロからのトレーニングを、多様なデータセットで評価する。
• 事前学習済みチェックポイントの選択とトレーニング戦略に関する実践的推奨を特定する。

提案手法

• 複数サイズ（Ti、S、B、L）およびResNet+ViT ハイブリッドを ImageNet-1k および ImageNet-21k で統一的な設定の下に TPU/JAX/Flax で事前学習する。
• 正則化（Dropout、 stochastic depth）と augmentation（Mixup、 RandAugment）を 28 構成のハイパーパラメータスイープで適用する。
• 下流データセットで SGD とコサイン学習率スケジュールを用いてファインチューニングを行い、ゼロからの学習と上流モデルからの転移を比較する。
• データ規模、計算予算、 augmentation の強度を体系的に変化させ、性能の景観をマッピングする。

実験結果

リサーチクエスチョン

• RQ1 データサイズ、 augmentation、正則化が ViT の性能をモデルサイズ別でどのように影響するか？
• RQ2 AugReg を用いた通常の下流タスクで、上流データセットの転移学習はゼロからのトレーニングよりコスト効果が高いか？
• RQ3 事前学習データを増やすことと augmentation の増加が VTAB タスク全体で転移性能にどう影響するか？
• RQ4 計算制約の下で事前学習済みチェックポイントとトレーニング戦略を選択する際の実用的な指針は何か？

主な発見

• AugReg は、同等のタスクでのトレーニングデータを 10 倍増やす効果に匹敵する。
• 実用的なデータセットでは、事前学習済みモデルからの転移学習の方が、ゼロからの AugReg 学習より総じてコスト効率が高く、より良い結果をもたらす。
• より多くの事前学習データは、さまざまな VTAB タスクでより転移しやすい一般的なモデルを生み出す。
• augmentation は中規模データセットでより効果的だが、非常に大規模な上流データ（例：ImageNet-21k、固定計算）では AugReg は計算を増やさない限りしばしば害となる。
• 転移の際には上流の検証精度で上流モデルを選ぶことが通常は効果的である。ImageNet-21k のチェックポイントを使用することが推奨される。
• パッチサイズを 32 に増やすと、同程度の速度の薄いモデルよりも優れた性能を発揮できる可能性があり、パッチサイズはパラメータ数を超える重要な容量要因である。