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

要約

複雑な数学的問題を解決するためには、効果的な推論が重要です。
最近の大規模な言語モデル（LLMS）は、長い考え方の推論を通じてテスト時間計算をスケーリングすることにより、パフォーマンスを向上させました。
ただし、トランスベースのモデルは、2次計算の複雑さと線形メモリ要件により、コンテキストの長さの拡張が本質的に制限されています。
この論文では、メモリ効率の高い推論を可能にするMambaアーキテクチャに基づいて構築された、新しいハイブリッド線形RNN推論モデルM1を紹介します。
当社のアプローチは、既存の推論モデルから蒸留プロセスを活用し、RLトレーニングを通じてさらに強化されています。
AIMEおよび数学のベンチマークでの実験結果は、M1が以前の線形RNNモデルよりも優れているだけでなく、同様のスケールで最先端のDeepSeek R1蒸留推論モデルのパフォーマンスにも一致することを示しています。
また、生成速度を非常にパフォーマンスの高い汎用推論エンジンVLLMと比較し、同じサイズの変圧器と比較して3倍以上のスピードアップを観察します。
スループットスピードアップにより、自己整合性の投票を使用して、固定世代の時間予算でDeepSeek R1蒸留トランス推論モデルと比較して、より高い精度を達成することができます。
全体として、ハイブリッドマンバ推論モデルを導入し、自己整合性または長い思考推論を使用して、テスト時間生成をスケーリングするためのより効果的なアプローチを提供します。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google