文章摘要
【关 键 词】 MoR、Transformer、推理速度、内存优化、LLM
KAIST、Mila和谷歌DeepMind团队近期发布了一种名为Mixture-of-Recursions（MoR）的全新大型语言模型（LLM）架构，该架构被认为有潜力成为Transformer的替代者。MoR在推理速度上提升了2倍，训练FLOPs减少，KV缓存内存直接减半，显著优化了模型性能。在135M到1.7B的参数规模下，MoR展示了新的帕累托前沿，即在相同的训练FLOPs下，困惑度更低、小样本准确率更高，并且吞吐量提升超过2倍，全面超越了传统的Transformer架构。
MoR的核心创新在于其递归混合体设计，通过小型路由器为每个token的隐藏状态打分，仅高分token继续循环，其余则提前退出。这种设计使得模型在处理文本时能够根据token的复杂性进行动态调整，避免了传统Transformer对所有token一视同仁的处理方式，从而显著提升了计算效率。此外，MoR还采用了共享块设计，每个token最多循环4次，进一步减少了内存和算力的消耗。
在预训练和推理过程中，MoR通过路由机制和KV缓存策略实现了动态调整递归步骤。路由机制包括专家选择路由和token选择路由，前者通过分层筛选机制模拟早退出机制，后者则在一开始就决定每个token的递归深度。KV缓存策略则包括按递归层缓存和跨递归共享，前者实现了局部化计算，显著提升了内存使用效率，后者则通过重复使用KV对减少了计算量。
实验结果表明，在相同训练计算预算下，MoR以更少参数优于基线模型，并且在相同数据量下，MoR用更少计算量仍优于基线模型。此外，MoR在推理吞吐量评估中也表现出色，通过连续深度批处理技术显著提升了推理速度。例如，在最大批设置下，MoR-4速度可提升2.06倍。
总体而言，MoR架构展示了其在模型规模和计算量增长时的良好可扩展性和高参数效率，有望成为未来LLM架构的重要发展方向。谷歌DeepMind通过这一研究，展示了在AI模型设计中，高效分工调度和节省算力的重要性，为未来的AI发展提供了新的思路。
原文和模型
【原文链接】 阅读原文 [ 2819字 | 12分钟 ]
【原文作者】 新智元
【摘要模型】 deepseek-v3
【摘要评分】 ★★★★★
相关文章
