[論文レビュー] Replacing softmax with ReLU in Vision Transformers

Mitchell Wortsman, Jaehoon Lee+2 more|arXiv (Cornell University)|Sep 15, 2023

Advanced Neural Network Applications被引用数 8

ひとこと要約

本論文は Vision Transformer の attention における softmax を ReLU に置換することを調査し、シーケンス長スケーリングを用いることで ReLU-attention が ImageNet-21k でのスケーリング性能で softmax-attention に匹敵しうることを示し、並列化の改善も提供する。

ABSTRACT

Previous research observed accuracy degradation when replacing the attention softmax with a point-wise activation such as ReLU. In the context of vision transformers, we find that this degradation is mitigated when dividing by sequence length. Our experiments training small to large vision transformers on ImageNet-21k indicate that ReLU-attention can approach or match the performance of softmax-attention in terms of scaling behavior as a function of compute.

研究の動機と目的

• vision transformers における softmax attention の点ごとの代替案を探索する。
• ReLU ベースの attention をシーケンス長でスケーリングした場合、softmax の性能に匹敵するかを評価する。
• 精度と訓練効率に対するシーケンス長スケーリングと qk-layernorm の影響を評価する。
• ReLU-attention におけるゲーティングと代替活性化の効果を分析する。
• 大規模な非-softmax attention を用いた Vision Transformers の訓練に関する実践的指針を提供する。

提案手法

• phi = L^{-alpha} h として attention を定式化し、α ~ 1 の ReLU に焦点を当てる。
• ReLU-attention を phi = L^{-1} ReLU と定義し、softmax attention と比較する。
• 精度を評価するために様々な非線形性とシーケンス長スケーリングを実験する。
• ハイパーパラメータを変更せず BigVision の ImageNet-21k および ImageNet-1k の訓練設定を用いる。
• qk-layernorm の除去の影響とゲーティッド attention ユニットの影響を評価する。

実験結果

リサーチクエスチョン

• RQ1 vision transformer の attention において softmax を ReLU に置換し、シーケンス長でスケーリングした場合、同等の精度を得られるか。
• RQ2 シーケンス長スケーリング（L^{-alpha}）は、異なる非線形性でどのように性能に影響するか。
• RQ3 大規模での ReLU-attention における安定性と精度における qk-layernorm の役割は何か。
• RQ4 gating メカニズムを追加することで ReLU-attention のシーケンス長スケーリングの必要性が減るか。
• RQ5 これらの置換は訓練計算量と並列化にどのような影響を与えるか。

主な発見

• L^{-1} スケーリングを用いた ReLU-attention は ImageNet-21k の訓練において softmax-attention のスケーリング傾向と一致できる。
• alpha が概ね 1 に近いシーケンス長スケーリングは、 tested モデル全体で一般に最良の結果を示す。
• qk-layernorm は試験されたスケールでは大きな影響を与えないが、スケールに応じて影響が異なる可能性がある。
• ゲーティングはシーケンス長スケーリングの利点を排除しない；最良の精度は依然として L^{-alpha} スケーリングで得られる。
• ReLU-attention を用いると、従来の softmax attention よりも少ない gather 操作で sequence length に対する並列化が可能になる。