实现线性化的关键：分析驱动的 Transformer 线性化

论文信息

标题: The Key to Going Linear: Analysis-Driven Transformer Linearization

作者: Anna Kuzina, Paul N. Whatmough, Babak Ehteshami Bejnordi

发布日期: 2026-07-08

PDF 链接: 下载 PDF

背景与动机：突破 Transformer 长上下文瓶颈

基于 Transformer 的大型语言模型在长上下文场景中面临严峻的计算挑战。因果自注意力机制的复杂度随序列长度 $T$ 呈二次增长，同时键值（KV）缓存的内存占用也线性膨胀，这使得长上下文推理成本高昂。一种极具吸引力的解决方案是事后线性化（post hoc linearization）：将已预训练的完整注意力模型转换为线性时间架构，而无需从头重新训练。近年来涌现出大量此类方法，但它们通常同时采用多种干预手段——如低秩适配（LoRA）、滑动窗口注意力（SWA）、混合路由和蒸馏——这些混合因素使得一个根本性的设计问题变得模糊不清：

哪种线性状态更新机制能够最好地逼近预训练的 softmax 注意力？

高通 AI 研究院的这篇论文正是围绕这一问题展开。作者提出了一种分析驱动的线性化范式，在最严格的条件下（完全冻结预训练骨干，仅训练新引入的线性注意力参数）隔离不同线性替换机制的能力，并首次从理论上揭示了 delta 风格的更新自然匹配 softmax 的关键几何结构。

核心方法：将线性注意力统一在投影视角下

研究比较了三类主流的线性注意力机制替代方案，所有实验均保持 Transformer 原有的 Q、K、V 投影权重和 RMSNorm 参数完全冻结，只训练各机制特有的小规模参数（见图 1）。这种严格设定剥离了其他干预的混淆效应，让替换机制的固有近似能力赤裸裸地暴露出来。

三种机制分别是：

归一化核注意力（以 Hedgehog 为代表） 使用可学习的特征映射 $\phi(q)$ 和 $\phi(k)$ 近似指数核： $\exp(q^\top k) \approx \phi(q)^\top \phi(k)$ ，进而通过递推方式计算注意力输出，维持一个归一化项 $Z_t$ 。该形式天然保留了 softmax 的归一化属性，但表达能力受限于特征映射的设计。
门控线性注意力（GLA） 采用无归一化的状态递推： $S_t = G_t S_{t-1} + k_t^\top v_t$ ，其中 $G_t = \operatorname{diag}(\alpha_t)$ 是一个与输入相关的对角遗忘门。这种机制具备硬件高效的实现，但其遗忘因子仅依赖于当前 token 的隐状态，不直接与键交互，因此在捕捉键之间的结构性关系时存在盲区。
门控 Delta 网络（GDN / Kimi Linear） 在 GLA 的基础上引入了 Delta 规则： $S_t = (I - \beta_t k_t k_t^\top) G_t S_{t-1} + \beta_t k_t^\top v_t$ 。这里的核心是一个秩 1 修正 $(I - \beta_t k_t k_t^\top)$ ，其效果是保留与当前键正交的成分，而衰减对齐的成分。这使得每个新 token 能够根据自身的键方向有选择性地抑制旧记忆。

理论突破：Softmax 的一阶近似揭示 Delta 更新的天然优势

论文的核心理论贡献是一组引理，它们清晰地揭示了为什么 delta 式更新优于纯门控累积。作者从 softmax 注意力权重出发：

P_i^{\text{full}} = \frac{\exp(q^\top k_i)}{\sum_j \exp(q^\top k_j)}

通过一阶泰勒展开（在 $\bar{x} = \frac{1}{t}\sum_j q^\top k_j$ 处线性化），并经过键归一化和均值投影近似，最终推导出：

P_i^{\text{full}} \approx \frac{1}{t\sqrt{d}} q^\top \prod_j (I - \beta_j k_j k_j^\top) k_i + \frac{1}{t}

这一形式与 GDN 的隐含注意力权重（公式 9）完美吻合。推导过程中关键步骤包括：

将键的均值减法替换为投影： $k_i - \bar{k} \approx (I - \Pi_{\bar{k}}) k_i$ ，误差在键方向集中时很小。
将单一正交投影近似为一串秩 1 更新的乘积： $\prod_{j} (I - \beta_j k_j k_j^\top)$ 在键高度聚集于均值方向时，行为与 $I - c \Pi_{\bar{k}}$ 一致。

这向量的乘积结构表明，GDN 中的每个 $I - \beta_j k_j k_j^\top$ 操作相当于对先前的键进行方向性抑制，这正是 softmax 注意力中查询 - 键交互产生动态重加权的一种线性模拟。相比之下，GLA 的对角遗忘矩阵仅依赖时间先后而不具备这种键依赖的几何修正能力，因此逼近软注意力时必然会丢失重要信息。

基于这一洞察，作者提出了一个 GLA 的变体 kGLA（键门控线性注意力），将对角遗忘矩阵与近似秩 1 项 $\operatorname{Diag}(k_j^2)$ 结合，虽不如完整秩 1 形式灵活，但保留了对角结构带来的硬件效率，并试图缩小与 GDN 的差距。

实验发现：从严格线性化到实用补偿

严格设置下的性能断层

在 LLaMA 3.1 8B 和 Qwen 3 8B 上使用纯线性替换（无 SWA、无投影适应）训练 10M token 后，GDN 全面优于 GLA、kGLA 和 Hedgehog（见表 1）。以 LLaMA 3.1 为例，GDN 在常识推理（CR）上达到 63.08，而 GLA 仅 61.81，Hedgehog 更是只及 36.13；MMLU 和 Lambada 上 GDN 也优势明显。值得注意的是，核方法的完全失败表明，在严格冻结骨干的条件下，仅靠特征映射学习注意力模式是不可行的。

补偿措施的引入与效果

为弥合剩余的性能差距，论文逐步重新引入常见的结构补偿：

滑动窗口注意力（SWA）：将输入分为线性路径（长程）和 SWA 路径（局部），两者独立计算后求和，避免了混合权重。这种设计让每个 token 要么进入线性缓存，要么进入局部窗口，清晰分离了职责。
Sink tokens：理论分析和 MSE 曲线显示，序列开头的少量 token 最难被近似（因为泰勒展开中的均匀偏置 $1/t$ 在小 $t$ 时显著）。作者将固定的 64 token 缓存预算划分为 56 token 滑动窗口和 8 个 sink token（始终由完整注意力处理的初始 token）。这一简单修改带来了明显的性能提升（表 2），完全弥补了常识推理上的差距，并在 MMLU 上大幅缩小了与原始模型的差距。
Q/K/V 投影适应：添加短卷积（核尺寸 8）或 LoRA（秩 8）到 Q、K、V 投影中。结果显示短卷积普遍带来更好且更稳定的提升，而在知识密集型的 MMLU 上差距仍维持在约 6%，但在上下文依赖型 Lambada 任务上几乎追平原始模型（表 3）。

最终系统对比与扩展

在最终配置（GDN + 64 token 预算 SWA + sink + 短卷积）下，模型在常见推理基准上与全注意力模型匹敌，并在 5-shot MMLU 上显著优于 LoLCaTs 和 Liger-GLA 等先前的线性化方法（表 4）。更难得的是，即使将缓存预算扩大到 128 以匹配其他自适应 token 缓存方法（Lizard、LoLA、STILL），该方法仍展现出竞争力，特别是在 MMLU 上领先。

长上下文评估方面，在 S-NIAH 和扩展 RULER 测试中，该方法在较大缓存预算（896 token）下性能强劲，但暴露出在多并行记忆检索任务上的短板（表 9），表明未来可结合内容感知的 token 选择策略。

此外，该方法展现出优异的模型规模扩展性：在 Qwen 3 系列上从 0.6B 到 32B 参数均保持线性化后性能与原始模型的平行缩放（图 4）。在混合设置中，若根据层的关键浓度和投影残差信号选择少量层保留全注意力，可比均匀选择更有效地提升上下文依赖类任务（图 5），为实际部署提供了灵活的精度 - 效率折中。

实践启示与未来方向

这项研究为希望线性化大模型的研究者和工程师提供了清晰的方法论：首选 Delta 更新机制（如 GDN），并配合滑动窗口加少量 sink token 的结构。在进行线性化时，应避免在黑色盒中同时使用多种技术，而应优先理解哪些结构能弥补线性机制的固有缺陷。短卷积是一种性价比极高的增强，而 LoRA 需更长的训练才能释放潜力。对于混合推理场景，测量各层的 “线性化难度”（如键浓度、投影残差）可以指导保留全注意力的层选择。

长远来看，该工作将线性注意力近似误差溯源至键的几何分布，为开发更高级的线性化策略提供了理论武器。结合自适应 token 选举机制对齐这些误差分布，有望进一步解锁长上下文线性推理的全部潜力。

总结

本文通过严格的分析和实验，确定了 Delta 风格线性更新（具体为 GDN）是后 hoc 线性化 Transformer 的最佳核心组件。一阶理论深刻解释了 softmax 注意力中依赖键的秩 1 动态，为选择线性注意力机制提供了指导原则。通过重新引入简洁的补偿结构——sink tokens、无加权混合的 SWA 和短卷积——最终方案在保持几乎可忽略的参数增量的同时，超越了先前的后线性化方法，并能稳定扩展到数十亿参数模型。这一分析驱动的方法不仅得出了一套强大且可解释的线性化配方，更为未来长上下文高效推理的系统设计奠定了坚实基础。