arXiv:2504.10449cs.LG2025-04被引 21

用Mamba架构实现高效推理,突破传统模型算力瓶颈。

M1: Towards Scalable Test-Time Compute with Mamba Reasoning Models

  • 基于Mamba的线性RNN结构,支持长序列推理且内存高效。
  • 在AIME和MATH上性能媲美顶尖蒸馏Transformer模型。
  • 生成速度超同规模Transformer三倍,适合高并发推理场景。

有效推理对解决复杂数学问题至关重要。近期大语言模型通过扩展测试时计算量的长链思维推理显著提升了性能,但基于Transformer的模型因二次计算复杂度和线性内存需求,难以扩展上下文长度。本文提出一种新型混合线性RNN推理模型M1,基于Mamba架构,实现高效推理。该方法通过从现有推理模型中进行知识蒸馏,并进一步结合强化学习训练。在AIME与MATH基准上的实验表明,M1不仅优于此前线性RNN模型,还达到了与同规模深度求索R1蒸馏推理模型相当的性能。我们还将生成速度与高性能通用推理引擎vLLM对比,发现其相比同规模Transformer提升超过3倍。在固定生成时间预算下,利用自一致性投票策略,M1可实现比DeepSeek R1蒸馏Transformer更高的准确率。综上,本文提出一种混合Mamba推理模型,为通过自一致性或长链思维推理实现测试时生成的可扩展性提供了更优方案。

原文摘要 · Abstract (English)

Effective reasoning is crucial to solving complex mathematical problems. Recent large language models (LLMs) have boosted performance by scaling test-time computation through long chain-of-thought reasoning. However, transformer-based models are inherently limited in extending context length due to their quadratic computational complexity and linear memory requirements. In this paper, we introduce a novel hybrid linear RNN reasoning model, M1, built on the Mamba architecture, which allows memory-efficient inference. Our approach leverages a distillation process from existing reasoning models and is further enhanced through RL training. Experimental results on the AIME and MATH benchmarks show that M1 not only outperforms previous linear RNN models but also matches the performance of state-of-the-art Deepseek R1 distilled reasoning models at a similar scale. We also compare our generation speed with a highly performant general purpose inference engine, vLLM, and observe more than a 3x speedup compared to a same size transformer. With throughput speedup, we are able to achieve higher accuracy compared to DeepSeek R1 distilled transformer reasoning models under a fixed generation time budget using self-consistency voting. Overall, we introduce a hybrid Mamba reasoning model and provide a more effective approach to scaling test-time generation using self-consistency or long chain of thought reasoning.

推理模型Mamba测试时计算数学推理

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。