[論文レビュー] Mega: Moving Average Equipped Gated Attention

Xuezhe Ma, Chunting Zhou+6 more|arXiv (Cornell University)|Sep 21, 2022

Topic Modeling被引用数 36

ひとこと要約

Megaは指数移動平均を組み込んだゲーティッド注意機構を単一ヘッドに統合し、線形複雑度のMega-chunk変種を提供して、複数モダリティにおける長距離系列タスクで強力な性能を発揮します。

ABSTRACT

The design choices in the Transformer attention mechanism, including weak inductive bias and quadratic computational complexity, have limited its application for modeling long sequences. In this paper, we introduce Mega, a simple, theoretically grounded, single-head gated attention mechanism equipped with (exponential) moving average to incorporate inductive bias of position-aware local dependencies into the position-agnostic attention mechanism. We further propose a variant of Mega that offers linear time and space complexity yet yields only minimal quality loss, by efficiently splitting the whole sequence into multiple chunks with fixed length. Extensive experiments on a wide range of sequence modeling benchmarks, including the Long Range Arena, neural machine translation, auto-regressive language modeling, and image and speech classification, show that Mega achieves significant improvements over other sequence models, including variants of Transformers and recent state space models.

研究の動機と目的

• 強い帰納的バイアスと柔軟な注意機構を組み合わせることで、長距離系列モデリングの改善を動機づける。
• 指数移動平均（EMA）によって強化された単一ヘッドのゲーティッド注意機構を開発する。
• スケーラブルな系列処理のための線形時間のMega-chunk変種を導入する。
• 多様なモダリティとベンチマークにおけるMegaの有効性を示す。
• 単一ヘッドのゲーティッド注意の表現力に関する理論的根拠を提供する。

提案手法

• 局所的で時間とともに減衰する依存性を捉えるため、多次元の減衰EMAを導入する。
• EMAを、EMA強化表現に由来するクエリ、キー、値を用いたGAUベースの単一ヘッドゲーティッド注意フレームワークに組み込む。
• さまざまなモダリティに対して、softmaxの代替としてラプラスベースの注意関数を採用する。
• Megaブロック内でゲート付き残差接続と正規化手順を用いて訓練を安定化させる。
• 固定長チャンクをEMAがチャンク間の文脈を提供することで処理することにより、線形時間/空間計算量を達成するMega-chunkを提案する。
• 言語、視覚、音声のタスク間で注意関数（softmax、relu^2、laplace）を比較し、適切なオプションを選択する。

実験結果

リサーチクエスチョン

• RQ1 標準の注意機構と比較して、ゲーティッド注意にEMAを組み込むことで局所および長距離の依存性のモデリングが改善されるか？
• RQ2 EMAを組み込んだ単一ヘッドのゲーティッド注意は、多頭注意に匹敵する表現力を達成できるか？
• RQ3 Mega-chunkは計算コストを線形複雑度に削減しつつ性能を維持できるか？
• RQ4 Megaは多様なモダリティ（テキスト、画像、音声）とタスク（LRA、MT、LM、Vision、Speech）でどのように性能を発揮するか？
• RQ5 ドメインを超えてMegaを最適に促進する注意関数はどれか？

主な発見

Model	LRA (Acc. ↑)	WMT16 (BLEU ↑)	WT103 (PPL ↓)	ImageNet (Acc. ↑)	SC (Acc. ↑)	Avg (↑)	Speed	Mem.
XFM	59.24	–	18.66	81.80	✗	–	–	–
XFM ${\u001d}$	37.11	65.21	79.14	42.94	71.83	59.24	–	–
Reformer	50.67	–	20.95	–	–	–	–	–
Linformer	51.36	–	52.27	–	–	–	5.5×	0.10×
BigBird	55.01	–	59.29	–	–	–	1.1×	0.30×
Performer	51.41	–	53.82	42.77	77.05	–	5.7×	0.11×
Luna-256	61.95	–	79.56	47.86	78.55	–	4.9×	0.16×
S4-v1	80.48	–	87.09	87.26	86.05	80.48	–	–
S4-v2	86.09	–	90.90	88.65	94.20	86.09	–	–
S4-v2 ${\u001d}$	85.86	–	90.94	88.48	94.01	85.86	4.8×	0.14×
Mega	88.21	29.18	18.07	82.31	97.30	88.21	2.9×	0.31×
Mega-chunk	85.66	90.19	90.97	85.80	94.41	85.66	5.5×	0.13×

• MegaはすべてのLong Range Arenaタスクおよび関連ベンチマークでTransformerおよびS4系を上回る。
• MegaはLRAで平均精度88.21を達成。S4系は86.09–86.35、XFMは59.24。線形時間変種で強い長距離モデリングを実証。
• Mega-(chunk)は大幅な速度とメモリ利得を提供。例：LRA Textタスクで通常のTransformerと比較して5.5倍速く、メモリ使用量は13％。
• 言語タスクではMega内のsoftmax注意が高い性能を示し、laplace注意は視覚・音声タスクで優れる。
• Mega-baseはTransformer-baseと比較してWMT’16英独翻訳でBLEUを1.1ポイント超改善。
• ImageNet-1kではMegaはDeiT-Bより約0.5％高いtop-1精度を達成。