News

Attention Residuals

原文链接：https://arxiv.org/abs/2603.15031 作者：Kimi Team (Moonshot AI) 机构：Moonshot AI（月之暗面）发布日期：2026-03-16 代码：https://github.com/MoonshotAI/Attention-Residuals

速查卡

项目	内容
一句话总结	用 softmax 注意力替代固定权重的残差累加，让每层可以选择性地聚合之前各层的输出
大白话版	传统 Transformer 每一层只能看到”前面所有层混在一起的结果”，就像把所有颜料搅在一起。AttnRes 让每层可以从之前的各层中”点菜”——需要哪层的信息就多用哪层的
核心数字	Block AttnRes 等效于 1.25× 计算量的性能提升；48B 模型 GPQA-Diamond +7.5；推理延迟增加 <2%
评级	A — 对 Transformer 基础架构的优雅改进，理论清晰、实验全面、工程可行
代码	https://github.com/MoonshotAI/Attention-Residuals（开源）
关键词	Residual Connection, Depth-wise Attention, PreNorm Dilution, Scaling Law, Block AttnRes, MoE

核心 Insight

残差连接（Residual Connection）是现代 LLM 的基石。从 2015 年 ResNet 提出以来， $h_l = h_{l-1} + f_{l-1}(h_{l-1})$ 这个更新规则几乎从未被根本性地改变过。Kimi 团队观察到一个深刻的对偶性：

残差连接在”深度”维度上扮演的角色，与 RNN 在”时间”维度上扮演的角色完全一致。

展开残差递推式： $h_l = h_1 + \sum_{i=1}^{l-1} f_i(h_i)$ ——每一层接收的是之前所有层输出的均匀加权求和。这就像 RNN 把整个序列历史压缩成一个隐藏状态一样。

在序列建模中，Transformer 用注意力机制替代了 RNN 的固定递推，让每个位置可以选择性地访问历史位置。AttnRes 对”深度”做了完全相同的事情。

为什么这个想法 work？

PreNorm（当前主流范式）的残差累加导致隐藏状态的幅值随深度 $O(L)$ 增长。这产生了三个问题：

无选择性访问：不同类型的层（注意力层 vs MLP 层）接收到同样的混合状态，无法获取针对性信息
信息不可逆丢失：早期层的信息被淹没在累加的信号中，深层无法选择性地恢复
输出幅值膨胀：深层必须学习越来越大的输出才能在累积的残差中”发出声音”

AttnRes 通过引入 softmax 注意力让每层用学到的、输入依赖的权重来选择性聚合前面各层的输出，优雅地解决了所有三个问题。

方法详解

整体架构

标准残差:   h_l = h_{l-1} + f_{l-1}(h_{l-1})     【固定权重累加】
                        ↓ 对偶变换
Full AttnRes: h_l = Σ α_{i→l} · v_i             【softmax 注意力加权】
                        ↓ 工程优化
Block AttnRes: 分块累加 + 块间注意力             【O(Nd) 内存/通信】

关键技术组件

组件 1: Full Attention Residuals

做什么： 用 softmax 注意力替代固定权重的深度方向累加。

怎么做： 对第 $l$ 层，定义：

Query: $q_l = w_l$ （每层一个可学习向量， $w_l \in \mathbb{R}^d$ ）
Key/Value: $k_i = v_i$ ，其中 $v_0 = h_1$ （token embedding）， $v_i = f_i(h_i)$ （第 $i$ 层的输出）
注意力权重： $\alpha_{i \to l} = \frac{\phi(q_l, k_i)}{\sum_{j=0}^{l-1} \phi(q_l, k_j)}$
其中 $\phi(q, k) = \exp(q^\top \text{RMSNorm}(k))$

关键公式：

$h_l = \sum_{i=0}^{l-1} \alpha_{i \to l} \cdot v_i$

直觉解释： 每层有一个”偏好向量” $w_l$ ，通过与前面各层输出做点积来决定”想要什么信息”。RMSNorm 防止输出幅值大的层主导注意力权重。

开销分析： 对每个 token，Full AttnRes 需要 $O(L^2 d)$ 计算和 $O(Ld)$ 内存。因为深度 $L$ 远小于序列长度，计算开销微不足道。在标准训练中（无激活重计算），这些层输出本来就要为反向传播保存，所以无额外内存开销。

组件 2: Block Attention Residuals（工程核心）

做什么： 在大规模分布式训练（流水线并行 + 激活重计算）中使 AttnRes 实用。

问题是什么： 在流水线并行中，Full AttnRes 需要将所有 $L$ 个层输出跨阶段传输，通信量 $O(Ld)$ 不可接受。

怎么做： 将 $L$ 层分成 $N$ 个块（block），每块 $S = L/N$ 层：

块内：使用标准残差累加，将所有层输出求和为一个块表示 $b_n = \sum_{j \in B_n} f_j(h_j)$
块间：对 $N$ 个块表示 + token embedding 应用 softmax 注意力

数值例子： 假设 48B 模型有 54 层（27 个 Transformer 块），每 6 层一个 AttnRes 块：

$N = 9$ 个块 + 1 个 token embedding = 10 个深度方向的”注意力源”
内存/通信从 $O(54d)$ 降到 $O(10d)$ ，约 5.4× 减少

关键发现： 实验表明 $N \approx 8$ 就能恢复 Full AttnRes 的绝大部分收益。

组件 3: 工程基础设施优化

跨阶段缓存（Cross-stage Caching）： 流水线并行中，每个物理阶段处理多个虚拟阶段。缓存策略让已接收的块表示在后续虚拟阶段中被重用，通信从 $O(C^2)$ 降到 $O(P^2)$ （ $V$ 倍改进， $V$ 为虚拟阶段数）。

两阶段推理（Two-phase Computation）： 利用 pseudo-query $w_l$ 与层计算解耦的特性：

Phase 1：批量计算一个块内所有层的块间注意力（并行化）
Phase 2：顺序计算块内注意力，通过 Online Softmax 合并

结果： 训练开销 <4%（流水线并行下），推理延迟 <2%。

与现有方法的关键区别

维度	标准残差	Highway/ReZero	mHC	AttnRes
权重类型	固定（全 1）	可学习/门控	可学习矩阵	输入依赖（softmax）
信息源	仅前一层 $h_{l-1}$	仅前一层	$m$ 个并行流	所有前面的层
每层 I/O	$3d$	$3d$	$34d$ （ $m$ =4）	$5.5d$ （Block）
选择性	无	有限	有限	完全选择性

实验结果

Scaling Law 实验

在 5 个模型规模（194M 到 528M 激活参数）上的拟合结果：

方法	Scaling Law 拟合	5.6 PFLOP/s-days 处的 Loss
Baseline (PreNorm)	$L = 1.891 \times C^{-0.057}$	1.714
Block AttnRes	$L = 1.870 \times C^{-0.058}$	1.692
Full AttnRes	$L = 1.865 \times C^{-0.057}$	1.692

解读： Block AttnRes 相当于 1.25× 的等效计算量提升。这是一个”免费午餐”级别的收益——几乎不增加计算成本就获得了显著的性能提升。Full AttnRes 和 Block AttnRes 在最大规模上的差距缩小到仅 0.001。

48B 模型主实验

将 Block AttnRes 集成到 Kimi Linear 架构（48B 总参数 / 3B 激活参数），在 1.4T token 上预训练。

类别	基准测试	Baseline	AttnRes	提升
通用	MMLU	73.5	74.6	+1.1
通用	GPQA-Diamond	36.9	44.4	+7.5
通用	BBH	76.3	78.0	+1.7
数学	Math	53.5	57.1	+3.6
数学	GSM8K	81.7	82.4	+0.7
代码	HumanEval	59.1	62.2	+3.1
代码	MBPP	72.0	73.9	+1.9
中文	CMMLU	82.0	82.9	+0.9
中文	C-Eval	79.6	82.5	+2.9

解读：

在所有评估的任务上均有提升，无一退步
提升在多步推理任务上最为显著（GPQA-Diamond +7.5, Math +3.6），这与”改进深度方向信息流有利于组合式任务”的假设一致
知识导向的基准（MMLU, TriviaQA）也有稳定提升

消融实验

变体	Val Loss	与 Full AttnRes 差距	说明
Full AttnRes	1.737	—	完整方法
+ 输入依赖 query	1.731	-0.006	更好但增加 $d \times d$ 投影
输入无关混合（去掉 query/key）	1.749	+0.012	输入依赖性是关键
用 sigmoid 替换 softmax	1.741	+0.004	softmax 的竞争性归一化更好
去掉 RMSNorm	1.743	+0.006	RMSNorm 防止大幅值层主导
多头注意力（H=16）	1.752	+0.015	深度方向最优混合跨通道一致
Block AttnRes (S=4)	1.746	+0.009	绝大部分收益保留
DenseFormer	1.767	+0.030	固定系数不够，需要输入依赖
mHC	1.747	+0.010	AttnRes 用更少 I/O 达到更好效果
Baseline (PreNorm)	1.766	+0.029	基线

关键发现：

输入依赖权重是核心：从固定系数（DenseFormer, 1.767）到可学习系数（mHC, 1.747）到输入依赖（AttnRes, 1.737），每一步都有显著提升
单头优于多头：深度方向的最优混合在不同通道间是一致的——“当一层的输出相关时，它整体都相关”
RMSNorm 不可或缺：在 Block AttnRes 中尤其关键，因为块级表示在多层上累加后幅值差异更大
Block 大小优雅退化： $S = 2, 4, 8$ 的效果接近， $N \approx 8$ 块是实践中的最佳折衷

最优架构分析

在固定计算量和参数量下的 25 种架构配置搜索中：

Baseline 最优点： $d_{model}/L_b \approx 60$ （较宽较浅）
AttnRes 最优点： $d_{model}/L_b \approx 45$ （较窄较深）

含义： AttnRes 能更有效地利用深度，使得最优架构偏向更深的网络。这与”改进深度方向信息流”的设计目标完美一致。

复现评估

维度	评分(1-5)	详细说明
数据可得性	⭐⭐⭐	论文使用内部训练数据，但方法对数据不敏感，可用公开数据集验证
代码可得性	⭐⭐⭐⭐⭐	GitHub 开源，包含 PyTorch 伪代码（论文 Figure 2 直接可用）
算力需求	⭐⭐	Scaling law 实验需要多规模训练；48B 模型需要大量 GPU
工程复杂度	⭐⭐⭐⭐	Block AttnRes 是 drop-in replacement，集成复杂度低
预期收益	⭐⭐⭐⭐⭐	1.25× 等效计算量提升，几乎无额外成本，对所有使用标准残差连接的模型都适用

复现建议： Block AttnRes 的实现非常简洁（论文 Figure 2 的伪代码可以直接用），建议先在小模型上验证 scaling law 趋势，然后集成到现有训练流程中。关键要注意 pseudo-query 必须零初始化。

批判性分析

局限性

论文自述的局限：

48B 模型实验只训练了 1.4T token——与完整训练（通常 10T+）相比仍是早期
最优架构搜索在固定预算下进行，不同预算下结论可能不同

我们额外发现的问题：

长序列场景未充分验证：论文提到了 32K 上下文扩展，但没有在超长上下文（128K+）下系统评估。Block 表示在 prefill 阶段占用 $N \cdot T \cdot d$ 内存，对超长序列可能成为瓶颈。
与其他架构创新的交互：论文只在 Kimi Linear（MoE + KDA/MLA 混合）上验证。AttnRes 与其他架构创新（如 State Space Models、线性注意力）的交互效果未知。
训练稳定性：论文提到 pseudo-query 必须零初始化以避免训练不稳定，但没有详细分析为什么其他初始化策略会失败，以及在更大规模下是否还有其他稳定性陷阱。

改进方向

自适应 Block 大小： 当前 Block 大小是超参数。可以让模型自动学习在哪里放置 Block 边界——不同深度区域可能需要不同粒度的跨层访问。
与 Post-Training 的结合： AttnRes 只在预训练中验证。它对 SFT、RLHF、DPO 等后训练阶段的影响值得探索——更好的深度信息流可能也有利于对齐训练。
推广到其他”固定累加”场景： 论文揭示的”时间-深度对偶”可能在更多场景中适用，例如 MoE 中的专家聚合、或多模态融合中的模态混合。

独立观察

深层理论意义： 论文通过统一的结构化矩阵分析（Section 6.2）证明了标准残差和所有现有变体都对应深度方向的”线性注意力”，而 AttnRes 是第一个实现深度方向”softmax 注意力”的方法。这完成了与序列维度相同的”线性→softmax”转换——这个类比非常优雅。
对模型压缩的启示： AttnRes 学到的注意力权重（Figure 8）显示强烈的对角主导模式（局部性）+ 选择性跳跃连接 + 持续的 embedding 权重。这意味着模型确实需要选择性跨层访问——传统的层剪枝可能丢失了重要信息。
对 Transformer 架构演进的意义： 这是继 RoPE（改进位置编码）、GQA/MLA（改进注意力效率）、MoE（改进前馈网络）之后，对 Transformer “最后一个未被改进的核心组件”——残差连接——的首次成功大规模改进。

对领域的影响

这篇论文可能会改变 LLM 架构设计的默认选项。Block AttnRes 的”drop-in replacement”特性意味着它可以几乎无成本地集成到任何现有训练流程中。如果社区广泛验证了 1.25× 等效计算量的收益，这可能成为未来所有 LLM 的标准配置——就像 RMSNorm 取代 LayerNorm 一样。