1 个月前

递归混合:学习动态递归深度以适应性地进行 token 级计算

Sangmin Bae, Yujin Kim, Reza Bayat, Sungnyun Kim, Jiyoun Ha, Tal Schuster, Adam Fisch, Hrayr Harutyunyan, Ziwei Ji, Aaron Courville, Se-Young Yun
递归混合:学习动态递归深度以适应性地进行 token 级计算
摘要

扩展语言模型可以解锁令人印象深刻的能力,但随之而来的计算和内存需求使得训练和部署变得昂贵。现有的效率提升方法通常针对参数共享或自适应计算中的某一方面,而如何同时实现这两方面的优化仍然是一个开放的问题。我们引入了递归混合(Mixture-of-Recursions, MoR)框架,该框架在一个单一的递归变压器中结合了这两个效率维度。MoR 在递归步骤之间重用共享的层堆栈以实现参数效率,而轻量级路由器则通过动态分配不同的递归深度给各个标记来实现自适应的标记级处理。这使得 MoR 只需在特定递归深度下仍然活跃的标记之间进行二次注意力计算,并通过选择性缓存这些标记的关键值对进一步提高内存访问效率。除了这些核心机制外,我们还提出了一种关键值共享变体,该变体重用第一次递归中的关键值对,旨在减少预填充延迟并降低内存占用。在从 1.35 亿到 17 亿参数的不同模型规模上,MoR 形成了一个新的帕累托前沿:在相同的训练浮点运算次数和更小的模型规模下,它显著降低了验证困惑度并提高了少样本准确性,同时相比普通模型和现有的递归基线模型提供了更高的吞吐量。这些改进表明,MoR 是一条有效的路径,可以在不承担大型模型成本的情况下达到大型模型的质量。