[Moonshot AI] Attention Residuals: 打破残差连接的“内容无关”禁锢

TL;DR

在现代 LLM 架构中，残差连接（Residual Connections）几乎是不可撼动的基石。然而，Kimi 团队发现，这种自 ResNet 以来沿用至今的 $h_l=h_{l-1}+f(h_{l-1})$ 模式存在一个致命缺陷：它是内容无关且权重固定的。

本文提出的 Attention Residuals (AttnRes) 彻底重构了这一逻辑。它将深层模型看作一个“深度维度的序列”，利用 Softmax Attention 让每一层可以从前面的层中“点菜式”地挑选最相关的特征。实验发现，在 48B 规模的模型上，这不仅稳定了训练动态，更在逻辑推理和代码任务上实现了跨越式提升。

痛点深挖：被遗忘的“深度稀释”问题

为什么标准残差连接（尤其是 PreNorm 变体）在超深网络中会遇到瓶颈？

1. Uniform Aggregation（均一聚合）：每一层都被强制赋予相同的权重（unit weight）。这意味着第 1 层的 embedding 和第 50 层的特征被一锅端地揉在一起，模型缺乏筛选机制。
2. O(L) Magnitude Growth：由于不断加和，隐藏状态的量级随深度线性增长。为了在如此巨大的“背景音”中产生影响，深层网络被迫学习极大的输出，这导致了训练的不稳定性。
3. 信息不可逆丢失：一旦特征在残差流中被错误地合并或抵消，后续层无法再将其召回。

核心机制：深度上的注意力机制

作者提出了一个极具启发性的直觉：深度其实就是另一种维度的序列。 在序列建模中，我们已经从 RNN（固定递推）进化到了 Transformer（全局注意力）；那么在深度维度上，为什么我们还要坚守残差（固定递推）呢？

从 Full AttnRes 到 Block AttnRes

核心公式非常简单直观： $h_l={\sum }_{i=0}^{l-1}{\alpha }_{iol}\cdot v_i$ 其中 $\alpha$ 是由该层特有的伪查询向量 (Pseudo-query) 与前层特征计算出的权重。

为了在大规模训练中降低计算和通信开销（避免每一层都要传输前面所有层的激活值），作者设计了 Block AttnRes：

• 将层分为 $N$ 个块（通常 $N=8$）。
• 块内保留普通残差。
• 块间通过全局 Attention 进行特征检索。

这种设计不仅将内存开销从 $O(L)$ 降到了 $O(N)$，还巧妙地利用**跨阶段缓存（Cross-stage Caching）**解决了流水线并行中的通信瓶颈。

实验与结果：推理能力的阶跃

在 1.4T Tokens 的预训练实验中，AttnRes 展示了惊人的稳定性。

1. 训练动态的改善

观察上图（中间与右侧），标准基线（Baseline）的输出量级随深度失控增长，而 AttnRes 呈现出周期性的平稳状态。更重要的是，AttnRes 的梯度分布更加均匀，解决了底层权重更新过慢的“梯度消失”隐患。

2. 关键战绩

• Scaling Law 优势：Block AttnRes 的训练效率相当于基线的 1.25 倍计算量投入。
• 深度任务表现：在 GPQA-Diamond 推理任务中，AttnRes 相比基线提升了 +7.5 分；在 Math 任务中提升了 +3.6 分；HumanEval 提升了 +3.1 分。

深度洞察：为什么它有效？

通过对学习到的注意力权重可视化（见下表），我们发现了有趣的现象：

• Embedding Persistence：模型底层（h1）的信息会被许多深层持续关注，特别是 Self-Attention 层。
• Content-dependent Selection：模型确实学会了根据当前 Token 的需求，跳过不相关的中间层，直接提取早期特征。

总结与局限

Takeaway: Attention Residuals 正式宣告了“深度维度”也可以被注意力化。它通过极小的参数代价（每层仅一个向量），显著增强了模型对自身内部特征的调度能力。

局限性: 虽然 Block 变体极大缓解了压力，但 Full AttnRes 在硬件互联带宽不足的情况下仍有挑战。此外，对于推理延迟敏感的超轻量模型，AttnRes 引入的额外部署复杂性需要审慎评估。

这篇论文标志着大模型架构优化的一个新方向：不仅要在水平方向（序列）做长，更要在垂直方向（深度）做深、做活。