Transformerの本質的欠陥を解決する：復旦大学らが提案した線形複雑性SOFT

[[437909]]

Visual Transformer (ViT) は、パッチ単位の画像トークン化と自己注意メカニズムの助けを借りて、さまざまな視覚認識タスクで SOTA を達成しました。ただし、自己注意モジュールを使用すると、Transformer のようなモデルの空間計算量と時間計算量は両方とも O(n^2) になります。自然言語処理の分野の研究者は、自己注意計算を線形複雑度に近づけるためにさまざまな試みを行ってきました。

最近、復旦大学、サリー大学、Huawei Noah's Ark Lab の研究者らが研究で詳細な分析を行い、これらの試みは理論上欠陥があるか、実験では視覚認識に効果がないことを示しました。さらに、これらの方法の限界は、近似プロセス中にソフトマックス自己注意が維持されることであることがわかりました。具体的には、従来の自己注意は、ラベル付けされた特徴ベクトル間のスケールされたドット積を正規化することによって計算されます。このソフトマックス演算を維持すると、Transformer の線形化の複雑さが軽減されます。これを踏まえて、本研究では初めてソフトマックスフリートランスフォーマー（SOFT）を提案した。

自己注意におけるソフトマックスを除去するために、ドット積類似度の代わりにガウスカーネル関数が使用され、それ以上の正規化は行われません。これにより、低ランク行列分解を介して完全な自己注意行列を近似することができます。近似の堅牢性は、ニュートン・ラプソン法を使用してムーア・ペンローズ逆行列を計算することによって実現されます。 ImageNet での広範な実験により、SOFT が既存の ViT バリアントの計算効率を大幅に向上させることが実証されています。重要なのは、線形複雑度の場合、SOFT ではより長いトークン シーケンスが許可され、精度と複雑度の間で優れたトレードオフが得られることです。

• 論文アドレス: https://arxiv.org/abs/2110.11945
• プロジェクトアドレス: https://github.com/fudan-zvg/SOFT

Transformer モデルには、計算とメモリ使用量の 2 次複雑性というボトルネックがあります。これは、自己注意メカニズムの本質的な機能です。トークンのシーケンス (単語や画像パッチなど) を入力として与えると、自己注意モジュールは 1 つのトークンを他のすべてのトークンに関連付けることで、特徴表現を反復的に学習します。これにより、推論中に n × n の注意行列を計算して保存する必要があるため、長さ n のトークン シーケンスに対して計算 (時間) とメモリ (スペース) の両方で O(n2) の二次複雑度が発生します。この問題は視覚において特に深刻です。適度な空間解像度であっても、2D 画像のトークン化によって NLP よりもはるかに長いシーケンスが生成されることがあります。したがって、この二次の複雑さにより、ViT モデルは、視覚認識タスクにとって非常に重要な、高い空間解像度を持つ画像をモデル化することができません。

自然な解決策は、近似によって自己注意計算の複雑さを軽減することです。実際、NLPでは多くの試みがなされてきました[33、5、18、38]。例えば[33]は学習可能な予測を通じてキーと値の長さを短くするという単純なアプローチを採用した。この大まかな近似は必然的にパフォーマンスの低下につながります。対照的に、[5, 17]は両方ともカーネルメカニズムを利用してソフトマックス正規化を近似し、自己注意における計算を線形化します。 [18]は代わりにハッシュ戦略を採用し、最も類似したペアを選択的に計算した。最近、[38]はNyström行列分解を用いて多項式反復を通じて完全な注意行列を再構築し、ランドマーク行列の擬似逆行列を近似した。

ただし、ソフトマックス正規化は、行列分解プロセス中に単純に繰り返されるため、理論的には不適切です。この研究では、これらのアプローチはいずれも視覚に適用した場合効果的ではないことが実験的に判明しました (セクション 4.2 を参照)。この研究では、既存の効率的なTransformerの限界がソフトマックス自己注意の使用によって引き起こされていることを発見し、初めてソフトマックスフリーのTransformerを提案しました。より具体的には、既存のすべてのトランスフォーマー（線形化の有無にかかわらず）では、トークン特徴ベクトル間のスケールされたドット積の上にソフトマックス正規化が必要です。このソフトマックス操作を維持すると、その後の線形化が困難になります。

この障害を克服するために、本研究では、空間と時間における線形複雑度がO(n)であるSOFTという新しいソフトマックスフリーの自己注意メカニズムを提案しました。具体的には、SOFT はガウスカーネルを使用して類似度 (自己注意) 関数を定義し、その後のソフトマックス正規化を必要としません。このソフトマックスフリーの注意行列を用いて、この研究ではさらに、それを近似するための新しい低ランク行列分解アルゴリズムを導入しました。ニュートン・ラプソン法を用いて行列のムーア・ペンローズ逆行列を確実に計算することにより、近似の堅牢性を理論的に保証できます。

この研究の主な貢献は次のとおりです。

• 線形空間および時間計算量を持つ新しいソフトマックスフリー Transformer が提案されています。
• この研究における注目行列近似は、理論的保証を備えた新しい行列分解アルゴリズムによって達成されます。
• 視覚認識タスクにおけるこの方法のパフォーマンスを評価するために、SOFT をコア自己注意コンポーネントとして使用し、さまざまな機能を備えた一連の一般的なバックボーン アーキテクチャを設計しました。広範囲にわたる実験により、SOFT モデルは線形複雑度でより長い画像トークン シーケンスを入力として受け取ることができることが示されています (図 1b)。その結果、SOFTは、同じモデルサイズでの精度と複雑さのトレードオフの点で、ImageNet [9]分類における最先端のCNNおよびViTバリアントよりも優れています（図1a）。

下の図 2 にモデルの概略図を示します。

図 2: 提案されたソフトマックスフリー自己注意 (SOFT) 法の概略図。 PE: 位置埋め込み。破線: 線形投影。 dh: 各アテンション ヘッドの非表示の dim。 ◦ 行列のドット積を表します。

著者らは2つの実験設定を採用した。最初の設定では、公平な比較のために、この研究ではすべての方法に同じ Tiny (表 2) アーキテクチャを使用します。つまり、SOFT のコアとなる自己注意ブロックは、各ベースライン独自の注意ブロックに置き換えられますが、アーキテクチャの残りの部分は変更されません。 [35]のスペース削減モジュールはLinformer [34]の特殊なケースであることに注意する。研究者らは、削減率をこの方法と同じに設定した。この研究では、同じ統一されたサンプリングの考え方を使用して、Nyströmformer の 1D ウィンドウ平均化 (NLP タスク用) を 2D 平均プーリング (画像用) に置き換えました。ダウンサンプリング比は、本研究の方法と一致した。また、Reformer[19]には公式のコードリリースがなく、Locality Sensitive Hashing（LSH）モジュールには入力トークンの長さに関する厳しい要件があるため、この方法は本研究の比較には含まれていないことも言及する価値がある。

下の表 1 から、次のことがわかります。

• Tiny アーキテクチャの Transformer と比較すると、Linear Transformer アプローチでは、同様のパラメータ サイズを維持しながら、メモリと FLOP が大幅に削減されます。
• SOFT 法は、すべての線形化方法の中で最も優れた分類精度を達成しました。
• この方法の推論速度は他の線形トランスフォーマーと同等ですが、トレーニング速度は Nystromformer よりもわずかに遅く、Performer および Linformer よりも遅くなります。

研究者らは、モデルのトレーニング速度が遅いのは主にニュートン・ラプソン反復法によるもので、ムーア・ペンローズ逆行列の精度を確保するためには、これを順番に適用することしかできないと指摘している。要約すると、同等の推論速度を考慮すると、トレーニング コストの増加は優れた精度の価値があると研究者は考えています。

この研究では、最先端の代替手段と比較し、ImageNet-1K 検証セットでトップ 1 の精度を報告しています。 FLOP 計算バッチ サイズは 1024 です。図1aと表3から、次のことがわかります。(i) 全体的に、ViTとそのバリアントはCNNよりも優れた分類精度を生み出します。 （ii）本研究は、ViT [11]やDeiT [31]、最先端のCNN RegNet [26]などの最近の純粋な視覚Transformerベースの手法の中で最高のパフォーマンスを達成しました。 （iii）SOFTは、アーキテクチャ構成において最も類似したTransformerのPVT [35]をすべてのバリエーションで上回っています。注意モジュールが主な違いであるため、これによりモデルの有効性が直接検証されます。 (iv) 提案された方法は、ViT の効率性の限界に対処しながら、より少ないパラメータと浮動小数点計算を必要とする最先端の ViT バリアントである Twins よりも優れています。

SOFT と代替方法を使用して注意がどのように学習されるかについての洞察を提供するために、図 3 はさまざまな比較モデルの注意マスクを示しています。各モデルについて、この論文では最初の 2 つのアテンション ヘッドの出力を示します。 SOFT は、ピクセル間のローカルな関係と長距離の関係の両方をキャプチャする際に堅牢性と汎用性を発揮することは明らかです。興味深いことに、SOFTはImageNet[9]のオブジェクト分類データセットでトレーニングされているにもかかわらず、同じカテゴリ内のインスタンス間で共有される意味概念とインスタンス固有の特徴の両方を学習できるようです。

興味のある読者は、研究の詳細について原著論文を読むことができます。