# 《Attention Residuals》深度解读

**论文信息**
- 标题：Attention Residuals
- 作者：Kimi Team (Moonshot AI)
- 发表时间：**2026年3月**（arXiv:2603.15031）
- arXiv：[2603.15031](https://arxiv.org/abs/2603.15031)
- GitHub：[MoonshotAI/Attention-Residuals](https://github.com/MoonshotAI/Attention-Residuals)

---

## 一句话总结

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

---

![alt text](atten_residuals.png)

## 背景与动机

### 标准残差连接的局限性

现代大语言模型（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$ 块输出的聚合表示（块内所有层输出的求和）
- 注意力只在块级别进行，显著降低复杂度

#### 复杂度分析

| 方法 | 每层内存开销 | 通信开销 | 计算复杂度 |
|------|-------------|----------|------------|
| 标准残差 | $O(d)$ | $O(d)$ | $O(1)$ |
| Full AttnRes | $O(Ld)$ | $O(Ld)$ | $O(L)$ |
| **Block AttnRes** | $O((S+N)d)$ | $O((S+N)d)$ | $O(S+N)$ |

其中 $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）：**

| 激活参数 | Tokens | Baseline Loss | Block AttnRes | Full AttnRes |
|----------|--------|---------------|---------------|--------------|
| 194M | 38.7B | 1.931 | 1.909 | 1.899 |
| 241M | 45.4B | 1.895 | 1.875 | 1.869 |
| 296M | 62.1B | 1.829 | 1.809 | 1.804 |
| 436M | 87.9B | 1.766 | 1.746 | 1.737 |
| 528M | 119.0B | 1.719 | 1.693 | 1.692 |

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

### 4.2 消融实验 (Ablation Studies)

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

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

| 变体 | Loss |
|------|------|
| Baseline (PreNorm) | 1.766 |
| DenseFormer | 1.767 |
| mHC | 1.747 |
| **AttnRes Full** | **1.737** |
| w/ input-dependent query | 1.731 |
| w/ input-independent mixing | 1.749 |
| w/ sigmoid | 1.741 |
| w/o RMSNorm | 1.743 |
| SWA (W=1+8) | 1.764 |
| **Block (S=4)** | **1.746** |
| w/ multihead (H=16) | 1.752 |
| w/o RMSNorm | 1.750 |

**关键发现**：
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）：**

**通用能力任务：**

| 任务 | Baseline | AttnRes | 提升 |
|------|----------|---------|------|
| MMLU | 73.5 | 74.6 | +1.1 |
| MMLU-Pro | 52.2 | 52.2 | 0.0 |
| GPQA-Diamond | 36.9 | 44.4 | +7.5 |
| BBH | 76.3 | 78.0 | +1.7 |
| ARC-Challenge | 64.6 | 65.7 | +1.1 |
| HellaSwag | 83.2 | 83.4 | +0.2 |
| TriviaQA | 69.9 | 71.8 | +1.9 |

**数学与代码任务：**

| 任务 | Baseline | AttnRes | 提升 |
|------|----------|---------|------|
| GSM8K | 81.7 | 82.4 | +0.7 |
| MGSM | 64.9 | 66.1 | +1.2 |
| Math | 53.5 | 57.1 | +3.6 |
| CMath | 84.7 | 85.1 | +0.4 |
| HumanEval | 59.1 | 62.2 | +3.1 |
| MBPP | 72.0 | 73.9 | +1.9 |

**中文任务：**

| 任务 | Baseline | AttnRes | 提升 |
|------|----------|---------|------|
| CMMLU | 82.0 | 82.9 | +0.9 |
| C-Eval | 79.6 | 82.5 | +2.9 |

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

| 指标 | 标准残差 | AttnRes | 改进 |
|------|----------|---------|------|
| 验证集最终 Loss | 1.766 | 1.737 | -1.6% |
| 输出幅度（深层） | 大幅增长 | 保持平稳 | 显著改善 |
| 梯度分布 | 深层梯度稀释 | 更均匀分布 | 显著改善 |

**关键发现**：
- 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 工具，但在正文中展示的分析相对有限，读者可能希望看到更多关于如何使用该工具诊断模型的细节。

---

## 延伸阅读

### 相关工作对比

| 方法 | 核心思想 | 与本文的关系 |
|------|----------|--------------|
| **ReZero** (Bachlechner et al., 2021) | 用可学习标量缩放残差分支 | AttnRes 扩展了这一思想，用 attention 权重替代单一标量 |
| **mHC** (Xie et al., 2026) | 用混合矩阵实现跨层连接 | 论文中与 AttnRes 对比，AttnRes 性能更优 (1.737 vs 1.747) |
| **DenseFormer** (Pagliardini et al., 2024) | 加权聚合前面所有层 | 论文中对比，AttnRes 显著优于该方法 (1.737 vs 1.767) |
| **Highway Networks** (Srivastava et al., 2015) | 门控残差连接 | 与 AttnRes 正交，可以结合使用 |
| **Deepspeed's Curriculum Learning** | 自适应深度训练 | 与 AttnRes 的深度选择机制有相似之处 |
| **Dynamic Depth Networks** | 动态选择网络深度 | AttnRes 可以看作是一种软性的、可学习的动态深度 |

### 推荐资料

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

2. **Deep Residual Learning for Image Recognition** (ResNet, He et al., 2016)
   - 残差连接的奠基性工作

3. **On Layer Normalization in the Transformer Architecture** (PreNorm, Xiong et al., 2020)
   - PreNorm vs PostNorm 的分析

4. **ReZero is All You Need** (Bachlechner et al., 2021)
   - 可学习残差缩放的相关工作

**技术博客**：
- 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](https://xrollout.org/blog/attn_residuals.zh)
>
> 我会在网站持续更新 MemoryHub 的工程实现细节和开源进展，欢迎收藏关注。