3月16日,Kimi团队发布了一篇名为《Attention Residuals》的论文,引来了一种大佬的转发和评论。

马斯克称这篇论文令人印象深刻(Impressive work from Kimi),卡帕西甚至反思“随机梯度下降(SGD)本身也是一种 ResNet(它的模块就是前向 + 反向传播),残差流就是权重…… 我们是不是还没把《注意力机制就是一切》这句话按字面意思理解透?”(SGD is a ResNet too (the blocks of it are fwd+bwd), the residual stream is the weights so... 🤦 We're not taking the Attention is All You Need part literally enough?)

这次我们不贩卖焦虑,实打实地看这篇论文究竟说了件什么事。

1 摘要(看论文必看摘要)

结合前置归一化(PreNorm)的残差连接是现代大语言模型(LLM)的标准组件,然而这类残差连接以固定的单位权重累加所有层的输出。这种无差别的聚合方式会导致隐藏状态的量级随模型深度出现不可控的增长,进而逐步稀释每一层的贡献。我们提出了注意力残差机制(Attention Residuals,简称 AttnRes),该机制将固定累加模式替换为针对前序层输出的 softmax 注意力聚合,让每一层都能通过可学习的、随输入动态变化的权重,选择性地整合早期层的表征。为了解决大规模模型训练中,对所有前序层输出做注意力计算带来的内存与通信开销问题,我们进一步提出了块注意力残差(Block AttnRes):将模型层划分为多个块,仅对块级表征做注意力计算,在保留全量注意力残差大部分性能收益的同时,大幅降低了内存占用。结合基于缓存的流水线通信机制与两阶段计算策略,块注意力残差可以作为标准残差连接的即插即用替代方案,仅带来极低的额外开销。

缩放定律实验证实,该方法的性能提升在不同尺寸的模型中均稳定存在,消融实验也验证了基于内容的深度维度选择机制带来的收益。我们进一步将注意力残差机制集成到 Kimi Linear 架构中(总参数量 48B / 激活参数量 3B),并在 1.4T tokens 的语料上完成了预训练。实验结果显示,注意力残差机制缓解了前置归一化带来的层贡献稀释问题,让不同深度的输出量级与梯度分布更均匀,同时在所有评估的下游任务中均实现了性能提升。

2 解决了什么问题?

2.1 问题1:传统残差连接的固定权重聚合导致层贡献稀释和隐藏状态失控增长

2.1.1 通俗化解释

先搞懂两个最基础的背景:

我们先把大语言模型类比成一个100层的食品加工厂。

你输入的一句话,就是最开始送进工厂的原材料。模型的每一层,就是工厂里的一个「加工车间」,每个车间只干一件事:给手里的食材做一步加工,再传给下一个车间;最终从最后一个车间出来的,就是AI给你的回答。

传统残差连接 就是这个加工厂定的一条铁规矩:

每个车间加工完,必须把上一个车间给你的所有食材,原封不动、一分不少地全部保留,再把自己新加工出来的食材加进去,一起传给下一个车间。绝对不能扔前面的任何东西,也不能挑挑拣拣,必须全留。

什么是 固定权重聚合

就是上面这条规矩,是完全死板、没有任何弹性的:

前面每一个车间的食材,传到后面的时候,分量占比永远是100%全加进去,不能多也不能少;

不管前面某个车间做的东西是关键核心,还是没用的废料,都必须完整保留。

用更通俗的话说:这就像你写一篇作文,要求你每加一段新内容,必须把之前写的所有段落,一字不差地全抄一遍,再写新的。不管之前写的是废话还是金句,都必须全留,所有段落的地位完全平等。


什么是 隐藏状态失控增长

隐藏状态 可以理解成每个车间拿到手的食材总重量。

类比到这里时,你应该已经看出来了:因为每一个前序环节的结果都必须被100%保留,所以越到后面隐藏状态越会发了疯地增长。

因为铁规矩要求 “前面所有东西必须全留,再加新的”,那食材的总重量,就会像滚雪球一样,越滚越大,完全刹不住车:

第1个车间:拿到初始原材料100克,加工完传给下一个车间;

第2个车间:必须保留这100克,再加自己新做的1克,总重量101克;

第3个车间:必须保留101克,再加自己新做的1克,总重量102克;

……

到第100个车间:手里的食材总重量已经到了199克,是最开始的快2倍。

层数越多,总重量涨得越离谱,这就是 失控增长——模型越往深层走,要处理的数据体量就越大,完全不受控制。


什么是 层贡献稀释

这是上面那个问题带来的最致命的后果,一句话说就是:越靠后的车间,自己辛辛苦苦做的新东西,在整个食材包里的占比越来越小,几乎没存在感,做了跟没做一样。

还是用上面的例子看:

第2个车间,自己加的1克新食材,在总重量101克里,还能占约1%的比例,好歹能被下一个车间注意到;

到了第100个车间,自己加的1克新食材,在总重量199克里,只占0.5%都不到,几乎被前面99个车间的旧食材完全淹没了。

哪怕这个第100层的车间,加工出了能决定最终成品好坏的关键食材,也会被前面几十层的旧内容盖过去,根本起不了作用。这就是 贡献稀释 —— 每层的价值,随着模型深度增加,被越冲越淡

更麻烦的是,它会形成一个死循环:

深层车间发现自己加的东西没人在意,为了让自己的内容能被看到,就只能拼命把自己的1克食材,做成10克、100克往里面猛加。这又反过来让食材总重量涨得更快,下一个车间的稀释更严重。

最终就是:隐藏状态越滚越大,每层的贡献却越来越淡,模型深度加了,但效果根本没跟上,甚至还会变得不稳定

而这一切的根源,就是最开始那条「固定权重、必须全量累加、不能选择」的死规矩。

2.1.2 概念讲解:传统残差连接(Standard Residuals)

原文在 2.1 Training Deep Networks via Residuals 和 6.2 Residual Connections as Structured Matrices 两个章节,给出了三层递进的数学定义

核心层间递推公式(最基础的残差更新规则)

展开后的全量累加公式(直接体现聚合本质)

原文对上述递推公式做了展开,明确:

结构化矩阵形式(全局权重的矩阵化表达)

原文6.2节将传统残差连接抽象为深度混合矩阵M,给出了全局数学表达:

这是原文最核心的传统残差连接架构示意图,用流程图直观展示了传统残差的逐层累加机制,和 Full/Block AttnRes 形成直接对比
图中专门给出了传统残差连接对应的全1下三角矩阵可视化,和上述矩阵公式完全对应,直观呈现了传统残差的权重分布特征

2.1.3 概念讲解:固定权重聚合(Fixed Unit Weights / Uniform Aggregation)

这个概念是传统残差连接的核心属性,原文有明确的数学定义、权重公式,并通过对比公式和可视化图片强化了该概念

1. 核心权重定义

原文摘要、引言、2.1节均明确:传统残差连接对所有前序层输出,采用固定的单位权重1进行累加,没有任何可学习、动态调整的系数。

对应数学形式为:

2. 对比式数学印证

原文3节给出了AttnRes的动态权重公式,反向明确了传统残差的固定权重本质:

该图展示了AttnRes学习到的动态权重分布,反向印证了传统残差没有动态选择、全层固定权重的核心特征

2.1.4 概念讲解:隐藏状态失控增长(Uncontrolled Hidden-state Growth)

核心量级增长公式

原文摘要、引言、5.2节均给出了明确的数学刻画:its unweighted accumulation causes hidden-state magnitudes to grow as O(L) with depth

含义:传统残差的无权重累加,会导致隐藏状态的量级随模型深度L呈线性增长,这是对「失控增长」最直接的数学定义。

增长规律的数学推导

该结论直接来自传统残差的累加公式:

上面展示的 Figure 5 是原文专门针对该概念的核心可视化图片,也是最直观的讲解图。图(b)中 Baseline 曲线(传统残差连接)的输出量级,随 Transformer Block Index(模型深度)呈单调增长,完美对应了O(L)的失控增长规律;同时对比了 Block AttnRes 的有界周期性波动曲线,直观凸显了传统残差隐藏状态随深度持续增长的问题。

补充:Figure 5 (a) 的 Validation Loss 曲线,也间接展示了隐藏状态失控增长带来的训练效果差距——传统残差的验证损失始终显著高于AttnRes。

2.2 问题2:跨层信息访问的内存与通信开销过高,难以支持大规模模型训练

还是用之前100 层食品加工厂的类比,先搞懂核心概念,再讲开销高的问题,最后说为啥搞不定大规模模型训练。

2.2.1 什么是跨层信息访问?

这是为了解决第一个问题(层贡献稀释) 想出来的最直接的办法:

既然传统残差只能拿上一个车间的东西,导致后面车间挑不到前面的有用内容,那干脆打破规矩,让每个车间能直接看、直接用前面所有车间的加工成果;比如第 50 层车间,不光能拿 49 层的,还能直接找第 10 层、第 20 层的核心食材用,不用再被前面的内容稀释。

这个 跳过相邻车间,直接找前面任意车间要东西 的操作,就是 跨层信息访问。它本身是个好想法,能解决层贡献稀释,但问题是实现这个想法的成本太高,需要占超多内存,还要花超多时间传数据,也就是所谓的“内存与通信开销过高”。

2.2.2 为什么跨层信息访问会让内存开销高?

内存就相当于加工厂的仓库,模型训练时,每一层的加工成果(隐藏状态)都要存在仓库里,才能被后面的层调用。

传统残差因为只拿上一层的东西,仓库只需要临时存上一个车间的成果就行。比如第50层车间干完,只需要把自己的成果传给51层,同时把49层的成果清掉,仓库永远只留上一个,特别省空间。

可是一旦要跨层访问,规则就变了:前面所有车间的成果,都不能清,必须一直存在仓库里!

因为你永远不知道后面哪个车间会突然找前面某一层要东西。比如第100层车间可能要找第10层的成果,那第10层的东西从被加工出来开始,就要在仓库里存到第100层干完,中间一点都不能丢。

举个直观的例子:

传统残差:仓库最多只存1份加工成果,全程省空间;

跨层访问:100层的工厂,仓库要同时存99份成果(第1到99层),层数越多,要存的东西就越多,仓库直接被堆满。

对应到模型里就是:大规模模型少则几百层、多则上千层,跨层访问需要把每一层的隐藏状态都永久保存,内存占用会跟着层数呈倍数暴涨,普通的服务器内存根本扛不住,这就是内存开销过高

2.2.3 为什么还会让通信开销高?

这里需要先补一个常识:超大模型的训练,根本不是一台服务器能搞定的(就像超大型的1000层加工厂,不可能挤在一个厂区里,会拆成10个厂区,每个厂区管100层)。

实际训练中,会把模型的几十层/上百层,分给多台服务器(机器) 分别计算,比如机器A管1-20层、机器B管21-40层、机器C管41-60层……这就是大规模模型的分布式训练

而 跨层信息访问 的要求是:任意一层能找前面所有层要东西,这就意味着:不同机器之间要疯狂传数据!

比如机器C的41层车间,要找机器A的5层车间要成果,那就要把机器A里第5层的成果,通过网络传到机器C里;如果机器C的每一层,都要找机器A、B里的任意层要东西,那机器之间的网络就会被这些数据传输出占满,就像厂区之间的运货车堵满了所有道路。

这种机器之间为了跨层访问,需要疯狂传数据、占满网络资源的情况,就是通信开销过高。传数据的速度,甚至赶不上机器计算的速度,模型训练会被拖得特别慢。

2.2.4 为什么这会导致难以支持大规模模型训练?

大模型的核心特点就是:层数多、参数量大,必须用多台机器分布式训练。

而跨层信息访问的两个问题,刚好戳中了大规模模型的痛点:

层数越多,需要保存的层成果就越多,单台机器的内存直接不够用,就算加内存,成本也会高到离谱;

机器数量越多,跨机器传数据的量就越大,网络通信会彻底堵死,训练速度变得极其慢(比如原本训练一次要1天,现在要10天),甚至直接跑不起来。

简单说就是:这个解决办法本身是对的,但实现的成本太高,撑不起大规模模型的训练需求。


补充一个小结论(衔接论文方案)

论文里的 Full AttnRes 其实就是这种纯跨层信息访问。确实能解决层稀释,但也确实有内存和通信开销高的问题;而 Block AttnRes 就是为了解决这个开销问题想的妙招:把车间分成几个大车间组,只允许跨组访问,不允许跨单个车间访问,既保留了跨层选内容的核心优势,又大幅减少了要存的内容和要传的数据。

2.3 问题3:深度模型中梯度分布不均,早期层梯度占比过大,影响训练稳定性

还是用100层食品加工厂的类比,先把梯度这个核心词讲透,再一步步说清梯度分布不均、早期层梯度占比大的问题,以及为什么会让训练不稳定。

2.3.1 什么是梯度?(训练中的调整指令)

模型的训练过程,本质就是AI自己练手艺、调流程的过程。

AI第一次做出来的成品肯定不好,会根据成品的糟糕程度,给工厂里的每一个车间发调整指令:告诉车间“你的加工步骤哪里错了、要怎么改、改多少,才能让最终成品更好”。

这个从最后一个车间往回传的、指导各车间改流程的调整指令,就是梯度

核心特点:梯度是从后往前传的(先告诉最后一个车间怎么改,再一层层往回告诉前面的),只有各车间都按指令精准调整,模型的手艺才能越练越好。

2.3.2 传统残差下,为什么会梯度分布不均,早期层梯度占比过大?

传统残差连接对梯度的传递规则,和对食材的传递规则一样,是死板的全量回传,没有任何过滤和调整。

就像加工食材时,前面的内容会全量传给后面;梯度回传时,后面的调整指令也会全量传给前面,甚至因为层层叠加,越往前的车间,收到的指令强度越大

举个直观的100层加工厂例子:

最后一层做出成品,发现甜度太高了,收到的调整指令是:把甜度调小10%(这是初始梯度);

第100层按规则,把这个“调小10%”的指令原封不动全传给99层,99层又全传给98层……一直往回传;

因为传统残差的“全量累加”特性,这个指令在回传过程中不会被削弱,反而会一点点叠加,传到第50层时,指令变成了“调小20%”,传到第10层时变成“调小50%”,传到第1层时,指令直接变成了“调小80%”;

最终导致的结果是:越靠前的早期车间,收到的调整指令越强、改的幅度要求越大;越靠后的深层车间,收到的指令越弱、改的幅度要求越小

这就是 梯度分布不均,早期层梯度占比过大

2.3.3 核心问题:为什么这会影响训练稳定性?

训练的稳定性,指的是模型的手艺能平稳提升,不会忽好忽坏、越练越差。而早期层梯度占比过大,会让整个工厂的调整节奏彻底乱套,具体体现在3个方面:

早期车间乱改基础,后面所有车间全跟着错

早期层(1-10层)是工厂的基础加工环节(比如洗食材、切食材、调基础味道),这些步骤是后面所有车间加工的前提。

现在早期车间收到的是“调小80%”的强指令,相当于让它把基础加工流程大改特改。比如原本切丁的改成切丝,原本淡盐的改成无盐。基础步骤乱了,后面90层车间就算按自己的弱指令精准调整,最终成品也会一塌糊涂。

深层车间白忙活,模型永远练不出核心手艺

深层层(90-100层)是工厂的核心加工环节(比如调最终口味、做造型、定成品风格),模型的核心能力(比如理解语义、做推理、生成通顺回答)都是靠深层车间练出来的。

但深层车间只收到“调小10%”的弱指令,就算按要求改了,幅度太小,几乎没变化——相当于厨师被要求「微调半勺盐」,根本尝不出差别,深层的核心手艺永远练不出来,模型训练就是白忙活。

整个工厂的调整忽左忽右,永远定不下来

因为早期层的强指令是叠加出来的、不稳定的,这次回传是“调小80%”,下次可能因为成品的一点小变化,变成“调大70%”。

早期车间一会往这大改,一会往那大改,导致整个加工流程的波动特别大——这次做的甜,下次做的咸,再下次做的淡,模型永远没法练出稳定的手艺,训练过程忽好忽坏,这就是训练不稳定的核心表现。


最后补个关联:这个问题也是「固定权重聚合」的连锁反应

其实这个梯度问题,和第一个问题(层贡献稀释、隐藏状态增长)是一根藤上的两个瓜,根源都是传统残差的固定权重全量累加:

正是因为食材传递时全量累加、没有筛选,才导致梯度回传时也全量叠加、没有调整,最终既出现了前向的层贡献稀释,又出现了反向的梯度分布不均。

而论文里的 AttnRes,解决这个问题的核心逻辑也很简单:不仅在食材传递时做动态筛选(注意力加权),在梯度回传时也会根据注意力权重做过滤和调整——把早期层的强指令调温和,把深层层的弱指令调明显,让所有车间收到的调整指令幅度差不多,梯度分布均匀了,训练自然就稳定了。

2.4 问题4:不同类型层(如注意力层、MLP 层)对前序层信息的需求差异无法被满足

还是用100层食品加工厂的老类比,先把模型里的“注意力层”、“MLP层”翻译成工厂里的不同车间类型,再讲清它们的信息需求差异,以及传统残差为啥满足不了。

2.4.1 先铺垫两个关键:模型里的两类核心车间

大模型里的每一层,其实不是随便的加工车间,主要分两种功能完全不同的核心车间(就像食品加工厂里,有的车间管选食材,有的车间管深加工,干的活不一样,需要的原材料也完全不同):

注意力层(选食材车间):核心工作是从前面所有加工成果里,挑出和当前任务相关的关键内容。比如你问AI“红烧肉怎么做”,注意力层就专门从前面的信息里,挑出“五花肉、冰糖、生抽”这些关键食材,忽略无关内容。

这类车间的需求:需要看前面很多层的信息,挑出有用的,不用全要,但要能选

MLP层(深加工车间):核心工作是把注意力层挑好的关键内容,做精细化加工、提炼规律。比如注意力层挑出了“五花肉、冰糖”,MLP层就负责研究“五花肉炒到什么程度、冰糖放多少能炒出糖色”。

这类车间的需求:只需要注意力层刚挑好的最新关键内容,不用看更早的层,要的是精准、新鲜,不是多

简单说:注意力层要广撒网选重点,MLP层要盯着最新的精加工,二者对前序层信息的需求类型、需求范围完全不一样,这就是不同类型层的信息需求差异。

2.4.2 核心问题:传统残差连接为啥满足不了这种差异?

传统残差的核心毛病还是那个死板的一刀切规则:不管你是什么类型的车间,不管你需要什么信息,一律只给你上一个车间的全部成果,没有任何选择余地。

这个规矩直接让两类车间的需求都被憋住了,全都干不了自己的本职工作。

注意力层(选食材车间):想选重点,却没的选

注意力层的本职是【从前面所有层挑关键信息】,但传统残差只给它上一个车间的成果——相当于选食材的车间,本来想从仓库里挑五花肉、冰糖,结果工厂只给它递过来上一个车间切好的土豆丝,它根本没机会挑到真正需要的关键内容,最后只能硬着头皮用土豆丝干活,选食材的作用完全失效。

MLP层(深加工车间):想要精准料,却收到一堆杂料

MLP层的本职是【盯着最新的关键内容精加工】,但上一个车间如果是注意力层(被规矩憋住没选好料),传过来的就是一堆没筛选的杂料。相当于深加工车间想加工五花肉,结果收到的是土豆丝、青菜、萝卜混在一起的杂料,根本没法做精细化加工,最后只能随便炒炒,做不出合格的成品。

2.4.3 更糟的连锁反应:两类车间互相拖后腿,模型越练越差

因为规矩的一刀切,注意力层和MLP层不仅自己干不好活,还会互相拖后腿,形成恶性循环:

注意力层没的选,只能把杂料传给下一个MLP层;

MLP层拿着杂料,只能做出劣质加工成果,再传给下一个注意力层;

下一个注意力层拿到劣质杂料,更挑不出关键内容,再传给下一个MLP层…

最终从最后一层出来的成品(AI的回答),要么答非所问,要么逻辑混乱,模型根本练不出核心能力。


补个通俗的小例子:比如你问AI:李白的《静夜思》是在哪写的?

正常情况下:注意力层该从前面的信息里,挑出“李白、静夜思、创作地点”这些关键信息,MLP层再基于这些信息,提炼出“扬州”这个答案;

但在传统残差下,注意力层只能拿到上一层的模糊信息,根本挑不出“静夜思创作地点”这个关键内容,只能把“李白、唐诗、月亮、思乡”这些杂料传给MLP层;MLP层拿着这些杂料,要么答“不知道”,要么乱答成“长安”。

2.4.4 最后衔接论文方案:AttnRes是怎么解决的?

论文里的 AttnRes(Attention Residuals,注意力残差),本质就是给不同类型的车间松绑和定制化供料;打破「只能拿上一层成果」的死规矩,让每一层(不管是注意力层还是MLP层),都能根据自己的需求,主动从前面任意一层里挑信息。

注意力层想选前面10层、20层的关键内容,就自己挑;

MLP层只想用刚从注意力层挑好的最新内容,就只挑这一层的;

而论文里的 Block AttnRes,就是在自主选料的基础上,给工厂加了个合理的仓库管理规则,既让车间能选料,又不会让仓库堆太多料、厂区之间传太多料(解决之前的内存/通信开销问题)。

3 研究过程是怎样的?

3.1 核心方案设计

前面讲的 4 个核心问题,根源都是传统残差「固定权重、只能用上一层成果、一刀切」的死规矩,而 AttnRes(注意力残差)就是专门打破这个规矩的解决方案。

它分两个递进的版本:先做了能完美解决所有问题的 Full AttnRes(全注意力残差),再针对它的落地成本问题,优化出了可大规模商用的 Block AttnRes(块注意力残差)。

下面我们全程用之前的「100 层食品加工厂」类比,同时对应论文里的数学公式和图表,一步步讲透。

3.1.1 Full AttnRes:全注意力残差

通俗解释:

我们给加工厂的每一个车间,都配了 1 个专属的「选料采购员」。每个采购员都有自己专属的选料标准,可以把前面所有车间的加工成果全部看一遍;采购员会给每个前序车间的成果打分,再把分数换算成占比(所有占比加起来刚好 100%);最后按占比,把各个车间的成果混合起来,作为当前车间的专属加工原料。

简单说:从传统残差的“必须全量收下上一层的所有东西”,变成了“我自己挑前面所有层的好东西,按重要性混合着用”,完美解决了之前的 4 个核心问题。

对应论文里的数学公式拆解:

论文在【3 Attention Residuals: A Unified View of Time and Depth】部分中给出了 Full AttnRes 的完整数学定义,我们将每个符号和加工厂的类比对应上,公式就会一目了然。

核心公式1:选料打分规则(对应公式 (2) 和 (3))

核心公式2:混合原料的最终规则(对应论文公式 (4))

论文中的 Figure 1 (b) Full Attention Residuals,就是这个方案的直观流程图。

图里每个层(Attention/MLP)都标注了一个$$w$$,就是我们说的「采购员的选料标准」。箭头从所有前面的层,都连到了当前层,代表当前层可以访问所有前序车间的成果。和左边 (a) 的传统残差对比:传统残差只有上一层到下一层的箭头,而 Full AttnRes 是全连接的箭头,代表可以自由选择所有前序层的内容。

Full AttnRes 的优点是完美解决了之前的 4 个核心问题,层贡献不稀释、梯度分布均匀、不同类型层可以按需选料。但缺点也很明显:内存和通信开销太高 ——100 层的工厂,每个车间都要存前面 99 个车间的成果,仓库(内存)直接被堆满,跨厂区(多机器训练)传数据的量也暴涨,没法用在大规模模型上。

这就引申出接下来的Block AttnRes(块注意力残差)。

3.1.2 Block AttnRes:块注意力残差

这是论文的核心落地优化方案,本质是在几乎不损失效果的前提下,把 Full AttnRes 的开销降到可商用的水平

我们把100个车间,分成8个大车间组(也就是“块(Block)”),比如组1是1-12层、组2是13-24层… 以此类推。

每个大组干完所有车间的活,会出一个「本组的总成果包」,只需要存这个成果包,不用存组里每个车间的单独成果。

每个车间选料的时候,不用看前面 99 个单独车间的成果,只需要看「前面所有大组的成果包」,再加上自己组里前面车间的半成品就行。

简单说:从「看前面所有单个车间」,变成了「看前面所有大组的总成果」,要存的内容从 100 份变成了 8 份,内存直接省了 90%,还保留了 99% 的效果。

对应论文里的数学公式拆解:

核心公式 1:块成果包的定义

核心公式 2:块注意力的选料范围(对应论文第 3 章的V矩阵定义)

剩下的打分、算占比、混合原料的规则,和 Full AttnRes 完全一致,只是选料范围从「所有单个层」变成了「块成果包 + 组内半成品」。

核心架构示意图:论文 Figure 1 (c) Block Attention Residuals

图里把所有层分成了一个个 Block(大车间组),箭头不再是从所有单个层连过来,而是从前面的 Block 连到当前层。

论文中第5页还给出了落地伪代码:

block_attn_res函数:就是「采购员选料混合」的核心逻辑。输入blocks:前面所有大组的成果包;partial_block:自己组里的半成品;proj:采购员的选料标准;核心步骤:把成果包和半成品堆起来 --> 统一称重(RMSNorm) --> 用选料标准打分 --> 换算成占比 --> 按占比混合成原料,返回给车间。

forward函数:就是每个车间的完整加工流程。先调用block_attn_res拿到混合原料 --> 做注意力 / MLP 加工 --> 如果到了大组的边界,就把本组的总成果包存起来,给后面的组用。

3.2 实验设计与执行

所有的理论都是为了解决实际问题。所以这部分内容我们先明确论文中所有实验的核心总目标:验证 AttnRes 方案。

是不是真的能解决最开始的 4 个核心问题?在小模型到大模型里都好用吗?哪个设计是核心?能不能真的落地商用?

所有实验都严格遵循「控制变量」原则:除了残差连接这个零件,模型的其他所有配置、训练数据、总算力、学习率等完全一致,就像两个完全一样的食品加工厂,只换了「原料传递规则」,保证对比出来的结果绝对公平。

3.2.1 实验基础配置(对应论文第5章中的 Architecture Details)

论文中用来做测试的「标准加工厂」,是 Kimi 团队已经成熟商用的 Kimi Linear 大模型架构

类比成:一个成熟的现代化食品加工厂,核心是「混合专家(MoE)」模式 —— 工厂有 256 个备用车间,每次只开 9 个最适合当前任务的车间干活,省算力还高效;

最大测试模型:48B 总参数(工厂总设备数),3B 激活参数(每次实际开动的车间数),共 54 个加工层(27 个 Transformer 块);

唯一改动:只把传统残差连接,换成我们设计的 Full/Block AttnRes,其他所有零件完全不动,保证对比公平。

3.2.2 实验 1:缩放定律验证 —— 是不是所有大小的模型都能用?

实验目的:验证 AttnRes 的效果是不是通用的:是只在小玩具模型里有用,还是从几十M的小作坊到几百M的中型工厂,全场景都比传统残差好用?

论文中做了5种不同规模的模型(从194M到528M激活参数,对应从小作坊到中型工厂),每个模型都做3个版本:

  1. 基线版(Baseline):传统残差连接
  2. Full AttnRes 版:全注意力残差
  3. Block AttnRes 版:8个块的块注意力残差

所有版本用完全一样的训练数据、算力、学习率,训练完成后看验证损失(validation loss) —— 类比成「成品的不合格率,数值越低,说明手艺越好、成品越合格」。

实验结果对应的原文图表和结果解读:

Table 2:详细实验数据

Table 2 中每一行对应一个模型大小,核心结论一眼可见:

所有模型尺寸下,Full AttnRes 和 Block AttnRes 的不合格率(Val. Loss)都比传统残差基线更低;

最大的528M模型:基线损失 1.719,Block AttnRes 损失 1.693,Full AttnRes 损失 1.692;Block和Full的效果几乎一模一样,没有差距。

Figure 4:缩放定律拟合曲线

Figure 4 中横坐标是训练用的总算力,纵坐标是成品不合格率(损失),线越靠下越好。

可以看到,两条 AttnRes 的线,全程都在传统基线的下方,说明不管用多少算力训练,AttnRes 的效果都更好;

核心亮点:Block AttnRes 用1倍的算力,就能达到传统基线 1.25 倍算力才能做到的效果,相当于直接省了25%的训练成本。

实验1的结论:AttnRes 在所有模型尺寸下都稳定优于传统残差,Block AttnRes 几乎完全保留了 Full AttnRes 的效果,还把内存开销降了一个数量级,完美解决了「跨层访问内存开销高」的问题。

3.2.3 实验 2:大规模落地验证——商用大模型里真的好用吗?

实验目的:这是论文最核心的落地实验:验证在真正要上线的48B超大商用模型里,AttnRes 能不能解决最开始的4个核心问题?训练稳不稳定?实际答题能力强不强?

用 48B 的 Kimi Linear 大模型,给它装上 Block AttnRes,用 1.4T tokens 训练(类比成用1.4万亿份食材,让工厂练了完整一轮手艺),和同配置的传统残差模型对比两个维度:

训练过程的动态变化:看看隐藏状态还会不会疯涨?梯度分布均不均匀?训练稳不稳定?

下游任务考试成绩:让两个模型去考各种类型的「能力测试」,看谁的分数更高。

实验结果对应的原文图表和结果解读:

Figure 5:训练动态对比图(3个子图对应3个核心问题)

Figure 5(a) 验证损失(成品不合格率):横坐标是训练步数,纵坐标是不合格率,AttnRes 的线全程都在基线下方,越往后训练,差距越大;说明 AttnRes 的工厂,手艺进步更快,最终成品的合格率更高。

Figure 5(b) 输出量级(隐藏状态大小):横坐标是工厂的车间序号(模型深度),纵坐标是食材总重量(隐藏状态量级)。传统基线的线:一路向上疯涨,完美对应了「隐藏状态失控增长」的问题;AttnRes的线:呈现有界的波浪形,一直在固定范围内波动,完全不会疯涨,完美解决了隐藏状态失控增长、层贡献稀释的问题

Figure 5(c) 梯度量级(调整指令强度):横坐标是车间序号,纵坐标是调整指令的强度(梯度大小):传统基线的线:前半段特别高,后半段特别低,对应「早期层梯度占比过大,深层梯度几乎为0」的问题;AttnRes 的线:几乎是平的,所有车间收到的调整指令强度都差不多,完美解决了梯度分布不均、训练不稳定的问题

Table 3:下游任务考试成绩对比

论文中让两个模型(Baseline 和 AttnRes 分别对应的两个模型版本)考了三大类能力测试,可以看出:AttnRes 在所有测试里,分数都等于或高于传统基线,尤其是复杂任务提升巨大。

实验2的结论:Block AttnRes 在商用超大模型里,完美解决了最开始的4个核心问题,训练更稳定,模型的知识、推理、代码、中文能力全方面提升,而且推理延迟只增加了不到2%,完全可以直接落地商用。

3.2.4 实验3:消融实验——哪个零件是核心?少了哪个就不行?

实验目的:找到 AttnRes 里真正起作用的核心设计,避免加了没用的花架子;可以类比成我们给新流水线装了好几个零件,现在把零件一个个拆掉,看看拆了哪个之后,流水线就不好用了。

用16层的固定模型,只改动 AttnRes 里的一个设计,其他所有配置完全不变,看成品不合格率(损失)的变化:损失越高,说明这个改动让效果变差了,这个零件就是核心。

实验结果对应的原文图表和结果解读:

Table 4:核心组件消融结果 和 Figure 6:块大小对效果的影响

从 Table 4 中可以看出:

基线损失 1.766,Full AttnRes 最优损失 1.737;

如果把「动态选料的权重」改成固定不变的 DenseFormer,损失直接回到1.767,和基线一样差。说明「输入依赖的动态权重」是 AttnRes 的核心,没有它就没有效果提升

把 softmax 换成 sigmoid,损失涨到 1.741;去掉给食材统一称重的 RMSNorm,损失涨到 1.743。说明这两个设计也是必须的,少了效果就打折扣。

从 Figure 6 中可以看出:

横坐标是每个块的大小(块越大,越接近传统残差;块越小,越接近 Full AttnRes),纵坐标是损失。块大小为1的时候,就是 Full AttnRes,损失最低;块大小为2、4、8的时候,损失几乎和 Full AttnRes 差不多,没有明显上涨;块大小超过16之后,损失快速上涨,越来越接近传统基线。

实验3的结论:AttnRes 的核心是「输入依赖的动态权重+softmax归一化+RMSNorm,少了任何一个效果都会变差;块大小选8左右,是效果和内存开销的最优平衡,这也是论文里选8个块的原因。

3.2.5 实验4:架构最优性分析——AttnRes更喜欢什么样的模型?

实验目的:之前传统残差有个问题:模型做太深了,效果反而会变差(因为层贡献稀释)。那AttnRes解决了这个问题之后,会不会改变模型的最优设计?它更喜欢「宽而浅」的模型,还是「深而窄」的模型?

在论文原文中关于这部分的实验,相当于是固定工厂的总预算、总工人数量(总算力、总参数完全不变),试了25种不同的「深度-宽度」组合:有的是宽而浅(车间少,每个车间大),有的是深而窄(车间多,每个车间小),看哪种组合的成品不合格率最低。

实验结果对应的原文图表和结果解读:

Figure 7:架构热力图,颜色越深,损失越低,效果越好

3.2.6 实验5:注意力权重可视化——它真的学会按需选料了吗?

实验目的:我们之前说,AttnRes 能满足「注意力层要广撒网、MLP层要精加工」的不同需求,那它真的学会了吗?直接把每个车间的「选料占比」画出来,一眼就能看清。

在原文中,将训练好的模型里,每个层给前序层打的权重(也就是采购员给每个前序车间的原料占比)画成热力图:颜色越深,代表给这个车间的原料占比越高,越关注这个车间的成果

实验结果对应的原文图表和结果解读:

Figure 8:注意力权重热力图

从 Figure 8 中可以看到3个核心点:

  1. 保留了本地连贯性:每层最关注的还是上一个相邻层(对角线颜色最深),不会乱选料,保证了加工的连贯性;
  2. 不会忘本:最开始的初始原材料(词嵌入),全程都有不低的权重,说明模型一直记得最开始的问题,不会答着答着就跑题;
  3. 完美匹配不同层的需求:左边的注意力层(选食材车间)的热力图:颜色分布更广,会关注前面很多层的内容,符合「广撒网选重点」的需求;右边的MLP层(深加工车间)的热力图:颜色几乎全集中在对角线上,只关注最近的内容,符合「盯着最新内容精加工」的需求。

实验5的结论:AttnRes真的学会了根据层的类型,按需选择前序层的信息,完美解决了「不同类型层的信息需求差异无法被满足」的第4个核心问题;而且 Block AttnRes 和 Full AttnRes 的选料逻辑几乎完全一样,说明块结构没有破坏核心的选料能力。

3.3 所有实验的最终总结

五组实验层层递进,完整验证了:AttnRes完美解决了最开始提出的4个核心问题,Block AttnRes在几乎不增加训练、推理开销的前提下,实现了模型效果的全方面提升,完全可以作为传统残差连接的替代方案,用在大规模商用大模型里。

写在最后

按逻辑顺下来好像应该写点什么这篇论文对行业内模型训练的影响,但我也不是某家公司的研究员,甚至所有理论和实操都是自学的。读完论文后也只能直观地感觉到以后模型训练的效率会提高,成本会下降。

至于更深远的影响,网上有很多大佬都已提出~

0
0