来源:AI寒武纪
kimi k3为什么可以取得突破,我们可以从4个月前(3月21)月之暗面ceo 杨植麟在英伟达GTC大会上的演讲回忆一下。
演讲的主题是Kimi最新模型K2.5背后的扩展方法论。整场演讲几乎全程在讲技术细节,涉及优化器、长上下文架构、智能体训练范式,以及一个刚刚公布的新架构预告。
但从这次演讲中可以窥见杨植麟的LLM扩展的三大核心技术野心:1.主要Adam优化器诞生于2014年,团队现在做出了一个可以直接替换Adam的开源版本Muon Clip,用它训练Transformer大语言模型,效果会明显更好;2.注意力机制诞生于8年多之前,团队现在推出了它的线性版本Kimi Linear,不必在每一层都用全注意力,线性注意力在长短上下文任务上都能表现得更好;3.残差连接,团队推出了开源版本的注意力残差架构。
现在大家看到了,仅仅过去了4个月,K3就是最好的答案
扩展的三个维度
整场演讲的主线,是团队在三个维度上做扩展。
第一个维度是最经典的scaling law:横轴是训练token数量,纵轴是loss,token越多,loss越低。团队关注的重点不是单纯堆token数量,而是提升token的利用效率,让同样数量的token能换来更低的loss,这靠的是更好的架构和优化器。
第二个维度是扩展上下文长度。上下文越长,模型在给定位置预测下一个token的准确率就越高,这意味着模型能完成更复杂的任务。
第三个维度是智能体数量。团队提出了一种新的训练范式,叫智能体蜂群,让模型不再单纯依赖单个智能体处理任务,转而编排一群智能体并行完成子任务,从而提升处理复杂任务的能力。
这三个维度可以直接映射到智能体的语境里:token效率对应更强的先验知识,让智能体在做强化学习搜索答案时更高效;长上下文对应能持续运行更久的智能体,可能连续跑上几天、几周甚至几个月去完成复杂任务;智能体蜂群则是叠加在前两者之上的又一个维度。最终团队想要的,是一群拥有超长上下文、又具备强先验的智能体,在整个强化学习系统里协同搜索答案。
维度一:怎么让每个token更值钱
杨植麟先展示了一张机器学习史上很经典的图,来自Kaplan等人的scaling law论文:只要按比例扩大训练token数量、模型参数量和算力,loss就会持续下降,这是过去几年整个行业的重要基础。
但团队真正在意的是提升token效率。杨植麟特别强调,token效率不只是效率问题,它直接关系到智能的上限。他举了个例子:假设手上有50万亿条高质量token,应用一个新优化器之后效率提升两倍,效果上就相当于凭空多出了100万亿条token。眼下整个行业正在逼近所谓的数据墙,高质量数据总量有限,如果把数据量看作一个常数,提升token效率就意味着能榨取出更强的智能,而不只是省了算力这么简单。
团队为此重点投入的是Muon优化器,一种二阶优化器。它的核心做法是,每一次梯度更新都会被转换,使更新的每一项彼此正交,这和传统的Adam优化器差异很大。如果实现得当,可以带来两倍的token效率提升。团队是首个证明Muon优化器可以扩展用于大语言模型训练的团队,论文里提出了两项关键技术:一是权重衰减,这对扩展到更大模型至关重要;二是保证更新幅度和Adam保持一致的水平,团队引入了一个可调系数应用在每次更新上,让最终的更新幅度和Adam保持可比。为了在英伟达GPU集群上做到内存高效,团队还专门开发了一套分布式Muon优化器实现,把优化器状态切分到不同的数据并行组里。在相同参数量和相同训练token数的对照实验里,把AdamW换成Muon优化器之后,各项指标全面明显提升。
不过把Muon扩展到一万亿参数规模时,团队遇到了新的麻烦:训练不稳定。具体表现是,注意力机制里的最大logit值会迅速爆炸,超过1000,而正常情况下这个值一般在50甚至不到100。与此同时训练会发散,loss先往下走一段,随后突然爆炸,没法按预期收敛。
团队的解法叫QK clip。做法是对网络里的每一个注意力头,在前向传播过程中计算出当前的最大logit值,再据此算出一个缩放系数,分别应用到key和query的投影上,把query和key的数值限制在一个给定范围内,从而避免爆炸。效果很直接:应用clip前后两条训练loss曲线几乎完全重合,说明这项技术不会拖累loss下降的速度;而如果单独看最大logit这个中间指标,可以看到它一开始照常上升,在数值100左右被卡住维持一段时间,之后随着训练步数增加自然回落。网络自己找到了办法把最大logit值控制住,同时完全不影响收敛。团队把这项技术用到了K2模型的训练上,成功把参数量扩展到一万亿,这也是机器学习历史上第一次把Muon优化器用在这种规模的训练上。
维度二:把上下文拉得更长
第二个维度的引子,是一张相对少被提及但很关键的图,对比的是Transformer和LSTM。左边的图显示,在相同参数量和相同训练token数量下,Transformer的训练loss更低,这也是Transformer会成为如今主流架构的原因。右边的图更有意思:横轴是上下文里的token位置,可以看到Transformer的训练loss会随着上下文变长持续下降,而LSTM的loss在某个token数量之后就不再下降,进入饱和状态。这说明Transformer在捕捉长上下文信息上明显更强。
这项能力在智能体时代格外重要,因为任务正在变得越来越复杂,需要的上下文也越来越长,比如理解一整个代码仓库,或者跑很长的智能体轨迹去完成任务,像是从零开始写一个操作系统内核,这些都是LSTM做不到的。团队的研究目标,是做出一种架构,既能高效扩展到更长的上下文,又能在更靠后的token位置上取得更低的单token loss。这正是团队推出Kimi Linear架构的出发点。
Kimi Linear里包含一种新的线性注意力变体,叫Kimi Delta Attention,是在原始delta rule的基础上,通过改进循环记忆机制得到的。同时团队把线性注意力层和全注意力层按1比3的比例混合使用,兼顾长上下文能力和实现效率。
原始线性注意力的问题在于,它的记忆是全局性的,只有一个统一的衰减因子在起作用,导致模型要么几乎遗忘掉所有历史信息,要么几乎把所有信息都保留下来,但没办法在长上下文里有选择地丢掉不重要的部分。Kimi Delta Attention引入了细粒度的衰减因子,把原本的标量衰减系数换成了一个对角矩阵,分别控制每个通道各自的衰减速度。一部分通道可以衰减得很慢,从而在很长的范围内保留信息,另一部分通道则可以很快遗忘掉旧信息,及时接收新信息,整体上提升了模型的表达能力。
为了适配现代GPU,团队采用了分块的方式来做并行计算。由于衰减项从标量变成了矩阵,不能像以前那样简单地提出来单独处理,团队因此重写了整套计算公式,引入了矩阵求逆运算以及累积衰减因子,在不牺牲效率的前提下实现并行计算。更重要的是,这次重写在数学上和原始公式完全等价,属于精确改写,效率提升的同时也不会带来任何性能损失。
从结果看,在偏短上下文的任务比如MMLU上,Kimi Linear的表现优于MLA和GDN;在偏长上下文的任务比如ruler上,Kimi Linear同样更好,而且比MLA更省算力。当上下文进一步扩展到一百万token甚至更长时,Kimi Linear相比基线效率提升更加明显,这也是第一个能在短上下文任务、长输入任务和长输出任务上全面超过全注意力架构的方案。
维度三:一群智能体一起干活
团队为智能体蜂群设计了一个编排者,也可以理解为主智能体,负责协调整体任务。它可以生成一批子智能体并给它们分配任务,也可以收集子智能体返回的结果,通过反复迭代这个过程,最终完成比单个智能体能力范围更大的复杂任务。
杨植麟拿人类社会做了类比:办一家公司需要不同的角色分工,需要一个类似CEO的角色把任务拆解、分配给不同岗位,最终整个组织朝着同一个目标推进。放到智能体蜂群里,可能同时存在负责AI研究的智能体、负责网页开发的智能体、负责物理研究的智能体,各自研究不同的课题,最后再生成一批事实核查员、网页开发者和文件下载员,把所有结果汇总成一份完整报告。
从另一个角度看这个范式:横轴是任务复杂度,用一组模型在该任务上的准确率来衡量,纵轴是执行时间。用上智能体蜂群之后,执行时间比单智能体大幅缩短,效率明显更高。这意味着可以把蜂群规模扩展到比如100个甚至1000个子智能体,在可接受的时间范围内完成复杂任务,从而真正产生经济价值。团队认为这种扩展可以体现在多个方向:扩展输入,比如并行下载、阅读成百上千份资料;扩展输出,比如并行写出一百页的文献综述;扩展执行的动作,比如同时对十个不同任务做数据分析;也可以扩展编排本身,也就是学会怎么拆分子任务、怎么把结果汇总起来。
技术层面,团队为智能体蜂群系统专门设计了新的奖励目标,相比传统的单智能体强化学习,这里一共考虑三项奖励。
第一项叫实例化奖励,用来鼓励模型生成子智能体,防止出现退化成单智能体串行执行的现象,尤其是在训练早期要鼓励并行执行,这项奖励的权重会随训练推进逐渐衰减。
第二项叫完成奖励。团队观察到,有些子任务被创建出来之后从来没有被完成过,模型几乎是靠不断生成一堆子智能体去刷第一项奖励,而这些任务本身可能太复杂,或者压根没有意义。因此加入完成奖励,鼓励每个子任务都保持相对较高的完成比例,确保生成出来的子任务是有意义的,同样采用训练前期权重高、后期权重低的衰减策略。
第三项是标准的结果奖励,衡量整个任务最终是否完成。
三项奖励加总起来,构成了整个强化学习系统的优化目标。与此同时,团队也搭建了配套的基础设施,用来支持并行执行、支持不同的奖励函数,并尽可能提升整个智能体蜂群系统的输入输出效率。
三条线合到一起,就有了K2.5
把token效率、长上下文、智能体蜂群这三条线放在一起,就诞生了Kimi K2.5
杨植麟提到,团队从K2.5身上观察到不少有意思的能力,比如视觉能力和编程能力融合之后,会涌现出很多新玩法,例如看一段视频,就能生成一个复现或者风格迁移自这段视频的网页,这些能力都源于预训练阶段成功且稳定的训练过程。
他还展示了K2.5基础模型的训练曲线,称这是自己见过最漂亮的曲线之一。模型经历了超过15万亿token的训练,在K2.5阶段又额外训练了15万亿token,整个训练过程非常稳定,尤其是引入新的Muon优化器之后,没有出现任何loss尖峰。这种平滑稳定的训练过程,产出了一个很强的基座模型,团队可以在此基础上继续微调,获得视频里展示的那些新能力。整个训练在英伟达H800 GPU集群上完成,集群里每个节点配备2TB内存,GPU之间通过NVLink连接。
K2.5的另一项关键创新,是它成为首个原生实现视觉和文本联合训练的开源模型。以往开源模型通常的做法是先在纯文本上训练,比如先用20万亿token训练出一个文本模型,再额外用2万亿token左右的后训练,把视觉能力叠加上去。K2.5的做法不一样,团队从训练的第一天就把视觉和文本训练过程融合在一起,这被称为早期融合,从进度0%开始就把视觉token和文本token合并训练。初步实验显示,这种方式效果优于先文本后视觉的后期融合做法,团队观察到的一些新能力也正是来自这种训练方式。举个例子,如果想让模型看图生成代码,必须把视觉和文本放进同一个分支里训练才能实现,如果两个分支各自独立训练是做不到的,必须把两种模态对齐进同一个共享表示空间里才行。
另一个有意思的发现是,这两种模态其实可以互相增强。行业里此前普遍的认知是,给文本模型加上视觉能力往往会拖累文本本身的表现,但团队发现,只要训练方式得当,两种模态反而能够相互促进,这是这次训练过程里的一个重要发现。
具体来看,视觉训练能提升文本表现。团队做了一组对比实验,只用纯视觉任务做强化学习,不涉及任何文本任务,比如教模型数数、回答视觉问答,不包含任何数学或编程题目,结果发现即便是偏重推理的文本任务,表现也跟着提升了。反过来,文本训练也能提升视觉表现。如果文本基础足够扎实,训练过程中其实不需要额外准备视觉方向的SFT数据。团队采用的方法叫zero vision SFT,也就是视觉方向几乎不用SFT数据,唯一使用的SFT数据是文本SFT数据,再对文本和视觉做联合强化学习。结果显示,在完全没有专门视觉数据的情况下,视觉类任务的表现依然能做到接近业内最好水平。这说明只要在预训练阶段把两种模态对齐进同一个共享空间,足够强的文本基座本身就能带动视觉能力一起提升。视频里展示的强视觉设计和前端编程能力,同样来自这套视觉文本联合训练方式。
下一代架构可能长什么样
演讲快结束时,杨植麟顺带介绍了团队前一天刚刚放出技术报告的新架构,叫注意力残差,算是对下一代架构方向的一次预告。
出发点很简单:团队此前在时间维度上用到的一些技术思路,能不能搬到深度这个维度上来用一遍。
这要从残差连接说起。杨植麟回忆起十年前在ICML 2016的一场分享上第一次听到这个想法,称它是个很精彩的点子。在ResNet出现之前,没人能训练很深的网络,一旦层数增加就会遇到梯度爆炸、梯度消失之类的训练稳定性问题;而ResNet提出之后,理论上可以堆叠任意多的层数,不用再担心训练稳定性。
他引用了Ilya两年前一次演讲里的说法:残差连接本质上是LSTM的一个变体,只是把方向旋转了90度。这句话可以这样理解:LSTM是循环网络的一种,处理过程是循环的,每一步都会把上一步的隐藏状态拿过来,通过某种门控机制生成当前状态;而残差连接在深度这个维度上做的其实是同一件事,把上一层的输出拿过来,通过某个函数处理后生成当前层的输出,只是具体的计算形式不一样。比如残差连接用的是固定的加法,把上一层的隐藏状态和当前层的输出直接相加,本质上还是把循环网络那套逻辑用在了深度维度上。
由此团队想到一个新方向:既然深度维度上可以套用LSTM的逻辑,那能不能换成用注意力机制来做同一件事,替代原本类似LSTM的门控函数,毕竟注意力机制已经在Transformer时代证明了自己足够强大。具体做法是,每一层的输出不再只依赖上一层的隐藏状态,而是把之前所有层的隐藏状态都纳入考虑,通过注意力操作对这些历史隐藏状态做聚合,计算出当前层的状态。团队把这个思路称为把注意力旋转90度,可以看作是残差连接在这套类比之下的一种自然延伸。
为了控制通信和内存开销,团队还设计了一个效率更高的变体,叫分块注意力残差。做法是把神经网络的所有层划分成若干个块,比如每块16层,或者每块4层,只在每个块的输出上应用注意力残差,块内部依然沿用标准的残差连接。这样能大幅降低开销,同时训练精度几乎不受影响。
效果上,这套新架构在scaling law层面把token效率提升了24%,原本50万亿条高质量token能达到的效果,现在相当于凭空多出了60多万亿条token的效果。验证集loss也持续低于原来的曲线,说明这套优化过程更加稳定。在编程、数学和推理密集型任务上提升最明显,团队在GPQA、数学和HumanEval等基准测试上都观察到了相应的提升。
杨植麟认为,这个时代做研究的思路和过去不太一样。十年前的研究更偏重发表新想法,但很难做到足够严谨的实验验证,难以得出扎实的结论。而现在有了所谓的扩展阶梯,团队有足够的资源在不同规模上训练和验证模型,配合一整套基准测试来衡量进展,因此更容易得出有把握的结论。这也是为什么很多看起来已经很成熟的老技术,最近又重新出现新进展的原因之一。
最后他总结,团队会继续在token效率、长上下文、智能体数量这三个维度上扩展模型,未来还会看到不同的架构和优化器同时在这三个维度上做优化,也会不断发现新的扩展维度,智能体蜂群不会是终点。
k3将于7月27之前开源,到时候各个推理商提供k3后你再看A厂的叙事还能不能撑得住。
阅读最新前沿科技趋势报告,请访问21世纪关键技术研究院的“未来知识库”
未来知识库是 “21世纪关键技术研究院”建 立的在线知识库平台,收藏的资料范围包括人工智能、脑科学、互联网、超级智能,数智大脑、能源、军事、经济、人类风险等等领域的前沿进展与未来趋势。目前拥有超过8000篇重要资料。每周更新不少于100篇世界范围最新研究资料。 欢迎扫描二维码或访问https://wx.zsxq.com/group/454854145828进入。
截止到2月28日 ”未来知识库”精选的百部前沿科技趋势报告
(加入未来知识库,全部资料免费阅读和下载)
热门跟贴