Gated Linear Attention Transformers with Hardware-Efficient Training

发表时间: 2023-12 · arXiv:2312.06635

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

作者/机构: Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, Yoon Kim

速读

一句话结论 本文提出了一种带有数据依赖门控机制的线性注意力模型 GLA Transformer 及其硬件高效的训练算法 FLASHLINEARATTENTION，在保持线性时间推理的同时，实现了媲美 LLaMA 和 Mamba 的语言建模性能与更快的实际训练速度。

要解决什么问题 现有线性注意力模型面临性能与实际硬件效率的双重卡点。在性能层面，标准线性注意力由于缺乏衰减项或遗忘门机制，难以像传统 RNN 那样根据输入动态丢弃无用信息，这被认为是导致其在处理长上下文任务时存在不稳定性、语言建模效果显著落后于标准 Softmax 注意力的核心原因。尽管近期模型引入了衰减因子，但往往是全局且数据无关的。在效率层面，现有的线性注意力实现未能充分适配现代 GPU 的计算模型与内存层次结构：基础的循环形式需要将所有时间步的隐藏状态存储在容量大但速度慢的全局显存（HBM）中，导致极高的 I/O 成本，且逐元素操作算术强度低，无法利用张量核心加速；并行形式虽然能像 FLASHATTENTION 那样优化 I/O，但其计算复杂度随序列长度呈二次方增长，长序列训练成本高昂；而现有的分块并行形式虽然在理论上平衡了两者，但由于缺乏 I/O 感知设计，导致在中等序列长度（如 2K 到 4K）下的实际运行速度反而慢于经过高度优化的 Softmax 注意力实现。

怎么做的 核心思路是为线性注意力引入数据依赖的门控机制，并为其设计一套 I/O 感知的硬件高效分块训练算法。首先，作者提出了门控线性注意力（GLA），通过引入一个随时间变化且依赖当前输入的遗忘门，让模型能够动态保留或遗忘历史信息。为了在参数效率、状态大小和张量核心利用率之间取得平衡，区别于 Mamba 无法使用张量核心的设计，GLA 将门控矩阵参数化为向量形式，其循环更新公式定义为：$$S_t = (\alpha_t^T \mathbf{1}) \odot S_{t-1} + k_t^T v_t$$ 其中 $\alpha_t$ 是输入经过低秩线性层和 Sigmoid 激活后得到的门控向量。其次，为了绕开显存读写瓶颈，作者设计了 FLASHLINEARATTENTION 算法。该算法采用分块计算策略，将张量逐块加载到高速缓存（SRAM）中复用，避免频繁访问全局显存，并提供了非物化（节省显存）和物化（支持序列级并行）两个版本。针对 GLA 中门控累积乘积在对数空间计算无法使用半精度矩阵乘法的问题，算法引入了关键的二级分块技术：将大块进一步切分为子块，子块间的交互通过张量核心执行半精度矩阵乘法，而子块内部的局部计算则在全精度下进行以保证数值稳定，从而大幅减少了非半精度计算的开销。此外，为了解决反向传播时显存占用过大的问题，作者推导出了关于门控对数梯度的闭式解公式，彻底避免了在显存中实例化庞大的隐藏状态张量。最终，将多头 GLA 层与 SwiGLU 前馈网络交错，构建出参数量与标准 Transformer 相当的 GLA Transformer 模型。

效果如何 实验在 SlimPajama 数据集的 100B 词元子集上从头训练了 340M 和 1.3B 两种规模的模型，硬件基于 H100 GPU。对比基线包括代表标准注意力强基线的 Transformer++（LLaMA 架构）、代表数据无关衰减线性路线的 RetNet，以及代表状态空间模型路线的 Mamba。在语言建模零样本测试中，GLA Transformer 的平均准确率全面超越 RetNet，并达到了与 Transformer++ 和 Mamba 相当的水平。在合成的多查询关联召回任务中，GLA 优于 RetNet 和 Mamba，但在真实的密集召回任务上，所有次二次方模型仍显著落后于标准 Transformer。在长度泛化方面，在 2K 长度上训练的 GLA 模型能够平滑外推到 18K 甚至 20K 以上的序列而困惑度没有显著下降，表现优于在 4K 后性能断崖式下跌的 Mamba。消融实验表明，细粒度的向量门控优于标量门控，且设置 4 个注意力头是性能与显存的最佳折中。在训练效率上，得益于硬件感知算法，GLA 的吞吐量高于同等规模的 Mamba，且在序列长度超过 4096 时优势进一步扩大，其底层算子速度也超越了 FLASHATTENTION-2。作者承认的局限在于，受限于计算资源，该方法尚未在更大参数规模和更海量的数据集上进行验证，其在更大规模下的扩展性仍有待观察。

A1 主要贡献

本文致力于解决现有线性注意力模型相较于标准Softmax注意力在性能和实际运行速度上的不足。

核心问题:

  1. 性能差距: 尽管线性注意力通过其循环（RNN）形式实现了线性时间复杂度的推理，但在语言建模等任务上，其性能通常显著低于标准的Softmax注意力。
  2. 效率问题: 现有的线性注意力实现（包括循环、并行和分块并行形式）没有充分考虑硬件的I/O特性，导致在实际运行中，特别是在中等序列长度上，其速度慢于高度优化的Softmax注意力实现（如FLASHATTENTION）。
  3. 门控机制：现有的线性注意力变体虽然引入了衰减因子以提升性能，但这些因子通常是全局且数据无关的，而数据依赖的门控机制在传统RNN中已被证明至关重要。

研究目标与创新点:

  1. 提出硬件高效的线性注意力算法 (FLASHLINEARATTENTION)：
    • 本文设计了一种I/O感知的硬件高效算法，用于训练线性注意力模型。该算法通过在内存移动和并行性之间进行权衡，实现了显著的速度提升。
    • 该算法的实现（名为FLASHLINEARATTENTION）作为一个独立的层，即使在短序列（如1K）上，其速度也超过了FLASHATTENTION-2。

  1. 提出门控线性注意力 (Gated Linear Attention, GLA)：

    • 将FLASHLINEARATTENTION算法推广到一种更具表现力的、带有数据依赖门控机制的线性注意力变体。
    • 这种门控机制允许模型根据输入动态地“遗忘”或保留信息，解决了标准线性注意力因缺乏衰减项而难以处理长上下文任务的潜在不稳定性。

  2. 构建并验证GLA Transformer模型:

    • 将GLA层作为标准注意力层的替代品，构建了GLA Transformer模型。
    • 在中等规模的语言建模实验（在15B/100B tokens上训练340M/1.3B参数模型）中，GLA Transformer的性能与强大的LLaMA架构以及最新的线性时间推理模型（如RetNet和Mamba）相当。
    • GLA Transformer在长度泛化方面表现尤为出色，一个在2K长度上训练的模型能够泛化到超过20K的序列，而困惑度没有显著下降。
    • 在训练速度方面，GLA Transformer的吞吐量高于同等规模的Mamba模型。

A3 背景知识

2.1 并行与循环形式

标准Softmax注意力。自回归Transformer采用Softmax注意力机制，它接收一个输入序列 $X \in R^{L \times d}$（$L$为长度，$d$为隐藏维度），并通过以下公式计算输出 $O \in R^{L \times d}$：

公式
公式

其中 $W_Q, W_K, W_V \in R^{d \times d}$ 是可学习的矩阵，$M \in \{-\infty, 1\}^{L \times L}$ 是一个掩码，防止模型关注未来的词元。这种并行形式可以在给定完整输入 $X$ 的情况下并行计算 $O$，从而实现高效训练。然而，在推理过程中，Transformer必须使用循环形式，即维护一个不断增长的键（key）和值（value）集合（KV缓存），并对它们执行注意力操作。
公式
公式

线性注意力。线性注意力机制将 $exp(q_t k_i^T)$ 替换为一个核函数 $k(x,y)$，该函数有关联的特征图 $\phi$（即 $k(x,y) = \langle\phi(x), \phi(y)\rangle$）。这简化了输出 $o_t$ 的计算，因为求和操作可以与矩阵乘法交换顺序：

公式
公式

通过定义 $S_t = \sum_{i=1}^{t} \phi(k_i)^T v_i$ 和 $z_t = \sum_{i=1}^{t} \phi(k_i)^T$，其中 $S_t \in R^{d \times d}, z_t \in R^{d \times 1}$，上述计算可以重写为一个RNN：
公式
公式

尽管已有多种核函数被探索，但最近的研究发现，使用线性核（即设置 $\phi$ 为恒等函数）且不带归一化项在实践中效果很好【索引83，Retentive network: A successor to transformer for large language models，2023a】。这产生了一个（未归一化的）线性注意力层，其更新方程如下：
公式
公式

从公式1可以清楚地看到，线性注意力层本质上是一个线性循环层，其拥有矩阵值的隐藏状态 $S_t$，该状态通过外积 $k_t^T v_t$ 进行更新。因果线性注意力的并行形式的复杂度仍然是 $L$ 的平方，其形式为 $O = (QK^T) \odot M V$，其中 $M$ 是一个下三角掩码矩阵。由于掩码 $M$ 的存在，无法利用矩阵乘法的结合律将并行形式的复杂度从二次降为线性。

2.2 分块并行形式

平衡并行与循环。线性注意力的分块并行形式在并行形式和循环形式之间取得了平衡【索引38，Transformer quality in linear time，2022；索引83，Retentive network: A successor to transformer for large language models，2023a】，并允许进行次二次方、部分并行的训练。

分块计算过程。具体来说，假设输入 $X$ 被分成多个不重叠的块，每个块的长度为 $C$。令 $S[i] \in R^{d \times d}$ 是处理完 $i$ 个块后的块级隐藏状态，即 $S[i] := S_{iC}$。令 $Q[i], K[i], V[i]$ 分别是第 $i$ 个块对应的查询、键和值向量。我们有以下块间的递推关系：

公式
公式

这里 $S[0]$ 可以初始化为零或来自前一个段落的隐藏状态。一个块内所有RNN输入的总和（即 $K[i]^T V[i]$）可以并行计算，复杂度为 $O(C^2d)$。

块内并行计算。输出的块内并行计算由以下公式给出：

公式
公式

其中，“块内”分量 $O_{intra}[i+1]$ 的计算形式与标准并行形式完全相同，耗时 $O(C^2d + Cd^2)$。“块间”分量 $O_{inter}[i+1]$ 计算了来自前一个块的隐藏状态的贡献，耗时 $O(Cd^2)$。因此，总的训练复杂度为 $O(\frac{L}{C}(C^2d + Cd^2)) = O(LCd + Ld^2)$，当 $L > d$ 时，这小于 $O(L^2d)$。值得注意的是，设置 $C = L$ 恢复了并行形式，而 $C = 1$ 则恢复了循环形式。

A2 方法细节

3 硬件高效的线性注意力

本节描述了FLASHLINEARATTENTION，这是一种I/O感知、硬件高效的线性注意力算法，其精神类似于FLASHATTENTION【索引21，Flashattention: Fast and memory-efficient exact attention with io-awareness，2022b；索引18，Flashattention-2: Faster attention with better parallelism and work partitioning，2023】。

3.1 硬件高效算法的原则

高效算法的设计考量。一个高效的算法应该了解现代硬件的计算模型、内存层次结构和专用计算单元。

3.2 线性注意力的硬件考量

不同形式的硬件效率分析。本节讨论与不同形式线性注意力效率相关的硬件因素。

3.3 FLASHLINEARATTENTION: 采用分块形式的硬件高效线性注意力

I/O感知的实现。我们描述了分块形式的I/O感知、硬件高效实现。我们提供了两个版本，其前向和后向传播根据块级隐藏状态 $S[n]$ 是否物化到HBM中而有所不同。算法1和图1展示了前向传播过程。核心思想是使用分块（tiling）技术逐块加载张量，并在芯片上（SRAM）重复使用张量块，以尽可能避免多次HBM I/O。例如，当 $Q[n]$ 加载到SRAM时，可以在芯片上同时计算 $Q[n]S$ 和 $(Q[n]K^T[n] \odot M)V[n]$，避免了两次加载 $Q[n]$，从而节省了HBM I/O。

算法1：FLASHLINEARATTENTION：前向传播
算法1：FLASHLINEARATTENTION：前向传播

图1：(a) FLASHLINEARATTENTION 非物化版本。此版本更节省内存。(b-c) FLASHLINEARATTENTION 物化版本。此版本支持序列级的分块并行。
图1：(a) FLASHLINEARATTENTION 非物化版本。此版本更节省内存。(b-c) FLASHLINEARATTENTION 物化版本。此版本支持序列级的分块并行。

公式
公式

两个版本的对比。

性能评估。图2显示了我们实现的运行速度和内存占用。两个版本的FLASHLINEARATTENTION都明显快于FLASHATTENTION-2【索引18，Flashattention-2: Faster attention with better parallelism and work partitioning，2023】和纯PyTorch（即非I/O感知）实现的分块线性注意力，展示了I/O感知的优势。

图2：在单张H100 GPU上的速度比较，批次大小为32，头数量为16，头维度为64，块大小为64。x轴和y轴均为对数尺度。w/ m. 和 w/o m. 分别表示使用和不使用在HBM中物化隐藏状态的FLASHLINEARATTENTION。
图2：在单张H100 GPU上的速度比较，批次大小为32，头数量为16，头维度为64，块大小为64。x轴和y轴均为对数尺度。w/ m. 和 w/o m. 分别表示使用和不使用在HBM中物化隐藏状态的FLASHLINEARATTENTION。

4 门控线性注意力

引入门控机制的动机。标准线性注意力（公式1）缺少衰减项或遗忘门，这在RNN中被证明是至关重要的【索引35，Long short-term memory，1997；索引14，Learning phrase representations using rnn encoder-decoder for statistical machine translation，2014；索引86，The unreasonable effectiveness of the forget gate，2018】。缺少衰减项使得模型难以“遗忘”信息，并被认为是导致线性注意力在长上下文任务中不稳定的部分原因【索引11，Linear Transformers Are Faster After All，2024】。虽然最近的工作通过引入一个全局的、非数据依赖的衰减因子 $\gamma \in (0,1)$ 来提升性能【索引83，Retentive network: A successor to transformer for large language models，2023a；索引76，Scaling transnormer to 175 billion parameters，2023b】，但本文考虑使用一种数据依赖的门控机制。我们证明，尽管门控因子更具表现力，由此产生的门控线性注意力（GLA）层仍然支持硬件高效的分块形式以进行高效训练。

4.1 GLA的循环与并行形式

循环形式。GLA拥有一个随时间变化的二维遗忘门 $G_t \in (0,1)^{d_k \times d_v}$：

公式
公式

这种基于哈达玛积的循环形式非常通用，涵盖了许多近期的具有二维隐藏状态的RNN，如表1所示。
表1：近期模型在门控线性注意力框架下的表述，它们在Gt的参数化上有所不同。偏置项被省略。
表1：近期模型在门控线性注意力框架下的表述，它们在Gt的参数化上有所不同。偏置项被省略。

门控矩阵Gt的参数化。GLA设计的核心在于 $G_t$ 的参数化，需要在参数效率、状态大小和训练效率之间取得平衡。一个简单地将 $x_t$ 映射到 $G_t$ 的方法需要一个大小为 $d \cdot d_k \cdot d_v$ 的矩阵，参数效率低下。Mao (2022)【索引62，Fine-tuning pre-trained transformers into decaying fast weights，2022】提出了一种更高效的基于外积的低秩参数化 ($G_t = \alpha_t^T \beta_t$)。在Mamba中【索引28，Mamba: Linear-time sequence modeling with selective state spaces，2023】， $G_t$ 是通过将一个数据无关的可学习矩阵 $A$ 与一个数据依赖的向量 $\alpha_t$ 结合得到的，这使得矩阵可以是满秩的，但也阻止了张量核心的使用。Mamba-2【索引19，Transformers are ssms: Generalized models and efficient algorithms through structured state space duality，2024】通过一个更受限的门控机制 $G_t = \gamma_t \mathbf{1}^T \mathbf{1}$（其中 $\gamma_t$ 是一个标量）解决了这个问题，使其能够利用张量核心。

本文采用的参数化。本文在标量和完全低秩参数化之间采取了折中方案，使用 $G_t = \alpha_t^T \mathbf{1}$。这得到了以下循环形式：

公式
公式

其中 $\alpha_t$ 是通过一个低秩线性层后接sigmoid函数在 $x_t$ 上参数化的。这个形式是通用的，可以涵盖多个近期的RNN模型【索引44，Gateloop: Fully data-controlled linear recurrence for sequence modeling，2023；索引79，Hgrn2: Gated linear rnns with state expansion，2024b；索引71，Eagle and finch: Rwkv with matrixvalued states and dynamic recurrence，2024】。

并行形式。展开上述循环公式，可以得到一个并行形式：

公式
公式

令 $b_t := \prod_{j=1}^{t} \alpha_j$，我们可以重写为：
公式
公式

这可以表示为一个并行形式：
公式
公式

然而，这种形式在数值上不稳定，因为 $b_t$ 是门控值 $\alpha_j \in (0,1)^{1 \times d}$ 的累积乘积，当 $t$ 很大时可能变得非常小。为了解决这个问题，可以在对数空间中计算：
公式
公式

但与普通线性注意力不同，公式4不能通过标准矩阵乘法表示，因此无法利用张量核心上的半精度矩阵乘法。

4.2 GLA的分块并行形式

推导分块形式。我们推导了GLA的分块形式，类似于基础线性注意力的分块形式（§2.2）。块内操作在块级别上实现上述并行形式以获得 $O_{intra}$。对于块间计算，我们有：

公式
公式

直观上，$\Lambda[i+1]$ 编码了从块开始的累积衰减，用于传播前一个块的隐藏状态 $S[i]$；而 $\Gamma[i+1]$ 编码了到块结束的衰减，用于累积信息以添加到下一个隐藏状态 $S[i+1]$。
图3：用于说明GLA中分块计算的注意力风格图。块间依赖（灰色）在分块形式中不直接计算。块内依赖通过二级分块/平铺建模，其中子块间部分（橙色）由半精度矩阵乘法计算，而子块内部（粉色）在对数空间中以全精度计算。
图3：用于说明GLA中分块计算的注意力风格图。块间依赖（灰色）在分块形式中不直接计算。块内依赖通过二级分块/平铺建模，其中子块间部分（橙色）由半精度矩阵乘法计算，而子块内部（粉色）在对数空间中以全精度计算。

4.3 硬件高效的GLA

适配FLASHLINEARATTENTION。有了分块形式，我们可以将§3中提出的FLASHLINEARATTENTION算法应用于门控情况。这种适配还依赖于以下两个关键技术。

公式
公式

这个公式可以通过对公式4求导轻易得到，从而避免了大规模的物化操作。

4.4 GLA Transformer

模型架构。我们将GLA层推广到多头情况，并将其与前馈网络（FFN）交错构建成一个Transformer风格的模型。

公式
公式

公式
公式

公式
公式

公式
公式

通过这种设计，一个GLA层的总参数量（大约 $4d^2$）与常规的Softmax注意力层相当。

A4 实验环境

数据集:

模型架构:

硬件配置:

软件配置:

A4 实验结果

主要结果:

表2：GLA Transformer与Transformer++、RetNet和Mamba的对比结果。所有模型均在SlimPajama数据集的同一子集上使用Mistral分词器进行训练。340M/1.3B模型分别训练15B/100B词元。单个任务的性能是零样本测试。我们报告了与Gu & Dao (2023)相同的任务集上的主要结果。最后一列显示了所有使用（归一化）准确率作为指标的基准测试的平均值。
表2：GLA Transformer与Transformer++、RetNet和Mamba的对比结果。所有模型均在SlimPajama数据集的同一子集上使用Mistral分词器进行训练。340M/1.3B模型分别训练15B/100B词元。单个任务的性能是零样本测试。我们报告了与Gu & Dao (2023)相同的任务集上的主要结果。最后一列显示了所有使用（归一化）准确率作为指标的基准测试的平均值。

召回密集型任务 (Recall-intensive tasks):

图4：在合成MQAR任务上的准确率(%)。
图4：在合成MQAR任务上的准确率(%)。

表3：不同模型在Arora et al. (2024)测试的三个召回密集型任务中的比较。所有任务分数越高越好。
表3：不同模型在Arora et al. (2024)测试的三个召回密集型任务中的比较。所有任务分数越高越好。

长序列训练与长度外推:

图5：在SlimPajama和PG19测试集上的长度外推。我们从头开始在SlimPajama上使用不同的训练长度预训练1.3B模型100B词元。*表示模型使用截断式BPTT，超过12个段，每个段长度为2K。
图5：在SlimPajama和PG19测试集上的长度外推。我们从头开始在SlimPajama上使用不同的训练长度预训练1.3B模型100B词元。*表示模型使用截断式BPTT，超过12个段，每个段长度为2K。

消融研究:

表4：在340M模型上训练7B词元的消融研究结果。我们通过最后200个训练步骤的平均困惑度来评估模型变体。
表4：在340M模型上训练7B词元的消融研究结果。我们通过最后200个训练步骤的平均困惑度来评估模型变体。

训练效率:

图6：在H100上的训练吞吐量和内存占用。
图6：在H100上的训练吞吐量和内存占用。

A7 补充细节

5.4 局限性与未来工作

当前研究的局限性。由于计算资源有限，本文的实验未能进行更大规模的验证。因此，GLA如何扩展到更大的模型和数据集上尚不清楚。

未来工作的展望。预计随着模型规模的扩大，GLA的训练效率优势会比Mamba更加明显，因为它能更好地利用张量核心，并且与张量并行兼容。未来的一个有趣方向是将GLA应用于其他模态，特别是那些具有长程依赖性的数据，以利用线性注意力的效率优势。

6 相关工作

线性RNNs/SSMs/Transformers。传统的RNN由于隐藏状态之间的非线性依赖和昂贵的顺序更新而难以扩展。线性RNNs、状态空间模型（SSMs）和线性Transformer消除了非线性依赖，使得沿时间维度的训练可以并行化，成为Transformer架构有竞争力的次二次方替代方案。

数据依赖的衰减率。数据依赖的衰减率一直被认为是RNN的关键。与以往依赖于前一隐藏状态和当前输入的遗忘门不同，Martin & Cundy (2018)【索引63，Parallelizing linear recurrent neural nets over sequence length，2018】提出遗忘门应仅依赖于当前输入以实现并行训练，这一策略在HGRN【索引76，Scaling transnormer to 175 billion parameters，2023b】等模型中被证明有效。RWKV-v6【索引71，Eagle and finch: Rwkv with matrixvalued states and dynamic recurrence，2024】和Mamba【索引28，Mamba: Linear-time sequence modeling with selective state spaces，2023】也使用了类似的数据依赖衰减率。

隐藏状态维度的扩展。RNNs依赖于固定维度的隐藏状态来编码历史信息，其维度大小影响模型的表达能力。线性Transformer通过外积参数化来扩展隐藏维度。而线性SSM则通过单输入单输出（SISO）策略来扩展。当SSM参数是数据无关时，可以通过快速傅里叶变换（FFT）高效训练。但对于数据依赖的SSM参数，如Mamba【索引28，Mamba: Linear-time sequence modeling with selective state spaces，2023】，则无法使用FFT，需要使用并行扫描算法的自定义CUDA核进行训练。为了将所有隐藏状态放入SRAM，Mamba的隐藏状态扩展率最高只能到16。相比之下，本文提出的硬件感知训练算法为扩展隐藏维度提供了一种替代的高效方法，使其可以扩展到更宽的范围，这在召回密集型任务中被证明是有用的。

A5 结论

本文提出了一种用于训练带有数据依赖门控机制的线性注意力Transformer的高效算法。该算法通过在浮点运算量（FLOPs）和并行性之间取得平衡，同时允许使用现代GPU上的张量核心单元进行半精度矩阵乘法，从而实现了高效训练。在语言建模上的实验表明，门控线性注意力（GLA）Transformer的性能与强大的基线模型相当。

A6 附录

A 扩展相关工作

A.1 线性注意力

A.2 序列并行

A.3 硬件感知算法

B 分块（门控）线性注意力的细节

FLASHLINEARATTENTION的反向传播。算法2列出了线性注意力的反向传播伪代码。

算法2 FLASHLINEARATTENTION: 反向传播
算法2 FLASHLINEARATTENTION: 反向传播

GLA的伪代码。算法3和4展示了物化版本GLA的前向/后向传播；算法5和6展示了非物化版本。列表1以Pytorch风格代码展示了二级分块算法。

算法3：门控线性注意力的前向传播（带物化）
算法3：门控线性注意力的前向传播（带物化）

算法4：门控线性注意力的反向传播（带物化）
算法4：门控线性注意力的反向传播（带物化）

算法5：门控线性注意力的前向传播（不带物化）
算法5：门控线性注意力的前向传播（不带物化）

算法6：门控线性注意力的反向传播（不带物化）
算法6：门控线性注意力的反向传播（不带物化）

def gated_linear_attention_forward (Q , K , V , a , C , c ) :
    # Q/K/V: query / key / value
    # a: log forget gate
    # C/c: chunk size , subchunk size
    # L: sequence length , d: head dimension
    L , d_k = Q.shape
    d_v = V.shape[-1]
    S = torch.zeros(d_k, d_v)
    O = torch.empty_like(V)
    # cumsum of log decay within a chunk
    B = torch.empty_like(a)
    # local compute of cumulative product of decay within a chunk
    for i in range(0, L // C):
        b = torch.zeros(d_k)
        for j in range(0, C):
            b += a[i]
            B[i] = b
    for i in range(0, L // C):
        r = range(i*C ,(i +1) * C)
        # (C, d) chunking
        bq , bk , bv , bb = Q[r] , K[r] , V[r] , B[r]
        b = bb[-1, None]
        #inter-chunk w/ matmul
        q, k, g = bq * (bb.exp()), bk * ((b - bb).exp()), b.exp()
        o = q @ S
        # hidden state update
        S = g.t() * S + k.t() @ bv
        #intra-chunk (secondary chunking)
        for j in range(0, C // c):
            t = range(j *c , (j +1) * c)
            #(c, head_dim ) subchunking
            q , k , v , b = bq[t] , bk[t] , bv[t] , bb[t]
            p = torch.zeros(c, c)
            #intra-subchunk w/o matmul .
            for m in range(c):
                for n in range(m +1):
                    p[m, n] = torch.sum(q[m] * k[n] * ((b[m] - b[n]).exp()))
            o[t] += p @ v
            # inter-subchunk w/ matmul
            z = b[0, None]
            q = q * (b - z).exp()
            for u in range(0, j):
                y = range(u *c , (u +1) * c)
                p = q @ (bk[y] * (z - bb[y]).exp()).t()
                o[t] += p @ bv[y]
        O[r] = o
    return O

dlogαt的推导。通过展开循环公式，并对 $o_t$ 关于 $\log b_t$（其中 $b_t$ 是累积门控值）求导，可以得到 $d\log b_t$ 的表达式。由于 $\log b_t = \sum_{i=1}^{t} \log\alpha_i$，那么 $d\log\alpha_t$ 就是 $d\log b_i$ 从 $i=t$ 到 $L$ 的累积和。这个过程避免了直接计算与 $S_{t-1}$ 相关的梯度，从而实现了内存高效。

dlogαt的梯度形式
dlogαt的梯度形式

展开循环
展开循环

关于查询/键向量的梯度
关于查询/键向量的梯度

关于累积门控对数的梯度
关于累积门控对数的梯度

dlogbt的最终形式
dlogbt的最终形式

C 通用门控线性注意力

通用形式。在主论文中，我们使用了一个简化的参数化，其中 $G_t = \alpha_t^T \mathbf{1}$。这里探讨了更通用的形式 $G_t = \alpha_t^T \beta_t$。

公式
公式

尽管实验发现学习 $\beta$ 并没有带来性能提升，但该通用形式仍然支持并行和分块形式，这可能对未来的线性注意力模型发展有用。

C.1 并行形式。通过展开循环并利用克罗内克/外积的混合积性质，可以推导出该通用形式的并行形式。最终的并行形式可以表示为 $O = ( (Q \odot B) (K/B)^T \odot M ) (V \odot D)$，其中 $B$ 和 $D$ 分别是 $\alpha$ 和 $\beta$ 的累积乘积矩阵。

公式推导过程
公式推导过程

C.2 分块并行形式。该通用形式同样可以推导出分块并行形式，用于高效训练。块间递推和块内并行计算的公式被相应地推广，以包含 $\alpha$ 和 $\beta$ 两个门控向量。这个分块形式是现有多种线性注意力形式的泛化，例如，通过设置不同的门控值，它可以退化为普通线性注意力或RetNet的分块形式。

内存高效的梯度计算。对于通用形式，关于 $\alpha$ 和 $\beta$ 的梯度也存在闭式解，允许在不实例化隐藏状态 $S$ 的情况下计算 $d\alpha$ 和 $d\beta$，从而实现内存高效。

梯度公式
梯度公式

D 额外的实验结果

扩展性能结果。表5展示了在所有11个任务上的完整结果，包括1.3B模型的5-shot结果。这些结果进一步证实了GLA模型与强基线模型（如Transformer++和Mamba）相比具有竞争力。

表5：扩展的零样本和五样本性能结果。所有模型都在SlimPajama数据集的同一子集上使用Mistral分词器进行训练。340M/1.3B模型分别训练了15B/100B词元。最后一列显示了所有准确率的平均值。
表5：扩展的零样本和五样本性能结果。所有模型都在SlimPajama数据集的同一子集上使用Mistral分词器进行训练。340M/1.3B模型分别训练了15B/100B词元。最后一列显示了所有准确率的平均值。