Attention Residuals 深度解读该论文由 Kimi Team 提出，核心是为解决现代大语言模型（LLM）

该论文由 Kimi Team 提出，核心是为解决现代大语言模型（LLM）中标准残差连接的固有缺陷，将注意力机制引入深度维度，提出 Attention Residuals（AttnRes）及可规模化的 Block AttnRes，成为标准残差的即插即用替代方案，在 48B 参数量的 Kimi Linear 模型上验证了显著性能增益，以下从 7 个核心维度展开解读。

一、该论文解决了什么问题（Research Problem）？

问题的定义是什么？

现代 LLM 的核心构建块是PreNorm 结合的标准残差连接，其通过固定单位权重累加所有层输出 $h_l=h_{l-1} + f_{l-1}(h_{l-1})$ 实现梯度传播，但该设计存在两大核心缺陷：
- 深度维度的信息聚合为无差别的均匀累加，无输入依赖的可学习权重，无法选择性强调 / 抑制不同层的贡献，早期层信息易被掩埋，且大量层可被剪枝而性能损失极小；
- 无权重的累加导致隐藏状态幅度随深度呈 $O(L)$ 增长，引发PreNorm稀释问题：深层需学习更大的输出才能在累加中获得有效影响，同时梯度分布随深度严重失衡，早期层梯度过大，训练稳定性下降。
此外，现有残差改进方法要么受限于加法递归范式，要么跨层访问的内存 / 通信开销过高，难以在大规模 LLM 中落地。
应用场景是什么？

大规模预训练语言模型（ LLMs ）的训练与推理，尤其针对深层数、大参数量的混合专家（MoE）架构 LLM（如论文中的 Kimi Linear 48B 总参 / 3B 激活参模型），也适用于常规稠密 Transformer 架构。
为什么重要？
- 残差连接是 LLM 实现深度训练的基础，其缺陷直接限制模型的有效深度和表达能力，是制约 LLM 性能进一步提升的核心瓶颈；
- PreNorm 是当前 LLM 的主流范式，其稀释问题导致训练效率低、层贡献不均、下游推理能力受限，尤其对多步推理、代码生成等组合式任务影响显著；
- 大规模 LLM 的训练 / 推理对内存、通信、延迟要求严苛，现有改进方法无法兼顾性能提升和工程可落地性，亟需轻量、可规模化的残差连接替代方案。

二、以前是怎么解决的（Prior Work）？

前人针对残差连接的缺陷提出了三类改进方法，但均未从根本上解决问题，核心瓶颈集中在深度维度的信息聚合方式和规模化开销两方面：

三类主流解决方法

方法类别	代表工作	核心思路
单状态递归变体	Residual、Highway、DeepNorm	仅能访问前一层的聚合隐藏状态，通过门控、缩放、归一化优化累加规则，但无跨层直接访问
多状态递归	HC/mHC、SiameseNorm、DDL	用多并行流拓宽递归状态，通过流混合实现有限的层信息交互，但仍依赖前一时刻的聚合流状态
跨层连接方法	DenseFormer、MRLA、DenseNet	让每层直接访问前序所有层输出，但 DenseFormer 用固定标量权重（无输入依赖），MRLA 接近线性注意力而非 softmax 检索

核心瓶颈
- 聚合方式的本质缺陷：深度维度的信息聚合要么是固定权重（无自适应能力），要么是输入依赖但无softmax 的竞争性选择，无法实现层的精准选择性聚合；
- 跨层访问的局限性：要么仅能访问前一层的聚合状态（信息不可逆丢失），要么直接访问所有层但计算 / 内存开销过高，无法在大规模 LLM 的流水线并行、激活重计算场景中落地；
- PreNorm 稀释未被根本解决：所有方法均未脱离加法递归范式，隐藏状态随深度增长、梯度分布失衡的核心问题依旧存在；
- 工程适配性差：现有跨层方法未针对 LLM 的分布式训练 / 推理做基础设施优化，通信、延迟开销无法接受。

三、作者的核心思想是什么（Core Idea）？

论文的核心创新基于序列 - 深度的对偶性，将 Transformer 解决 RNN 序列递归问题的思路迁移到深度维度，从根本上重构残差连接的聚合方式，核心思想可概括为 3 点：

对偶性理论基础：残差连接在深度维度的固定递归聚合，与 RNN 在序列维度的递归问题高度相似 —— 二者均将前序信息压缩为单一聚合状态，丢失细粒度选择性访问能力；而 Transformer 用注意力机制解决了 RNN 的序列递归瓶颈，因此可将注意力引入深度维度解决残差的递归缺陷。
注意力替换固定累加：将深度维度的固定加法聚合 $h_l= \sum_iv_i$ 替换为输入依赖的softmax 注意力聚合 $h_l=\sum_i\alpha_{i→l}·v_i$ ，其中 $\alpha_{i→l}$ 是由每层可学习伪查询计算的注意力权重，实现层的精准、自适应选择性聚合，从根源上缓解 PreNorm 稀释。
规模化的折中设计：为解决 Full AttnRes 的 $O(Ld)$ 内存 / 通信开销，提出Block AttnRes，将层划分为块并仅在块级表示上做注意力，把开销降至 $O(Nd)$ （N 为块数），同时证明 N≈8 时可恢复 Full AttnRes 的大部分收益，实现性能与效率的平衡。
范式统一：从结构化矩阵视角证明标准残差及所有变体均为深度维度的线性注意力，而 AttnRes 将其推广为深度维度的softmax 注意力，完成了与序列维度一致的 “线性 - softmax” 范式转变。

四、方法结构是什么（Method）？

论文提出的方法体系包含核心注意力残差模块和规模化基础设施优化两部分，前者分为 Full AttnRes 和 Block AttnRes（即插即用，兼容现有LLM架构），后者针对大规模训练/推理做工程优化，确保方法的落地性。

1. Full Attention Residuals（全注意力残差）

为每层设计可学习伪查询 $w_l∈\mathbb{R}^$ ，以token嵌入 $h_1$ 和前序层输出 $f_i(h_i)$ 为Key/Value，通过RMSNorm归一化避免大幅度输出主导注意力权重，最终通过softmax计算注意力权重 $\alpha_{i→l}$ ，层输入为权重加权的前序表示和，核心公式：

$\alpha_{i \to l}=\frac{\exp(w_l^\top \text{RMSNorm}(v_i))}{\sum_{j=0}^{l-1}\exp(w_l^\top \text{RMSNorm}(v_j))}, \quad h_l=\sum_{i=0}^{l-1}\alpha_{i→l}·v_i$

其中 $v_0=h_1$ （嵌入）， $v_i=f_i(h_i)$ （i≥1，层输出）。

计算开销： $O(L^2d)$ ，内存开销： $O(Ld)$ ；
普通训练无额外内存开销：复用反向传播的激活值，仅在规模化训练（激活重计算、流水线并行）中存在开销。

2. Block Attention Residuals（块注意力残差）

为解决 Full AttnRes 的规模化开销，将L层划分为N个块，是Full AttnRes 的工程折中版本，核心步骤：

块内累积：每个块内通过标准残差将层输出累加为单一块级表示 $b_n=\sum_{j∈B_n}f_j(h_j$ ，块内后续层还会维护部分和 $b_n^{i-1}$ （前i-1层的累加值）；
块间注意力：块内第一层仅对前序块表示+token嵌入做注意力，后续层额外对块内部分和做注意力，注意力权重计算与Full AttnRes一致；
灵活插值：N=L时退化为 Full AttnRes，N=1时退化为标准残差，论文验证N≈8时可恢复Full AttnRes的大部分收益。

核心优势：将内存/通信开销从 $O(Ld)$ 降至 $O(Nd)$ ，计算开销从 $O(L^2d)$ 降至 $O(N^2d)$ 。

3. 规模化基础设施优化

针对大规模LLM的流水线 并行训练和长上下文推理做专项优化，使Block AttnRes的开销可忽略：

训练端：跨阶段缓存：消除流水线并行下块表示的冗余传输，将通信开销从 $O(C)$ （C为总块数）降至 $O(P)$ （P为物理阶段数），训练端整体开销<4%；
推理端：两阶段计算：先批量计算块间注意力（ amortize内存访问），再串行计算块内注意力，结合在线softmax实现无误差融合，推理延迟开销<2%；
长上下文优化：序列分片预填充：将块表示沿序列维度分片到多个张量并行设备，将长上下文（如128K）的内存占用从15GB降至1.9GB/设备，结合分块预填充可进一步降至0.3GB/设备。

论文还提供了PyTorch风格的Block AttnRes伪代码，确保方法的可复现性。

五、实验验证（Experiment）

实验基于Kimi Linear（MoE Transformer架构，48B总参/3B激活参），预训练1.4T tokens，从Scaling Law、训练动态、下游性能、消融实验四个维度验证，核心实验设置和结果如下：

1. 数据集

预训练数据集：Kimi Linear的1.4T tokens通用语料，采用4096→32K的渐进式上下文长度训练；
下游评估数据集：覆盖3大方向，共14个基准：
- 通用语言理解与推理：MMLU、MMLU-Pro、GPQA-Diamond、BBH、ARC-Challenge、HellaSwag、TriviaQA；
- 数学与代码生成：GSM8K、MGSM、Math、CMath、HumanEval、MBPP；
- 中文语言理解：CMMLU、C-Eval。

2. 对比方法

baseline：PreNorm结合的标准残差连接；
经典残差改进方法：DenseFormer（固定跨层权重）、mHC/mHC-lite（多流递归）；
AttnRes变体：Full AttnRes、Block AttnRes（N≈8）。

3. 核心评价指标

预训练/验证：验证损失（拟合幂律曲线 $L=A×C^{-\alpha}$ ）、层输出幅度、梯度幅度分布；
效率指标：内存访问成本（每token每层层的读/写）、训练/推理开销（内存、延迟、通信成本）；
下游性能：各基准的准确率/分数；
架构分析：深度-宽度-注意力头数的架构偏好（ $d_{model}/L_b$ 、 $H/L_b$ ）。

4. 核心实验结果

Scaling Laws：AttnRes在所有模型尺度上均显著优于baseline，Block AttnRes可达到baseline1.25×的计算优势（相同损失下，Block AttnRes所需计算量仅为baseline的80%）；Full AttnRes性能最优，且与Block AttnRes的性能差距随模型尺度增大而缩小。

训练动态：Block AttnRes彻底缓解PreNorm稀释——输出幅度随深度有界且周期性分布（而非单调增长），梯度幅度在所有层均匀分布（而非早期层梯度过大），验证损失全程低于baseline且差距随训练推进扩大。
下游性能：Block AttnRes在所有下游任务中均优于baseline，组合式任务增益显著：GPQA-Diamond（+7.5）、Math（+3.6）、HumanEval（+3.1），中文任务C-Eval（+2.9）、CMMLU（+0.9），通用任务MMLU（+1.1）。
效率对比：Block AttnRes的每token每层层内存访问成本仅为5.5d，远低于mHC的34d，训练开销<4%，推理延迟开销<2%，为标准残差的即插即用替代方案。

5. 消融实验（16层模型）

验证AttnRes各模块的有效性，核心结论：

输入依赖混合是核心：移除输入依赖，用固定标量权重替换注意力，性能从1.737降至1.749；
softmax与RMSNorm不可替代：softmax替换为sigmoid（1.741）、移除RMSNorm（1.743）均导致性能下降，softmax的竞争性归一化是精准选择的关键；
块大小的权衡：块大小S=2/4/8时性能接近Full AttnRes（1.737），S=16/32时性能向baseline退化，论文建议N≈8为最优折中；
单头注意力更优：多头深度注意力（H=16）性能降至1.752，说明深度维度的最优混合在通道上是均匀的；
跨层访问的价值：滑动窗口聚合（SWA，仅访问最近8层）性能为1.764，远差于Full/Block AttnRes，证明远距离层的选择性访问比近距离层的大量访问更重要。

六、论文的贡献

论文的贡献分为理论贡献、方法贡献、应用贡献三大类，实现了理论创新-方法设计-工程落地-实际验证的闭环：

1. 理论贡献

首次揭示了序列-深度的对偶性，将残差连接的深度递归与RNN的序列递归建立形式化关联，为深度维度的模型设计提供了新的理论视角；
从结构化矩阵视角，提出深度混合矩阵 $M$ 的分析框架，证明标准残差及所有变体均为深度维度的线性注意力，AttnRes是首个将其推广为深度维度softmax注意力的方法，完成了序列/深度维度的注意力范式统一；
分析了深度混合矩阵的半分离秩，为后续残差连接的设计提供了统一的量化分析标准。

2. 方法贡献

提出Attention Residuals（AttnRes） ，用输入依赖的softmax注意力替换残差连接的固定累加，从根源上解决了PreNorm稀释问题，实现了深度维度的精准、自适应信息聚合；
提出可规模化的Block AttnRes，将内存/通信开销从 $O(Ld)$ 降至 $O(Nd)$ ，实现了性能与效率的最优权衡，且与现有LLM架构完全兼容；
设计了适配大规模LLM的基础设施优化策略（跨阶段缓存、两阶段推理、序列分片预填充），使Block AttnRes的训练/推理开销可忽略，成为标准残差的即插即用 替代方案。

3. 应用贡献

将AttnRes集成到Kimi Linear 48B MoE模型，在1.4T tokens预训练中实现了更均匀的输出/梯度分布，验证了方法在大规模模型中的有效性；
在14个下游基准中实现全任务性能提升，尤其在多步推理、代码生成、数学推理等组合式任务中增益显著，为LLM的性能提升提供了新路径；
验证了AttnRes在不同模型尺度的一致性收益，Block AttnRes可达到baseline1.25×的计算优势，为大规模MoE LLM的设计和训练提供了新范式；
揭示了AttnRes的架构偏好（更深更窄的网络），为后续LLM的架构设计提供了重要指导。

七、论文的局限

尽管AttnRes在性能和工程落地性上表现优异，但论文仍明确指出了方法的局限性，同时也为后续研究指明了方向：

Full AttnRes的规模化受硬件限制
- Full AttnRes虽性能最优，但内存/通信开销为 $O(Ld)$ ，受限于当前硬件的流水线并行和激活重计算机制，无法在大规模LLM中直接落地，只能依赖Block AttnRes做折中，未来需等待硬件互连技术的改进。
块大小的选择缺乏形式化准则
- 论文中块数N≈8是经验性选择，尚无形式化的块大小选择准则，更大的块会导致性能向baseline退化，更小的块则增加内存/通信开销，如何根据模型尺度、硬件环境自适应选择块大小仍需研究。
架构偏好的落地限制
- AttnRes偏好更深更窄的网络（固定参数量下， $d_{model}/L_b$ 更小），但更深的模型因推理的串行计算特性会带来更高的推理延迟，该架构偏好难以直接在实际部署中落地，需结合模型压缩、推理加速技术。
注意力机制的拓展性有限
- 当前AttnRes采用的是基础softmax注意力，未探索更高效的注意力变体（如线性注意力、稀疏注意力、长程注意力），在超深层数（L>1000）的模型中，块级注意力的开销仍可能上升。
验证场景的泛化性待补充
- 实验主要基于MoE 架构的Kimi Linear模型，未在纯稠密Transformer、小参数量模型、非语言模型（如CV、语音）中做充分验证，AttnRes的跨架构、跨领域泛化性仍需进一步测试。
伪查询的设计存在优化空间
- 当前采用的是每层固定的可学习伪查询，消融实验中发现输入依赖的动态伪查询性能更好（损失从1.737降至1.731），但因引入 d×d 投影和串行内存访问问题未采用，如何设计轻量的动态伪查询是后续重要方向。
未探索与其他残差改进的结合
- AttnRes未与现有残差改进方法（如Highway、mHC、SiameseNorm）结合，如何实现AttnRes与这些方法的互补，进一步提升性能仍需研究。