M1 : Vers un calcul évolutif au moment du test avec les modèles de raisonnement Mamba

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

Résumé

Un raisonnement efficace est crucial pour résoudre des problèmes mathématiques complexes. Les récents modèles de langage de grande taille (LLMs) ont amélioré les performances en augmentant le calcul au moment du test grâce à un raisonnement en chaîne de pensée longue. Cependant, les modèles basés sur les transformers sont intrinsèquement limités dans l'extension de la longueur du contexte en raison de leur complexité computationnelle quadratique et de leurs besoins en mémoire linéaire. Dans cet article, nous introduisons un nouveau modèle de raisonnement hybride RNN linéaire, M1, construit sur l'architecture Mamba, qui permet une inférence efficace en mémoire. Notre approche tire parti d'un processus de distillation à partir de modèles de raisonnement existants et est encore améliorée par un entraînement par renforcement (RL). Les résultats expérimentaux sur les benchmarks AIME et MATH montrent que M1 surpasse non seulement les précédents modèles RNN linéaires, mais égalise également les performances des modèles de raisonnement distillés Deepseek R1 de pointe à une échelle similaire. Nous comparons également notre vitesse de génération avec un moteur d'inférence généraliste très performant, vLLM, et observons une accélération de plus de 3x par rapport à un transformer de même taille. Avec cette accélération du débit, nous sommes en mesure d'atteindre une précision plus élevée par rapport aux modèles de raisonnement transformer distillés DeepSeek R1 sous un budget de temps de génération fixe en utilisant un vote d'auto-cohérence. Globalement, nous introduisons un modèle de raisonnement hybride Mamba et proposons une approche plus efficace pour mettre à l'échelle la génération au moment du test en utilisant l'auto-cohérence ou un raisonnement en chaîne de pensée longue.

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.