M1: 使用Mamba推理模型迈向可扩展的测试时计算

有效的推理对于解决复杂的数学问题至关重要。最近的大型语言模型 (LLM) 通过长链思考推理扩展测试时计算，从而提高了性能。然而，基于 Transformer 的模型由于其二次计算复杂度和线性内存需求，在扩展上下文长度方面受到固有的限制。在本文中，我们介绍了一种新颖的混合线性 RNN 推理模型 M1，它建立在 Mamba 架构之上，可以实现内存高效的推理。我们的方法利用了来自现有推理模型的蒸馏过程，并通过 RL 训练进一步增强。在 AIME 和 MATH 基准测试上的实验结果表明，M1 不仅优于以前的线性 RNN 模型，而且在相似规模下也与最先进的 Deepseek R1 蒸馏推理模型的性能相匹配。我们还将我们的生成速度与高性能通用推理引擎 vLLM 进行了比较，并观察到比相同大小的 Transformer 快 3 倍以上。通过吞吐量加速，我们能够在固定的生成时间预算下，使用自洽性投票获得比 DeepSeek R1 蒸馏 Transformer 推理模型更高的准确率。总的来说，我们介绍了一种混合 Mamba 推理模型，并提供了一种更有效的方法，使用自洽性或长链思考推理来扩展测试时生成。