
新智元报道
新智元报道
【新智元导读】Transformer杀手来了?KAIST、谷歌DeepMind等机构刚刚发布的MoR架构,推理速度翻倍、内存减半,直接重塑了LLM的性能边界,全面碾压了传统的Transformer。网友们直呼炸裂:又一个改变游戏规则的炸弹来了。
就在刚刚,KAIST、Mila和谷歌DeepMind团队等放出重磅炸弹——
一个名为Mixture-of-Recursions的全新LLM模型架构。
这个崭新的架构,被业内认为有潜力成为Transformer杀手!
它的推理速度提升2倍,训练FLOP减少,KV缓存内存直接减半。
最终,在135M到1.7B的参数规模下,MoR直接划出了一个新的帕累托前沿:相同的训练FLOPs,但困惑度更低、小样本准确率更高,并且吞吐量提升超过2倍。
全面碾压传统的Transformer!

论文链接:https://arxiv.org/abs/2507.10524
其实,学界很早就发现,Transformer复杂度太高,算力需求惊人。
比如最近CMU大牛、Mamba架构作者Albert Gu就表示,Transformer模型能力的局限太大,所谓token就是胡扯。

而谷歌产品负责人Logan Kilpatrick公开指出了注意力机制的缺陷——不可能实现无限上下文,还强调必须要在核心架构层进行全面创新。
今天谷歌DeepMind的这项研究,和这些大牛的观点不谋而合了。
对此,网友们纷纷表示实在炸裂。
有人预测,潜在空间推理可能会带来下一个重大突破。

显然,对于代码、数学、逻辑这类分层分解问题的任务,MoR都是一个改变游戏规则的重磅炸弹。


甚至还有人评论道:看起来像是Hinton的胶囊网络重生了。


LLM发展到如今,接下来该怎样做?靠堆参数、加层数,让它更聪明吗?
这项研究告诉我们:真正的高手,从来都不是靠堆料,而是靠设计的艺术。
这次他们做出的MoR全新架构,直译出来是「递归混合体」,直接让LLM推理速度噌噌翻倍!
所以,MoR究竟做了什么?
简而言之,它做了以下两点。
1. 不对所有token一视同仁
LLM在处理文本时,会把句子拆成一个个token,不过,像「的」「是」「在」这种词,并不需要多高深的推理,只需要一次前向传播就够了。而复杂的token,则需多次经过同一层栈。
MoR的聪明之处就在于,因token而异。
MoR的秘密武器是小型路由器,会为每个token的隐藏状态打分,仅高分token的会继续循环,其余的则提前退出。

2. 循环复用:一个模块搞定全部
传统Transformer的思路就是不断「堆层」,堆得越高,处理能力越强。但这样的代价,就是内存和算力:模型会越来越慢,越来越贵。
而MoR则反其道而行之,专门设计了共享块,每个token最多循环4次,只要路由器说「完成」,就提前跳出循环。
总之,如果说Transformer是一个庞大的工厂流水线,那MoR就更像一支高效的特种部队。未来的AI,恐怕不会再比拼谁更重,而是谁更会分工调度、节省力气。
而谷歌DeepMind,已经敏锐地把握到了这一点,给我们演示了这一趋势的早期范本。

只靠Scaling law,把语言模型做大,确实能让它能力暴涨,但训练、部署所需的算力和成本也跟着暴涨。
现在常见的「瘦身」招数,要么是把参数共享(省显存),要么是按需计算(省算力)。
但目前仍缺乏一种能将两者有机融合的架构。
「递归混合」(Mixture-of-Recursions, MoR),充分发挥了递归Transformer的潜力(见图1),成功融合了两者。

图1:递归混合(Mixture-of-Recursions,MoR)概览
(左)每个递归步骤包含固定层堆栈和决定token是否继续递归的路由器(中间灰框区域)。
(中)完整模型结构,其中共享递归步骤根据路由决策对每个token最多应用𝑁𝑟次。
(右)展示token级递归深度的路由模式示例,颜色越深表示该token在递归块中的计算越活跃。底部数字以不同颜色标注各文本token的递归步数:1步、2步和3步 。
在统一的架构中,MoR同时实现了三种效率优化:
通过共享权重压缩参数量;通过动态路由减少冗余计算;通过智能缓存减少内存开销。

在预训练和推理过程中,MoR为每个token动态调整递归步骤,依赖两大组件:
路由机制和KV缓存策略。
受top-k门控机制的启发,研究人员提出了专家选择路由(Expert-choice routing)(见图2a)。
在这种模式下,递归深度可看作「专家」,每轮递归时这些专家会挑选它们认为最值得处理的top-k个token。

为了让递归更具一致性,团队还引入了分层筛选机制(hierarchical filtering):只有在第r层中被选中的token,才有资格参与第r+1层的评估。
这一设计模拟了早退出机制,让模型在训练初期就能自动「筛选」出需要深度处理的token,把算力集中到最困难的token上。
与前者不同,token选择路由(token-choice routing)(见图2b),在一开始就决定了每个token要走多少次递归流程。
具体来说,根据第1层的隐藏状态,模型会计算每个专家的打分(如通过softmax或sigmoid)。
假设有𝑁𝑟个专家,每个专家对应一次递归,那么模型将token分配给得分最高的专家。 token将被送入前i层递归,每层都按顺序进行处理。
在这种方式中,token在进入网络时就确定好递归深度,同时避免了每一层的重新选择,提升了推理效率。
表2左比较了两种方法:
expert-choice路由的优点在于,它可以实现理想的计算负载均衡。然而,它容易信息泄露。
相比之下,token-choice路由天然不会泄露信息。但这种方式负载分配不均。

表2:路由策略与键值缓存策略的比较。(左)两种路由策略总结:专家选择与令牌选择;(右)缓存策略相对于普通Transformer的相对成本效率

图2:混合递归(MoR)的架构组件。(a)专家选择路由;(b)token自主选择路由;(c)KV缓存策略
针对MoR模型,研究人员提出了两种KV缓存策略:按递归层缓存和跨递归共享。
1.按递归层缓存(见图2c上)是「选择性缓存」:只有被路由到某一递归层的Token,才会在该层生成并存储它的KV对。
注意力计算仅在当前递归层的缓存内进行,这种设计有助于实现局部化计算,显著提升了内存使用效率,并减少I/O负担。
2.跨递归共享(见图2c):只在第一个递归层生成并缓存KV对,然后在之后所有层中重复使用。这种机制下,每一层参与注意力计算的Query数量可能会减少。
也就是说,所有Token无论在后续层是否继续参与计算,都可以完整地访问历史上下文,无需重新计算。
表2右对比了两种缓存策略:
-
按递归层缓存:KV内存与I/O负担,被压缩为原来的一半左右。
-
跨递归共享:只能线性压缩注意力计算量,而且KV的读写次数较高,可能会成为性能瓶颈。

表3:在等计算量与等token数条件下,MoR、递归Transformer、普通Transformer的比较

研究者从零开始预训练模型,采用基于Llama的Transformer架构,参考了SmolLM开源模型的配置,在FineWeb-Edu的验证集和六个few-shot基准测试集上进行了评估。
在相同训练计算预算下,MoR以更少参数优于基线模型
在相同的训练预算(16.5e18 FLOPs)下,研究者将MoR模型与标准Transformer和递归Transformer进行了对比。

在四种模型规模(135M、360M、730M和1.7B参数)下,不同计算预算对应的验证损失对如图
如表3所示,MoR模型采用专家选择路由和两次递归(Nr=2),不仅在验证损失上更低,在few-shot平均准确率上也优于标准基线。
这得益于MoR更高的计算效率,使其在相同FLOPs预算下能处理更多的训练token。
在相同数据量下,MoR用更少计算量仍优于基线模型
为了隔离架构差异的影响,研究者在固定训练token数量(20B)的前提下进行分析。
结果证实,在少了25%训练FLOPs的情况下,MoR模型(𝑁𝑟=2)仍然实现了更低的验证损失和更高的准确率,超越了标准和递归基线。
与标准基线相比,MoR模型的训练时间减少了19%,峰值内存使用量降低了25%。
这就要归功于专门设计的分层过滤机制和按递归进行的注意力机制。
此外,MoR的性能也会受路由与缓存策略的影响。
评估一种新模型架构设计的核心标准之一,是其在模型规模和计算量增长时,性能是否能持续提升。
因此,研究团队全面对比了MoR与标准Transformer(Vanilla)和递归Transformer。
实验设置
实验的模型规模有四种:135M、360M、730M 和1.7B 参数。
对于递归Transformer和MoR配置,递归次数统一设为3。
在三个不同的计算预算下,进行预训练:2e18、5e18和16.5e18 FLOPs。
MoR架构:可扩展且参数高效
如图3所示,在所有参数规模和算预算力下,MoR始终优于递归基线模型。
尽管在最小规模(135M)时,MoR表现略逊于标准Transformer,但随着模型规模扩大,这一差距迅速缩小。
当参数规模超过360M时,MoR不仅能够与标准Transformer持平,甚至在低计算量和中等计算预算下,表现更加优越。
总体而言,这些结果表明,MoR具备良好可扩展性和高参数效率,可替代旧架构。
通过参数共享,MoR能利用连续深度批处理技术,在推理阶段显著提升了吞吐量。
这种机制在解码过程中,旧序列完成后立刻填入新tokens,持续保持了GPU的高利用率。
实验设置
在360M参数规模下,在不同递归深度(2、3和4)下,团队测试了MoR模型。
利用深度批处理,MoR显著提升推理吞吐量
如图4a所示,在两种设置下,MoR变体的推理吞吐量都超过了普通Transformer。
递归深度越高,越多tokens会提早退出,从而减少KV缓存的使用,进一步大幅提升了推理速度。例如,在最大批设置(𝐵=Max)下,MoR-4速度可提升2.06倍。
实验表明,结合深度批处理机制与提前退出策略,可大幅加速MoR模型在实际的推理速度。

消融实验等更多内容和细节,请参阅原文。
(文:新智元)