Attention Residuals:让深层网络学会按需回看历史层,而不是把所有层一股脑叠上去
Attention Residuals让深层网络学会按需回看历史层而不是把所有层一股脑叠上去论文标题Attention ResidualsarXivhttps://arxiv.org/abs/2603.15031代码仓库https://github.com/MoonshotAI/Attention-Residuals作者Kimi TeamGuangyu ChenYu ZhangJianlin Su提交时间2026 年 3 月 一句话看懂这篇论文做的事其实可以用一句很朴素的话概括别再让每一层都被迫平等地背着全部历史往前走了应该让当前层自己决定究竟要更信最近几层还是该回头找更早的表示。传统 PreNorm 残差连接像一个只会“累加”的记账员前面每层的输出不管有用没用都以权重 1 加进来。层数一深隐藏状态幅值一路涨后面新层的声音反而被稀释掉。Attention Residuals的做法很直接——把“深度上的残差累加”改成“深度上的注意力选择”。如果你愿意把 Transformer 的时间维注意力想成“当前 token 去翻前文”那这篇论文的核心类比就是当前层也该能翻前面的层。 这篇论文到底在打什么靶子现代大模型几乎都爱用PreNorm。原因很现实训练稳、扩展方便、工程上成熟。可 PreNorm 有个长期潜伏的问题论文里把它说得很清楚所有残差分支都按固定单位权重累加会带来 hidden state 随深度失控增长进而稀释单层贡献。把这个过程想成开会就很好懂。一个会议室里每个发言人说完都不关麦后面的人还得继续叠在前面人的声音上。层数少时还凑合层数一多房间里的背景噪音会越来越大。后面那位想让别人听到自己只能把嗓门越提越高。于是你会看到两件事一起发生输出幅值越往深层越大因为新层得用更大的输出去盖过前面积累下来的“旧账”。梯度分布失衡早层会拿到异常重的梯度负担整条深度路径不好调。这也是论文里反复强调的PreNorm dilution。我觉得这个问题以前不是没人意识到而是大家更多把精力放在 token 维注意力、MoE、KV cache、长上下文这些更显眼的地方。残差连接因为太“基础设施化”反而长期被默认成一个不值得再动的组件。这篇论文最漂亮的地方就在这它没有再去发明一个更复杂的 block而是回过头去重做“层与层之间怎么连接”这件老问题。 从“时间注意力”类比到“深度注意力”作者提出的核心直觉叫Time-Depth Duality。RNN 在时间维上把所有历史压进一个状态里Transformer 用注意力替代它让当前位置能按需访问过去 token。那在网络深度上残差连接不也在做类似的压缩吗传统残差是hlhl−1fl−1(hl−1)h_l h_{l-1} f_{l-1}(h_{l-1})hlhl−1fl−1(hl−1)它等价于一种很粗暴的历史聚合前面所有层的影响都被塞进一个不断膨胀的单一状态里。AttnRes 改成hlα0→lh1∑i1l−1αi→lfi(hi),∑i0l−1αi→l1h_l \alpha_{0\to l} h_1 \sum_{i1}^{l-1} \alpha_{i\to l} f_i(h_i), \quad \sum_{i0}^{l-1}\alpha_{i\to l}1hlα0→lh1i1∑l−1αi→lfi(hi),i0∑l−1αi→l1这里最关键的变化就两个不再固定权重为 1而是学一个深度方向的 softmax 权重。当前层可以从更早层直接取信息不必只能吃上一层的混合结果。这就像把“只有一根接力棒”的比赛改成了“当前跑者可以回头从多个补给站拿水”。你不必盲信刚刚那一棒也不用把一路上所有东西都背在身上。更具体一点Full AttnRes 里每一层有一个可学习的伪 query 向量wlw_lwl键和值来自更早层的表示。作者用 RMSNorm 处理 key再做 softmax。这个小设计很关键因为如果不做归一化幅值大的层会天然在注意力里占便宜等于把“谁声音大谁赢”又偷偷带回来了。我很喜欢这个设计里的克制感。它没有把 query 做成输入相关的复杂投影而是先用每层一个可学习向量。这样做看上去“弱”但带来两个现实好处推理时不需要额外的d×dd \times dd×d投影开销同一 block 内多个层的 query 可以提前并行批处理后面基础设施优化全靠这个性质撑起来。这不是为了论文好看而是明显考虑过大模型训练和推理里的真账单。️ Full AttnRes 很优雅但工程上会卡在哪如果完全照理想方案来当前层要看所有更早层那就是Full AttnRes。它在概念上最纯粹也的确给出了最好 loss。问题也很直接需要保存所有层输出带来O(Ld)O(Ld)O(Ld)的存储与通信负担深度上全连接后虽然计算量O(L2d)O(L^2d)O(L2d)对L1000L \lt 1000L1000的网络还不算离谱但在大规模训练里真正难受的往往不是算力而是跨 stage 通信和激活保活。论文这部分说得很坦白。小规模训练时AttnRes 几乎不额外占内存因为反向传播本来就要留激活。可一旦上到分布式训练、激活重计算、pipeline parallel这些原本能释放再重算的中间状态现在得一直带着还得跨 stage 传。账就变了。所以作者没有停留在“理论最优”而是继续往下做了个更能落地的版本Block AttnRes。 Block AttnRes把“全层回看”压缩成“按块回看”Block AttnRes 的思路很像做摘要把LLL层切成NNN个 blockblock 内部的层输出先累加成一个 block 表示当前层主要对历史 block 表示做注意力对当前 block 内部还保留一个逐层更新的 partial sum避免块内信息彻底糊掉。于是Full AttnRes 里“看所有历史层”的问题被替换成“看历史 block 当前 block 的局部累积”。效果是什么存储和通信从O(Ld)O(Ld)O(Ld)降到O(Nd)O(Nd)O(Nd)计算从O(L2)O(L^2)O(L2)降到O(N2)O(N^2)O(N2)当N≈8N \approx 8N≈8时已经能吃到大部分收益。这点非常重要。很多架构论文的坏毛病是理论方案很好压缩版一做性能就塌。AttnRes 没塌。Block 版本和 Full 版本在大尺度下只差一点点这就有工程价值了。官方仓库给了一张很直观的总览图图1左边是标准残差所有历史只能被均匀加和中间是 Full AttnRes每层对所有历史层做深度注意力右边是 Block AttnRes把层分组后在块级别做选择把内存从O(Ld)O(Ld)O(Ld)压到O(Nd)O(Nd)O(Nd)。这张图基本把整篇论文的骨架讲完了。如果只看这张图我会把这篇工作的贡献归纳成一句话把“深度维度的跳连”从手工写死升级成了可学习的稀疏回看。⚙️ 真正让我高看一眼的是它把基础设施也补齐了很多论文做到这里就停了提出 Block 版本然后说“我们工程上也能跑”。这篇没有糊弄它把通信、prefill、解码延迟这些麻烦账都摊开了。1. 训练阶段跨 stage 缓存如果 pipeline parallel 里每次都把完整 block 历史重新发一遍通信会非常浪费。作者做了一个cross-stage caching前一轮虚拟 stage 收到的 block本地先缓存下一轮只传增量块。论文给出的结论很明确峰值单次传输代价从O(C)O(C)O(C)降到O(P)O(P)O(P)在 steady-state 1F1B 下这个改动可以把额外通信尽量盖到计算下面端到端训练开销在 pipeline 并行场景下低于 4%。2. 推理阶段两阶段计算Block AttnRes 还有一个麻烦每一层都回看历史 block会带来反复读取 KV 风格缓存的开销。作者把它拆成两段Phase 1同一 block 内所有层的 query 一次性并行打到历史 block 表示上Phase 2块内再顺序处理 partial sum并用 online softmax 把两路结果合并。这个设计本质上是在做“把随机小读合并成一次批量大读”。论文给出的每层 I/O 也很亮眼标准残差3dFull AttnRes24dBlock AttnRes5.5dmHC34d这组数字挺有说服力。Block AttnRes 的代价确实比标准残差高但远没高到不能接受而且明显比同类复杂残差方案省得多。3. 长上下文 prefill分片缓存长上下文时缓存 block 表示也会涨。论文举了一个例子128K 上下文、8 个 block 时如果直接存要到15 GB。作者进一步做了 sequence sharding按序列维切到PPP个 TP 设备上每卡内存可降到约 1.9 GB再叠加 chunked prefill比如 16K chunk单卡额外开销能压到0.3 GB 以下。这部分让我觉得作者不是只在做“paper architecture”而是真把它当一个要进大规模训练栈的组件在设计。 扩展规律实验它不是只在一个点上赢而是一路都赢论文先做了五个不同规模的 scaling law对比基线、Full AttnRes、Block AttnRes以及 mHC-lite。主表里的关键数字如下激活参数TokensBaselineBlock AttnResFull AttnResmHC-lite194M38.7B1.9311.9091.8991.906241M45.4B1.8951.8751.8741.869296M62.1B1.8291.8091.8041.807436M87.9B1.7661.7461.7371.747528M119.0B1.7191.6931.6921.694从趋势看有三件事Full AttnRes 基本一直最好Block AttnRes 跟得很紧到最大尺度只差0.001优势不是偶然点状提升而是整条 scaling curve 下移。对应的拟合曲线也很漂亮BaselineL1.891×C−0.057\mathcal{L} 1.891 \times C^{-0.057}L1.891×C−0.057Full AttnResL1.865×C−0.057\mathcal{L} 1.865 \times C^{-0.057}L1.865×C−0.057Block AttnResL1.870×C−0.058\mathcal{L} 1.870 \times C^{-0.058}L1.870×C−0.058图放在这里最直观图2三条曲线斜率差不多但 AttnRes 整体更低说明它不是改变了“怎么随算力增长”而是把同等算力下的损失整体往下拉。作者据此估算在 5.6 PFLOP/s-days 时Block AttnRes 相当于给了基线大约 1.25 倍的训练算力优势。我对这组结果的判断是这篇论文最硬的证据不是某个 benchmark 上多几点而是它让 scaling 曲线整条平移。这是架构工作里非常稀缺的信号。因为很多局部 trick 只在某个尺寸、某个训练 recipe、某个验证集上灵光规模一变就不稳。AttnRes 看起来没这个毛病。 放到 Kimi Linear 48B 上真实下游也有肉眼可见的涨幅作者没有满足于小模型 loss还把它接到Kimi Linear这套 MoE 架构上跑了完整 recipe。模型配置是27 个 Transformer blocks也就是 54 层256 个专家里每 token 路由 8 个再加 1 个共享专家总参数 48B激活参数 3BBlock AttnRes 采用每块 6 层也就是 9 个 block再加 token embedding一共 10 个深度源训练 recipe 与 Kimi Linear 1.4T token 方案保持一致更具体一点先在 1T tokens 上做 WSD 预训练再接约 400B 高质量 token 的中训后面继续把上下文长度拉到 32K全局 batch size 是 8M tokens。最终下游成绩如下评测BaselineAttnRes提升MMLU73.574.61.1MMLU-Pro52.252.20GPQA-Diamond36.944.47.5BBH76.378.01.7ARC-Challenge64.665.71.1HellaSwag83.283.40.2TriviaQA69.971.81.9GSM8K81.782.40.7MGSM64.966.11.2Math53.557.13.6CMath84.785.10.4HumanEval59.162.23.1MBPP72.073.91.9CMMLU82.082.90.9C-Eval79.682.52.9这里最抓眼球的是两类任务多步科学推理GPQA-Diamond 直接涨了7.5代码与数学Math3.6HumanEval3.1这很符合 AttnRes 的机制直觉。知识型任务有涨但最猛的地方出现在需要后层反复调用早层中间表示的复合推理任务上。你可以把它理解成模型不是只有更大记忆而是更会在深度方向上检索自己。我会把这部分看成论文最有产品味道的结果。因为如果一个残差改造只能让 validation loss 好看一点但下游不动那价值会打折。现在这些 benchmark 的提升说明深度信息流这件事真的会传导到实际能力上。 训练动态这篇论文最“显眼”的证据图官方仓库里另一张很关键的图是训练动态图3左图是验证损失中图是各层输出幅值右图是各层梯度幅值。基线模型里输出幅值随着深度一路涨梯度更多压在浅层Block AttnRes 把这种失衡压平了。这张图的价值在于它不是只告诉你“结果更好”还告诉你为什么会更好。输出幅值基线模型里越往后层输出 magnitude 越夸张到了末尾 block 几乎像冲天而起。这就是典型的 residual accumulation 失控。Block AttnRes 则呈现出一种周期性、有界的模式。每到 block 边界选择性聚合就像把残差流“清一次缓存”不会再让幅值无止境堆高。梯度幅值基线里浅层梯度特别大深层越来越弱很像一条失衡的水管。AttnRes 的 softmax 权重让不同深度源之间出现竞争梯度更平均训练也更舒服。这里我有个主观看法比起把它理解成“更强 skip connection”我更愿意把它理解成“给深度方向加了一套流量调度器”。残差连接原来像不限流的总线谁都往上写。AttnRes 给它装了个带竞争的路由规则结果是信息流和梯度流都没那么拥堵了。 消融实验真正起作用的到底是什么这部分是论文里我最关心的因为很多架构论文到收尾时会暴露“真正有效的其实是别的因素”。AttnRes 的消融比较干净。1. 和旧方案比赢点不在“多连几层”而在“按内容选层”16-layer 模型上Baseline1.766DenseFormer1.767mHC1.747Full AttnRes1.737Block AttnRes1.746DenseFormer 本质上也是让层看更多历史但权重是静态的、输入无关的结果几乎没赢。这说明一个要害光让历史连接变稠密没用关键是选择机制必须是内容相关的。这很像搜索系统。把更多文档都塞进候选集不代表答案就更好有用的是排序器而不是候选池本身。2. softmax 比 sigmoid 更合适把 softmax 换成 sigmoidloss 从1.737变成1.741。数字不算巨大但方向很稳定。作者的解释我认同softmax 有竞争关系能迫使模型在深度来源之间做更尖锐的选择sigmoid 更像“都给一点票”会把注意力重新搞成“人人有份”的旧残差味道。3. 输入相关 query 虽然更强但作者没贪这个便宜把 query 做成从当前 hidden state 投出来loss 可以到1.731比默认 Full AttnRes 的1.737更好。可作者没把这个版本当默认方案因为它会增加每层一个d×dd \times dd×d投影破坏 query 与前向解耦的性质让推理时块内并行失效内存访问也更难看。我很认同这个取舍。论文没有为了多刷 0.006 loss硬把工程代价翻上去。这类克制往往比单纯追 SOTA 更说明作者清楚自己在做什么。4. Block size 不必太细Block size 的扫描结果也有意思Full AttnRes也就是S1S1S11.737Block size 21.746Block size 41.746Block size 81.748再变大就逐步回到基线附近这说明一个很务实的点你没必要为了一点点理论完整性把块切得非常细。大约 8 个 block 就够用了。 架构扫描透露了一个更深的信号AttnRes 更吃深度论文还做了一组我觉得相当有味道的分析固定算力和参数预算扫不同的深度、宽度、头数配比。结果显示基线最优点在dmodel/Lb≈60d_{model}/L_b \approx 60dmodel/Lb≈60validation loss 为1.847AttnRes 最优点移到了dmodel/Lb≈45d_{model}/L_b \approx 45dmodel/Lb≈45validation loss 为1.802在固定参数预算下更低的dmodel/Lbd_{model}/L_bdmodel/Lb意味着更深、更窄。换句话说AttnRes 让模型更会利用深度了。这点我觉得很关键因为它不只是“给现有架构补点血”而是可能在更长期里影响架构选型如果残差连接不再强行把深度压成一个单状态那么“深一点到底有没有用”这个问题的答案可能会重新变。当然论文也很老实地说了偏深的网络在部署时通常延迟更高所以这不是简单的产品推荐。但从研究角度看这是一条很有后劲的线索有些时候我们不是缺深度而是旧残差让深度没法被好好消费。 我对这篇论文的判断不是花哨小 trick而是残差连接的一次重估价如果只用一句评价我会说这篇工作很像“把注意力从 token 维再复制一遍到 depth 维”但落地细节比这个口号扎实得多。我喜欢它的三个地方问题抓得准盯住了 PreNorm 里一个大家默认忍着用、但确实越来越痛的结构性问题。形式很干净没有发明一堆分支和矩阵而是把残差统一写进一个“深度注意力”框架里。工程补完度高训练通信、prefill 内存、解码 I/O 都给了方案和实测开销。如果非要挑毛病我觉得也有三处边界现在的主验证平台还是 Kimi Linear 系列。这当然很有说服力但还想看它在更广的 dense Transformer、不同 optimizer、不同数据配方上的迁移稳定性。Full AttnRes 依旧是更优上界。Block 版本已经很好可毕竟还是一种压缩近似。未来如果硬件和互联更强作者自己也承认会更倾向更细粒度 block甚至回到 full。这个方法更像预训练架构改动不是便宜的后装插件。你很难指望把一个已训练好的大模型直接打补丁替换残差收益大概率要靠从头训练才能出来。不过这些都不算硬伤。架构论文最怕的是“想法新鲜但证据链不完整”。AttnRes 的证据链是完整的。️ 如果把它迁到自己的模型里哪些经验最值得抄如果你做的是大模型预训练或架构研究我觉得这篇论文至少给了 5 个能直接拿走的启发残差不是铁板一块注意力、MoE、归一化都在被重构残差连接本身也值得重新设计。先处理深度信息流再谈更深网络很多“深层不如加宽”的经验规律可能部分只是旧残差的副作用。输入相关选择很重要静态加权跨层连接收益可能有限动态选择才真正改变表达路径。做架构时别回避系统账单如果一个结构不能讲清通信、缓存、prefill 和延迟那它离生产很远。用 block 压缩换工程可行性是很通用的范式不光能用于残差很多跨层或跨步记忆结构都能借鉴这个做法。如果你问我这篇论文最适合谁看我会给三个答案做大模型架构的人尤其是在琢磨残差、归一化、深度扩展的人做长训练栈和分布式系统的人因为里面的 pipeline cache 和 prefill 设计很实做推理与代码能力方向的人因为它在 GPQA、Math、HumanEval 这些任务上的收益很说明问题。 一个简化版伪代码Block AttnRes 到底怎么跑下面这段不是论文原文代码而是把它的核心流程压成了一个更容易看懂的伪代码# blocks: 历史 block 表示 [b0, b1, ..., b(n-1)]# partial_block: 当前 block 内已经累加的表示# w_l: 当前层的 learned pseudo-querydefblock_attn_res(w_l,blocks,partial_blockNone):sourceslist(blocks)ifpartial_blockisnotNone:sources.append(partial_block)keys[rmsnorm(x)forxinsources]scores[dot(w_l,k)forkinkeys]alphasoftmax(scores)h_lsum(a*vfora,vinzip(alpha,sources))returnh_lforblockinmodel.blocks:# Phase 1对历史 block 并行算# Phase 2块内顺序更新 partial_blockpartial_blockNoneforlayerinblock:hblock_attn_res(layer.w,blocks,partial_block)ylayer.forward(h)partial_blockyifpartial_blockisNoneelsepartial_blocky blocks.append(partial_block)真正在高性能实现里Phase 1 和 Phase 2 会拆得更细还会用 online softmax 合并统计量。但从算法直觉上你抓住一句话就够了当前层的输入不再是“上层状态 当前层输出”而是“对历史深度表示做一次选择性重组”。 一个更偏研究者视角的结论我会把Attention Residuals看成这样一篇论文它不是在重新发明 Transformer而是在重新定义 Transformer 里最习以为常的一根线——残差线。过去大家默认这根线越简单越好最好是常数 1、无脑相加、训练稳定。可当模型深度、专家规模、训练 token、长上下文都一路往上长时这种“简单”开始变成一种隐形束缚。AttnRes 的价值不只是把 loss 压低一点而是给了一个更一般的观点深度本身也可以被注意力化。一旦接受这个观点后面可以展开的空间就很大深度维是否能做更稀疏的路由不同 token 是否该用不同的 depth retrieval 策略残差、记忆、层间检索会不会最终长成统一框架这篇论文还不是这个方向的终点但很可能会变成一个明确的起点。 参考资料论文https://arxiv.org/abs/2603.15031官方仓库https://github.com/MoonshotAI/Attention-Residuals觉得有启发的话欢迎点赞、在看、转发。跟进最新AI前沿关注我的微信公众号机器懂语言