谷歌与KAIST联手推出MoR架构：用更低内存实现更快推理，能否成为Transformer的新对手？

谷歌 DeepMind 与韩国科学技术院（KAIST）的研究人员近日联合发布了一种名为“Mixture-of-Recursions”（MoR）的新型语言模型架构。这一新架构旨在降低大型语言模型的计算和内存需求，同时保持模型的高性能。传统的 Transformer 架构在处理大规模数据时面临高昂的计算和内存资源负担，MoR 通过一系列创新技术解决了这些问题。 MoR 的核心创新点在于结合了递归计算和动态路由机制，在一个统一框架内实现了参数共享、计算优化和内存管理的多目标优化。具体来说，传统的 Transformer 模型中，每个输入文本的 token 都要经过相同数量的计算层处理，而 MoR 允许不同的 token 根据自身复杂度接受不同深度的处理。这种方式不仅提高了处理效率，还能更好地分配计算资源。 在参数共享方面，MoR 采用了“Middle-Cycle”策略，即只在模型的中间层之间共享权重，保持第一层和最后一层使用独立参数。这种策略在参数效率和模型表达能力之间取得了较好的平衡。与此同时，MoR 还提出了两种新的 KV 缓存管理策略：“递归式缓存”和“递归共享”。前者只为需要特定递归步骤的 token 存储 KV 数据，从而减少内存占用和数据读写量；后者则利用所有 token 都经过第一个递归块的特点，只在第一步缓存 KV 数据，供后续递归步骤重复使用，最大程度地节省内存。 研究团队在不同规模的模型上进行了广泛测试。结果显示，MoR 模型在参数量比基准 Transformer 模型少近一半的情况下，仍能提升少样本学习任务中的准确率，从 42.3% 提高到 43.1%。此外，MoR 的高效计算使其在相同计算预算内能够处理更多训练数据，进而提升模型的整体性能。在一个固定训练数据量的对比实验中，MoR 模型在使用少 25% 计算量的情况下，性能依然超过基准模型，同时训练时间减少了 19%，峰值内存使用下降了 25%。 在推理性能方面，MoR 的优势尤为突出。它采用了连续深度批处理技术，可以将处于不同计算阶段的 token 组合到同一批次中进行处理。这一技术与模型的早期退出机制相结合，大幅提升了处理吞吐量。以 3.6 亿参数规模的模型为例，在特定配置下，MoR 实现了高达 2.06 倍的推理加速。 研究还观察到，MoR 模型在处理不同类型 token 时表现出与语义重要性相关的分配模式。内容丰富的重要 token 会获得更多的递归次数（如三次），而功能性词汇则只需较少的递归次数。这说明模型能够智能地分配计算资源，更好地理解文本内容的内在意义。 尽管 MoR 在实验中展现出巨大的潜力，但是否能彻底取代现有的 Transformer 架构仍有待进一步观察。MoR 的技术创新为未来的语言模型设计提供了多维度优化的思路，尤其是对于降低大型语言模型的部署和使用成本具有重要意义。 业内专家对 MoR 的发布给予了高度关注，认为这是 AI 领域效率优化的重要突破。谷歌 DeepMind 一直以来都是语言模型研究的前沿机构，MoR 的发展也是基于其先前在 MoD 等领域的研究成果。未来，MoR 可能会成为推动大语言模型应用普及的关键技术之一。