💻 computer science

Multi-Head Attention Residuals

本文引入了多头注意力残差（Multi-Head Attention Residuals, MHAR），这是一种无参数的架构修改，它通过多个路由头使不同的特征子空间能够独立地关注深度历史，从而解决了子空间不一致问题，并与标准 Transformer 相比，在各种模型规模上显著提升了验证损失和下游任务性能。

原作者： Cheng Luo, Zefan Cai, Junjie Hu

发布于 2026-07-31

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Cheng Luo, Zefan Cai, Junjie Hu

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正在建造一座宏伟的多层图书馆，每一层新楼都增加了一个新的知识房间。在人工智能的世界里，这些“图书馆”被称为 Transformer，它们是支撑我们今天使用的许多智能工具背后的“大脑”。为了理解一个故事或一道数学题，这座图书馆需要在处理第十层时，依然能记住第一层发生了什么。传统上，这些图书馆通过一条单一且狭窄的走廊连接每一层。当第十层的房间需要回溯时，它只能窥视紧邻其下方的那个房间。它必须信任下方的房间已经完美地总结了从第一层到当前层的所有内容。这就像试图通过询问坐在你正前方的人来回忆一整部电影的内容；如果那个人漏掉了一个细节，你也就会错过它。

为了解决这个问题，研究人员最近尝试了一个新想法，叫做“注意力残差”（Attention Residuals）。与其仅仅看向下方的房间，研究人员给每个房间配备了一个神奇的望远镜，使其可以观察任何之前的楼层。然而，这里有一个限制：望远镜只有一个透镜。房间里的每一个细节——无论是颜色、声音还是数字——都必须通过这同一个单透镜来决定应该关注哪一层。这就像是强迫一群口味迥异的朋友（一个喜欢动作片，一个喜欢纪录片，还有一个喜欢动画片）去达成一致，共同观看同一个频道。他们不得不做出妥协，而这种妥协往往意味着没有人能看到他们真正需要看到的内容。

这篇论文介绍了一个聪明的升级方案，叫做多头注意力残差（Multi-Head Attention Residuals, MHAR）。研究人员意识到，正如电影院有多个银幕同时播放不同类型的电影一样，图书馆的望远镜也应该有多个透镜。他们将那个单一的、共享的透镜拆分成了几个更小的、独立的透镜。现在，“动作片”部分可以观察充满动作感的楼层，而“纪录片”部分可以观察历史类楼层，这一切都不需要争吵或妥协。

团队通过构建不同规模的图书馆来测试这个想法，规模从 1 亿参数的小型模型一直到庞大的 10 亿参数模型。他们发现，对于小型图书馆，旧的单透镜系统表现尚可。但随着图书馆变得越来越大、越来越复杂，单透镜成为了一个瓶颈，甚至导致 AI 的表现比标准的走廊系统还要差。然而，新的多透镜系统则随着图书馆规模的扩大而变得越来越好。在最大规模下，与标准模型相比，新方法将 AI 的困惑度（以验证损失衡量）降低了约 0.063，而旧的单透镜方法实际上让情况恶化了 0.140。

研究人员还展示了这项升级并不需要制造更大的望远镜，也不需要增加图书馆的重量；它只是对现有玻璃进行的一种更聪明的排列方式。他们甚至开发了特殊的超快速工具，以确保图书馆不会因为额外的观察工作而变得迟缓，使速度几乎与原始设计一样快。最后，他们证明了即使是在一个巨大的、预建的 80 亿参数图书馆上，这个技巧依然有效，能够提升其解决数学问题和回答复杂科学问题的能力，而无需从头重建。核心发现是：随着 AI 模型变得更宽、更复杂，它们不再对该记住什么达成共识，因此，赋予它们多种独立的方式进行回溯，是释放其全部潜力的关键。

技术摘要：多头注意力残差 (Multi-Head Attention Residuals, MHAR)

1. 问题陈述

标准 Transformer 通过单一的加性残差流在深度方向上传播信息，其中每个子层仅读取紧接的前一状态（ $h_\ell = h_{\ell-1} + f_\ell(h_{\ell-1})$ ）。虽然早期层计算出的有用特征会保留在运行总和中，但它们变得无法被单独寻址。

注意力残差 (Kimi, 2025) 通过将所有先前的子层输出视为一种记忆来放宽这一限制，允许每个子层检索整个深度历史中经过学习的、经 Softmax 加权的组合。然而，该机制存在一个结构性瓶颈：它使用一个单一共享的路由查询 ( $q \in \mathbb{R}^d$ ) 来对深度历史进行评分。因此，所有 $d$ 个特征维度都被迫通过同一个分布 $\alpha$ 以相同的比例从相同的层中读取信息。

论文指出，这造成了一种**“被迫妥协” (forced compromise)**。随着模型宽度的增加，不同的特征子空间可能会对哪些历史层最相关产生分歧。单一查询无法同时满足这些不同的偏好，从而将截然不同的深度偏好压缩为单一的读取过程。这种妥协的代价随模型宽度的增加而增长，在更大规模下，单查询路由可能会变得有害。

2. 方法论：多头注意力残差 (MHAR)

作者提出了多头注意力残差 (MHAR)，这是对注意力残差的一种严格泛化，它在不增加参数或显著计算量的情况下消除了这种被迫妥协。

核心机制

MHAR 不再使用单一查询 $q$ ，而是将路由查询重塑为 $H$ 个针对每个子空间的头，每个头都有其独立的对深度历史的 Softmax。

重塑 (Reshaping)： 将单一查询向量 $q \in \mathbb{R}^d$ 重塑为形状为 $(H, d/H)$ 的张量。
独立路由 (Independent Routing)： 对于每个头 $h$ ，模型根据对应于该头的源向量切片计算一个独特的深度分布 $\alpha^{(h)}$ 。
块对角读取 (Block-Diagonal Read)： 检索到的混合物是块对角的；每个子空间独立地读取深度历史。
恢复 (Recovery)： 当 $H=1$ 时，可以完全恢复为原始的单头注意力残差。

实现细节

参数效率： 该变换是一种重塑操作，相对于单头基准增加了零参数。
计算效率： 开销微乎其微，仅涉及在深度轴上进行的 $H$ 个小型 Softmax（长度 $\le 2L+1$ ），而非一个大型 Softmax。
系统优化： 意识到深度路由是内存受限的（需要 $O(L^2 B T d)$ 的流量），作者提供了融合 Triton 内核。这些内核执行带有融合 RMSNorm 的在线 Softmax，消除了中间键 (keys) 和堆叠源 (sources) 的物化过程。这使得路由时间比 torch.compile 快 $2\times$ ，比 eager 参考实现快 $6.4\times$ ，将训练吞吐量提升至基准水平的 $0.55\text{--}0.88\times$ 。

超参数设置

路由头的数量 $H$ 设置为与模型中的 KV 头 数量相等。作者发现这是一个近乎最优且无需超参数的默认设置，使路由粒度与模型的消费粒度保持一致。

3. 核心贡献

MHAR 架构： 一种无参数修改的方法，通过将路由查询拆分为独立的头，消除了单查询深度路由中的强制妥协。
高效内核： 融合 Triton 内核使深度路由在大规模应用中变得可行，显著降低了内存开销和延迟。
恒等保持转换 (Identity-Preserving Conversion)： 一种将 MHAR 嫁接到预训练模型上的方法，通过使用“增量注意力残差”（将路由项添加到现有流中而非替换）来实现，从而在不引起损失突跳的情况下进行中途训练。
机制洞察： 对已训练查询的直接探测证实，子空间之间的分歧会随宽度增加，这是单头路由性能下降的主要驱动因素。

4. 实验结果

从头预训练 (Pretraining From Scratch)

在三个规模（100M, 350M, 1B）上基于 FineWeb-Edu 进行训练：

验证损失： MHAR 在验证损失上始终优于标准 Transformer 和单头注意力残达。
- 100M： 比基准降低 $-0.049$ nats。
- 350M： 比基准降低 $-0.080$ nats。
- 1B： 比基准降低 $-0.063$ nats。
缩放行为： MHAR 相对于单头路由的优势随规模扩大而增加。虽然单头路由在 100M 时有所帮助（$-0.039 $），但在 350M 时开始表现出负面影响（$ +0.055 $），并在 1B 时变得有害（$ +0.140$）。
计算等效性： MHE 实现了相对于基准 $1.3\text{--}1.5\times$ 的计算等效收益。

中途训练 (Mid-Training, 8B)

在 8B 模型 (Marin-8B) 上使用匹配调度控制的方案，并在面向能力的语料库上进行退火：

恒等保持： 转换过程未引入优化冲击；前 100 步的损失曲线与基准一致。
下游收益： MHAR 在推理任务中带来了显著的改进：
- GSM8K： $+3.2$ 分（配对 McNemar $p=0.004$ ）。
- GPQA： $+3.1$ 分（配对 McNemar $p=0.038$ ）。
- 在 MMLU、MATH 和代码任务中的收益在该特定设置下未达到统计学显著。

消融研究

头数 (Head Count)： 最优头数 $H$ 与 KV 头数一致。在欠训练阶段，更多的头总是会有所帮助；在收敛阶段，$H=KV$ 是曲线的“拐点”。
宽度 vs. 深度： 对已训练查询的探测显示，“宽度分歧”（不同子空间之间发散的路由偏好）随模型宽度单调增加，验证了强制妥协的假设。
头对齐 (Head Alignment)： 将路由头与注意力 KV 组对齐相比于任意子空间路由并无明显收益，这表明收益来自于头的多样性，而非其特定的对齐结构。

5. 重要性与主张

论文声称，多头注意力残差 (MHAR) 解决了 Transformer 在访问深度历史方面的一个根本架构限制。通过允许不同的特征子空间独立地关注不同的层，MHAR 在单查询路由失效的情况下实现了有效的扩展。

作者强调：

改进并非源于增加参数或计算量，而是由于移除了一个结构性瓶颈。
由于采用了融合内核，该方法在规模化应用时是实用的，缓解了深度路由的内存受限问题。
该方法具有通用性，既可以从头开始训练，也可以通过恒等保持的中途训练进行应用。
多头路由的必要性是模型宽度的直接结果；随着模型增长，单查询带来的“强制妥协”会成为负担，使得 MHAR 成为扩展深度感知架构的关键组件。