[論文レビュー] Early Convolutions Help Transformers See Better

Tete Xiao, Mannat Singh+4 more|arXiv (Cornell University)|Jun 28, 2021

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 より劣る理由を動機付けて理解する。
• 初期視覚処理（patchify stem）が ViT の最適化問題を引き起こしているかを調査する。
• 最小限の畳み込み stem をドロップイン置換として評価し、元の ViT stem とのトレーニングダイナミクスを比較する。
• 多様なモデルサイズとデータ規模での収束速度、オプティマイザの適合性、ハイパーパラメータの安定性、ピーク性能を評価する。

提案手法

• ViT patchify stem（stride-p, p×p）を約 5 層の軽量畳み込み stem に置換する。
• FLOPs と実行時間の parity を保つために1つのトランスフォーマーブロックを削除する。
• ViT P（元の patchify stem）と ViT C（畳み込み stem）を複数サイズ（1G–36G FLOPs）および ImageNet スケール（1k–21k）で訓練する。
• エラー分布関数（EDF）を用いて、トレーニング長の安定性、オプティマイザの安定性（AdamW vs SGD）、学習率・ウェイトデカイ安定性を評価する。
• EMA 重みと標準的なデータ拡張レシピを用いた制御されたトレーニング設定下で、ピーク性能を比較する。

実験結果

リサーチクエスチョン

• RQ1 パッチファイ stem を畳み込み stem に置換することは、ViT の最適化安定性と収束速度に影響を与えるか？
• RQ2 畳み込み stem を持つ ViT は、精度の大きな低下を招くことなく SGD でも AdamW でも効果的に訓練できるか？
• RQ3 畳み込み stem は、patchify 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 より top-1 誤差を約 1–2 ポイント低くする。
• 公正なピーク性能の比較の下で、ViT C は中程度の複雑さの範囲で ViT P および競合する CNN を上回る；ImageNet-21k の事前訓練では ViT C がすべてを上回る。