M1:面向可扩展测试时间计算的 Mamba 推理模型
机器学习
2025-09-10 v3
摘要
有效的推理是解决复杂数学问题的关键。最近的大型语言模型 (LLM) 通过扩展测试时间计算(通过长链式思考推理)提升了性能。然而,基于 Transformer 的模型因其二次计算复杂度和线性内存需求,在扩展上下文长度方面存在固有限制。本文引入一种新型混合线性 RNN 推理模型 M1,基于 Mamba 架构构建,实现内存高效推理。我们的方法通过从现有推理模型蒸馏过程获得优势,并通过强化学习训练进一步提升。AIME 和 MATH 基准测试结果表明,M1 不仅超越了以往的线性 RNN 模型,还在规模相似的情况下与最新 Deepseek R1 蒸馏推理模型保持一致的性能。我们还与高度性能通用推理引擎 vLLM 进行生成速度比较,发现与相同规模的 Transformer 模型相比,速度提升超过 3 倍。通过吞吞量提升,我们能够在固定生成时间预算下使用自洽投票获得更高的准确率。总体而言,我们引入一种混合 Mamba 推理模型,为使用自洽或长链式思考推理扩展测试时间生成提供了更有效的方法。
引用
@article{arxiv.2504.10449,
title = {M1: Towards Scalable Test-Time Compute with Mamba Reasoning Models},
author = {Junxiong Wang and Wen-Ding Li and Daniele Paliotta and Daniel Ritter and Alexander M. Rush and Tri Dao},
journal= {arXiv preprint arXiv:2504.10449},
year = {2025}
}
备注
Code is available https://github.com/jxiw/M1