[論文レビュー] SOFT: Softmax-free Transformer with Linear Complexity

Jiachen Lu, Jinghan Yao+7 more|arXiv (Cornell University)|Oct 22, 2021

Advanced Neural Network Applications参考文献 45被引用数 61

ひとこと要約

SOFTは、ソフトマックスなしの自己注意機構をガウスカーネルとNyströmベースの低ランク近似、 Newton-Raphsonを介したムーア・ペンローゼ逆行列を用いて導入し、線形時間/空間計算量を達成するとともに、他の線形トランスフォーマーよりImageNetの精度を改善します。

ABSTRACT

Vision transformers (ViTs) have pushed the state-of-the-art for various visual recognition tasks by patch-wise image tokenization followed by self-attention. However, the employment of self-attention modules results in a quadratic complexity in both computation and memory usage. Various attempts on approximating the self-attention computation with linear complexity have been made in Natural Language Processing. However, an in-depth analysis in this work shows that they are either theoretically flawed or empirically ineffective for visual recognition. We further identify that their limitations are rooted in keeping the softmax self-attention during approximations. Specifically, conventional self-attention is computed by normalizing the scaled dot-product between token feature vectors. Keeping this softmax operation challenges any subsequent linearization efforts. Based on this insight, for the first time, a softmax-free transformer or SOFT is proposed. To remove softmax in self-attention, Gaussian kernel function is used to replace the dot-product similarity without further normalization. This enables a full self-attention matrix to be approximated via a low-rank matrix decomposition. The robustness of the approximation is achieved by calculating its Moore-Penrose inverse using a Newton-Raphson method. Extensive experiments on ImageNet show that our SOFT significantly improves the computational efficiency of existing ViT variants. Crucially, with a linear complexity, much longer token sequences are permitted in SOFT, resulting in superior trade-off between accuracy and complexity.

研究の動機と目的

• 視覚における長いトークン列を扱う際の二乗計算量を回避するため、効率的なTransformerの必要性を動機づける。
• ソフトマックスなしの自己注意機構を提案し、線形計算量を実現する。
• Newton-Raphsonで求めたムーア・ペンローゼ逆行列を用いたNyströmベースの低ランク近似を開発し、堅牢な注意を実現する。
• ImageNet上でSOFTベースのバックボーンを設計・評価し、精度–計算量のトレードオフを評価する。

提案手法

• ソフトマックスベースの注意を、対称で [0,1] にあるガウスカーネルアテンション S = exp(Q ⊖ K) に置換する。
• 小さなボトルネック m を用いたNyström分解により、完全なアテンション行列を線形時間/空間計算量で近似し、Ŝ = Pᵀ A† P とする。
• 数値的なロバスト性を確保するため、Newton–Raphson iterations を用いてムーア・ペンローゼ逆 A† を計算する（A₀ = αA, A_{k+1} = 2A_k − A_k A A_k）。
• 畳み込み、平均プーリング、その他のサンプリング法を用いてボトルネックトークンをサンプルする。安定性と効率の観点から平均プーリングを好む。
• SOFTをピラミッド型ビジョントランスフォーマーのバックボーン内の層として実装し、d_e, h, n, m, sp などの特定ハイパーパラメータを用いてSOFTの派生を構築する。

実験結果

リサーチクエスチョン

• RQ1 Vision Transformersにおいてsoftmax-freeなガウスカーネルアテンションはsoftmaxベースのアテンションと同等の精度を達成できるか。
• RQ2 Nyströmベースの低ランク近似とムーア・ペンローゼ逆行列は視覚タスクで安定した学習と線形計算量を提供するか。
• RQ3 SOFTの精度–効率のトレードオフを最適化する設計選択（ボトルネックサイズ m、サンプリング戦略）は何か。
• RQ4 ImageNetおよびNLPベンチマークでSOFTは他の線形/効率的なトランスフォーマーとどう比較されるか。
• RQ5 SOFTを視覚認識のためのピラミッド型トランスフォーマー骨格に組み込むことによって、アーキテクチャ上の利点は何か。

主な発見

• SOFTはアテンションの計算量を線形時間・空間に抑え、長いトークン列を扱える。
• Newton–Raphson によるムーア・ペンローゼ逆を用いたNyström分解は堅牢なアテンション近似を提供する。
• ImageNet 上でSOFTベースのバックボーンは、精度–計算量のトレードオフにおいていくつかのCNNおよびViT系統より優れた性能を示す。
• SOFTはLinformer、Performer、Nyströmformerと同程度の設定下で競争力があるか、あるいはそれを上回る精度を示す。
• アブレーションによりボトルネックサイズ m ≈ 49 が精度と計算量の良いバランスを提供することが分かった。平均プーリングサンプリングが tested methods の中で最良の性能を示す。