HyperAI超神经

M1:通过Mamba推理模型实现可扩展的测试时计算

Junxiong Wang, Wen-Ding Li, Daniele Paliotta, Daniel Ritter, Alexander M. Rush, Tri Dao
发布日期: 4/16/2025
M1:通过Mamba推理模型实现可扩展的测试时计算
摘要

有效的推理对于解决复杂的数学问题至关重要。近期,大规模语言模型(LLMs)通过扩展测试时的计算,采用长链推理方法显著提升了性能。然而,基于变压器的模型由于其计算复杂度呈二次增长和内存需求呈线性增长,本质上在扩展上下文长度方面存在局限性。在本文中,我们介绍了一种新型的混合线性RNN推理模型M1,该模型基于Mamba架构构建,能够实现高效的内存推理。我们的方法利用了从现有推理模型中提取知识的过程,并通过强化学习(RL)训练进一步增强。在AIME和MATH基准测试中的实验结果表明,M1不仅超越了先前的线性RNN模型,而且在相似规模下与最先进的Deepseek R1蒸馏推理模型的性能相当。我们还将其生成速度与高性能通用推理引擎vLLM进行了比较,发现在相同规模的变压器模型中,M1的速度提高了3倍以上。借助于吞吐量加速,在固定生成时间预算下使用自一致性投票方法时,我们能够实现比DeepSeek R1蒸馏变压器推理模型更高的准确率。总体而言,我们引入了一种混合Mamba推理模型,并提供了一种更有效的测试时生成扩展方法,无论是通过自一致性还是长链推理。