[论文解读] Transformer Quality in Linear Time

Weizhe Hua, Zihang Dai+2 more|arXiv (Cornell University)|Feb 21, 2022

Topic Modeling被引用 45

一句话总结

FLASH 引入门控注意单元（Gated Attention Unit，GAU）和混合块注意力，以实现线性时间注意力，同时保持 Transformer-like 质量，在长上下文语言模型任务上显著提升训练速度。

ABSTRACT

We revisit the design choices in Transformers, and propose methods to address their weaknesses in handling long sequences. First, we propose a simple layer named gated attention unit, which allows the use of a weaker single-head attention with minimal quality loss. We then propose a linear approximation method complementary to this new layer, which is accelerator-friendly and highly competitive in quality. The resulting model, named FLASH, matches the perplexity of improved Transformers over both short (512) and long (8K) context lengths, achieving training speedups of up to 4.9$ imes$ on Wiki-40B and 12.1$ imes$ on PG-19 for auto-regressive language modeling, and 4.8$ imes$ on C4 for masked language modeling.

研究动机与目标

• 通过实现可扩展、高效的注意力来解决 Transformer 的长上下文限制。
• 开发一种新的层（门控注意单元，Gated Attention Unit）在较弱的注意力下保持高质量。
• 提出一种线性时间注意近似（混合块注意力），将局部二次注意力与全局线性分量整合。
• 证明 FLASH 在大规模数据集（Wiki-40B、PG-19、C4）上可匹配或接近 Transformer++ 的质量，同时提供显著的训练加速。
• 提供广泛的消融与与现有线性和二次注意力模型的对比分析。

提出的方法

• 引入门控注意单元（GAU），其中 O = (U ⊙ ĤV)W_o 且 ĤV = AV，A 由门控机制计算，可使用简单的单头注意力。
• 表明 GAU 即使在较弱的注意力下也能通过用注意力获取的 V 来门控 U，使其能达到 Transformer 的质量。
• 开发混合块注意力：将序列分割为不重叠的块；在每个块内应用局部二次注意力，在块之间应用全局线性注意力（因果或非因果），并通过门控将结果结合。
• 在 FLASH 中表示 GAU，结合 GAU 块与混合块注意力，以实现线性复杂度且质量损失最小。
• 提供伪代码并讨论利于加速器实现的实现，重点在于最小化内存重新格式化并在实践中保持线性扩展。
• 在 C4、Wiki-40B 和 PG-19 上进行广泛的双向和自回归语言建模实验，模型规模约为 1.10 亿到 5.00 亿参数。

实验结果

研究问题

• RQ1 GAU 是否能够使用较弱的单头注意力实现 Transformer 级别的质量？
• RQ2 混合块注意力是否能够在保持高质量表现的同时实现线性时间复杂度以适应长序列？
• RQ3 在 MLM 和 LM 任务中，FLASH-Quad（二次）与 FLASH（线性）与 Transformer++ 等线性注意力方法在不同上下文长度下有何比较？
• RQ4 在大规模数据集上将上下文长度从 512 扩展到 8K 时，速度提升与质量权衡如何？
• RQ5 在下游长上下文任务（如 TriviaQA）的微调中，基于 GAU 的 FLASH 相对于 Transformer++ 与基线表现如何？

主要发现

• GAU 使用较小的单头注意力即可达到 Transformer 级别的质量。
• FLASH 实现线性可扩展性并带来显著的训练加速：在 Wiki-40B 的自回归语言模型上最高可达 4.9×，在 PG-19 上达到 12.1×，在 MLM 的 C4 上达到 4.8×。
• FLASH-Quad 与 FLASH 在 Wiki-40B 的上下文长度 512–8192 下提供更低的延迟和有竞争力的困惑度，对于较长上下文（如 PG-19）有更大提升。
• 与 Transformer++ 相比，FLASH 在 PG-19 上将训练成本降低最多 12.1×，并在各任务上实现有利的质量-速度权衡。
• 消融研究表明局部二次与全局线性分量都至关重要；GAU 相较于 MC-TFM++（Transformer++ 的混合块变体）在显著性上具有优势。
• 块大小与设计选择（如先验到全 token 的选项）影响微调性能与速度，一般而言较大块对长上下文更有利。