价值千万的文章，拜读一下

摘要

带有前置归一化（PreNorm）[60]的残差连接[12]是当前大语言模型（LLM）的标准配置，但这类残差连接会以固定单位权重累加所有层的输出。这种均匀聚合方式会随着网络层数加深，导致隐藏状态数值无节制增大，进而逐渐稀释每一层的特征贡献度[27]。 本文提出**注意力残差机制（AttnRes）**：摒弃固定累加方式，对前序网络层的输出采用softmax注意力聚合，让每一层能够根据输入，通过可学习权重选择性融合早期特征表示。 针对大规模模型训练中，对所有前序层输出做注意力计算带来的显存与通信开销问题，我们进一步提出**分块注意力残差（Block AttnRes）**。该方法将网络层划分为多个模块，仅对模块级特征进行注意力聚合，在保留完整注意力残差大部分性能增益的同时，降低显存占用。 结合基于缓存的流水线通信策略与两阶段计算策略，分块注意力残差能够低成本、无缝替代传统残差连接，实用性更强。 缩放定律实验表明，该方法在不同规模的模型中均能稳定提升性能；消融实验验证了**基于内容的跨层自适应筛选**具备显著优势。此外，我们将注意力残差集成至 Kimi Linear 架构[69]（总参数量480亿，激活参数量30亿），并在1.4万亿词元数据集上完成预训练。实验结果显示，注意力残差有效缓解了前置归一化带来的特征稀释问题，使不同网络层的输出幅值与梯度分布更加均衡，并全面提升了模型在各类下游任务上的效果。

1 引言

标准残差连接[12]是当前大语言模型（LLM）的基础核心组件[35, 51, 9]。其更新公式为 ，通常被视作梯度高速公路：通过恒等映射让梯度绕过变换层，从而保证深层网络能够稳定训练。但残差连接还有一个较少被关注的作用。 对递推公式展开可以发现，每一层都会以**等权重**累加所有前置层的输出；也就是说，残差连接决定了信息在纵向深度上的聚合方式。现如今，序列混合、专家路由等机制都采用了可学习、与输入相关的权重策略[53, 20, 9]，但纵向跨层聚合仍然依赖固定的单位权重，无法自适应增强或抑制不同网络层的特征贡献。 在实际应用中，前置归一化（PreNorm）[60]已成为主流范式。然而，无权重的累加方式会使隐藏状态幅值随层数 $L$ 线性增长（复杂度为 $O(L)$），不断稀释各层的相对特征贡献[27]。浅层信息容易被淹没，且无法被选择性提取；实验表明，大语言模型中相当一部分网络层可以被剪去，且模型性能几乎不受损失[11]。

现有改进方法，例如缩放残差路径[54]、多分支递推结构[72]，仍局限于加法递推形式；而支持跨层信息交互的方法[36, 56]，又难以实现大规模部署。这一困境，与注意力机制问世之前循环神经网络（RNN）在序列维度遇到的难题高度相似。

Transformer 在 token 维度已经进入“attention时代”，但在 layer（depth）维度还停留在“RNN时代”。

我们发现：**纵向跨层聚合**与循环神经网络的**序列递推计算**存在形式对偶关系。基于这一对偶性，本文提出**注意力残差机制（AttnRes）**。将固定累加形式 改进为 ，其中 为通过每层单独学习的伪查询向量 计算得到的 Softmax 注意力权重。该轻量化方法仅为每层增加一个 d 维向量，即可实现跨层、基于内容的自适应信息筛选。事实上，传统残差连接及各类递推改进方法，本质上都属于**纵向线性注意力**；而 AttnRes 将其拓展为**纵向 Softmax 注意力**。如同注意力机制革新了序列建模一样，该方法完成了深度维度上从线性注意力到 Softmax 注意力的升级（见第6.1、6.2节）。

第一类，single-state。所有的递推公式化简到最后都可以得到： 都是线性的。
第二类，multi-state，就是多分支，列入siamesenorm：

类似一种在postnorm和prenorm中间做的一种巧妙地tradeoff。但是化简到最后，说白了还是依旧线性。
类似的HC/mHC也是一样，把每次layer传播只考虑上一层变成考虑全部了，除了weight选取方法变了，化简到最后仍然是wihi的形式，因此本质依旧线性。
 

第三类，每次都要访问之前所有隐藏层，但是卷积等操作本质仍然是线性

在常规训练场景下，完整注意力残差（Full AttnRes）几乎不会增加额外开销，因为反向传播本身就需要缓存各层输出。但在大规模训练中，模型通常采用激活重计算与流水线并行策略，此时跨阶段缓存和传输各层激活值会带来开销。为此，我们进一步提出**分块注意力残差（Block AttnRes）**：将网络层划分为 $N$ 个模块，每个模块通过传统残差聚合得到单一特征表征，仅对模块级特征执行跨块注意力计算。该方法将显存与通信开销降至 $O(N d)$。结合底层工程优化（见第4节），分块注意力残差可**无缝替换**传统残差连接，训练额外开销极低，推理时延几乎无增加。

缩放定律实验表明：在不同算力预算下，AttnRes 均稳定优于基线模型；其中，Block AttnRes 的收敛效果，等同于基线模型提升1.25倍算力后的效果。我们将 AttnRes 集成到 Kimi Linear 架构[69]（总参数量480亿，激活参数量30亿），并在1.4万亿词元数据集上完成预训练。训练动态分析显示：AttnRes 有效缓解了前置归一化带来的特征稀释问题，使不同网络层的输出幅值保持稳定、梯度分布更加均匀。在下游评测任务中，本文模型全面超越基线模型。

主要贡献

1. **注意力残差机制** 提出 AttnRes，用可学习的纵向 Softmax 注意力替代固定残差累加，并给出可扩展变体 Block AttnRes，将显存与通信开销由 $O(Ld)$ 降至 $O(Nd)$。通过结构化矩阵统一分析，证明传统残差及其递推变体等价于纵向线性注意力，而 AttnRes 实现了纵向 Softmax 注意力。

2. **大规模工程优化方案** 设计适配大规模训练的系统优化策略：基于流水线并行的跨阶段缓存机制，消除冗余数据传输；两阶段推理策略，结合在线 Softmax 实现跨块注意力复用。最终模型训练额外开销极小，常规推理场景下时延增加低于2%。

3. **全面实验验证与机理分析** 基于1.4万亿词元预训练的480亿参数模型，通过缩放定律实验、组件消融实验、下游基准测试，全面验证 AttnRes 的有效性，证明其显著优于传统残差连接。训练动态分析进一步表明：AttnRes 缓解了 PreNorm 的特征稀释问题，实现隐藏状态幅值收敛、跨层梯度分布均衡。

2 动机 (Motivation)

符号说明。 考虑形状为 的一批输入序列，其中 $B$ 是批大小，$T$ 是序列长度，$d$ 是隐藏维度。为了清晰起见，我们针对单个 token 编写公式： 表示进入第 $l$ 层的隐藏状态，其中 是层索引，$L$ 是总层数。Token 嵌入为 。函数 $f_l$ 表示第 $l$ 层执行的变换。在 Transformer 模型中，我们将每个自注意力（Self-Attention）或 MLP 视为一个单独的层。

2.1 通过残差训练深层网络

残差学习 (Residual Learning)。 残差学习已被证明是训练深层网络的关键技术，因为它允许梯度绕过变换。具体而言，每一层按如下方式更新隐藏状态：

，展开这个递归式，第 $l$ 层的隐藏状态是嵌入项与之前所有层输出的总和：。残差连接背后的核心见解是恒等映射（Identity Mapping）：每一层都为信息和梯度的不失真流动保留了一条直接路径。在反向传播期间，针对中间隐藏状态的梯度为：

展开该乘积会得到单位矩阵 $I$ 加上涉及层雅可比矩阵的高阶项。恒等项始终存在，无论深度如何，都为从损失函数到任何层的梯度提供了直接路径。

推广残差。 尽管有效，但残差更新中的固定单位系数对每一层的贡献采取了均一化处理，缺乏跨深度调节混合比例的机制。Highway Networks 引入了可学习的元素级门控来放宽这一限制：

其中 在变换路径和恒等路径之间进行插值。更广义地说，两者都是加权递归 的实例，残差设置对应 ，Highway 设置对应 。

局限性。 无论是固定的还是门控的，这两种方法都存在一个根本性的约束：每一层只能访问其直接输入 $h_{l-1}$，这是一个合并了所有先前层输出的单一压缩状态，而不是这些输出本身。这导致了几个局限性：(1) 缺乏选择性访问：不同的层类型（如注意力层与 MLP 层）接收到的是相同的聚合状态，尽管它们可能受益于不同的权重分配；(2) 不可逆损失：通过聚合丢失的信息无法在更深的层中被选择性地恢复；(3) 输出膨胀：后面的层为了在累积的残差中获得影响力，会学习生成越来越大的输出，这可能会导致训练不稳定。这些局限性促使我们需要一种机制，让每一层能从之前所有层中选择性地聚合信息。

---

3 注意力残差：时间与深度的统一视图

上述局限性让人联想到序列建模中类似的瓶颈，这暗示我们可以为“深度”维度寻求类似的解决方案。

时间与深度的对偶性。 就像随时间演变的 RNN 一样，残差连接在深度上将所有先验信息压缩为单一状态 $h_l$。对于序列建模，Transformer 通过用注意力机制取代递归改进了 RNN，允许每个位置根据数据依赖的权重选择性地访问所有先前位置。我们提议对深度采用相同的方法论：

其中 是满足 的层特定注意力权重。与序列长度（可达数百万 token）不同，网络深度通常适中（$L < 1000$），这使得在深度上进行 $O(L^2)$ 的注意力计算在计算上是可行的。我们将此方法称为注意力残差 (Attention Residuals)，简写为 AttnRes。

3.1 全注意力残差 (Full Attention Residuals)

注意力权重可以写为 ，其中 $\phi$ 是核函数。我们采用带有归一化的，在深度上产生 softmax 注意力：

对于每一层 $l$，我们定义：

其中查询向量 是  中层特定的可学习向量。 内部的 RMSNorm 可防止输出幅值较大的层主导注意力权重。第 $l$ 层的输入随后为：

我们称这种形式为全注意力残差 (Full AttnRes)。对于每个 token，Full AttnRes 需要 $O(L^2d)$ 的算术运算和 $O(Ld)$ 的内存来存储层输出。由于深度远小于序列长度，算术成本适中。

开销。 $O(Ld)$ 的内存与反向传播已经保留的激活值完全重叠，因此 Full AttnRes 在普通训练中不会引入额外的内存开销。然而，在大规模训练中，广泛采用激活重计算和流水线并行：原本会被释放并重新计算的层输出现在必须为后续所有层保持活跃状态，且在流水线并行下，每个输出还必须跨越阶段边界进行传输。此时内存和通信开销都会增长为 $O(Ld)$。

分块优化。 Full AttnRes 中一个刻意的设计选择是：伪查询向量 $w_l$ 是一个与层的分前向计算解耦的可学习参数。这种独立性意味着任何层组的注意力权重都可以并行计算，而无需等待它们的顺序输出。特别是，这允许将 $L$ 层分组为 $N$ 个块（每个块 $S$ 层），并在每个块内批量进行注意力计算，将每层的内存 I/O 从 $O(Ld)$ 降低到 $O((S+N)d)$。然而，在当前的分布式训练方案下，主要成本不是本地内存带宽，而是流水线并行下的跨阶段通信。这促使我们引入下文的 Block AttnRes 变体，它将跨阶段表示的数量从 $L$ 减少到 $N$。我们预见未来互连技术的改进将使完整的 $O(Ld)$ 通信变得实用，从而充分释放 Full AttnRes 的潜力。

5.4.2 对学习得到的 AttnRes

分布模式分析 图8可视化了**完整注意力残差（full AttnRes）**和**分块注意力残差（block AttnRes，分块数 N=8）**在16头模型（见表2）中学习得到的权重 。

每张热力图展示：第 $l$ 层注意力层/多层感知机（MLP）层（行）对前面各来源层（列）的注意力分配情况，并将注意力前置层与MLP前置层分开展示。本文总结三点核心发现：

- **局部性得以保留** 每一层对紧邻上一层的注意力权重最高；同时热力图出现少量非对角线集中权重（例如：第4层关注浅层特征、分块设置下第15~16层回溯前文特征）。这说明模型自动学习出了传统残差之外的跳跃连接。

- **层级功能分化** 初始嵌入特征 $h_1$ 在所有层中始终保有可观权重，在注意力前置层中尤为明显。MLP前置输入更倾向于依赖邻近特征（对角线集中），而注意力前置输入拥有更大感受野。该规律符合固有特性：注意力模块负责跨层信息交互，MLP负责局部特征变换。

- **分块 AttnRes 保留原有结构** 完整版与分块版均保持对角线主导、嵌入特征持久生效、层级分化三大特性。这表明：分块压缩具备隐式正则化效果，同时完整保留核心信息传输路径。