Attention Residuals
Kimi用深度方向Softmax注意力重塑残差连接

🎯PreNorm稀释：深度方向的根本问题

自He等人提出残差连接（ResNet）和Xiong等人将PreNorm引入Transformer以来，"残差+PreNorm"已成为所有现代大语言模型的标准范式。其核心机制极其简单：每一层的输出以固定单位权重1累加到之前的隐状态上，即 h_l = h_{l-1} + v_l = sum_{i=0}^{l} v_i。这种设计使得梯度可以通过残差路径直接回流到早期层，有效缓解了深层网络的梯度消失问题。

然而，随着模型深度从早期的12层（BERT-Base）增长到今天的上百层（Llama、Kimi等），一个根本性问题日益突出：固定单位权重的累加导致隐状态幅度随深度线性增长（O(L)），每一层的相对贡献被逐步稀释。具体而言，第1层的输出v_1在h_1中占100%的权重，但在h_100中仅占1/100——早期层的信息被后续层的累加"淹没"了。这正是PreNorm的"稀释效应"（dilution effect）。

这一问题的经验证据是多方面的。首先，大量研究表明，现代LLM中相当比例的层可以在不显著影响性能的情况下被剪枝——这意味着这些层的贡献确实被稀释到几乎可忽略。其次，PreNorm的隐状态幅度随深度增长不受控，导致深层和浅层的梯度分布极不均匀。再者，信息瓶颈理论分析表明，固定等权累加使得深层难以选择性检索早期层的特定信息——所有早期信息被无差别地混合在一起。

此前的改进尝试包括缩放残差路径（scaled residual paths）和多流递推（multi-stream recurrences），但它们仍然束缚在加性递推框架内。另一些引入跨层访问的方法（如DenseNet风格的跨层连接）则难以扩展到大规模训练。Kimi Team敏锐地观察到：这个问题与RNN在序列方向面临的困境有着精确的对偶关系——RNN用固定递推处理序列，直到注意力机制提供了选择性检索的替代方案。那么，深度方向是否也可以经历从"线性递推"到"softmax注意力"的跃迁？

💡核心洞察：从线性递推到Softmax注意力

AttnRes的核心洞察可以通过一个统一的结构化矩阵分析来精确表述。考虑深度方向的隐状态累加 h_l = sum_{i=0}^{l-1} v_i，这可以看作是一个深度方向的"注意力"操作，其中每个层的权重固定为1（均匀分布）。更一般地，此前所有的残差变体——包括标准残差、缩放残差、多流递推——都可以统一表达为深度方向的线性注意力：权重是固定的、不依赖于内容的。

AttnRes将这一框架推进到深度方向的softmax注意力。具体而言，第l层的隐状态不再是所有先前层输出的等权求和，而是通过softmax注意力权重加权求和：

h_l = Σ_{i=0}^{l-1} α_{i→l} · v_i

其中注意力权重 α_{i→l} = softmax(w_l · RMSNorm(v_i))，由一个每层唯一的可学习伪查询（pseudo-query）w_l ∈ R^d 和RMSNorm归一化计算得出。这意味着每一层可以根据当前内容选择性地检索之前任意层的表示——某些层可能主要依赖第3层的特征，而另一些层可能更关注第50层的输出，权重完全由学习决定且输入相关。

这一设计的精妙之处在于其极简性：每层仅增加一个d维向量w_l作为参数，注意力计算仅涉及一个标量点积（w_l与归一化后的v_i），不引入任何额外的多头机制或复杂投影。与序列方向的注意力（O(T²d)复杂度）不同，深度方向的注意力复杂度仅为O(Ld)——因为深度L通常远小于序列长度T，这使得Full AttnRes在标准训练中几乎无额外开销（层输出本就为反向传播而保留）。

从"线性→softmax"的转变在序列方向已被证明是变革性的——它使得Transformer彻底取代了RNN。AttnRes在深度方向完成了同样的跃迁，这不仅是技术上的改进，更是一种架构哲学的对称性完成。

🏗️AttnRes与Block AttnRes的设计

Full AttnRes的概念简洁明了，但在大规模训练中面临实际挑战。当模型使用激活重计算（activation recomputation）和流水线并行（pipeline parallelism）时，原本为反向传播保留的层输出激活需要被显式保存和跨流水线阶段通信——这使得O(Ld)的内存和通信开销成为瓶颈。为此，作者提出了Block AttnRes。

Block AttnRes的分块策略：将L层划分为N个块（典型N≈8），每个块内部使用标准残差连接进行累加，将块内所有层的输出缩减为一个块级表示（block representation）。跨块注意力仅在这N个块级表示上执行，将内存和通信开销从O(Ld)降至O(Nd)。当N=8时，Block AttnRes恢复了Full AttnRes大部分的性能增益，同时作为标准残差连接的即插即用替代品，仅有边际训练成本和可忽略的推理延迟开销。

具体实现上，每个Transformer层包含注意力（Attention）和前馈网络（MLP）两个子层，Block AttnRes在两个子层之前各执行一次跨块注意力。块大小以"子层数"计（每个Transformer层=2个子层），当层号达到块边界时，当前块的部分和被追加到块表示列表中，开始新的块。伪代码中的关键逻辑如下：

块内累加：partial_block 在块内通过标准残差累加各子层输出。块边界时，partial_block 被添加到 blocks 列表，新块从零开始。

跨块注意力：将所有已完成块的表示加上当前部分块堆叠为V [N+1, B, T, D]，对V做RMSNorm后与投影权重w做点积得到logits，softmax(0)归一化后加权求和得到当前层的输入h。这个h随后送入注意力或MLP子层。

这一设计的关键洞见是：块内保留标准残差的简单性，块间引入内容相关的选择性。块内的标准残差确保了块内信息流的畅通和梯度传播；块间的softmax注意力则解决了跨块信息检索的选择性问题。这种分层策略在效果和效率之间取得了优雅的平衡。

⚙️工程实现：让分块注意力在大规模训练中实用

Block AttnRes的工程实现是使其在大规模训练中实用的关键。作者开发了两个核心基础设施优化。

跨阶段缓存（Cross-stage Caching）：在流水线并行中，模型的不同层被分配到不同的设备阶段（pipeline stages）。Block AttnRes需要跨阶段访问块表示，如果每次都重新传输将导致大量通信开销。跨阶段缓存策略将块表示缓存在目标阶段，避免冗余传输——只有首次生成块表示时才需要跨阶段通信，后续访问直接从缓存读取。这与Flash Attention的tiling策略有异曲同工之妙：通过巧妙的缓存管理将通信最小化。

两阶段推理策略（Two-phase Inference）：在推理时，Block AttnRes的跨块注意力需要所有块表示才能计算当前层的输入。如果按顺序计算，会导致推理延迟增加。两阶段策略利用online softmax技术将跨块注意力分摊到推理过程中：第一阶段预计算块表示的归一化键（keys），第二阶段在生成每个token时仅需做一次轻量的加权求和。最终推理延迟开销低于2%——对于生产环境而言，这是一个几乎无感的开销。

这两个优化的组合使得Block AttnRes在大规模训练和推理中都成为标准残差连接的实用替代品。训练开销是边际的（主要来自额外的RMSNorm和投影计算），推理开销可忽略——这意味着任何使用PreNorm的现有模型都可以以极低的成本升级到Block AttnRes。

📊实验结果：Scaling Law与下游任务

作者从三个层面验证了AttnRes的有效性：Scaling Law实验、消融实验和大规模预训练下游评估。

Scaling Law：在多个模型尺寸和计算预算下，AttnRes一致优于基线。最引人注目的结果是：Block AttnRes匹配了基线用1.25倍计算量训练的损失。换言之，Block AttnRes等效于"免费"获得25%的额外计算量——这在 Scaling Law 的指数衰减曲线尾部是极为显著的提升，因为在该区域每一点损失下降都需要大量额外计算。

消融实验：作者验证了内容相关的深度选择性的重要性。将AttnRes的softmax注意力替换为固定权重（即退化为线性注意力）后，性能增益大幅消失——这证明了"内容相关"而非"加权求和"本身是关键。此外，块数量N的消融显示N≈8是一个好的平衡点：更少的块损失太多选择性，更多的块增加开销但收益递减。

下游任务（Kimi Linear 48B / 3B激活，1.4T tokens预训练）：

最显著的提升出现在GPQA-Diamond（+7.5分）——这是一个需要多步推理的困难基准，表明AttnRes对需要深层推理的任务有特别大的帮助。代码生成（HumanEval +3.1）和数学（Math +3.6）的提升也印证了这一点。这些任务的共同特点是：它们需要模型在深层有效地组合来自不同层的信息——正是AttnRes选择性检索所擅长的。

训练动态分析：AttnRes显著改善了PreNorm稀释问题。基线模型的隐状态幅度随深度线性增长，而AttnRes的隐状态幅度在深度方向保持有界（bounded）。梯度分布也更为均匀——浅层和深层的梯度范数差异大幅缩小，这意味着所有层都能有效地参与学习，而非只有少数层主导梯度更新。

🏭生产验证：从48B到2.8T的Kimi K3

AttnRes不仅是一个学术验证——它已经被部署到生产级模型中。2026年7月，Kimi发布了K3模型，这是全球首个开放的3T级别模型，总参数量2.8万亿。Kimi K3的架构以Kimi Delta Attention（KDA）和Attention Residuals（AttnRes）为核心组件。

KDA负责序列方向的信息流扩展，AttnRes负责深度方向的信息流优化——两者共同构成了Kimi K3的架构骨干。在MoE层面，K3使用了Stable LatentMoE框架，有效激活896个专家中的16个（极端稀疏率1.8%）。配合Quantile Balancing路由、Per-Head Muon优化器、SiTU激活函数和Gated MLA，这些架构创新共同实现了相比Kimi K2约2.5倍的Scaling效率提升。

K3在多个基准上展现了前沿性能：GPQA-Diamond 93.5分、SWE Marathon 42.0分、BrowseComp 91.2分。K3还能在48小时自主运行中完成芯片设计——为基于自身架构的nano模型设计并在4mm²内闭合100MHz时序的芯片，这展示了AttnRes赋能的深层推理能力在实际agentic场景中的价值。

从48B实验模型到2.8T生产模型的跨越，验证了AttnRes在极端规模下的稳定性和有效性。特别是Block AttnRes的工程优化——跨阶段缓存和两阶段推理——在2.8T参数规模下仍然保持<2%的推理延迟开销，证明了其工程设计的前瞻性和鲁棒性。

🔧技术线索与实操Tips

基于AttnRes论文和Kimi K3的生产经验，以下是可直接应用于LLM架构研发的技术线索和实操建议。

Tip 1：将AttnRes视为标准残差连接的即插即用升级。Block AttnRes的设计目标是drop-in replacement——不需要修改注意力机制、MLP或归一化层。实现上仅需在每个子层之前添加跨块注意力逻辑，以及每层一个d维伪查询参数。如果你的模型已经有PreNorm架构，集成Block AttnRes的代码改动量约50-100行。GitHub上的官方仓库（MoonshotAI/Attention-Residuals）提供了PyTorch风格伪代码，可直接作为实现参考。

Tip 2：块数量N=8是经过验证的良好起点。消融实验表明，N在6-10范围内效果稳定，N=8是一个安全的默认值。过少的块（N<4）会损失太多选择性——块表示过于粗糙，跨块注意力无法有效区分不同层的信息。过多的块（N>16）增加内存和通信开销但收益递减——当块大小接近单层时，Block AttnRes退化为Full AttnRes，失去了分块的效率优势。对于层数L=80的模型，N=8意味着每块10层（5个Transformer层），这是一个平衡选择性和效率的甜蜜点。

Tip 3：在Scaling Law实验中用Block AttnRes匹配1.25x计算量的基线。如果你在做架构搜索，可以用Block AttnRes作为"计算效率乘数"——在相同计算预算下，Block AttnRes等效于1.25倍计算量的标准残差基线。这意味着在比较新架构与基线时，应将Block AttnRes作为更强的基线——否则你可能高估自己架构的优势。

Tip 4：RMSNorm是AttnRes注意力计算的关键组件，不要省略。伪代码中跨块注意力使用RMSNorm对块表示V进行归一化后再与伪查询做点积。如果省略RMSNorm，不同块的表示幅度差异（PreNorm稀释导致的O(L)增长）会使得点积被高幅度块主导，softmax退化为近似argmax——失去选择性加权的意义。RMSNorm确保所有块表示在相同幅度尺度上参与注意力计算。

Tip 5：推理优化的关键是两阶段策略和online softmax。如果你在部署AttnRes模型，不要为每个token重新计算所有跨块注意力。两阶段策略将块表示的归一化键预计算并缓存，推理时仅需做一次轻量加权求和。这使得推理延迟开销<2%——对于需要高吞吐量的API服务而言，这是可接受的。如果你的推理框架支持prefix caching，可以进一步将块表示的键缓存与KV cache统一管理。

Tip 6：在流水线并行中实现跨阶段缓存以避免冗余通信。Block AttnRes需要跨流水线阶段访问块表示。如果每次前向传播都重新传输块表示，通信开销会抵消AttnRes的效率优势。跨阶段缓存策略将块表示缓存在使用它的阶段，仅首次生成时传输。这与ZeRO的参数缓存策略类似——通过将通信摊销到训练过程中来最小化同步开销。

Tip 7：关注GPQA-Diamond等深度推理基准来评估AttnRes的收益。AttnRes的最大收益出现在需要多步推理的任务上（GPQA +7.5分）。如果你的评估套件主要包含浅层知识检索任务（如TriviaQA），可能低估了AttnRes的价值。建议在评估中包含GPQA-Diamond、MATH、HumanEval等需要深层信息组合的基准。

Tip 8：AttnRes与MoE架构天然兼容。Kimi K3的成功证明了AttnRes与极端稀疏MoE（896选16）的兼容性。MoE的专家路由发生在每个Transformer层内部，而AttnRes的跨块注意力发生在层之间——两者在计算流上不冲突。如果你在训练大型MoE模型，AttnRes是一个低风险的架构增强选项。

🔮对LLM架构设计的启示

AttnRes的意义超越了其作为残差连接改进的直接贡献。它揭示了一个更深层的架构设计原则：深度方向和序列方向的信息流应当对称地处理。

在序列方向，我们已经经历了从RNN的固定递推到注意力机制的选择性检索的范式转变。这一转变催生了Transformer，彻底改变了NLP。然而在深度方向，我们一直停留在固定递推——残差连接本质上就是深度方向的RNN。AttnRes完成了这一对称性：深度方向也拥有了选择性检索的能力。

这一对称性的完成暗示了更多可能性。如果深度方向可以受益于softmax注意力，那么序列方向的其他创新是否也有深度方向的对应物？例如，Flash Attention的tiling策略是否可以启发深度方向的高效注意力实现？Sparse attention（如Longformer的滑动窗口+全局token）是否可以启发深度方向的稀疏跨层连接？这些问题值得探索。

另一个启示是关于架构搜索空间的方向。当前的LLM架构创新主要集中在注意力机制的变体（GQA、MLA、线性注意力、Delta Attention等）和MoE路由策略上。AttnRes指出了一条不同的路径：不改变层内计算，而是改变层间信息流。这条路径的优势在于它与层内创新正交——你可以同时使用KDA（序列注意力创新）和AttnRes（深度信息流创新），两者互不干扰。这种正交性意味着架构创新的空间比单一维度的优化要广阔得多。

最后，AttnRes的成功也验证了一个实用的架构创新方法论：用最小的参数和计算开销解决最根本的问题。AttnRes每层仅增加一个d维向量，推理开销<2%，但它解决的是PreNorm稀释这一影响所有深层模型的基础问题。这种"最小干预、最大杠杆"的设计哲学，与LoRA在微调领域的成功有异曲同工之妙——不是所有创新都需要大规模的架构重写，有时一个精准的小改动就能撬动显著的收益。

📄参考信息