Linear Attention/zh

Article
Topic area	Deep Learning
Prerequisites	Transformer, Softmax Function

概述

线性注意力是一类注意力机制，其时间和内存开销随序列长度线性扩展，与原始 Transformer 架构中所用标准 softmax 注意力的二次开销形成对比。其核心思想是用一个可写成特征映射内积的核替代 softmax 注意力中的指数相似度函数，再利用矩阵乘法的结合律重排计算顺序。其结果是一种注意力层，对于自回归解码而言，可表示为具有固定大小隐藏状态的线性循环神经网络，因而在长上下文语言建模、流式推理和端侧部署中颇具吸引力。

线性注意力以牺牲部分 softmax 注意力的表达能力为代价换取渐近效率。这种取舍是否划算取决于工作负载：对于很长的序列与常数内存推理，它通常是明显的胜利；而对于短到中等长度的上下文，常数因子和质量差距可能抵消其理论上的优势。诸如门控线性注意力以及状态空间式模型等现代变体已显著缩小了质量差距，并构成了 2020 年代初期若干高效序列模型的基础。

背景与动机

标准的缩放点积注意力对查询 $ Q \in \mathbb{R}^{N \times d} $、键 $ K \in \mathbb{R}^{N \times d} $ 和值 $ V \in \mathbb{R}^{N \times d} $ 计算输出

$ {\displaystyle \mathrm{Attention}(Q, K, V) = \mathrm{softmax}\!\left(\tfrac{Q K^\top}{\sqrt{d}}\right) V.} $

矩阵 $ Q K^\top $ 的形状为 $ N \times N $，因此其时间和内存开销均为 $ \mathcal{O}(N^2 d) $。对于长文档、高分辨率图像或音频波形，这种二次扩展成为训练和推理的主要开销，催生了大量关于高效注意力的文献。线性注意力是其中最简单也最具影响力的方法之一。

第二个动机来自自回归解码。标准 Transformer 解码器必须在每一步关注所有先前的词元，这就需要存储一份随上下文长度增长的键值缓存。线性注意力解码器维护对过去的固定大小摘要，因此每词元的解码开销在序列长度上是常数，且内存不会增长。

核重构

线性注意力的出发点是把（未归一化的）softmax 相似度写成核 $ k(q, k) = \exp(q^\top k / \sqrt{d}) $，然后用一个允许显式特征映射 $ \phi : \mathbb{R}^d \to \mathbb{R}^{d'} $ 的正定核加以替换，使得

$ {\displaystyle k(q, k) = \phi(q)^\top \phi(k).} $

代入并写出对单个查询 $ q_i $ 的注意力输出，可得

$ {\displaystyle y_i = \frac{\sum_{j=1}^{N} \phi(q_i)^\top \phi(k_j)\, v_j}{\sum_{j=1}^{N} \phi(q_i)^\top \phi(k_j)} = \frac{\phi(q_i)^\top \sum_{j=1}^{N} \phi(k_j)\, v_j^\top}{\phi(q_i)^\top \sum_{j=1}^{N} \phi(k_j)}.} $

关键的一步是第二个等号，它利用矩阵乘法的结合律把 $ \phi(q_i) $ 提到求和号外。两个求和

$ {\displaystyle S = \sum_{j=1}^{N} \phi(k_j)\, v_j^\top \in \mathbb{R}^{d' \times d}, \qquad z = \sum_{j=1}^{N} \phi(k_j) \in \mathbb{R}^{d'}} $

并不依赖于查询，因此可在 $ \mathcal{O}(N d' d) $ 时间内预先计算并在所有查询间重用。完整的注意力输出此时按 $ \mathcal{O}(N d' d) $ 而非 $ \mathcal{O}(N^2 d) $ 计算，从而获得名副其实的线性扩展。

同样的技巧也适用于 softmax 的分母，它转化为对 $ \phi(q_i)^\top z $ 的归一化。一些实现完全省略分母，转而依靠层归一化来稳定输出幅度。

特征映射

$ \phi $ 的选择同时决定了该层的表达能力与开销。已有若干家族被提出：

• 逐元素的正映射。 Katharopoulos 等人引入了 $ \phi(x) = \mathrm{elu}(x) + 1 $，这是一种与输入维度相同的廉价非负映射，是规范的“线性 Transformer”基线。
• 随机特征映射。 Performer 模型用正随机特征 $ \phi(x) \propto \exp(W x - \|x\|^2 / 2) $ 近似 softmax 核，其中 $ W $ 包含随机正交投影。这在保持线性开销分解的同时，在期望意义上恢复了 softmax 注意力。
• 多项式特征映射。 取 $ \phi(x) = (1, x, x \otimes x, \dots) $ 得到一个多项式核；截断到低次给出一种可处理的、表达能力受控的线性注意力变体。
• 恒等映射。 令 $ \phi(x) = x $ 把该层退化为普通双线性注意力。这种做法很快，但除非配合归一化或门控，否则通常表现不佳。

在实践中，特征映射很少是质量的瓶颈；更重要的设计选择是归一化、门控以及循环状态如何更新。

循环形式与自回归推理

对于因果（自回归）建模，累积和 $ S_t $ 与 $ z_t $ 可增量更新：

$ {\displaystyle S_t = S_{t-1} + \phi(k_t)\, v_t^\top, \qquad z_t = z_{t-1} + \phi(k_t).} $

每个新词元都对矩阵值状态 $ S_t $ 施加一次秩一更新。在时刻 $ t $ 的输出为

$ {\displaystyle y_t = \frac{\phi(q_t)^\top S_t}{\phi(q_t)^\top z_t}.} $

这正好是线性 RNN 的形式，其固定大小的隐藏状态维度为 $ d' \times d $。每词元的解码开销为 $ \mathcal{O}(d' d) $，与序列长度无关，也不存在不断增长的键值缓存。推理时的内存消耗为常数，正是这一性质推动了对长上下文语言模型中线性注意力的重新关注。

训练：并行形式

若依赖循环形式进行训练会很慢，因为通过长时序循环进行反向传播难以沿时间轴并行化。所幸，给出循环形式的同一套代数也给出了并行形式：累积和可以用前缀扫描计算，或者更常见地，把整个 $ N \times N $ 的注意力矩阵以适中大小的分块物化，以便让 GPU 保持忙碌。诸如 FLA（Flash Linear Attention 库）的分块并行形式以及 RetNet 中所用的核函数等现代实现，正是利用这一点在长序列上以可与 softmax Transformer相媲美的速度训练，同时保持线性的渐近开销。

因果掩码带来一个微妙之处：累积和必须遵守词元的顺序，因此先对所有 $ j $ 求和再施加掩码的朴素向量化做法是错误的。正确的实现要么使用前缀扫描，要么把序列切分为分块，并把块内的 softmax 风格注意力与块间的线性更新结合起来。

变体与扩展

大量后续模型都建立在线性注意力的基本思路之上：

• Performer。 使用随机特征对 softmax 核进行逼近，提供无偏估计以及对逼近方差的理论保证。
• 带 elu+1 的线性 Transformer。 Katharopoulos 等人的原始公式，被广泛用作基线。
• RetNet。 用衰减因子 $ \gamma \in (0, 1) $ 取代无界的累积和，得到 $ S_t = \gamma S_{t-1} + \phi(k_t) v_t^\top $。该衰减赋予模型多尺度保持的性质，并被证明等价于一种分块计算。
• 门控线性注意力（GLA）。 用依赖于数据的门控替代标量衰减，在保留线性开销的同时恢复了更多 softmax 注意力的选择性行为。
• 选择性状态空间模型。 诸如 Mamba 的架构虽然严格来说并非线性注意力模型，却共享线性循环结构，并可在一个密切相关的框架中加以表述。在近期文献中，这两个家族已大幅趋同。

与 softmax 注意力的比较

相对于标准的 softmax 注意力，线性注意力提供：

• 渐近开销。 时间和内存上从 $ \mathcal{O}(N^2 d) $ 降为 $ \mathcal{O}(N d' d) $。这一优势在大约 $ N \approx 2{,}000 $ 到 $ 8{,}000 $ 词元处变得决定性，具体取决于硬件和常数因子。
• 常数推理内存。 没有不断增长的 KV 缓存；状态形状固定为 $ d' \times d $。
• 便于流式处理。 新词元可通过秩一更新加入。

代价是：

• 表达能力降低。 Softmax 注意力可以以非常尖锐的分布选择性地聚焦于少量词元；而线性注意力的有界秩状态无法重现任意尖锐的模式。这在经验上表现为更弱的关联式召回和复制能力。
• 对特征映射选择的敏感性。 质量随 $ \phi $、归一化和门控显著变化，因此朴素的替换往往表现不佳。

在实践中，把 softmax 与线性注意力层交错排列的混合架构是一种常见折中。

局限性与开放问题

线性注意力有界秩的循环状态既是其定义性特征，也是其主要局限。需要精确检索任意过去词元的任务，例如使用长示例集的上下文学习，最清楚地暴露了这一差距。若干工作试图弥补缺失的容量：门控、多尺度衰减、更大的头数，以及使用诸如 delta 规则等非线性更新规则。

第二个开放问题与硬件有关。循环形式紧凑但本质上是顺序的，而并行形式则需要精心的核函数工程才能匹配融合的 softmax 注意力核函数的吞吐量。Flash Linear Attention 等库已经大幅缩小了差距，但其实现层面的成熟度仍逊于标准注意力。

最后，线性注意力有时被批评为向 RNN 的倒退，确实，循环形式承袭了训练深层循环网络的全部经典困难。现代变体通过精心的初始化、层归一化和衰减参数化来缓解这一问题，但该问题真实存在，在采用线性注意力构建新系统之前值得认真理解。