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

扩大语言模型的规模可以释放出令人印象深刻的能力，但随之而来的计算和内存需求也使得训练和部署成本高昂。现有的效率优化方法通常要么专注于参数共享，要么采用自适应计算策略，但如何同时实现这两方面的优势仍然是一个开放问题。我们提出了“递归混合”（Mixture-of-Recursions，MoR）这一统一框架，将两种效率维度融合在一个统一的递归 Transformer 架构中。MoR 在递归步骤中复用一组共享的层级，从而实现参数效率；同时，通过轻量级路由机制，在每个 token 层面实现自适应的“思考”深度，动态地为不同的 token 分配不同的递归深度。这使得 MoR 仅在特定递归深度下仍处于活跃状态的 token 之间集中进行二次复杂度的注意力计算，并通过有选择地缓存这些 token 的键值对，进一步提升内存访问效率。除了这些核心机制之外，我们还提出了一种 KV 共享变体，复用了首次递归中的键值对，专门用于降低预填充阶段的延迟和内存占用。在从 1.35 亿到 17 亿参数的不同模型规模上，MoR 形成了一条新的帕累托前沿：在训练计算量相同甚至模型更小的情况下，MoR 显著降低了验证集困惑度，提升了少样本准确率，并且相比基础版本及其他现有的递归基线模型具有更高的吞吐性能。这些优势表明，MoR 是一条在不带来大模型成本的前提下实现大模型质量的有效路径。

论文试图解决语言模型在训练和部署过程中计算和内存需求过高的问题，同时兼顾参数效率与自适应计算的双重目标。当前的研究通常只关注其中一方面，而如何结合两者仍是一个开放问题。

论文提出了一种名为Mixture-of-Recursions (MoR) 的统一框架，在单一递归Transformer中结合了参数共享和自适应计算。MoR通过共享层堆栈实现参数效率，并利用轻量级路由机制动态分配不同token的递归深度，从而提高内存访问效率并降低预填充延迟。

{MoR在135M到1.7B参数范围内形成了新的Pareto前沿，在相同训练FLOPs和更小模型规模下显著降低了验证困惑度并提升了few-shot准确率。,提出了KV共享变体以减少预填充延迟和内存占用。,相比传统和现有递归基线模型，MoR实现了更高的吞吐量。,实验设计覆盖多个模型规模，展示了其在性能和效率上的综合优势。}

{"Efficient Training of Language Models using Parameter Sharing","Adaptive Computation Time for Transformer Models","Recursive Transformers for Long Sequence Modeling","Dynamic Sparse Attention Mechanisms in Neural Language Models","KV Cache Optimization Techniques for Fast Inference in Large Language Models"}