KAIST：LLM混合递归推理

标题：Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation
来源：arXiv, 2507.10524

摘要

缩放语言模型解锁了令人印象深刻的能力，但伴随的计算和内存需求使训练和部署都很昂贵。现有的效率工作通常针对参数共享或自适应计算，留下了如何同时实现两者的问题。我们引入了混合递归 (MoR)，这是一个统一框架，它结合了单个递归 Transformer 中两个效率轴。MoR在递归步骤中重用共享的层堆栈来实现参数效率，而轻量级路由器通过将不同的递归深度动态分配给单个令牌来实现自适应令牌级思维。这允许 MoR 仅在在给定递归深度仍然活跃的令牌之间关注二次注意力计算，通过仅选择性地缓存它们的键值对来进一步提高内存访问效率。除了这些核心机制之外，我们还提出了一个 KV 共享变体，它重用第一个递归中的 KV 对，专门设计用于减少预填充延迟和内存占用。在从135M到1.7B参数的模型尺度上，MoR形成了一个新的帕累托边界:在相同的训练FLOPs和更小的模型尺寸下，它显著降低了验证困惑度，提高了少镜头精度，同时与普通递归和现有的递归基线相比提供了更高的吞吐量。这些收益表明，MoR 是朝着大型模型质量迈出的有效途径，而不会招致大型模型成本。

️文章简介

研究问题：如何动态地为每个令牌分配递归深度，以提升大语言模型（LLM）的计算效率和性能？
主要贡献：提出了一种混合递归（MoR）框架，通过动态分配递归深度和高效的KV缓存策略，实现了语言模型的计算效率和速度的显著提升。

重点思路

设计了一种动态递归路由机制，为每个令牌分配特定的递归步骤，聚焦计算在更复杂的令牌上。
采用递归KV缓存策略，只有在递归步骤中被选择的令牌存储其键值对，从而降低内存使用并提高计算效率。
通过集成固定参数块的多次应用，实现一个统一的框架，同时结合参数共享和动态计算深度。

分析总结

MoR模型在与Vanilla和递归模型比较时，在同等训练计算预算下，表现出更低的验证损失和更高的少样本准确率。
MoR通过分动态分配令牌的递归深度，显著减少了训练中的FLOPs和内存使用，同时提高了推理吞吐量。
在推理方面，MoR利用连续深度批处理机制，进一步提升了推理速度，尤其是在递归深度增加时能够显著提高吞吐量。

个人观点

论文的创新点在于引入了混合递归框架（MoR），通过动态调整每个令牌的递归深度和高效的KV缓存策略，实现了模型在计算效率、内存使用和推理速度方面的重大优化，突破了传统递归模型的限制。

附录