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

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

论文试图解决大规模语言模型在训练和部署过程中高昂的计算和内存成本问题。具体来说，它旨在同时实现参数共享和自适应计算，以提高效率，这是一个在现有研究中尚未被充分探索的问题。

论文提出了一种名为Mixture-of-Recursions（MoR）的新框架，该框架结合了参数共享和自适应计算两个效率维度。MoR通过共享递归步骤中的层栈实现参数效率，并利用轻量级路由机制动态分配不同标记的递归深度，从而实现标记级别的自适应推理。这一思路相比现有研究具有显著的创新性。

1. MoR通过减少活跃标记的注意力计算和选择性缓存键值对来提升内存访问效率。 2. 提出了一种KV共享变体，进一步降低预填充延迟和内存占用。 3. 实验结果显示，MoR在训练FLOPs相同且模型更小的情况下，显著降低了验证困惑度并提高了少样本准确率。 4. MoR在吞吐量方面优于基线模型，展示了其在实际应用中的潜力。 5. 论文并未提及是否开源代码，但实验设计全面，覆盖了多个模型规模和任务。

1. Efficient Training of Language Models using Parameter Sharing 2. Adaptive Computation Time for Transformer Models 3. Sparse Mixture of Experts for Efficient Language Modeling 4. Recursive Transformer Architectures for Long Context Processing 5. Memory-Efficient Attention Mechanisms in Transformers