arXiv:2507.10524cs.CLcs.LG2025-07NeurIPS被引 113

动态调整每令牌计算深度,实现高效大模型推理。

Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation

  • 用共享层栈+轻量路由,让不同令牌自适应递归深度。
  • 在135M~1.7B参数模型上,验证困惑度更低、吞吐更高。
  • 适合追求高性价比大模型推理的开发者和研究者。

扩大语言模型能带来显著能力提升,但随之而来的计算与内存开销使训练和部署成本高昂。现有高效方法通常只关注参数共享或自适应计算,未能同时实现二者。我们提出混合递归(Mixture-of-Recursions, MoR),一种统一框架,将两种效率机制集成于单一递归Transformer中。MoR通过在递归步骤间复用共享层栈实现参数效率,同时利用轻量路由器为每个令牌动态分配不同的递归深度,仅对仍活跃的令牌进行二次注意力计算,并选择性缓存其键值对,进一步提升内存访问效率。此外,我们还提出一种KV共享变体,复用首次递归的键值对以进一步降低内存占用。在135M至1.7B参数规模的模型上,MoR实现了新的帕累托前沿:在相同训练浮点运算量且模型更小时,显著降低验证困惑度并提升少样本准确率,同时比标准和现有递归基线获得更高吞吐量。这些成果表明,MoR是实现大模型性能而不付出大模型成本的有效路径。

原文摘要 · Abstract (English)

Scaling language models unlocks impressive capabilities, but the accompanying computational and memory demands make both training and deployment expensive. Existing efficiency efforts typically target either parameter sharing or adaptive computation, leaving open the question of how to attain both simultaneously. We introduce Mixture-of-Recursions (MoR), a unified framework that combines the two axes of efficiency inside a single Recursive Transformer. MoR reuses a shared stack of layers across recursion steps to achieve parameter efficiency, while lightweight routers enable adaptive token-level thinking by dynamically assigning different recursion depths to individual tokens. This allows MoR to focus quadratic attention computation only among tokens still active at a given recursion depth, further improving memory access efficiency by selectively caching only their key-value pairs. Beyond these core mechanisms, we also propose a KV sharing variant that reuses KV pairs from the first recursion, specifically designed to further decrease memory footprint. Across model scales ranging from 135M to 1.7B parameters, MoR forms a new Pareto frontier: at equal training FLOPs and smaller model sizes, it significantly lowers validation perplexity and improves few-shot accuracy, while delivering higher throughput compared with vanilla and existing recursive baselines. These gains demonstrate that MoR is an effective path towards large-model quality without incurring large-model cost.

递归模型自适应计算内存优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。