KIMIK3采用AttentionResiduals技术,深度方向信息流从等权累加变为可学习注意力聚合

月之暗面(Moonshot AI)在 KIMI K3 模型中引入 Attention Residuals(AttnRes)机制,将传统 Transformer 中固定权重的残差连接升级为深度维度上的可...

人工智能 阅读量: 1.74万

月之暗面(Moonshot AI)在其最新大语言模型 KIMI K3 的架构设计中,系统性地引入了一项关键技术创新:Attention Residuals(AttnRes)。这项技术并非简单改进,而是对 Transformer 残差连接范式的结构性重构——将原本在深度维度上“等权相加”的标准残差操作,替换为由内容与深度共同决定的“可学习注意力加权”,从而在不显著增加计算开销的前提下,有效缓解深层网络中的信息瓶颈问题。

传统 PreNorm Transformer 的残差结构,其信息流可被形象理解为“所有历史层输出以固定权重 1 倒入同一容器”。随着网络层数加深,新层输入信号相对容器容量越来越微弱,导致梯度衰减与表征退化。这一现象在超深模型(如 60 层以上)中尤为明显,构成训练稳定性与表达能力的硬性约束。

AttnRes 的核心思想是:让当前层主动“询问”哪些历史层对其当前任务最有价值,并据此动态分配权重。具体实现中,每层引入一个可学习的伪查询向量(pseudo-query),用于计算其与各历史层输出之间的注意力分数,再经 softmax 归一化后加权聚合。这意味着深度方向的信息流动从“刚性累加”转变为“柔性选择”,本质上是对 RNN 到 Transformer 范式跃迁在深度轴上的复刻。

为兼顾理论效果与工程落地,KIMI K3 采用的是 Block AttnRes 变体。该版本不保存每一层输出,而是将若干连续层划分为一个块(block),仅维护块级表示。在推理与训练时,当前块的输出与历史块表示共同参与深度注意力计算。此举显著降低显存占用与通信开销,使该机制可在 vLLM、SGLang 等主流推理框架中高效部署,并支持流水线并行优化。

图中清晰展示了三类残差机制的结构差异:(a) 标准残差为线性累加;(b) Full AttnRes 对每层历史输出做注意力加权;(c) Block AttnRes 在块粒度上执行相同逻辑,牺牲部分细粒度分辨率换取系统级可行性。

实证数据支撑了该设计的有效性。在 16 层模型的消融实验中,Block AttnRes 相较 Baseline 显著改善训练动态:验证损失下降更快且更稳定;各块输出幅值波动减小,表明表征更均衡;梯度幅值衰减趋势明显缓解(见图1)。尤其值得注意的是,Block AttnRes 在训练后期仍能维持较低的梯度噪声水平,而基线模型则出现梯度震荡加剧现象,印证其对深层优化过程的稳定作用。

文章配图

进一步的缩放律分析(Scaling Law)显示,Block AttnRes 在同等算力预算下(PFLOP/s-day)可达到更低的验证损失,其拟合曲线斜率约为 -0.058,优于 Baseline 的 -0.057,相当于在相同计算量下实现约 1.25 倍的性能增益。该结果说明 AttnRes 不仅未带来额外算力惩罚,反而提升了模型参数与计算资源的利用效率。

验证损失与计算效率关系图:Baseline、Full AttnRes、Block AttnRes 的 Loss-PFLOP/s 曲线

在更大规模模型上的验证更具说服力。针对 48B 参数量级的模型,月之暗面团队在多个下游任务(如 MMLU、CMMLU、C-Eval 等中文与英文综合能力评测)中对比了 Baseline 与 AttnRes 版本。结果显示,引入 AttnRes 后,模型在多数任务上取得一致提升,尤其在需要长程依赖与多步推理的任务中优势更为突出。这表明深度方向的信息重加权机制确实增强了模型对历史上下文的甄别与整合能力。

从技术谱系看,AttnRes 属于“跨层访问”(Cross-layer access)范畴,与 DenseNet、MRLA 等方法同源但路径不同。区别在于:DenseNet 使用固定权重串联各层;MRLA 引入卷积式通道投影;而 AttnRes 则通过注意力机制实现动态、内容感知的深度聚合。表格中最后一行明确标注了 AttnRes 的更新规则:当前层输出为历史块表示与当前块部分和的注意力混合,权重由可学习参数动态生成,属于“Dynamic”类型且支持多流(m streams)访问。

值得指出的是,AttnRes 并非孤立创新。同期已有多种残差路由方案(如 DeepNorm、ReZero、LayerScale)试图缓解信息稀释问题,但大多聚焦于归一化或缩放系数调整,未触及“聚合方式本身”的根本变革。AttnRes 的独特性在于,它将注意力机制从序列维度(token-level)迁移至深度维度(layer-level),形成一种新的信息流拓扑结构。

残差机制分类汇总表:Single-state recurrence、Multi-state recurrence、Cross-layer access 三大类共 12 种方法对比

在工程实现层面,该机制已集成至 SGLang 与 vLLM 两大开源推理框架。其核心计算分为两阶段:第一阶段跨块并行计算各块表示;第二阶段在块内顺序执行注意力聚合与在线 softmax 归一化。这种设计既保障了流水线并行效率,又避免了全层存储带来的显存爆炸风险。

目前,AttnRes 已应用于 KIMI K3 模型,并可能延伸至后续版本如 PI0.5、PI0.7 等。尽管原始技术报告未公开完整训练细节,但其在 48B 模型上的下游任务表现,以及在 16 层模型中的训练动态改善,已足够说明该机制具备实用价值与推广潜力。

随着大模型向更深、更宽方向持续演进,如何高效管理深度维度的信息流将成为架构设计的关键战场。AttnRes 提供了一条不同于传统归一化调优的新路径:不是压制梯度或放大信号,而是重构信息整合逻辑本身。这一思路或将启发更多面向深度轴的结构创新。