M1: Rumo à Computação Escalável em Tempo de Teste com Modelos de Raciocínio 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

Resumo

O raciocínio eficaz é crucial para resolver problemas matemáticos complexos. Modelos de linguagem de grande escala (LLMs) recentes têm impulsionado o desempenho ao escalar a computação em tempo de teste por meio de longas cadeias de raciocínio. No entanto, modelos baseados em transformers são intrinsecamente limitados na extensão do comprimento do contexto devido à sua complexidade computacional quadrática e requisitos lineares de memória. Neste artigo, introduzimos um novo modelo híbrido de raciocínio RNN linear, M1, construído sobre a arquitetura Mamba, que permite inferência eficiente em termos de memória. Nossa abordagem aproveita um processo de destilação de modelos de raciocínio existentes e é ainda aprimorada por meio de treinamento com RL. Resultados experimentais nos benchmarks AIME e MATH mostram que o M1 não apenas supera modelos anteriores de RNN linear, mas também iguala o desempenho dos modelos de raciocínio destilados Deepseek R1 de última geração em uma escala similar. Também comparamos nossa velocidade de geração com um motor de inferência de propósito geral altamente performático, vLLM, e observamos um aumento de mais de 3x em comparação com um transformer de mesmo tamanho. Com o aumento da taxa de transferência, conseguimos alcançar maior precisão em comparação com os modelos de raciocínio transformer destilados DeepSeek R1 sob um tempo fixo de geração usando votação de autoconfirmação. No geral, introduzimos um modelo de raciocínio híbrido Mamba e fornecemos uma abordagem mais eficaz para escalar a geração em tempo de teste usando autoconfirmação ou longas cadeias de raciocínio.

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.