《Attention Residuals》深度解读

论文信息
- 标题：Attention Residuals
- 作者：Kimi Team (Moonshot AI)
- 发表时间：2026年3月（arXiv:2603.15031）
- arXiv：2603.15031
- GitHub：MoonshotAI/Attention-Residuals

一句话总结

本文提出 Attention Residuals (AttnRes)，用 softmax attention 机制 替代 Transformer 中固定的残差连接，使每一层能够选择性地聚合前面所有层的输出，从而解决 PreNorm 带来的隐藏状态膨胀和梯度稀释问题。

背景与动机

标准残差连接的局限性

现代大语言模型（LLM）广泛采用 PreNorm + 残差连接 的架构【原文 §1】：

$$h_l = h_{l-1} + f_l(\text{LayerNorm}(h_{l-1}))$$

这种设计有两个主要作用：
1. 梯度高速公路：通过恒等映射让梯度绕过变换层，实现深层网络的稳定训练
2. 逐层信息聚合：每一层接收前面所有层输出的均匀加权和

然而，作者指出这种均匀聚合存在根本性缺陷：

"Unlike sequence mixing and expert routing, which now employ learnable input-dependent weighting, this depth-wise aggregation remains governed by fixed unit weights, with no mechanism to selectively emphasize or suppress individual layer contributions."
（与现在使用可学习、输入依赖加权的序列混合和专家路由不同，这种深度方向的聚合仍然由固定的单位权重控制，没有机制来选择性地强调或抑制各个层的贡献。）

PreNorm 的隐藏状态膨胀问题

PreNorm 架构虽然比 PostNorm 更稳定，但带来了 "uncontrolled hidden-state growth"（不可控的隐藏状态增长） 问题【原文 §1】：

• 每一层都累加前面所有层的输出
• 随着深度增加，隐藏状态的幅度不断增大
• 这导致梯度分布不均匀，深层梯度被稀释

作者引用先前工作 [27] 指出这一问题，并提出用 Attention Residuals 来解决。

核心方法

3.1 Attention Residuals (AttnRes)

作者提出用 softmax attention 机制替代固定的残差连接【原文 §3】。

核心思想

标准残差连接：
$$h_l = h_{l-1} + f_l(\text{LN}(h_{l-1}))$$

所有前面层的影响是均匀累加的。

Attention Residuals：
$$h_l = \sum_{i=0}^{l-1} \alpha_{l,i} \cdot v_i$$

其中 $\alpha_{l,i}$ 是通过 softmax attention 计算出的输入依赖权重：
$$\alpha_{i \rightarrow l} = \frac{\phi(\boldsymbol{q}_l, \boldsymbol{k}_i)}{\sum_{j=0}^{l-1} \phi(\boldsymbol{q}_l, \boldsymbol{k}_j)}$$

注意力分数通过查询-键机制计算，采用 RMSNorm 归一化：
$$\phi(\boldsymbol{q}, \boldsymbol{k}) = \exp\left(\frac{\boldsymbol{q}^\top \text{RMSNorm}(\boldsymbol{k})}{\sqrt{d}}\right)$$

其中查询向量 $\boldsymbol{q}_l = \boldsymbol{w}_l$ 是每层学习的参数，键向量 $\boldsymbol{k}_i = \boldsymbol{v}_i$ 是第 $i$ 层的输出。

视觉对比（基于原文 Figure 1）

标准残差连接 (a)：

Embedding → [Layer 1] → [Layer 2] → [Layer 3] → Output
              ↑ +       ↑ +         ↑ +
             (uniform additive accumulation)

Full Attention Residuals (b)：

Embedding → [Layer 1] → [Layer 2] → [Layer 3] → Output
              ↓           ↓           ↓
              └─────[Attention]─────┘
                   (selective aggregation
                    with learned weights)

Block Attention Residuals (c)：

Embedding → Block 1 (Layers 1-4) → Block 2 (Layers 5-8) → Output
                          ↓                         ↓
                          └──────────[Attention]────┘
                           (block-level selective
                            aggregation)

3.2 Block Attention Residuals (BlockAttnRes)

Full Attention Residuals 需要对所有前面的层计算 attention，这在模型很深时会带来显著的内存和通信开销【原文 §3.2】。

解决方案：分层聚合

作者提出 Block Attention Residuals，将层分组为块（blocks）：

$$h_l = \sum_{b=0}^{B-1} \alpha_{l,b} \cdot \boldsymbol{b}_b$$

其中：
- 模型被划分为 $N$ 个块，每个块包含 $S = L/N$ 层
- $\boldsymbol{b}_b$ 是第 $b$ 块输出的聚合表示（块内所有层输出的求和）
- 注意力只在块级别进行，显著降低复杂度

复杂度分析

其中 $L$ 是层数，$d$ 是隐藏维度，$S$ 是每块层数，$N$ 是块数（$N \ll L$）。

实际实现优化

作者还引入了以下优化【原文 §4】：

1. Cache-based Pipeline Communication：利用缓存机制减少跨层通信开销
2. Two-phase Computation Strategy：分阶段计算，Phase 1计算块间attention，Phase 2计算块内attention

这些优化使得 BlockAttnRes 成为实用的、可直接替换标准残差连接的方案。

3.3 理论分析

作者分析了 Attention Residuals 如何解决 PreNorm 的问题【原文 §6.2】。

隐藏状态控制

标准 PreNorm：
$$\|h_L\| \approx \|h_0\| + \sum_{l=1}^{L} \|f_l(\text{LN}(h_{l-1}))\|$$

由于所有残差都是正向累加，隐藏状态随深度单调增长。

Attention Residuals：
$$\|h_L\| = \|\sum_{i=0}^{L-1} \alpha_{L,i} \cdot \boldsymbol{v}_i\|$$

由于 softmax 权重 $\alpha_{L,i}$ 的总和为 1（或可控值），隐藏状态的幅度可以被有效控制，避免无限制增长。

梯度流动改善

在标准残差连接中，深层梯度需要反向传播经过所有前面的层，导致梯度稀释。Attention Residuals 允许梯度通过 attention 权重选择性地流动，使得梯度可以更直接地从损失函数传递到相关的层，改善训练稳定性。

实验与结果

作者在多个维度上验证了 Attention Residuals 的有效性【原文 §5】。

4.1 缩放定律实验 (Scaling Law Experiments)

作者在多个模型规模下验证了改进的一致性【原文 §5.1, Table 2】。

实验设置：
- 模型规模：从 194M 到 528M 激活参数（5个规模点）
- 训练数据：Common Crawl 子集
- 训练 tokens：38.7B - 119.0B

论文真实数据（Table 2）：

关键发现：
1. 改进随规模增加：更大的模型从 Attention Residuals 中获益更多
2. 一致性：在所有测试规模下，AttnRes 都优于标准残差连接
3. Block vs Full：Block AttnRes 在大部分规模下接近 Full AttnRes 的性能

4.2 消融实验 (Ablation Studies)

作者验证了内容依赖的深度选择机制的贡献【原文 §5.3, Table 4】。

论文真实数据（Table 4，16层模型）：

关键发现：
1. 内容依赖机制的重要性：使用输入独立混合（1.749）比完整 AttnRes（1.737）性能差
2. Softmax 优于 Sigmoid：使用 sigmoid 代替 softmax 导致性能下降（1.741 vs 1.737）
3. RMSNorm 的重要性：移除 RMSNorm 会导致性能下降（Full: 1.743, Block: 1.750）

4.3 Kimi Linear 架构集成

作者在实际的 Kimi Linear 架构中集成了 Attention Residuals【原文 §5.2, Table 3】。

论文真实数据（Table 3）：

通用能力任务：

数学与代码任务：

中文任务：

训练动态改善【原文 Figure 5】：

关键发现：
- AttnRes 在 GPQA-Diamond（+7.5）、Math（+3.6）、HumanEval（+3.1）等复杂推理任务上提升最大
- 在 MMLU-Pro 上持平，在其他所有任务上都有提升
- 平均来看，AttnRes 在各领域任务上都有 1-4 个百分点的稳定提升

亮点与局限

亮点

1. 问题识别精准：准确指出了标准残差连接和 PreNorm 的根本问题——均匀累积导致的隐藏状态膨胀和梯度稀释。

2. 方法设计优雅：用 softmax attention 替代固定权重是一个自然的扩展，既保留了残差连接的梯度流动优势，又引入了内容依赖的选择机制。

3. 工程实现务实：BlockAttnRes 的提出体现了从理论到实践的务实态度，通过分层聚合解决了内存和通信瓶颈，使得方法可以扩展到大规模模型训练。

4. 实验验证全面：从缩放定律（194M-528M参数）、消融实验（16层模型），到实际模型（Kimi Linear 48B总参/3B激活参数）的集成，实验设计全面且有说服力。

5. 开源贡献：代码开源有利于社区验证和进一步开发。

局限（原文承认的）

1. 计算开销：虽然 BlockAttnRes 大幅降低了开销，但相比标准残差仍有额外计算（attention 计算、跨层通信）。

2. 超参数敏感性：块大小 $S$、块数量 $N$ 等超参数可能需要针对不同模型规模进行调优。论文发现 $N \approx 8$ 在大部分规模下效果较好。

3. 架构局限性：实验主要在 decoder-only 架构上进行，对于 encoder-decoder 或 state space models 的适用性需要进一步验证。

局限（解读者补充）

1. 与现有技术的对比可以更充分：论文主要与标准残差对比，但与其他残差变体（如 ReZero, mHC, DenseFormer）的对比仅在消融实验中进行，缺少更系统的比较。

2. 长序列场景：虽然 BlockAttnRes 解决了内存问题，但对于超长序列（>100K），attention 本身的计算复杂度仍然是瓶颈。

3. 可解释性分析有限：论文提到了 Persistency Map 工具，但在正文中展示的分析相对有限，读者可能希望看到更多关于如何使用该工具诊断模型的细节。

延伸阅读

相关工作对比

推荐资料

核心论文：
1. Attention Residuals (本文)
- https://arxiv.org/abs/2603.15031
- GitHub: https://github.com/MoonshotAI/Attention-Residuals

1. Deep Residual Learning for Image Recognition (ResNet, He et al., 2016)

2. 残差连接的奠基性工作

3. On Layer Normalization in the Transformer Architecture (PreNorm, Xiong et al., 2020)

4. PreNorm vs PostNorm 的分析

5. ReZero is All You Need (Bachlechner et al., 2021)

6. 可学习残差缩放的相关工作

技术博客：
- Kimi/Moonshot AI 官方博客（如有发布相关解读）
- Hugging Face Blog 关于 Transformer 架构的讨论

开源代码：
- https://github.com/MoonshotAI/Attention-Residuals
- 可结合 Megatron-LM 或 DeepSpeed 进行大规模训练实验

总结与思考

核心贡献回顾

1. 问题定义清晰：准确识别了标准残差连接 + PreNorm 架构的核心问题——均匀累积导致的隐藏状态膨胀和梯度稀释。

2. 方法设计优雅：用 softmax attention 替代固定权重，引入内容依赖的深度选择机制，既保留了残差连接的梯度流动优势，又实现了灵活的信息聚合。

3. 工程实践务实：BlockAttnRes 的提出解决了全量 AttnRes 的内存和通信瓶颈，通过分层聚合使得方法可扩展到大规模模型训练。

4. 实验验证充分：从缩放定律（194M-528M参数）、消融实验（16层模型），到 Kimi Linear (48B总/3B激活) 的实际集成，多维度验证了方法的有效性和实用性。

方法论的启示

从现象到方案的研究路径：
1. 识别问题：PreNorm 的隐藏状态膨胀是已知问题，但之前缺乏系统性解决方案
2. 扩展机制：将序列混合中成功的 attention 机制扩展到深度方向
3. 工程优化：通过分层设计解决可扩展性问题
4. 全面验证：从理论分析到实际部署的完整验证

架构设计的思考：
- 残差连接虽然是 Transformer 成功的关键，但并非不可改进
- 通过引入学习机制替代固定规则，可以在保持稳定性的同时增强灵活性
- 深度方向的聚合与序列方向的混合同等重要，值得更多研究关注

未来展望

短期方向：
1. 更广泛的架构验证：在 encoder-decoder、视觉 Transformer、多模态模型中验证 AttnRes
2. 超参数自动调优：论文发现 $N \approx 8$ 效果较好，但需要更系统的自动调优方法
3. 与现有技术结合：与 MoE、稀疏注意力、量化等技术结合

长期方向：
1. 自适应深度网络：让模型动态决定需要多少层，实现样本依赖的计算量
2. 记忆机制扩展：将 AttnRes 与外部记忆结合，处理超长序列
3. 硬件协同设计：为 AttnRes 设计专用硬件加速器

对社区的贡献：
- 开源代码和工具（如 Persistency Map）的进一步完善
- 在更多开源模型（如 Llama、Mistral、Qwen）上的预训练实验
- 与 DeepSpeed、Megatron-LM 等训练框架的深度集成

完整技术文档及后续更新： https://xrollout.org/blog/attn_residuals.zh

我会在网站持续更新 MemoryHub 的工程实现细节和开源进展，欢迎收藏关注。