[論文レビュー] An Attention Free Transformer

Shuangfei Zhai, Walter Talbott+5 more|arXiv (Cornell University)|May 28, 2021

Advanced Neural Network Applications参考文献 26被引用数 42

ひとこと要約

Attention Free Transformer (AFT) を紹介。dot-product self-attention を非-attention、線形複雑度モジュールに置換し、AFT-local および AFT-conv のような派生も全球的な連結性を局所性や畳み込み重み共有で維持します。

ABSTRACT

We introduce Attention Free Transformer (AFT), an efficient variant of Transformers that eliminates the need for dot product self attention. In an AFT layer, the key and value are first combined with a set of learned position biases, the result of which is multiplied with the query in an element-wise fashion. This new operation has a memory complexity linear w.r.t. both the context size and the dimension of features, making it compatible to both large input and model sizes. We also introduce AFT-local and AFT-conv, two model variants that take advantage of the idea of locality and spatial weight sharing while maintaining global connectivity. We conduct extensive experiments on two autoregressive modeling tasks (CIFAR10 and Enwik8) as well as an image recognition task (ImageNet-1K classification). We show that AFT demonstrates competitive performance on all the benchmarks, while providing excellent efficiency at the same time.

研究の動機と目的

• スケールする Transformers が必要になる理由を、二次の注意コストのため動機づける。
• ドット積アテンションなしでグローバルな相互作用を維持するプラグイン型のアテンションフリーモジュールを提案する。
• 局所性と空間的重み共有を活用しつつグローバルな連結性を維持する AFT-local および AFT-conv を導入する。
• 画像自回帰モデリング、言語モデリング、画像分類の分野で競争力のある性能と効率を示す。

提案手法

• 標準のマルチヘッド注意を、Q, K, V の間の相互作用に置換し、KとVは学習された位置バイアスと結合され、結果はクエリと要素ごとに乗算される。
• AFT を Y_t = sigma_q(Q_t) ⊙ [ sum_{t'} exp(K_{t'}+w_{t,t'}) ⊙ V_{t'} / sum_{t'} exp(K_{t'}+w_{t,t'}) ] に定式化する、w は学習されたペアワイズ位置バイアス。
• パラメータを削減するために w_{t,t'} = u_t^T v_{t'} の因子化パラメータ化を提供。
• 局所性を持つスライド窓を用いた AFT-local (local w with s-bounded window) および空間的重み共有を持つ AFT-conv (CNNs に類似) 変種を導入。
• 位置バイアスなしの非常に高効率な極端な変種である AFT-simple を提供 (s=0)。
• 訓練の安定性と性能を改善するためのパラメータ化と再パラメータ化（例: w- factorization, w normalization）について議論。

実験結果

リサーチクエスチョン

• RQ1 アテンションフリーの機構は自回帰画像モデリング、言語モデリング、画像分類のタスクで Transformers と同等またはそれを上回ることができるか。
• RQ2 局所性と重み共有の変種（AFT-local, AFT-conv）は、グローバルな連結性を保ちながら効率と精度で実用的な利点を提供するか。
• RQ3 位置バイアスパラメータの因子化はモデルの性能とパラメータ数にどのような影響を与えるか。
• RQ4 AFT の変種と標準および他の効率的な Transformer との比較における計算量と実用的なランタイムはどうか。
• RQ5 AFT 変種は従来の Transformer アーキテクチャにプラグインとして統合できるか。

主な発見

モデル	L	d	h	訓練損失	テスト損失	イテレーション/秒	GB/GPU
PixelCNN	-	-	-	3.08	3.14	-	-
PixelCNN++	-	-	-	-	2.92	-	-
PixelSNAIL	-	-	-	-	2.85	-	-
Sparse Transformer strided	128	256	2	-	2.80	-	-
Image Transformer local2d	12	512	4	-	2.90	1.61	22.3
Transformer	12	512	4	2.90	2.88	1.35	30.6
Transformer	24	256	2	2.90	2.86	1.36	30.4
AFT-local-256	12	512	1	2.78	2.80	1.68	11.4
AFT-local-256	24	256	1	2.75	2.74	1.67	12.8
AFT-simple	24	256	1	2.82	2.89	2.15	9.5

• AFT 変種は CIFAR10、Enwik8、ImageNet-1K において競争力のある性能を達成しつつ、効率性にも有利な結果を示す。
• AFT-local は CIFAR10 の画像自回帰モデリングにおいて、速度とメモリ使用の両面で複数の Transformer ベースラインを上回る。
• 因子化された位置バイアスパラメータ化 は、非因子化された w と比較してパラメータを大幅に削減し性能を改善。
• AFT-simple はさらに高い効率で強力な結果を示し、いくつかのビジョン Transformer ベースラインに近づくかそれを上回る。
• AFT-conv は空間的重み共有を CNN 的な形で行い、グローバルな連結性を維持しつつ ImageNet Top-1 精度をベースライン構成より大幅に改善。
• タスク全体で、AFT 変種はシーケンス長と特徴次元に対して線形またはほぼ線形の時間/空間挙動を示し、標準のアテンションの二次コストとは対照的である。