M1: Verso un calcolo scalabile al momento del test con modelli di ragionamento Mamba

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

April 14, 2025

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

cs.AI

Abstract

Il ragionamento efficace è cruciale per risolvere problemi matematici complessi. I recenti modelli linguistici di grandi dimensioni (LLM) hanno migliorato le prestazioni aumentando il calcolo al momento del test attraverso lunghi processi di ragionamento a catena di pensiero. Tuttavia, i modelli basati su transformer sono intrinsecamente limitati nell'estendere la lunghezza del contesto a causa della loro complessità computazionale quadratica e dei requisiti di memoria lineari. In questo articolo, introduciamo un nuovo modello ibrido di ragionamento RNN lineare, M1, basato sull'architettura Mamba, che consente un'inferenza efficiente in termini di memoria. Il nostro approccio sfrutta un processo di distillazione da modelli di ragionamento esistenti ed è ulteriormente potenziato attraverso l'addestramento RL. I risultati sperimentali sui benchmark AIME e MATH mostrano che M1 non solo supera i precedenti modelli RNN lineari, ma eguaglia anche le prestazioni dei modelli di ragionamento distillati Deepseek R1 all'avanguardia su scala simile. Confrontiamo inoltre la nostra velocità di generazione con un motore di inferenza generico ad alte prestazioni, vLLM, e osserviamo un'accelerazione di oltre 3x rispetto a un transformer delle stesse dimensioni. Con l'aumento della velocità di elaborazione, siamo in grado di ottenere una maggiore accuratezza rispetto ai modelli di ragionamento transformer distillati DeepSeek R1 con un tempo di generazione fisso utilizzando il voto di autoconsistenza. In sintesi, introduciamo un modello di ragionamento ibrido Mamba e forniamo un approccio più efficace per scalare la generazione al momento del test utilizzando l'autoconsistenza o lunghi processi di ragionamento a catena di pensiero.

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.