如今,大型语言模型(LLM)的能力令人惊叹,但其巨大的计算和内存需求也让许多研究者和开发者望而却步。我们不禁要问:有没有一种方法,既能拥有大模型的强大能力,又能摆脱其高昂的成本呢?
今天,我们将深入探讨一篇于2025年7月发表的论文——《Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation》。该论文提出了一种名为Mixture-of-Recursions (MoR) 的创新框架,巧妙地结合了参数共享和自适应计算两大效率提升策略,为我们指明了一条通往“物美价廉”的AI模型之路。
核心思想:当模型学会“按需思考”
想象一下我们人类在处理信息时的情景:对于简单词汇(如“的”、“和”),往往一扫而过;而对于关键复杂的概念,则会反复琢磨、深入思考。MoR的核心思想正是来源于此——让模型为输入文本中的每个词元(token)动态地分配不同级别的“思考深度”。
MoR建立在递归Transformer(Recursive Transformer) 的基础上。传统的Transformer模型由许多独立的层堆叠而成,而递归Transformer则重复使用一个共享的层块(Recursion Block),从而大幅减少模型的参数量。
MoR在此基础上更进一步,它引入了一个轻量级的路由器(Router),其作用就像一个交通指挥员。对于输入的每个词元,路由器会判断它需要经历多少次递归计算。简单的词元可能只递归1次,而复杂的词元则可能需要递归3次甚至更多。
图:递归混合(MoR)概述。(左图)每个递归步骤由一个固定的层堆和一个路由器组成,路由器决定每个标记是通过还是退出。这个递归块与中间的灰色方框相对应。 中)完整的模型结构,其中共享递归步骤根据路由器的决定对每个标记最多应用𝑁𝑟 次。(右图)路由模式示例显示了按标记排列的递归深度,其中较暗的单元格表示递归块的计算正在进行中。下图显示了每个文本标记的递归步数,颜色分别为 1(紫)、2(蓝) 和 3(红)
这种机制带来了三大优势:
-
参数高效:通过复用层块,模型参数量显著降低。
-
计算自适应:计算资源被智能地分配给最需要它的词元,避免了在简单信息上的算力浪费。
-
内存优化:后续我们将看到,这种动态机制还为KV缓存的优化提供了可能。
两大关键技术:路由策略与KV缓存
MoR的实现依赖于两个核心组件:路由策略(Routing Strategies)和键值缓存策略(KV Caching Strategies)。
路由策略:如何决定“思考”的深度?
论文探讨了两种主要的路由策略来决定每个词元的递归深度:
• 专家选择(Expert-choice):在每个递归步骤中,路由器会从当前所有活跃的词元中挑选出“最需要”继续计算的top-k个。这种方式保证了每一步的计算负载是固定的,易于管理。
• 词元选择(Token-choice):在一开始,路由器就为每个词元“指定”好它需要经历的总递归次数。这种方式更符合自回归模型的特性,但可能导致不同递归深度的计算负载不均衡。
实验表明,专家选择的路由方式在性能上表现更优,因为它能更灵活地在每个阶段动态分配计算资源。
KV缓存策略:优化推理速度的“记忆”机制
在语言模型生成文本时,需要一个名为KV缓存的“记忆”系统来存储历史信息,以避免重复计算。动态递归深度给KV缓存带来了挑战:如果一个词元提前“退出”了计算,它在更深层次的KV缓存就会缺失。
为此,MoR设计了两种KV缓存策略:
• 递归式KV缓存(Recursion-wise KV Caching):在每个递归深度,只缓存当前仍在计算的那些词元的键值对。这是一种“本地化”的缓存,可以有效减少内存占用和I/O开销。
• 递归式KV共享(Recursive KV Sharing):所有词元都在第一次递归时生成KV缓存,并在后续所有递归步骤中共享使用。这种方式可以最大化地节省内存,尤其是在处理长文本时,但可能会因为信息不匹配而牺牲一些性能。
图:图 2:混合递归(MoR)的架构组件。(a) 专家选择路由: 在每个递归步骤中,路由器都会选择最前的𝑘 标记继续递归,随着递归深度的增加,活动标记集的范围会逐渐缩小。(b) 标记选择路由: 每个令牌从一开始就通过一个路由决策被分配到一个固定的递归步骤,从而定义了它在模型中的完整计算路径。© KV 缓存策略: 矩阵中的每个方格代表一个标记(行)是否与另一个标记的缓存密钥(列)相关。在 “递归式 KV 缓存”(顶部)中,每个递归步骤只缓存当前选定(未删除)令牌的密钥(蓝色),并且只关注这些条目。在 “递归 KV 共享”(下图)中,前一个标记的所有键都会在第一个递归步骤中缓存(紫色),并在随后的递归步骤中共享,以便进行关注操作。
实验结果:更少参数,更强性能
论文通过大量实验,将MoR与传统的Transformer(论文中称Vanilla)和标准的递归Transformer模型进行了比较。结果令人振奋:
• 同等计算预算下,性能更优:在相同的训练计算量(FLOPs)下,MoR模型(尤其是采用专家选择路由的)在验证集上的困惑度更低,在少样本学习任务上的准确率更高。值得注意的是,达到这一性能时,MoR的参数量仅为传统模型的一半左右。
• 同等训练数据下,效率更高:在处理相同数量的训练词元时,MoR能节省约25%的训练FLOPs,并将训练时间缩短19%。
• 推理速度大幅提升:得益于参数共享和动态的提前退出机制,MoR可以利用连续深度批处理(continuous depth-wise batching)技术,在推理时实现高达2.06倍的吞吐量提升。
表:在固定 FLOPs (16.5e18) 和令牌 (20B) 设置下,MoR、递归和 Vanilla 变换器的比较。所有模型都在 FineWeb-Edu 上进行了训练,并通过验证负对数似然(NLL)和少拍准确率进行了评估。对于 isoFLOP 行,训练令牌数(𝑁𝑡𝑜𝑘)因模型效率而异。对于固定令牌行,我们报告消耗的有效 FLOP。对于模型大小,我们报告非嵌入式参数计数。对于 KV 机制,我们区分为 Cache(递归缓存)和 Share(递归共享)。在递归模型中,所有标记都经过固定的递归深度(𝑁𝑟),而不是自适应深度。
此外,论文还对不同模型尺寸(从135M到1.7B)进行了扩展性分析(IsoFLOP Analysis),结果表明,随着模型规模的增大,MoR的优势愈发明显,甚至在超过360M参数规模后开始全面超越同等计算预算下的传统Transformer模型。
图:四种模型规模下不同计算预算的验证损失: 135M、360M、730M 和 1.7B 参数。对于 MoR 模型,我们使用专家选择路由和递归缓存。尽管使用的参数明显较少(约三分之一是由于层绑定,𝑁𝑅 = 3),但 MoR 的性能始终优于递归基线,并在更大规模上与标准 Transformers 相匹配或超过后者。
结论与展望
Mixture-of-Recursions (MoR) 框架成功地将参数共享的效率与自适应计算的智能融为一体,为构建高性能、低成本的语言模型提供了一个全新的、极具前景的范式。它证明了,我们不必一味地追求模型的“大”,而是可以通过让模型学会“思考”的深浅,来更高效地利用计算资源。
论文作者也指出了未来的研究方向,包括:
• 将MoR扩展到更大的模型规模和更多样的数据集上。
• 探索其在多模态(如视频、音频)和长文本处理领域的应用潜力。
• 将其与模型稀疏化等其他效率提升技术结合,探索更高的优化上限。
总而言之,MoR不仅是一项技术创新,更是一种设计哲学上的启示:在通往更强人工智能的道路上,智慧地分配资源,可能比无尽地堆砌资源更为关键。
论文链接
《Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation》:https://arxiv.org/abs/2507.10524