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

• URL: http://arxiv.org/abs/2504.10449v1
• Date: Mon, 14 Apr 2025 17:38:25 GMT
• Title: M1: Towards Scalable Test-Time Compute with Mamba Reasoning Models
• Authors: Junxiong Wang, Wen-Ding Li, Daniele Paliotta, Daniel Ritter, Alexander M. Rush, Tri Dao,
• Abstract summary: We introduce a novel hybrid linear RNN reasoning model, M1, built on the Mamba architecture.<n> Experimental results 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.
• Score: 72.75501495786297
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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.

Related papers

• Thinking Slow, Fast: Scaling Inference Compute with Distilled Reasoners [72.37408197157453]
  Recent advancements have demonstrated that the performance of large language models (LLMs) can be significantly enhanced by scaling computational resources at test time.<n>This raises a fundamental question: can models with lower complexity leverage their superior generation throughput to outperform similarly sized Transformers for a fixed computational budget?<n>To address this question and overcome the lack of strong subquadratic reasoners, we distill pure and hybrid Mamba models from pretrained Transformers.
  arXiv  Detail & Related papers  (2025-02-27T18:08:16Z)
• Towards Thinking-Optimal Scaling of Test-Time Compute for LLM Reasoning [113.49074603075032]
  Recent studies have shown that making a model spend more time thinking through longer Chain of Thoughts (CoTs) enables it to gain significant improvements in complex reasoning tasks. We explore whether scaling with longer CoTs can indeed impair the reasoning performance of Large Language Models (LLMs) in certain domains.
  arXiv  Detail & Related papers  (2025-02-25T10:48:05Z)
• Revisiting the Test-Time Scaling of o1-like Models: Do they Truly Possess Test-Time Scaling Capabilities? [61.85289698610747]
  We study whether o1-like large language models (LLMs) truly possess test-time scaling capabilities.<n>We find that longer CoTs of these o1-like models do not consistently enhance accuracy.<n>We propose Shortest Majority Vote, a method that combines parallel scaling strategies with CoT length characteristics.
  arXiv  Detail & Related papers  (2025-02-17T07:21:11Z)
• O1-Pruner: Length-Harmonizing Fine-Tuning for O1-Like Reasoning Pruning [98.3430004984531]
  We propose Length-Harmonizing Fine-Tuning (O1-Pruner) to minimize reasoning overhead while maintaining accuracy. Our code is coming soon at https://github.com/StarDewXXX/O1-Pruner.
  arXiv  Detail & Related papers  (2025-01-22T01:35:11Z)
• The Mamba in the Llama: Distilling and Accelerating Hybrid Models [76.64055251296548]
  We show how to distill large Transformers into linear RNNs by reusing the linear projection weights from attention layers with academic GPU resources.<n>The resulting hybrid model achieves performance comparable to the original Transformer in chat benchmarks.<n>We also introduce a hardware-aware speculative decoding algorithm that accelerates the inference speed of Mamba and hybrid models.
  arXiv  Detail & Related papers  (2024-08-27T17:56:11Z)
• Mamba: Linear-Time Sequence Modeling with Selective State Spaces [31.985243136674146]
  Foundation models are almost universally based on the Transformer architecture and its core attention module. We identify that a key weakness of such models is their inability to perform content-based reasoning. We integrate these selective SSMs into a simplified end-to-end neural network architecture without attention or even blocks (Mamba) As a general sequence model backbone, Mamba achieves state-of-the-art performance across several modalities such as language, audio, and genomics.
  arXiv  Detail & Related papers  (2023-12-01T18:01:34Z)
• A Data-driven feature selection and machine-learning model benchmark for the prediction of longitudinal dispersion coefficient [29.58577229101903]
  An accurate prediction on Longitudinal Dispersion(LD) coefficient can produce a performance leap in related simulation. In this study, a global optimal feature set was proposed through numerical comparison of the distilled local optimums in performance with representative ML models. Results show that the support vector machine has significantly better performance than other models.
  arXiv  Detail & Related papers  (2021-07-16T09:50:38Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.