掘金 人工智能 07月17日 17:40
「有望成为Transformer杀手」,谷歌DeepMind新架构MoR实现两倍推理速度
index_new5.html
../../../zaker_core/zaker_tpl_static/wap/tpl_guoji1.html

 

谷歌提出的Mixture-of-Recursions(MoR)模型,是一种革命性的LLM架构,旨在解决大型语言模型在计算和内存方面的巨大开销。MoR通过在递归Transformer基础上引入参数共享和自适应计算,实现了「无需承担大模型成本的大模型质量」。该模型为每个token动态分配计算资源和递归深度,显著降低了内存带宽压力和冗余计算,从而大幅提升了推理吞吐量和效率。实验结果表明,MoR在显著减少参数量和计算量的同时,能够超越基线模型,为高效LLM的预训练和大规模部署提供了新的可能性。

🌟 MoR架构统一了参数共享与自适应计算:通过将动态token级路由集成到参数高效的递归Transformer中,MoR实现了一种协同架构,使得模型在拥有大模型质量的同时,大幅降低了计算成本。它通过为每个token分配专属的递归深度,决定共享参数模块的重复应用次数,从而将计算资源精准分配给最需要的部分。

🚀 递归级KV缓存优化显著提升效率:MoR引入了递归式KV缓存策略,根据每个token的递归深度有选择性地存储和检索键值对,大幅降低了内存带宽压力。这种定向缓存机制无需后处理即可提升推理吞吐量,是其高效性的关键原因之一。

💡 MoR在参数效率和计算效率上均优于基线模型:实验结果显示,MoR在相同的计算预算下,参数量减少近50%却能取得更优的性能。同时,在固定训练token数量下,MoR使用比基线模型少25%的计算量,仍实现了更低的验证损失和更高的准确率,训练时间减少19%,峰值内存使用量下降25%。

📈 MoR具备良好的可扩展性与推理吞吐能力:MoR在不同模型规模下均稳定优于递归基线模型,并能追平甚至超越原始Transformer,证明了其作为标准Transformer可扩展且高效替代方案的潜力。此外,结合深度方向批处理和提前退出机制,MoR在实际部署场景中能够显著提升推理吞吐能力,最高可达2.06倍。

🤔 MoR为研究更深层思考的AI系统提供基础:MoR在生成过程中能够按每个token自适应地分配「思考深度」,这与语言模型中新兴的潜在推理和内部思考研究相一致,为开发更具认知启发性的AI系统奠定了基础,使其能够更深入地思考复杂问题。

大型语言模型已展现出卓越的能力,但其部署仍面临巨大的计算与内存开销所带来的挑战。随着模型参数规模扩大至数千亿级别,训练和推理的成本变得高昂,阻碍了其在许多实际应用中的推广与落地。

这促使研究者们围绕两个主要方向积极探索高效化技术:一是通过权重共享提升参数效率,二是根据输入复杂度动态分配计算资源,实现自适应计算。

而一切的基础则是 Transformer 架构。这次谷歌又一次坐不住了,在递归 Transformer 的基础上再次进化,发布了名为 Mixture-of-Recursions (MoR)的新 LLM 模型架构,有学者认为这个新架构「有潜力成为 Transformer 的杀手」。

MoR 框架引入了一种统一方法,同时实现参数共享和自适应计算。与以往孤立处理参数减少或自适应计算的方法不同,MoR 将动态 token 级路由集成到参数高效的递归 Transformer 中,创建了一种协同架构,实现了「无需承担大模型成本的大模型质量」。

在本研究中,作者提出了 Mixture-of-Recursions(MoR),一个统一框架,旨在充分发挥递归 Transformer 的潜力(如图所示)。MoR 通过端到端训练轻量级的路由,为每个 token 分配专属的递归深度:即根据每个 token 所需的「思考」深度,决定共享参数模块对其重复应用的次数,从而将计算资源精准地分配到最需要的部分。

这种基于 token 的动态递归机制,天然支持递归级的键值(KV)缓存,即根据每个 token 所分配的递归深度,有选择地存储和检索对应的键值对。该定向缓存策略显著降低了内存带宽压力,从而无需后处理操作就能提升推理吞吐量。

因此,MoR 在一个统一架构中同时实现了三项优化:

1. 参数共享:通过权重绑定显著减少参数量;

2. 计算路由:通过 token 路由避免冗余计算(减少无效的 FLOPs);

3. 递归级缓存:按需缓存键值对,降低内存访问开销,提升吞吐效率。

MoR 以此兼顾了性能、效率和计算资源利用率,是递归 Transformer 架构的重要演进。

MoR 总览

左图显示每个递归步骤由一组固定堆叠的层与一个路由(router)组成。路由会为每个 token 决定是继续通过递归块(进行下一轮递归),还是提前退出。图中灰色方框表示的就是该递归模块。

中图展示了完整的模型结构。对于每个 token,模型会根据路由的决策,对共享的递归模块最多重复应用 N_r 次,从而实现动态递归深度。

右图 是一个 token 级的路由示例,展示了每个 token 的实际递归次数。颜色越深的单元格表示该 token 在当前递归步骤中被激活参与了计算。底部显示了每个文本 token 实际经历的递归次数,使用不同颜色编码标注(如:紫色 1 次、蓝色 2 次、红色 3 次)。

核心架构与方法

递归 Transformer 模型基础

MoR 基于递归 Transformer,通过跨多个层共享权重来实现参数效率,而不是为每一层维护唯一参数。

多种参数共享模式

循环:递归块循环复用(例如,重复层 [0,1,2])

序列:每个块连续地重用相同的层(例如,[0,0,0],[1,1,1])

中周期 / 中序列:保留首层和尾层的独特参数,同时共享中间层权重的变体

中周期策略在经验上被证明最为有效,在参数效率和模型表达能力之间提供了最佳平衡。

动态路由机制

MoR 的核心创新在于其动态路由系统,该系统根据单个 token 的计算需求分配不同的递归深度。实现了两种主要的路由策略:

1、在每一个递归步骤 r,路由器 G 会对每个 token 当前的隐藏状态进行评估,并选择 top-k 的 token 继续进入下一层递归。

这会形成一个逐步收窄的计算漏斗:复杂的 token 会接受更深层的处理,而简单的 token 会更早退出计算。然而,这种方法在训练过程中会面临因果性挑战,因此需要采用一些专门的技术:

2、Token 选择路由(Token-Choice Routing):每个 token 在一开始就做出一次路由决策,以确定其完整的递归深度。这种方式天然满足因果性要求,但需要引入负载均衡机制:

KV缓存优化

MoR 引入了两种针对动态深度模型的专用 KV 缓存策略:

递归式 KV 缓存:在每个递归深度选择性地缓存 KV 对,注意力仅限于本地缓存的 token:

递归 KV 共享:所有 token 都会遍历第一个递归块,该块的 KV 对会被缓存并在后续递归中重用:


实验结果

MoR 在相同计算预算下以更少参数优于基线模型

在等效的训练预算(16.5 × 10¹⁸ FLOPs)下,作者将 MoR 模型与 原始 Transformer 及递归 Transformer 进行了比较。如表 3 所示,MoR 模型采用「专家选择路由」(expert-choice routing)和两层递归(N_r = 2),实现了更低的验证损失,并在少样本学习平均准确率上超越基线(43.1% 对比 42.3%)。值得注意的是,尽管参数量减少近 50%,MoR 依然取得了更优性能。这一优势归因于其更高的计算效率,使其在相同 FLOPs 预算下处理了更多训练 token。

在固定 FLOPs(16.5e18)与固定 token 数(20B)两种设定下,对 MoR、递归 Transformer(Recursive)和标准 Transformer(Vanilla)进行了对比。

MoR 在相同数据下以更少计算超越基线模型

为了单独分析架构差异,作者在固定训练 token 数量(20B)下进行评估。结果显示,MoR 模型(N_r = 2)在使用比基线模型少 25% 的 FLOPs 情况下,仍实现了更低的验证损失和更高的准确率。

这种理论上的效率优势在实践中带来了显著提升:相比于基线模型,MoR 的训练时间减少 19%,峰值内存使用量下降 25%。这些改进来源于 MoR 的层级过滤机制与递归注意力策略,有效缩短了序列长度,实现了更优的计算 - 准确率平衡,甚至在预训练阶段亦是如此。

MoR 性能受路由与缓存策略影响

作者还评估了 MoR 架构下的一些设计变体,特别是在轻量级的三层递归设置(N_r = 3)下仍具可比性的 Vanilla 模型。结果显示,采用「token 选择路由」时准确率较低(40.0%),而「专家选择路由」表现更好(42.6%),表明路由粒度对模型性能具有关键影响。

此外,采用 KV 缓存共享略微降低了性能,但显著提升了内存效率。在对内存占用有要求的部署场景中,这种性能与资源之间的权衡是值得接受的。

不同计算预算下的验证损失,在四种模型规模(135M、360M、730M 和 1.7B 参数量)下进行对比。

MoR 是一种可扩展且参数高效的架构

如图所示,MoR 在所有模型规模和计算预算下都稳定优于递归基线模型。尽管在最小规模(135M)下表现略低于原始 Transformer(可能由于递归容量瓶颈所致),但这种差距随着模型规模的扩大迅速缩小。在参数量超过 360M 时,MoR 不仅能够追平,而且在低至中等预算下常常超越原始 Transformer。

结果表明:MoR 是标准 Transformer 的一种可扩展且高效的替代方案。

它在使用显著更少参数的情况下,取得了优异的验证性能,因而非常适合用于预训练以及大规模部署。

MoR 通过连续深度批处理提升推理吞吐量

如图 4a 所示,在上述两种批处理设置下,所有 MoR 变体的推理吞吐量均高于 Transformer 基线模型。

随着递归深度的增加,更多 token 可在早期阶段退出,从而减少 KV 缓存的使用。这一机制显著提升了推理速度:

例如,MoR-4 在最大批量配置下的吞吐量提升可达 2.06 倍。

虽然这伴随着轻微的性能下降,但考虑到吞吐量的大幅提升,这种权衡是值得的。

实验结果表明:将深度方向批处理范式与提前退出机制相结合,能够显著提升 MoR 在实际部署场景中的推理吞吐能力。

图中 (a) 在固定批量和最大批量设置下,MoR 与原始 Transformer 在推理吞吐量与对数似然(log-likelihood)之间的帕累托前沿(Pareto frontier)表现对比。

(b) 参数共享策略对递归 Transformer 的影响:展示了递归次数为 𝑁𝑟 = 3 的递归 Transformers 在四种不同参数共享策略下的负对数似然(NLL)表现。这些模型在 100 亿 token 上预训练。图中红色虚线表示完整规模的原始 Transformer,黑色虚线表示参数量约为其三分之一的基线模型。

(c) 四种不同架构在启用 KV 缓存共享时的 NLL 表现对比。其中:

另有消融实验,请参照原论文。

总结与讨论

MoR 通过证明参数效率和自适应计算可以协同结合,为高效的 LLM 架构建立了一种新的范例。该框架实现大型模型质量的同时显著降低了计算需求的能力,解决了语言建模中的基本可扩展性挑战。

最重要的是,MoR 为开发更具认知启发性的 AI 系统提供了基础。其在生成过程中能够按每个 token 自适应地分配「思考深度」的能力,与语言模型中新兴的潜在推理和内部思考研究相一致。这表明 MoR 可以作为一个平台,用于研究模型如何学会在保持常规任务效率的同时,更深入地思考复杂问题。

该方法的确提高了推理速度,降低了 KV 缓存,但其是否能被称为「Transformer 杀手」,仍然存在很多质疑。

更多细节内容,请参考原论文。

Fish AI Reader

Fish AI Reader

AI辅助创作,多种专业模板,深度分析,高质量内容生成。从观点提取到深度思考,FishAI为您提供全方位的创作支持。新版本引入自定义参数,让您的创作更加个性化和精准。

FishAI

FishAI

鱼阅,AI 时代的下一个智能信息助手,助你摆脱信息焦虑

联系邮箱 441953276@qq.com

相关标签

MoR 大型语言模型 Transformer AI架构 计算效率
相关文章