杨植麟带队发论文，Kimi下一代模型架构曝光，推理之父惊呼：深度学习2.0要来了！

智东西
作者｜陈骏达
编辑｜云鹏

“一切都需要被重新思考，深度学习2.0时代即将来临。”在读完下方这篇来自月之暗面的最新论文后，前OpenAI大牛、“推理模型之父”Jerry Tworek发出感叹。

智东西3月16日报道，今天，月之暗面发布论文，提前预览了下一代模型的关键模块——注意力残差（Attention Residuals，简称AttnRes）。论文的核心亮点在于对大模型中最基础、但长期被忽视的结构——残差连接（Residual Connection）进行了重新设计。

在传统Transformer架构中，每一层都会将前一层的输出与当前层的计算结果直接相加，这种“等权累加”的方式虽然稳定了深度网络训练，但随着模型层数不断增加，所有历史层的表示被简单叠加，容易导致信息被逐渐稀释，深层网络的表达能力受到限制。

AttnRes用注意力机制替代固定的残差相加方式，让每一层能够从所有历史层表示中动态选择最有用的信息。每一层相当于一个查询（Query），在之前层的表示中进行检索并分配权重，使模型能更灵活地利用不同层级的特征。

在实现层面，论文还提出了块级注意力残差（Block AttnRes）来解决大模型训练时的计算和通信开销问题。

实验中，月之暗面将AttnRes集成到Kimi Linear架构（总参数量48B/激活参数量3B）中，在1.4T token上进行预训练。AttnRes缓解了预归一化的稀释问题，产生更均匀的跨深度输出幅度和梯度分布，从而缓解深层模型训练中的不稳定问题。

在相同计算预算下，引入AttnRes的模型在多个基准测试中均取得稳定提升，尤其是在数学推理、科学问答以及代码生成等需要复杂推理能力的任务上。

此外，Block AttnRes的验证损失为1.692，而基线为1.714，相当于约1.25倍的计算效率优势。

这篇论文在大模型圈引发广泛讨论，月之暗面官方在X平台发布的帖子阅读量已逼近60万次。除了Jerry Tworek之外，OpenAI联合创始人Andrej Karpathy也评论道：“随机梯度下降（常用的优化算法之一）也是一种ResNet，残差流就是权重……看来我们还没把‘Attention is All You Need’ 这句话完全按字面意思理解透。”

月之暗面的数十名研究员参与了这项研究，对项目贡献最大的三位研究员包括Guangyu Chen、Yu Zhang（Kimi-Linear一作）、Jialin Su。月之暗面的三位联合创始人杨植麟、吴育昕、周昕宇带领团队完成了这项研究。

一、标准残差连接，面临“等权累加”之困

由知名华人AI学者何恺明等人提出的标准残差连接，最早在ResNet中系统化提出，并迅速成为现代深度神经网络的核心结构之一。

随着深度学习的发展，这一设计不仅在计算机视觉中取得巨大成功，也逐渐成为大型语言模型和各类Transformer系统中的基础构建模块。

其核心思想非常直观：让网络学习“残差”（即输入与输出之间的差值），而不是直接拟合完整映射。标准残差连接中，每一层的输入等于所有之前层输出的累积和。残差连接实际上是一种跨深度的信息聚合机制，所有层的贡献权重都是1，也就是均匀加权。

在现代大模型中，残差结构通常还会与PreNorm（预归一化）结合使用。具体做法是在每个子层计算之前先进行层归一化，然后再进入注意力或前馈网络模块，最后再通过残差连接将输出加回原始表示。

这种结构在Transformer模型中被广泛采用，因为它能够进一步改善训练稳定性，使梯度在深层网络中更容易传播，并避免训练后期出现数值不稳定的问题。

标准残差连接之所以能获得广泛采用，在于其解决了深层网络训练困难的问题，允许梯度绕过复杂变换，缓解了梯度消失的问题。同时，残差使每层对表示的修改是渐进的，不会破坏已有表示。模型训练更稳定、更容易优化。

然而，这一方式也存在一些局限性。标准残差连接中所有层的权重一致，但现实情况是不同层信息的重要性不同。理想情况下应该采用可学习的权重。

此外，学习时，每层网络只能看到压缩后的状态，导致许多早期信息被混合压缩，无法单独访问某一层输出，导致信息不可逆的丢失。同时，后期层在模型的贡献中变大，早期层被逐渐淹没。

模型剪枝（prune）等方式对模型性能影响不大，这一现象也从侧面证明，当前广泛采用的残差结构，使很多层贡献很小。

正是在这样的背景下，月之暗面的研究团队提出了一种新的思路：用习得的softmax注意力机制来替代固定的残差累积，也就是AttnRes。

二、用注意力替代固定权重，让每一层学会“检索”历史

在AttnRes中，每一层在计算时不再仅仅依赖上一层的输出，而是可以对之前多个层的表示进行一次注意力检索。模型通过学习得到的一组权重，对不同层的输出进行加权组合，从而生成当前层的输入表示。

由于这些权重通过softmax归一化，它们能够反映不同层在当前上下文中的相对重要性，使信息聚合过程更加灵活和自适应。

这种机制带来的一个直接好处是，模型能够更有效地利用早期层的信息。由于每一层都可以直接访问历史层的表示，早期信息不再被完全压缩到单一状态中，而是可以在需要时被重新调用。

这种“跨层检索”的能力，使得深层网络在信息利用上更加接近于Transformer在时间维度上的注意力机制。

对于每个token，完整AttnRes需要O(L²d)的算术运算和O(Ld)的内存来存储层输出。由于深度远小于序列长度，算术成本适中。

当然，如果让每一层都对所有历史层进行注意力计算，也就是使用全注意力残差，计算成本会迅速增加。为了解决这一问题，研究团队进一步提出了Block AttnRes（块注意力残差）的结构设计。

该方法将网络层划分为若干块，每个块内部保持常规计算方式，而跨块之间则使用压缩后的表示进行注意力交互。

通过这种方式，模型在保持灵活信息检索能力的同时，也能有效控制计算复杂度，使这一机制在大规模模型中仍然具有可行性。

▲从左到右：传统残差、全注意力残差、块注意力残差

三、基础设施协同优化，计算效率达1.25倍

不过，Block AttnRes也给基础设施带来新的挑战。

在训练阶段，Block AttnRes带来更多的流水线并行中通信开销。标准残差结构在流水线阶段之间只需传输固定大小的隐藏状态，而Block AttnRes需要在每个阶段访问此前所有块的表征。

如果采用朴素实现，就必须在阶段转换时重复传输完整的历史块表示，通信量会随着块数量累积呈平方增长。为了解决这一问题，月之暗面引入跨阶段缓存（cross-stage cache）机制，大幅减少冗余通信。

同时，每个块只需在虚拟阶段中存储一次，结合激活检查点机制后，每层激活内存占用与标准Transformer基本一致，因此整体训练的额外开销很小，在实际系统中端到端训练时间增加不到4%。

在推理阶段，Block AttnRes的主要挑战来自跨层块表征的重复访问和长上下文缓存带来的内存压力。为此他们采用了两阶段计算策略，首先在阶段一中对缓存的块表征执行一次批量查询，为同一块内的所有层同时计算块间注意力，从而将原本每层都需要访问历史块的操作合并为一次矩阵计算，显著降低内存读取次数。

随后在阶段二中按顺序计算块内各层的注意力，并通过在线softmax将结果与阶段一的输出合并，这一阶段可以与周围算子进行内核融合，从而进一步减少I/O开销。

此外，为缓解长上下文预填充时块表征缓存过大的问题，系统将这些表征沿序列维度在张量并行设备之间进行分片，使每个设备只存储部分序列。通过两阶段计算和序列分片的结合，BlockAttnRes在推理时仅带来不到2%的延迟开销，同时显著降低了长上下文场景下的显存需求。

在实验方面，研究团队将这一结构集成到Kimi Linear架构中进行验证。

实验表明，在相同参数规模下，引入Block AttnRes后，模型在多个基准测试中都取得了性能提升。其中提升最多的是多步推理基准测试GPQA-Diamond，采用AttnRes的模型得分提升幅度超过了20%。

与此同时，该方法还表现出良好的计算效率，Block AttnRes在损失函数上与使用1.25倍计算量训练的基线模型相当。

结语：将注意力引入“网络深度维度”，有望成为下一代模型关键架构

通过将注意力机制引入深度维度，Attention Residuals为深层网络提供了一种更加灵活的信息聚合方式，使模型能够动态选择历史表示，而不是被动接受均匀累积的结果。这一思路有望进一步提升深层模型的信息利用效率。

从架构演进的角度来看，这项工作也提出了一个值得关注的思路：在过去，注意力机制主要用于解决序列维度上的信息交互问题，而AttnRes则将注意力引入到了“网络深度维度”，让模型能够在不同层之间建立更加灵活的信息通路，可能为下一代大模型架构演进提供新的方向。