Transformer危！谷歌MoR架构发布：内存减半推理速度还翻倍

鹭羽 发自 凹非寺
量子位 | 公众号 QbitAI

超越Transformer，谷歌推出全新底层架构——
Mixture-of-Recursions（MoR），注意不是MoE，它能推理速度提高2倍，而KV内存直接减半！
而且All in One，首次在单一框架中实现，用同一组参数处理不同任务的同时，进行动态分配计算资源。
就像给LLM开了个双层增强buff，模型性能和效率全都要。


谷歌DeepMind联合KAIST AI、Mila人团队通过统一参数共享、自适应递归深度和高效KV缓存，在保持大模型性能的同时降低计算和内存成本，形成新的效率最优解。
不少网友甚至将它形容为Transformer Killer。


更有甚者表示，该架构的出现或许能代表，潜在空间推理也许将会成为下一个LLM突破所在。


具体MoR创新在哪些方面？下面一一详解。
MoR：首次统一参数共享与自适应计算
Transformer的出现虽然带来了优秀的少样本泛化和推理能力，但随之而来庞大的计算和内存需求还是让训练和部署成为难题。
目前相关优化方法主要是参数共享和自适应计算，但往往只能二选一，无法同时兼顾。
于是研究人员提出了递归混合模型MoR，可以在单一递归Transformer中同时融合两个效率维度。
首先采用的递归Transformer，相比通过多个独特层构建token的标准Transformer，其直接将模型划分为递归块，复用一组共享参数池。


主要包含三种参数共享策略：
[ol]Cycle：循环复用层。Sequence：连续复用同一层。Middle变体：保留首尾层独特参数，仅共享中间层。[/ol]
利用参数共享，可以减少独特参数数量，提升分布式训练效率，并通过连续深度批处理消除计算 “气泡”，提高推理吞吐量。


然后MoR采用动态路由机制，通过轻量级路由器为每个token分配不同递归深度，并集中计算在复杂token上，可以分为两种：
Expert-choice路由：将每个递归步骤视作 “专家”，基于隐藏状态计算分数，通过阈值选择合适的token继续计算，采用层级过滤，优先为复杂token分配计算。Token-choice路由：初始阶段为每个token分配固定递归深度，通过softmax/sigmoid确定专家，然后token按分配深度依次完成递归。


另外，MoR本身还借助KV缓存策略，管理键值的存储与使用，同时确保内存效率的提升：
Recursion-wise缓存：仅缓存当前递归步骤中活跃token的KV对，将注意力计算限制在本地缓存，降低内存和IO需求。Recursive KV共享：复用首次递归的KV对供后续步骤使用，确保所有token可访问历史上下文，减少预填充操作，而此时注意力计算量下降幅度较小。
在三种策略共同作用下，MoR通过在解码每个token时直接进行潜在思考，路由机制让模型能进行自适应推理，突破了先前固定思考深度的限制，实现了参数效率与自适应计算的统一。
性能超越Transformer
研究人员在135M到1.7B不同参数规模的模型上，就原始Transformer、递归基线模型和MoR进行对比实验。


实验表明，在相同的16.5e18 FLOPs的训练预算下，MoR使用了将近50%的更少参数，但实现了更低的验证损失和更高的平均少样本准确率43.1%。
而vanilla模型的少样本准确率此时是42.3%，说明MoR拥有更高的计算效率，可以在相同FLOPs预算下处理更多训练token。
在训练固定的20B token时，MoR也减少了25%的训练FLOPs，训练时间还缩短了19%，峰值内存减少25%。
另外，通过分析路由策略，发现Expert-choice路由的性能会在一定程度上优于Token-choice路由，说明路由粒度会对性能产生重要影响。
研究人员还对MoR进行了IsoFLOP分析，发现在135M、360M、730M和1.7B参数规模，以及2e18、5e18、16.5e18的FLOPs预算下，MoR始终优于递归基线模型。


虽然受递归容量瓶颈限制，在135M时略逊于vanilla模型，但在360M及规模的进一步扩大，MoR性能接近甚至超过Vanilla模型，且参数仅为后者的1/3，验证了MoR的可扩展性。
在推理吞吐量评估上，360M规模的MoR模型在固定批大小和最大批大小设置下，均优于vanilla。


递归深度的增加让更多token提前退出，KV缓存占用减少，吞吐量显著提升，验证了深度批处理与早期退出结合对部署效率的提升。
谷歌对底层架构的再思考
这已经不是谷歌第一次对底层架构进行重新思考，甚至可以说，谷歌始终希望利用架构创新重构计算范式，重新寻找AI新的平衡。
例如混合专家模型（MoE）正是这一理念的集中体现。
最早是在2017年，谷歌首次将MoE引入LSTM层，通过稀疏门控机制，仅激活部分专家网络从事输入处理，但仍能让参数量高达137B的模型保持高效训练。


后面推出的GShard将MoE与Transformer结合，可实现动态负载均衡，2021年的Switch Transformer又进一步地简化了路由机制。
而Gemini 1.5 Pro就是采用的分层MoE架构，将专家网络与多模态处理深度结合，可以处理更为复杂的多模态任务，训练和服务效率也得到显著提升。


MoE的底层逻辑设计突破了传统全连接模型的计算缺陷，现在已经成为了许多超大规模模型的优先选择，为应对算力瓶颈提供了新范式。
另外还有像TokenFormer等可扩展架构，将模型参数视为可学习的token，通过增量训练无缝扩展模型规模，为未来千亿级模型的低成本迭代提供了可能。
所以有网友认为，关于谷歌如今推出的MoR，会在未来彻底改变AI世界规则，会超越Transformer吗？一切交给时间验证。


参考链接：
[1]https://x.com/deedydas/status/1945313404958466519
[2]https://www.alphaxiv.org/abs/2507.10524
[3]https://x.com/reza_byt/status/1945498424536862841
[4]https://arxiv.org/abs/1701.06538