告别残差加法，Kimi 给神经网络换了个 “智能引擎”

来源算法进阶 本文约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-Residuals01 时间与深度的对偶要理解 Attention Residuals 是在做什么我们得先看传统的残差连接出了什么问题。在大模型向更深、更强演进的过程中这种残差的加法机制带来了两个副作用1. 信息稀释 残差连接采用固定单位权重的均匀聚合导致浅层特征在向深层传递时其相对贡献度随深度线性衰减。这种「信息稀释」现象限制了深层网络对底层原始表示的直接利用能力。随着层数增加第一层的信息传到第一百层时已经被后面九十九层的信息层层冲淡。2. 隐藏状态爆炸为了在不断累加的残差流中维持信号强度深层模块往往需要输出模长更大的激活值。这种隐状态的无序扩张不仅破坏了数值稳定性还导致梯度分布不均增加了超大规模模型训练收敛的难度直接导致了训练的不稳定性。本文的天才之处在于发现模型的「深度」其实就是另一种形式的「时间」。论文作者之一的 Yulun Du 老师道出了该论文的核心思想将注意力旋转 90°。Attention Residuals (AttnRes) 由此诞生为每一层配备了一个「智能筛选器」。每一层都会发出一个 Query去之前的所有层里寻找最相关的特征并按需分配权重进行聚合。正如 RNN 在时间维度上的表现残差连接在深度维度上也将所有先验信息压缩到了单一状态中。在序列建模领域Transformer 通过将递归替换为注意力机制超越了 RNN使每个位置都能通过数据依赖的权重选择性地访问之前的所有位置。研究团队针对「深度」提出了同样的方法其中是层特定的注意力权重且满足。与动辄达到数百万 Token 的序列长度不同网络的深度通常较浅L1000这使得深度方向上OL2复杂度的注意力机制在计算上具有可行性。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%获得一个 「更聪明」 的深层网络。图 1Attention 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