M1: Skalierbare Testzeit-Berechnungen mit Mamba Reasoning Models

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

April 14, 2025

Autoren: Junxiong Wang, Wen-Ding Li, Daniele Paliotta, Daniel Ritter, Alexander M. Rush, Tri Dao

cs.AI

Zusammenfassung

Effektives Denken ist entscheidend für die Lösung komplexer mathematischer Probleme. Aktuelle große Sprachmodelle (LLMs) haben die Leistung durch die Skalierung der Berechnungen zur Testzeit mittels langer Ketten von Gedankengängen gesteigert. Allerdings sind transformerbasierte Modelle aufgrund ihrer quadratischen Rechenkomplexität und linearen Speicheranforderungen inhärent in der Erweiterung der Kontextlänge begrenzt. In diesem Artikel stellen wir ein neuartiges hybrides lineares RNN-Denkmodell, M1, vor, das auf der Mamba-Architektur basiert und speichereffiziente Inferenz ermöglicht. Unser Ansatz nutzt einen Destillationsprozess bestehender Denkmodelle und wird durch RL-Training weiter verbessert. Experimentelle Ergebnisse auf den AIME- und MATH-Benchmarks zeigen, dass M1 nicht nur frühere lineare RNN-Modelle übertrifft, sondern auch die Leistung der state-of-the-art Deepseek R1 destillierten Denkmodelle in ähnlichem Maßstab erreicht. Wir vergleichen auch unsere Generierungsgeschwindigkeit mit einer hochleistungsfähigen allgemeinen Inferenz-Engine, vLLM, und beobachten eine mehr als dreifache Beschleunigung im Vergleich zu einem Transformer gleicher Größe. Mit dem Durchsatzgewinn können wir unter einem festen Generierungszeitbudget mithilfe von Selbstkonsistenz-Abstimmungen eine höhere Genauigkeit im Vergleich zu den DeepSeek R1 destillierten Transformer-Denkmodellen erreichen. Insgesamt führen wir ein hybrides Mamba-Denkmodell ein und bieten einen effektiveren Ansatz zur Skalierung der Generierung zur Testzeit mithilfe von Selbstkonsistenz oder langen Ketten von Gedankengängen.

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.