[论文解读] 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 融入基于 GAU 的单头门控注意力框架，其中查询、键和值来自经 EMA 增强的表示。
• 采用 Laplace 基于的注意力函数作为对 softmax 的替代，以适应不同模态。
• 在 Mega 模块内使用门控残差连接和归一化步骤以稳定训练。
• 提出 Mega-chunk，通过对固定长度的区块进行处理来实现线性时间/空间复杂度，EMA 提供跨区块的上下文。
• 在任务间对比注意力函数（softmax、relu^2、laplace），以为语言、视觉和语音选择合适选项。

实验结果

研究问题

• RQ1 将 EMA 集成到门控注意力中是否能相比标准注意力更好地建模局部和长期依赖？
• RQ2 带 EMA 的单头门控注意力是否能够达到与多头注意力相当的表达能力？
• RQ3 Mega-chunk 变体在降低计算成本到线性复杂度的同时是否保持了性能？
• RQ4 Mega 在文本、图像、语音等多模态及 LRA、MT、LM、Vision、Speech 等任务上的表现如何？
• RQ5 哪些注意力函数最有利于跨领域的 Mega？

主要发现

模型	LRA（准确度↑）	WMT16（BLEU↑）	WT103（PPL↓）	ImageNet（准确度↑）	SC（准确度↑）	平均值（↑）	速度	内存
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 在 top-1 准确度方面大约比 DeiT-B 高出0.5%。