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

有效的推理能力对于解决复杂的数学问题是至关重要的。近期的大型语言模型（LLMs）通过扩展测试时的计算，利用长链推理显著提升了性能。然而，基于Transformer的模型由于其计算复杂度为二次方且内存需求为线性增长，在扩展上下文长度方面存在固有限制。在本文中，我们提出了一种新型的混合线性RNN推理模型M1，该模型基于Mamba架构构建，能够实现高效的内存推理。我们的方法通过从现有的推理模型中提取知识进行蒸馏，并进一步通过强化学习（RL）训练加以改进。在AIME和MATH基准上的实验结果表明，M1不仅超越了之前的线性RNN模型，而且在相同规模下与最先进的Deepseek R1蒸馏推理模型表现相当。此外，我们将生成速度与高性能通用推理引擎vLLM进行了比较，发现相较于同规模的Transformer模型，M1的速度提升了超过3倍。借助这种吞吐量加速，我们在固定生成时间预算内，通过自一致性投票机制，相比DeepSeek R1蒸馏Transformer推理模型实现了更高的准确性。总体而言，我们引入了一种混合Mamba推理模型，并提供了一种更有效的方法来扩展测试时生成能力，无论是通过自一致性还是长链推理方式。

该论文试图解决大型语言模型在扩展推理上下文长度时面临的计算复杂度和内存限制问题。这是一个现有问题，但通过引入新的模型架构，论文试图提供一种更高效的解决方案。

论文提出了一种基于Mamba架构的混合线性RNN推理模型（M1），结合了蒸馏过程和强化学习训练。相比传统的Transformer模型，M1通过降低计算复杂度和内存需求，实现了更高效的推理能力，同时支持长链推理任务。这种方法在不牺牲性能的前提下显著提高了生成速度。

实验设计包括在AIME和MATH基准上的测试，并与Deepseek R1等最先进的模型进行了比较。结果显示M1不仅在性能上匹配甚至超越了现有模型，还在固定生成时间预算下通过自一致性投票进一步提升了准确性。此外，M1在生成速度上比同规模的Transformer快3倍以上。论文未提及代码是否开源，但其高效推理方法值得深入研究，特别是在资源受限环境中的应用。

近期相关研究包括Deepseek团队发布的R1推理模型、vLLM高性能推理引擎以及利用链式思维推理的研究。例如，《Chain of Thought Prompting Elicits Reasoning in Large Language Models》探讨了如何通过提示工程提升推理能力，而《Long-Context Transformers》则研究了扩展Transformer上下文长度的技术。