Kimi 团队提出 Attention Residuals:用注意力机制革新深度学习残差连接

Kimi用注意力替代传统残差加法,提升深层网络效率、稳定性和推理表现。

原文标题:告别残差加法,Kimi 给神经网络换了个 “智能引擎”

原文作者:数据派THU

冷月清谈:

Kimi 团队发布论文 Attention Residuals,尝试用输入相关的注意力机制替代神经网络中沿用多年的传统残差加法。传统残差连接虽然简单有效,但在模型不断加深时会出现两类问题:一是浅层信息在层层累加中被逐步稀释,二是为了维持影响力,深层模块往往产生更大的激活值,带来隐藏状态膨胀与训练不稳定。

这项工作的核心视角是“深度可以看作另一种时间维度”。类似 Transformer 用注意力替代 RNN 在时间上的递归压缩,Attention Residuals 让每一层不再机械地接收统一权重的历史信息,而是通过 Query 对前面各层进行选择性读取,用 Softmax 权重聚合最相关的表示。这样,深层可以直接调用早期层的重要特征,而不必被中间层“冲淡”。

考虑到全量深度注意力在超深网络中的成本问题,团队又提出 Block AttnRes:在块内保留简单累加,在块间用注意力做选择性聚合,从而把开销控制在可落地范围。实验显示,该方法在多个规模上都带来更低损耗,折算约有 1.25 倍计算优势;在数学、代码、多步推理等任务上提升更明显,同时改善了梯度分布和隐藏状态稳定性。文章也进一步引申出一个值得关注的方向:如果深度与时间可以类比,未来基础架构和优化器设计是否也能借鉴类似思路。

怜星夜思:

1、问题1:如果 Attention Residuals 真能替代传统残差,你觉得它最先会在哪类模型里大规模落地?是大语言模型、视觉模型,还是别的方向?
2、问题2:传统残差加法这么多年一直是主流,你觉得它未来会被真正“替代”,还是更可能和 Attention Residuals 长期共存?
3、问题3:文章把“深度看作时间”这个类比讲得很强,你觉得这种视角除了残差机制,还可能启发哪些新的模型或优化器设计?
4、问题4:从论文结果看,Attention Residuals 在推理、数学、代码任务上更亮眼。你觉得这是巧合,还是说明这种机制更适合“需要跨层调信息”的任务?

原文内容

图片
来源:算法进阶
本文约2800字,建议阅读6分钟
本文介绍了 Kimi 团队用 Attention Residuals 替代传统残差机制的成果。


只要接触深度学习神经网络的读者们对「 」一定不会陌生。自从 2015 年 ResNet 诞生以来,这种「将输入直接加到输出上」的简单逻辑,统治了几乎所有神经网络架构。


但就在刚刚,沿用了十年的残差机制「升级」了。随橙想呢,替代方法竟然是「注意力机制」。背后的直觉是:模型的「深度」其实就是另一种形式的「时间」。


就连 OpenAI 「推理模型之父」,主导了 o1/o3 系列、Codex 编程模型及 GPT-4 的 STEM 能力开发的 Jerry Tworek 都深受这一论文启发,认为应当重新思考之前的一切,「深度学习 2.0」的时代即将到来。



这篇颠覆传统残差连接机制的工作来自 Kimi 团队,发布了一项重磅技术报告:Attention Residuals ,该方法旨在通过对前序层进行学习到的、依赖输入的注意力机制,来取代标准的深度递归。


  • 论文标题:Attention Residuals

  • 论文链接:https://github.com/MoonshotAI/Attention-Residuals/blob/master/Attention_Residuals.pdf

  • 项目链接:https://github.com/MoonshotAI/Attention-Residuals


01 时间与深度的对偶


要理解 Attention Residuals 是在做什么,我们得先看传统的残差连接   出了什么问题。


在大模型向更深、更强演进的过程中,这种残差的加法机制带来了两个副作用:


1. 信息稀释: 残差连接采用固定单位权重的均匀聚合,导致浅层特征在向深层传递时,其相对贡献度随深度线性衰减。这种「信息稀释」现象限制了深层网络对底层原始表示的直接利用能力。随着层数增加,第一层的信息传到第一百层时,已经被后面九十九层的信息层层冲淡。


2. 隐藏状态爆炸:为了在不断累加的残差流中维持信号强度,深层模块往往需要输出模长更大的激活值。这种隐状态的无序扩张不仅破坏了数值稳定性,还导致梯度分布不均,增加了超大规模模型训练收敛的难度,直接导致了训练的不稳定性。


本文的天才之处在于,发现模型的「深度」其实就是另一种形式的「时间」。



论文作者之一的 Yulun Du 老师道出了该论文的核心思想:将注意力旋转 90°。


Attention Residuals (AttnRes) 由此诞生:为每一层配备了一个「智能筛选器」。每一层都会发出一个 Query,去之前的所有层里寻找最相关的特征,并按需分配权重进行聚合。


正如 RNN 在时间维度上的表现,残差连接在深度维度上也将所有先验信息压缩到了单一状态   中。在序列建模领域,Transformer 通过将递归替换为注意力机制,超越了 RNN,使每个位置都能通过数据依赖的权重,选择性地访问之前的所有位置。研究团队针对「深度」提出了同样的方法:



其中   是层特定的注意力权重,且满足   。与动辄达到数百万 Token 的序列长度不同,网络的深度通常较浅(L<1000),这使得深度方向上O(L2)复杂度的注意力机制在计算上具有可行性。


02 注意力残差


1. 理论重构:完整的注意力残差

传统的残差连接(ResNet)本质上是深度递归:它像 RNN 一样,把过去所有层的信息死板地 「压缩」进一个求和状态中。


注意力权重可以表示为   。本文采用了带有归一化的指数核函数,即在深度维度上执行 Softmax 注意力:



  • 核心创新: 既然 Transformer 用注意力机制取代了 RNN,解决了长序列的遗忘问题;那么 AttnRes 就在深度上取代了残差累加。

  • 数学实现: 每一层不再是简单地加上前一层,而是发出一个可学习的 Query,去和之前所有层产生的 Key 做匹配。

  • Softmax 权重: 通过 Softmax 归一化,模型可以 「挑选」 出对自己最有用的某几层。比如第 50 层可以直接提取第 2 层的特征,权重占比可以高达 0.8,而不用担心被中间的 48 层稀释。


2. 工程落地:Block AttnRes 的分块策略

全量注意力(Full AttnRes)虽然完美,但在超深模型里会导致显存和通信量爆炸(  复杂度)。为了让模型跑得通,研究团队设计了块结构。


局部求和(Intra-Block): 把模型分成 N 个块。在块内部,各层输出依然进行简单的累加,缩减为一个 「块代表」(Representation):



全局调度(Inter-Block): 每一层在进行残差聚合时,不再盯着 「每一层」 看,而是盯着 「每一个块」 看。对于第 n 块中的第 i 层,其 Value 矩阵 定义为:




在这种设计下,网络的第一层接收 Token 嵌入;每个块的第一层接收之前所有的块表示及 Token 嵌入;块内的后续层则额外关注当前块内已产生的累加结果。最终的输出层聚合所有 N 个块表示。


  • 效率奇迹: 实验发现,即便模型有上百层,只要划分成 N≈8 个块,就能获得绝大部分性能增益。

  • 复杂度骤降: 内存开销从随层数 L 增长,降到了随块数增长。这意味着你可以用极小的代价(推理延迟增加 < 2%),获得一个 「更聪明」 的深层网络。


图 1:Attention Residuals 概览:(a) 标准残差(Standard Residuals): 采用均匀加法累加的传统残差连接方式。(b) 全量注意力残差(Full AttnRes): 每一层都通过学习到的注意力权重,有选择地聚合之前所有层的输出。(c) 块注意力残差(Block AttnRes): 将各层划分为若干个「块」,将内存开销从 O (Ld) 降低至 O (Nd)。


03 战果:1.25 倍的「计算杠杆」


根据论文信息,实验架构与 Kimi Linear 完全一致,这是一种遵循 Moonlight / DeepSeek-V3 设计的混合专家模型(MoE) Transformer。唯一的修改是在残差连接中加入了 AttnRes;模型深度、隐藏维度、专家路由和 MLP 结构等其他组件均保持不变。


研究团队测试了五种模型规模,并为每种规模训练了三个变体:PreNorm 基准模型、全量 AttnRes 以及约 8 个块的 Block AttnRes。



下图展示了拟合后的规模化曲线。



三个变体的斜率相似,但 AttnRes 在整个计算范围内一致实现了更低的损耗(Loss)。基于拟合曲线,在 5.6 PFLOP/s-days 的计算量下,Block AttnRes 的损耗为 1.692,而基准模型为 1.714,这相当于1.25 倍的计算优势(Compute Advantage) 。随着模型规模增大,Full 与 Block 变体之间的差距在缩小。


研究团队的最大模型基于 Kimi Linear 48B 配置:27 个 Transformer 块(共 54 层),在 256 个路由专家中激活 8 个,外加 1 个共享专家,总参数 48B,激活参数 3B。该模型采用 Block AttnRes,每块 6 层,共产生 9 个块外加 1 个 Token 嵌入,形成 10 个深度方向的来源。



上图展示了模型在 1T token 训练过程中的动态变化:


  • 验证损耗: AttnRes 在整个训练过程中始终保持较低的验证损耗,尤其在衰减(Decay)阶段差距进一步拉大。

  • 输出量级: 基准模型遭受 PreNorm 稀释问题:随着隐状态量级随深度单调增长,深层网络被迫从固定缩放的归一化输入中学习越来越大的输出,以维持影响力。而 Block AttnRes 将这种增长限制在每个块内,通过块边界的选择性聚合重置了累加过程,呈现出有界的周期性模式。

  • 梯度量级: 在所有残差权重固定为 1 的基准模型中,梯度流在深度上的分布极不均匀,导致早期层梯度过大。Block AttnRes 的可学习 Softmax 权重引入了来源之间的竞争,从而实现了显著更均匀的梯度分布。



下游性能表现: 如上表所示,Block AttnRes 在所有评测任务中均达到或超过了基准模型。

  • 提升显著的任务: 在多步推理任务中提升尤为突出,如 GPQA-Diamond (+7.5) 、Minerva Math (+3.6) 以及代码生成 b

  • 知识类任务: MMLU (+1.1)  TriviaQA (+1.9) 也展现了稳健的提升。


数据给出了最有力的证明:


  • 计算效率: 达到同样的性能,AttnRes 相比传统残差节省了约 20% 的计算量(1.25x 优势)。

  • 逻辑推理: 在数学、代码等硬核任务上提升显著。例如,在极难的 GPQA-Diamond 测试中,性能提升了 7.5 分。

  • 稳定性: 成功抑制了隐藏状态的数值爆炸,让深层网络依然能保持「冷静」和「高效」。


04 总结:Rethink & Imagine


用更高维的视角看基础架构的研究,时间和空间都是相通的。


这篇论文「将注意力旋转 90°」的思想,似乎带给 Karpathy 一些启示和思考。



ResNet 的残差流是信息在不同空间深度上的传递。SGD (随机梯度下降)的权重流是信息在不同时间维度上的传递。


研究团队觉得 ResNet 的加法太朴素了,所以提议用 Attention 来筛选过去每一层的输出。 既然 SGD 也是 ResNet,「Attention is All You Need」,那我们为什么不能在优化器里也加上 Attention?


架构的生命力,往往来自于对惯性的反思。


当我们回过头去审视那些基础架构,或许就能在过去的故纸堆中,发现更多通往未来的巧妙结合。


编辑:于腾凯

校对:林亦霖



关于我们

数据派THU作为数据科学类公众号,背靠清华大学大数据研究中心,分享前沿数据科学与大数据技术创新研究动态、持续传播数据科学知识,努力建设数据人才聚集平台、打造中国大数据最强集团军。




新浪微博:@数据派THU

微信视频号:数据派THU

今日头条:数据派THU

我觉得可以将注意力机制应用到特征维度上。不同的特征对于不同的任务,重要性肯定是不一样的。我们可以用注意力机制来动态地调整不同特征的权重,让模型更加关注重要的特征。这就像“把好钢用在刀刃上”,让模型更加高效。

可以考虑用一些近似注意力机制来降低计算复杂度。比如,稀疏注意力、线性注意力等。这些方法可以在保证一定性能的前提下,大幅减少计算量。将这些方法应用到 AttnRes 中,应该可以在性能和效率之间取得更好的平衡。

这个问题很有意思!图像处理中,可以考虑将图像的深度(比如卷积神经网络的层数)视为一种空间上的“时间”,用 AttnRes 来选择性地聚合不同层级的特征,理论上会比简单堆叠卷积层更好。长文本建模的话,Attention 本身就已经很强大了,但可以尝试用 AttnRes 来优化 Transformer 的层间连接,是不是能解决梯度消失问题?

从根源上解决问题可能是更好的思路,隐藏层爆炸本质上是信息过载在深度网络中的体现,是不是也可以考虑从信息论的角度入手,设计一些能够更有效地压缩和传递信息的网络结构?比如引入一些稀疏连接或者信息瓶颈。

会不会是因为 Attention 机制更不容易受到“信息稀释”的影响?在数学和代码任务中,一些早期的关键信息可能对最终结果至关重要。传统的残差连接可能会冲淡这些信息,而 Attention Residuals 可以让模型更灵活地提取和利用这些信息。

我觉得这个思路很赞!把深度看作时间,本质上是将空间关系转化为时间序列关系来处理。很多领域都有类似的结构,比如社交网络、知识图谱等等,都可以借鉴这种方法来优化信息传递和整合。

我觉得还可以试试混合精度训练(Mixed Precision Training)。用半精度(FP16)来训练模型,可以有效降低显存占用和计算时间。当然,需要仔细调整超参数,避免梯度消失等问题。

块大小的选择应该是一个trade-off。如果块太小,那么每个块内的信息聚合可能不够充分,模型无法有效地利用历史信息。如果块太大,那么计算复杂度会增加,并且可能引入更多的噪声。所以,需要根据具体的任务和数据进行实验,找到一个平衡点。

这个问题很有意思!Attention机制在NLP里已经证明了其价值,现在拓展到神经网络的深度方向,确实让人眼前一亮。我觉得这个思路很有潜力,也许可以尝试用类似的方式优化循环神经网络或者其他序列模型,让模型更好地关注序列中不同时间步的信息。

选择合适的块大小 N 值是一个trade-off问题。N 值越大,每个块内的计算量越小,但块间的交互信息可能损失更多;N 值越小,每个块内的计算量越大,但可以更好地保留块内的局部信息。我认为最佳的 N 值应该根据具体的任务和数据集进行调整,可以通过实验来确定。

谢邀,人在ICU,刚下飞机。我觉得楼上说的都有道理,但还可以从另一个角度看:残差连接本质上是加法,而加法是线性操作。在深度学习中,我们希望引入更多的非线性,增强模型的表达能力。Attention 机制引入了 Softmax 等非线性操作,可以使模型更好地学习特征之间的复杂关系。

大胆预测一下:Attention is all you need!既然 Transformer 已经在各个领域都取得了成功,那么 Attention Residuals 肯定也可以。说不定以后我们看到的网络架构都是基于 Attention 的!手动狗头。

这个问题问得好!我之前在跑实验的时候也遇到了类似的困惑。我的经验是,先从小N值开始尝试,然后逐渐增大N值,观察性能的变化。当性能提升不再明显时,就可以停止增大N值了。另外,还可以考虑使用一些自动化调参的工具,例如Bayesian Optimization,来寻找最佳的N值。

数值爆炸是个大问题啊!如果你训练过超大模型就知道了,经常会遇到NaN的问题,就是因为数值溢出了。Attention Residuals可以限制隐藏状态的增长,让模型训练更稳定,更容易收敛。尤其是对于MoE这种参数量巨大的模型,稳定性更重要。

有没有可能出现过拟合问题?注意力机制很强大,但如果数据量不够大或者分布不均衡,会不会让模型过度关注某些特定的层,反而降低了泛化能力?

这个角度很有意思!在控制理论里,系统的状态演化就既可以看作是在时间上的演变,也可以看作是在状态空间中的转移。某种程度上,时间序列分析和状态空间模型就是这种对偶性的体现叭。

楼上说得对,计算资源确实是个问题。另外,我觉得注意力机制的可解释性可能也会下降。ResNet 好歹还有个明确的梯度传递路径,这玩意儿一层套一层,debug 起来估计要命。

从更哲学的角度来说,世界的变化既可以看作是时间的流逝,也可以看作是事物状态的改变。关键在于我们要用动态的眼光去看待问题,不要只局限于某个维度。

我不觉得会是革命性的突破,更像是对现有模型的优化。ResNet 已经很成熟了,要完全替代它难度很大。但 Attention Residuals 或许能在某些特定领域找到自己的用武之地,比如需要更精细特征提取的任务。