M1: Hacia un cómputo escalable en tiempo de prueba con modelos de razonamiento Mamba

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

April 14, 2025

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

cs.AI

Resumen

El razonamiento efectivo es crucial para resolver problemas matemáticos complejos. Los modelos de lenguaje de gran escala (LLMs) recientes han mejorado el rendimiento al escalar el cómputo en tiempo de prueba mediante cadenas de razonamiento largas. Sin embargo, los modelos basados en transformadores están inherentemente limitados en la extensión de la longitud del contexto debido a su complejidad computacional cuadrática y requisitos de memoria lineal. En este artículo, presentamos un nuevo modelo híbrido de razonamiento RNN lineal, M1, construido sobre la arquitectura Mamba, que permite una inferencia eficiente en memoria. Nuestro enfoque aprovecha un proceso de destilación de modelos de razonamiento existentes y se mejora aún más mediante entrenamiento con RL. Los resultados experimentales en los benchmarks AIME y MATH muestran que M1 no solo supera a los modelos RNN lineales anteriores, sino que también iguala el rendimiento de los modelos de razonamiento destilados Deepseek R1 de última generación a una escala similar. También comparamos nuestra velocidad de generación con un motor de inferencia de propósito general de alto rendimiento, vLLM, y observamos una aceleración de más de 3x en comparación con un transformador del mismo tamaño. Con la aceleración del rendimiento, logramos una mayor precisión en comparación con los modelos de razonamiento transformadores destilados DeepSeek R1 bajo un tiempo fijo de generación utilizando votación de autoconsistencia. En general, introducimos un modelo de razonamiento híbrido Mamba y proporcionamos un enfoque más efectivo para escalar la generación en tiempo de prueba utilizando autoconsistencia o cadenas largas de razonamiento.

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.