Mixture-of-Recursions (MoR)：大型语言模型的统一效率范式#
摘要#
本文介绍了Mixture-of-Recursions (MoR)，一种创新的 Transformer 架构范式，通过统一参数共享、自适应令牌级计算和高效 KV 缓存机制，显著提升了大型语言模型的计算效率与性能。MoR 的核心创新在于引入轻量级路由器，使模型能够为每个令牌动态分配递归深度，实现自适应计算。实验表明，在同等训练 FLOPs 下，MoR 在减少 19%训练时间和 25%内存使用的同时，实现了 43.1%的少样本准确率（优于基线 42.3%）。此外，MoR 通过递归 KV 缓存策略将推理吞吐量提升了 2.06 倍，为大型模型质量提供了经济高效的实现途径。
1. 引言#
1.1 背景与挑战#
大型语言模型(LLMs)在少样本泛化和复杂推理任务中展现出卓越能力，但面临严峻挑战：
计算资源需求巨大：训练和部署需要极高的计算资源和内存开销
效率与性能的权衡：现有方法通常专注于参数效率或自适应计算的单一维度
递归 Transformer 的限制：动态递归尝试面临训练复杂和部署困难的问题
1.2 MoR 的核心目标#
MoR 提出统一框架以解决上述挑战：
在单个递归 Transformer 中同时实现参数效率和自适应计算
在同等训练 FLOPs 和更小模型尺寸下实现：
更低的验证困惑度
更高的少样本准确率
显著的吞吐量提升
2. MoR 架构设计#
2.1 三大核心机制#
MoR 整合了三个关键效率维度的创新：
机制 |
核心技术 |
主要优势 |
|---|---|---|
参数共享 |
递归层绑定技术 |
减少参数量，提高参数效率 |
自适应令牌级计算 |
轻量级动态路由器 |
按需分配计算资源 |
高效 KV 缓存 |
递归感知缓存策略 |
减少内存访问开销 |
2.2 参数共享机制#
MoR 采用创新的层绑定策略：
核心思想：在递归步骤中重用共享层堆栈而非使用独立层
参数共享策略：
Cycle 策略
Sequence 策略
Middle-Cycle 策略（最优）
Middle-Sequence 策略
Middle-Cycle 策略：保留第一层和最后一层的独特参数，中间层循环使用
效率优势：
分布式训练中相同参数在递归步骤中重复使用
连续深度批处理消除空闲期，提高吞吐量
2.3 自适应令牌级计算#
MoR 的核心创新在于动态分配计算资源：
路由机制：
轻量级路由器评估每个令牌的复杂度
为每个令牌分配最优递归深度
路由策略比较：
策略
优点
挑战
专家选择路由
完美负载均衡，固定计算预算
信息泄露，需辅助损失
令牌选择路由
无信息泄露
需平衡损失解决负载不均
计算优化：
仅对活跃令牌执行二次注意力计算
减少 25%训练 FLOPs
2.4 高效 KV 缓存#
针对动态深度模型的 KV 缓存挑战：
递归感知 KV 缓存：
仅缓存路由到特定递归步骤的令牌键值对
显著减少 KV 缓存大小和 IO 需求
递归 KV 共享：
在首个递归步骤缓存 KV 对并在后续递归中重用
减少预填充延迟
作为正则化技术提升性能
3. 实验与结果#
3.1 实验设置#
模型配置：
参数量：120M-1.3B
训练 FLOPs：16.5e18（与基线相同）
数据集：C4, Pile 等大规模文本语料
评估指标：
验证困惑度
少样本准确率（11 项任务平均）
训练时间/内存消耗
推理吞吐量
3.2 性能提升#
准确率与困惑度：
MoR 平均少样本准确率：43.1% (vs 基线 42.3%)
验证困惑度降低 5-8%
训练效率：
训练时间减少19%
峰值内存使用降低25%
模型可扩展性：
所有尺寸下优于递归基线
360M 参数模型超越 Vanilla Transformer
参数量仅为 Vanilla 的约三分之一
3.3 推理吞吐量优化#
连续深度批处理：
利用参数共享架构优势
立即填充已完成序列的空闲槽位
保持高 GPU 利用率
早期退出机制消除计算"气泡"
吞吐量提升：
MoR-4 最大批次加速：2.06 倍
所有 MoR 变体优于 Vanilla 基线
增加递归深度进一步减少 KV 缓存使用
3.4 深度分析#
令牌语义重要性：
内容丰富的令牌("People"、"Drugs")经历更多递归步骤
功能词("and"、"---")经历较少步骤
测试时扩展：
增加递归深度提高生成质量
更深递归专注于细化令牌表示
计算最优扩展：
在 isoFLOPs 约束下，优先增加模型尺寸而非训练长度
更大模型容量带来更显著性能提升
4. 讨论#
4.1 设计选择分析#
递归深度分配：
复杂令牌自动获得更多计算资源
简单令牌提前退出优化计算效率
路由机制选择：
令牌选择路由更适用于无信息泄露场景
专家选择路由在负载均衡场景表现更佳
参数共享策略：
Middle-Cycle 策略平衡参数效率与性能
4.2 与传统方法对比#
方法 |
参数量 |
自适应计算 |
KV 缓存效率 |
训练速度 |
|---|---|---|---|---|
Vanilla Transformer |
高 |
无 |
低 |
基准 |
静态递归模型 |
中 |
无 |
中 |
提高 10% |
专家混合 |
高 |
有 |
低 |
降低 15% |
MoR |
低 |
有 |
高 |
提高 19% |
5. 结论与未来工作#
5.1 主要贡献#
提出首个统一参数效率与自适应计算的 Transformer 框架
设计轻量级路由器实现令牌级动态计算分配
开发递归感知 KV 缓存策略解决内存瓶颈
实验验证 MoR 在效率与性能上的双重优势
5.2 未来研究方向#
推理优化：探索路由器如何适应思维链需求
规模扩展：将 MoR 应用于 3B+参数模型
自适应容量控制：开发更精细的模型设计策略
稀疏算法兼容：结合剪枝与量化技术
多模态扩展：应用于视觉、语音等领域
6. 参考文献#
Dehghani et al. "Universal Transformers", ICLR 2019
Fedus et al. "Switch Transformers", JMLR 2022
Clark et al. "ELECTRA", ICLR 2020
Lepikhin et al. "GShard", arXiv 2020
Roller et al. "Hash Layers for Large Sarchines", NeurIPS 2021
结论：Mixture-of-Recursions (MoR) 通过创新的统一框架，在参数效率、自适应计算和内存优化三个维度实现了突破性进展。实验证明，MoR 能够以更小的模型尺寸、更低的计算成本实现优于传统架构的性能，为大型语言模型的高效训练和部署提供了新的范式。随着未来研究的深入，MoR 有望成为实现高效大型语言模型的主流架构。