Transformer Quality in Linear Time

发表时间: 2022-02 · arXiv:2202.10447

原文: https://arxiv.org/abs/2202.10447

速读

一句话结论 本文提出了名为 FLASH 的线性复杂度模型家族，它通过结合门控注意力单元与混合分块注意力机制，在保持与完全增强版 Transformer 同等生成质量的前提下，在现代加速器上实现了对长上下文的真正线性扩展与数倍的训练加速。

要解决什么问题 尽管 Transformer 在处理长序列时面临二次方复杂度的卡点，但现有的线性注意力机制在实际应用中普遍存在三个致命缺陷。首先是质量不佳，近似方法的性能通常远逊于经过简单调整的增强版 Transformer（如引入门控线性单元的 Transformer++）。其次是工程开销大，许多稀疏或局部注意力方法高度依赖切片、拼接等不规则的内存重格式化操作，这在 TPU 或 GPU 等高度并行的加速器上极不友好，导致理论提速无法兑现。最后也是最关键的卡点，自回归训练效率极低。在语言建模任务中，传统线性注意力为了满足因果约束，需要进行类似 RNN 的顺序累积求和。这引入了与序列总长度等同的顺序依赖步数，不仅需要极其频繁的内存访问，还彻底锁死了硬件的并行计算能力，导致其训练阶段甚至比直接计算完整的二次方注意力还要慢。

怎么做的 核心思路是将注意力与门控机制深度融合以降低计算门槛，并采用分治策略将长序列分块，块内使用精确注意力，块间使用线性注意力。该方法由两个关键部件构成。第一个部件是门控注意力单元（GAU），负责减轻注意力机制的负担。研究发现引入门控机制后，模型对注意力精确度的依赖会大幅降低。GAU 将门控线性单元与注意力机制合二为一，使用通过注意力从全局检索到的表示作为门控信号。其核心更新式为： $$ O = (U \otimes \hat{V}) W_o $$ 其中 $U$ 是输入经过线性变换的表示，$\hat{V}$ 是通过注意力矩阵聚合后的表示，$\otimes$ 代表逐元素相乘。得益于门控补偿，GAU 能安全使用极简的单头、无 softmax 的弱注意力，大幅削减参数量，为后续近似计算打下基础。第二个部件是混合分块注意力，负责实现硬件友好的线性复杂度。为绕开自回归训练中的顺序依赖卡点，该机制将长度为 $T$ 的序列划分为 $G$ 个大小为 $C$ 的不重叠块。在每个块内部，独立计算精确的局部二次方注意力捕捉短距离依赖： $$ \hat{V}_{quad_g} = A_{quad_g} V_g $$ 在块与块之间，采用全局线性注意力捕捉长程交互： $$ \hat{V}_{lin_g} = Q_{lin_g} \sum_{h=1}^g K_{lin_h}^T V_h $$ 最终将两部分相加代入 GAU。这种设计将传统的逐词顺序累积求和提升到了块级别，使顺序依赖步数从总长度 $T$ 骤降到块数量 $T/C$，极大释放了加速器的并行能力。同时，坚持使用不重叠的块划分彻底避免了内存重格式化操作，确保了几行代码的更改就能兑现真正的线性提速。

效果如何 实验在 64 个 TPU-v4 核心上进行（附录验证了单张 Nvidia V100 GPU），评估了 110M 和 500M 两种参数规模的模型。训练数据包括用于掩码语言建模的 C4、用于自回归语言建模的 Wiki-40B，以及平均文档长达 69K 词的 PG-19。对比基线涵盖代表精确二次方路线的原始 Transformer、Transformer+（加旋转位置编码）、最强基线 Transformer++（加门控线性单元），以及代表线性路线的 Performer 和代表分块路线的 Combiner。结果显示，FLASH 是唯一能在困惑度上媲美 Transformer++ 的线性复杂度模型。在 Wiki-40B 的自回归任务中，为达到与 Transformer++ 相同的困惑度，FLASH 将训练时间缩短了 1.2 到 4.9 倍。在极度依赖长程信息的 PG-19 数据集（8K 上下文、500M 参数）上，FLASH 更是以 12.12 倍的惊人加速比达到了 Transformer+ 的最终困惑度。在 TriviaQA 的微调任务中，FLASH 取得了与 Transformer+ 几乎一致的 F1 分数（仅差 0.3），同时速度提升 2.7 倍。关于代价与局限性，作者坦承，虽然让相邻块重叠能在理论上提升质量，但这会引入大量内存重格式化操作严重拖慢速度，因此被舍弃。此外消融实验表明，局部二次方和全局线性注意力缺一不可，去掉任一个都会导致性能断崖式下跌；且模型对超参敏感，处理更长上下文时必须相应增大块大小（如上下文超 1024 时，块大小 512 优于 256）才能维持最佳效果。

A1 主要贡献

核心问题: 尽管Transformer模型【35, Vaswani, A. 等人, Attention is all you need, NIPS 2017】在深度学习领域取得了巨大成功，但其在处理长序列时面临着二次方复杂度的限制。这一限制阻碍了模型处理长期信息的能力，而这在许多应用中至关重要。

现有方法缺陷: 许多旨在通过高效注意力机制加速Transformer的现有技术，虽然在理论上可能具有线性复杂度，但在实践中存在一个或多个以下缺陷：
1. 质量不佳 (Inferior Quality): 与经过简单调整后的增强版Transformer（本文称为Transformer++）相比，许多高效注意力方法存在显著的性能下降，这种下降超过了其效率优势。
2. 实践中的开销 (Overhead in Practice): 这些方法通常会使Transformer层变得复杂，并需要大量的内存重格式化操作，导致其在GPU或TPU等加速器上的理论复杂度与实际速度之间存在巨大差距。
3. 自回归训练效率低下 (Inefficient Auto-regressive Training): 大多数注意力线性化技术虽然在推理时解码速度快，但在自回归任务（如语言建模）的训练中却非常缓慢。这主要是因为它们需要在大量步骤上进行类似RNN的顺序状态更新，无法充分利用现代加速器的并行计算能力。

图1：在Wiki40B上进行自回归语言建模时，FLASH相对于vanilla Transformer (TFM)和增强型Transformer (TFM++)的TPU-v4训练加速比。所有模型大小均在110M左右，使用2^18个token的批次大小训练125K步。
图1：在Wiki40B上进行自回归语言建模时，FLASH相对于vanilla Transformer (TFM)和增强型Transformer (TFM++)的TPU-v4训练加速比。所有模型大小均在110M左右，使用2^18个token的批次大小训练125K步。

研究目标与创新点: 本文旨在开发一个名为FLASH的新模型家族，该模型首次在质量上与完全增强的Transformer相媲美，同时在现代加速器上真正实现了对上下文大小的线性可扩展性。其开发分为两步：

  1. 提出门控注意力单元 (Gated Attention Unit, GAU):

    • 首先，本文设计了一种名为GAU的新层，如图2所示。GAU通过引入门控机制来减轻自注意力的负担。
    • 相比Transformer层，GAU层计算成本更低，且其质量对注意力机制的精确度依赖性更小。
    • 实验证明，使用一个较小的单头、无softmax的注意力机制的GAU，其性能与标准Transformer相当。
    • GAU虽然仍具有二次方复杂度，但它削弱了注意力的作用，为后续进行低质量损失的近似计算创造了条件。

  2. 提出混合分块注意力 (Mixed Chunk Attention):

    • 其次，本文提出了一种高效的方法来近似GAU中的二次方注意力，从而得到一个具有线性复杂度的层变体。
    • 核心思想是将token分组为块（chunks），然后在块内使用精确的二次方注意力，在块间使用快速的线性注意力，如图4所示。
    • 这种方法可以自然地转化为加速器友好的实现，仅需几行代码更改即可在实践中实现线性可扩展性。

图2：(a)一个增强的Transformer层，由门控线性单元(GLU)和多头自注意力(MHSA)两个块组成。(b)我们提出的门控注意力单元(GAU)。(c)门控注意力单元的伪代码。为了简洁，(a)和(b)中省略了残差分支上的跳跃连接和输入归一化。
图2：(a)一个增强的Transformer层，由门控线性单元(GLU)和多头自注意力(MHSA)两个块组成。(b)我们提出的门控注意力单元(GAU)。(c)门控注意力单元的伪代码。为了简洁，(a)和(b)中省略了残差分支上的跳跃连接和输入归一化。

主要成果: 实验表明，在512到8K的上下文长度范围内，FLASH在质量上与完全增强的Transformer (Transformer++)相当，同时在Wiki-40B上的自回归语言建模任务中实现了高达4.9倍的训练加速，在PG-19上实现了高达12.1倍的加速，在C4上的掩码语言建模任务中实现了4.8倍的加速。

A2 方法细节

2. 门控注意力单元 (Gated Attention Unit)

本节介绍了门控注意力单元（GAU），一个比Transformer层更简单但性能更优的层。尽管GAU对上下文长度仍具有二次方复杂度，但它为第3节将要介绍的近似方法提供了更好的基础。首先介绍相关的层：


其中$\otimes$代表逐元素相乘。在GLU中，每个表示$u_i$都由与同一token关联的另一个表示$v_i$进行门控。

图3：GAU与Transformer在短上下文长度(512)下自回归和掩码语言建模的对比
图3：GAU与Transformer在短上下文长度(512)下自回归和掩码语言建模的对比

表1：各种修改对GAU的影响
表1：各种修改对GAU的影响

表2：各种修改对MHSA的影响
表2：各种修改对MHSA的影响

3. 使用GAU实现快速线性注意力

从第2节中得到的两个观察结果，启发我们将GAU扩展到长序列建模：
- 首先，GAU中的门控机制允许使用更弱的（单头、无softmax）注意力而不损失质量。如果将这一直觉进一步应用于使用注意力建模长序列，GAU也可能提升近似（弱）注意力机制（如局部、稀疏和线性化注意力）的有效性。
- 此外，使用GAU时，注意力模块的数量自然增加了一倍——根据第2节，MLP+MHSA ≈ 2×GAU的成本。由于近似注意力通常需要更多层来捕捉完整的依赖关系【6, Dai, Z. 等人, Transformer-xl: Attentive language models beyond a fixed-length context, arXiv 2019】、【4, Child, R. 等人, Generating long sequences with sparse transformers, arXiv 2019】，这一特性也使得GAU在处理长序列时更具吸引力。

3.1. 现有的线性复杂度变体

3.2. 我们的方法：混合分块注意力 (Mixed Chunk Attention)

基于现有方法的优缺点，我们提出了混合分块注意力，它融合了部分注意力和线性注意力的优点。
- 准备工作: 输入序列首先被分块为$G$个大小为$C$的非重叠块。然后，根据GAU公式为每个块$g$生成$U_g \in R^{C \times e}$、$V_g \in R^{C \times e}$和$Z_g \in R^{C \times s}$。接下来，从$Z_g$生成四种注意力头：$Q_{quad_g}$、$K_{quad_g}$、$Q_{lin_g}$、$K_{lin_g}$。

代码1：混合分块注意力的伪代码。
代码1：混合分块注意力的伪代码。

3.2.1. 讨论

A4 实验环境

A4 实验结果

4.1. 双向语言建模 (MLM on C4)

图5：C4数据集上掩码语言建模的验证集结果 — 所有模型大小约110M，使用2^18个token的批次大小训练125K步。质量以负对数困惑度衡量。
图5：C4数据集上掩码语言建模的验证集结果 — 所有模型大小约110M，使用2^18个token的批次大小训练125K步。质量以负对数困惑度衡量。

4.2. 自回归语言建模 (LM on Wiki-40B & PG-19)

4.3. 微调 (Fine-tuning on TriviaQA)

4.4. 消融研究

A5 结论

本文提出了FLASH，一个旨在解决现有高效Transformer变体在质量和实际速度方面问题的实用解决方案。这一目标通过设计一个高性能层（门控注意力单元GAU）并将其与一个对加速器友好的近似策略（混合分块注意力）相结合来实现。在双向和自回归语言建模任务上的实验表明，FLASH在质量（困惑度）上与完全增强的Transformer相当，而训练速度则比现有最先进的模型快得多。未来的工作包括研究这个新模型家族的缩放定律以及其在更多下游任务上的表现。

A6 附录

A. 与Combiner的联系

B. 实验设置

B.1. 超参数

B.2. 模型规格

实验中使用的所有模型的详细规格总结在表8、9和10中。FLASH-Quad和FLASH使用SiLU/Swish作为非线性激活函数，因为它在我们的模型中略优于GELU。在一些掩码语言模型中使用了ScaleNorm，因为它在TPU-v4上比LayerNorm运行稍快且不影响模型质量。
- C4上的MLM模型配置 (表8)

表8：第4节中C4数据集上MLM实验的模型配置。
表8：第4节中C4数据集上MLM实验的模型配置。

- Wiki-40B上的LM模型配置 (表9)
表9：第4节中Wiki-40B数据集上LM实验的模型配置。
表9：第4节中Wiki-40B数据集上LM实验的模型配置。

- PG-19上的LM模型配置 (表10)
表10：第4节中PG-19数据集上LM实验的模型配置。
表10：第4节中PG-19数据集上LM实验的模型配置。

C. 额外的实验结果

C.1. GPU上的自回归训练

自回归训练的低效率问题不仅限于TPU。如表11所示，在单个Nvidia V100 GPU上，Performer由于需要对所有token进行顺序的累积求和，其延迟最高。相比之下，当上下文长度超过1024时，本文提出的FLASH延迟最低，证明了混合分块注意力机制的有效性。

表11：在单个Nvidia Tesla V100 GPU上，Wiki-40B自回归语言建模每个训练步骤的延迟比较 — 延迟以毫秒报告。OOM代表CUDA内存不足错误。Performer-Matmul使用矩阵乘法实现累积和(cumsum)。
表11：在单个Nvidia Tesla V100 GPU上，Wiki-40B自回归语言建模每个训练步骤的延迟比较 — 延迟以毫秒报告。OOM代表CUDA内存不足错误。Performer-Matmul使用矩阵乘法实现累积和(cumsum)。

C.2. MLM和LM的表格化结果

C.3. 块大小的消融研究

块大小的选择会影响FLASH的质量和训练成本。在极端情况下，当块大小等于上下文长度时，FLASH退化为FLASH-Quad；当块大小为1时，它变成一个低效的线性注意力模型。图8展示了在1K到8K的上下文长度下，四种不同块大小在质量和训练成本之间的权衡。

图8：FLASH在1K到8K上下文长度下块大小(C)的消融研究。(a) 上下文长度 = 1024 (b) 上下文长度 = 2048 (c) 上下文长度 = 4096 (d) 上下文长度 = 8192
图8：FLASH在1K到8K上下文长度下块大小(C)的消融研究。(a) 上下文长度 = 1024 (b) 上下文长度 = 2048 (c) 上下文长度 = 4096 (d) 上下文长度 = 8192

D. FLASH-Quad和FLASH的伪代码

下面展示了FLASH-Quad和FLASH的详细实现代码。
- 相对位置偏置伪代码 (Code 4)

def rel_pos_bias(n):
    """Relative position bias."""
    if n < 512:
        # Construct Toeplitz matrix directly when the sequence length is less than 512.
        w = tf.get_variable(
            &#39;weight&#39;,
            shape=[2 * n - 1],
            dtype=tf.float32,
            initializer=WEIGHT_INITIALIZER)
        t = tf.pad(w, [[0, n]])
        t = tf.tile(t, [n])
        t = t[..., :-n]
        t = tf.reshape(t, [n, 3 * n - 2])
        r = (2 * n - 1) // 2
        t = t[..., r:-r]
    else:
        # Construct Toeplitz matrix using RoPE when the sequence length is over 512.
        a = tf.get_variable(
            &#39;a&#39;,
            shape=[128],
            dtype=dtype,
            initializer=WEIGHT_INITIALIZER)
        b = tf.get_variable(
            &#39;b&#39;,
            shape=[128],
            dtype=dtype,
            initializer=WEIGHT_INITIALIZER)
        a = rope(tf.tile(a[None, :], [n, 1]), axis=0)
        b = rope(tf.tile(b[None, :], [n, 1]), axis=0)
        t = tf.einsum(&#39;mk,nk->mn&#39;, a, b)

    return t

def _get_scaledsin(embeddings):
    &quot;&quot;&quot;Create sinusoidal position embedding with a scaling factor.&quot;&quot;&quot;
    hidden_size = int(embeddings.shape[-1])
    pos = tf.range(tf.shape(embeddings)[1])
    pos = tf.cast(pos, tf.float32)
    half_d = hidden_size // 2
    freq_seq = tf.cast(tf.range(half_d), tf.float32) / float(half_d)
    inv_freq = 10000 ** -freq_seq
    sinusoid = tf.einsum(&#39;s,d->sd&#39;, pos, inv_freq)
    scaledsin = tf.concat([tf.sin(sinusoid), tf.cos(sinusoid)], axis=-1)
    scalar = tf.get_variable(
        &#39;scaledsin_scalar&#39;,
        shape=(),
        initializer=tf.constant_initializer(1 / hidden_size ** 0.5))
    scaledsin *= scalar
    return scaledsin