M1: К масштабируемым вычислениям во время тестирования с помощью моделей рассуждений Mamba

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

April 14, 2025

Авторы: Junxiong Wang, Wen-Ding Li, Daniele Paliotta, Daniel Ritter, Alexander M. Rush, Tri Dao

cs.AI

Аннотация

Эффективное рассуждение играет ключевую роль в решении сложных математических задач. Недавние крупные языковые модели (LLM) значительно улучшили производительность за счет масштабирования вычислений во время тестирования с использованием длинных цепочек рассуждений. Однако трансформерные модели имеют фундаментальные ограничения в увеличении длины контекста из-за их квадратичной вычислительной сложности и линейных требований к памяти. В данной статье мы представляем новую гибридную линейную модель рассуждений на основе RNN, M1, построенную на архитектуре Mamba, которая обеспечивает эффективное использование памяти при выводе. Наш подход использует процесс дистилляции из существующих моделей рассуждений и дополнительно улучшается за счет обучения с подкреплением (RL). Экспериментальные результаты на бенчмарках AIME и MATH показывают, что M1 не только превосходит предыдущие линейные RNN-модели, но и соответствует производительности современных дистиллированных моделей рассуждений Deepseek R1 аналогичного масштаба. Мы также сравниваем скорость генерации нашей модели с высокопроизводительным универсальным движком вывода vLLM и наблюдаем ускорение более чем в 3 раза по сравнению с трансформером того же размера. Благодаря увеличению пропускной способности, мы достигаем более высокой точности по сравнению с дистиллированными трансформерными моделями рассуждений DeepSeek R1 при фиксированном бюджете времени генерации, используя голосование с самосогласованностью. В целом, мы представляем гибридную модель рассуждений на основе Mamba и предлагаем более эффективный подход к масштабированию генерации во время тестирования с использованием самосогласованности или длинных цепочек рассуждений.

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.