⏶50
混合递归：学习动态递归深度以实现自适应的词元级计算
发表
由 
Sangmin Bae 提交
Sangmin Bae 提交作者: Sangmin Bae, 
Yujin Kim, Reza Bayat, Sungnyun Kim, Jiyoun Ha, Tal Schuster, Adam Fisch, Hrayr Harutyunyan, Ziwei Ji, Aaron Courville, Se-Young Yun
Yujin Kim, Reza Bayat, Sungnyun Kim, Jiyoun Ha, Tal Schuster, Adam Fisch, Hrayr Harutyunyan, Ziwei Ji, Aaron Courville, Se-Young Yun摘要
扩展语言模型能解锁卓越的能力，但随之而来的计算和内存需求使得训练和部署都变得十分昂贵。现有的效率优化工作通常要么针对参数共享，要么针对自适应计算，而如何同时实现两者仍然是一个悬而未决的问题。我们引入了混合递归（Mixture-of-Recursions, MoR），这是一个统一的框架，它在单个递归 Transformer 内部结合了这两个效率维度。MoR 在不同递归步骤中重用一个共享的层堆栈以实现参数效率，而轻量级路由器通过为单个 token 动态分配不同的递归深度，来实现自适应的 token 级思考。这使得 MoR 能够将二次方的注意力计算仅集中在给定递归深度下仍然活跃的 token 之间，并通过只选择性地缓存它们的键值对，进一步提高了内存访问效率。除了这些核心机制，我们还提出了一种 KV 共享变体，它会重用第一次递归的 KV 对，专门用于降低预填充延迟和内存占用。在从 1.35 亿到 17 亿参数不等的模型规模上，MoR 形成了一条新的帕累托前沿：在相同的训练 FLOPs 和更小的模型尺寸下，它显著降低了验证集困惑度并提高了少样本准确率，同时与标准基线和现有递归基线相比，提供了更高的吞吐量。这些成果表明，MoR 是在不产生大模型成本的情况下，实现大模型质量的一条有效路径。
扩展语言模型的规模可以释放出令人印象深刻的能力，但随之而来的计算和内存需求使得训练和部署都变得昂贵。现有的效率提升方法通常要么专注于参数共享，要么专注于自适应计算，而如何同时实现这两者仍然是一个悬而未决的问题。我们引入了递归混合模型（Mixture-of-Recursions, MoR），这是一个统一的框架，在单个递归 Transformer 内部结合了这两个效率维度。MoR 通过在不同递归步骤中复用一个共享的层堆栈来实现参数效率，同时，轻量级路由器通过为单个 token 动态分配不同的递归深度，实现了自适应的 token 级“思考”。这使得 MoR 能够仅在给定递归深度下仍然活跃的 token 之间集中进行二次方复杂度的注意力计算，并通过选择性地仅缓存这些 token 的键值对（key-value pairs），进一步提高了内存访问效率。除了这些核心机制，我们还提出了一种 KV 共享变体，它复用第一次递归的 KV 对，专门用于减少预填充（prefill）延迟和内存占用。在从 1.35 亿到 17 亿参数不等的模型规模上，MoR 形成了一条新的帕累托前沿（Pareto frontier）：在相同的训练 FLOPs 和更小的模型尺寸下，它显著降低了验证集困惑度（validation perplexity）并提高了少样本（few-shot）准确率，同时与普通（vanilla）和现有的递归基线模型相比，提供了更高的吞吐量。这些优势表明，MoR 是在不产生大模型成本的情况下，获得大模型质量的一条有效途径。代码将发布于 https://github.com/raymin0223/mixture_of_recursions。