AI热点 4 weeks ago 81 Views 10 Comments

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

量子位

Published 290 Articles

超越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,其直接将模型划分为递归块,复用一组共享参数池。

主要包含三种参数共享策略:

  1. Cycle:循环复用层。
  2. Sequence:连续复用同一层。
  3. Middle变体:保留首尾层独特参数,仅共享中间层。

利用参数共享,可以减少独特参数数量,提升分布式训练效率,并通过连续深度批处理消除计算 “气泡”,提高推理吞吐量。

然后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

本文来自微信公众号“量子位”,作者:鹭羽 ,36氪经授权发布。

量子位

量子位

290 Articles 41810 Views 9840 Fans

Comment (10)

User avatar

谷歌MoR架构太给力,LLM效率翻倍简直炸裂!

User avatar

“这MoR架构,我感觉它要统治世界,我们得赶紧学习!”

User avatar

“效率翻倍?这可能只是谷歌的营销手段,别太当真!”

User avatar

“炸裂!这简直是人类对人工智能的疯狂崇拜,我有点害怕!”

User avatar

“这MoR架构,我猜一定是谷歌在用外星人的技术,不然怎么解释得过这么快!”

User avatar

“效率翻倍,我感觉未来的世界,一切都快了,赶紧抓住这波潮流!”

User avatar

“谷歌这招,太狠了!就像给LLM喂了超级能量饮料,直接爆炸!”

User avatar

“翻倍效率,这说明人类的思考速度可能要被超越了,有点吓人!”

User avatar

“炸裂?炸裂得我都怀疑谷歌是不是在用炸药来训练模型!”

User avatar

“这MoR架构,简直是谷歌的神经病,效率翻倍?我感觉世界都要加速了!”

睡觉动画