M1: Naar schaalbare rekencapaciteit tijdens testen met Mamba-redeneermodellen

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

April 14, 2025

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

cs.AI

Samenvatting

Effectief redeneren is cruciaal voor het oplossen van complexe wiskundige problemen. Recente grote taalmodellen (LLMs) hebben de prestaties verbeterd door het schalen van testtijdberekeningen via lange ketens van gedachtegangen. Echter, transformer-gebaseerde modellen zijn inherent beperkt in het uitbreiden van de contextlengte vanwege hun kwadratische rekencomplexiteit en lineaire geheugenvereisten. In dit artikel introduceren we een nieuw hybride lineair RNN-redeneermodel, M1, gebouwd op de Mamba-architectuur, dat geheugenefficiënte inferentie mogelijk maakt. Onze aanpak maakt gebruik van een distillatieproces van bestaande redeneermodellen en wordt verder verbeterd door RL-training. Experimentele resultaten op de AIME- en MATH-benchmarks laten zien dat M1 niet alleen beter presteert dan eerdere lineaire RNN-modellen, maar ook de prestaties evenaart van state-of-the-art Deepseek R1 gedistilleerde redeneermodellen op een vergelijkbare schaal. We vergelijken ook onze generatiesnelheid met een zeer performante algemene inferentie-engine, vLLM, en observeren een meer dan 3x versnelling in vergelijking met een transformer van dezelfde grootte. Met deze doorvoersnelheidsverbetering kunnen we een hogere nauwkeurigheid bereiken in vergelijking met DeepSeek R1 gedistilleerde transformer-redeneermodellen onder een vast generatietijdbudget door gebruik te maken van zelfconsistentie-stemming. Over het geheel introduceren we een hybride Mamba-redeneermodel en bieden we een effectievere aanpak voor het schalen van testtijdgeneratie met behulp van zelfconsistentie of lange ketens van gedachtegangen.

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.