2025年5月,小米LLM核心团队发布了一项突破性研究,展示了他们全新的大语言模型MiMo-7B如何在数学和编程推理任务上取得卓越成绩。这项研究发表于arXiv(arXiv:2505.07608v1),论文链接为
https://github.com/xiaomimimo/MiMo。想象一下,如果你正在尝试教一个聪明的学生解决复杂的数学问题或编写高效的计算机程序。你会怎么做?你可能会从基础知识开始,然后逐步引导他们掌握更复杂的思维方式,最后通过反复练习和反馈来强化他们的能力。小米的研究团队正是采用了类似的方法来打造他们的MiMo-7B大语言模型,使它在复杂推理任务上表现出色。
一、为何MiMo-7B值得关注?
在人工智能领域,我们经常看到像OpenAI、DeepSeek和Anthropic这样的公司推出强大的推理型大语言模型,它们能够解决复杂的数学问题和生成高质量的代码。然而,这些模型通常体积庞大,参数量动辄达到320亿甚至更多。想象一下,这就像需要一台超级计算机才能运行的庞大软件。
而小米的研究人员提出了一个有趣的问题:我们能否创造一个体积小得多但推理能力同样出色的模型?他们的答案是肯定的。MiMo-7B只有70亿参数(相比之下,其他顶级推理模型可能是它的四倍大小),却在多项推理任务上超越了那些更大的模型,包括OpenAI的o1-mini。
这就像是在赛车世界中,一辆精心调校的紧凑型跑车击败了大排量的超级跑车。小米团队通过精心设计模型的"成长过程",实现了这一目标 - 从模型的"童年"(预训练阶段)到"青少年期"(后训练阶段)都进行了全面优化。
二、预训练:培养模型的基础思维能力
想象一下你正在教一个孩子阅读和理解世界。你会给他们提供各种各样的书籍、故事和知识,帮助他们建立广泛的知识基础。小米团队在模型预训练阶段就是这样做的,但他们采取了一些特别的方法。
更好的学习材料
首先,研究团队优化了数据预处理流程。就像为孩子挑选高质量的教材一样,他们改进了文本提取工具,特别注重保留网页中的数学公式和代码片段。传统的网页抓取工具通常会丢失这些重要的推理信息,就像从教科书中撕掉了所有的方程式和例题一样。
研究人员还开发了一个能在一天内完成的全局数据去重系统,这就像快速筛选掉重复的学习材料,确保模型接触到的每一段内容都带来新的学习价值。
多维度数据筛选
接下来,团队使用了多维度的数据筛选方法。想象一下,你不仅要考虑教材的主题是否合适,还要评估它的难度、连贯性和知识密度。传统的基于规则的过滤器可能会错误地过滤掉包含大量数学和代码内容的高质量网页,就像误把高级教材当作难以理解的内容而丢弃。
为解决这个问题,研究人员微调了小型语言模型作为数据质量评估器,对内容进行领域分类和多维度质量评估,确保模型能接触到丰富的推理模式。
人工合成的推理数据
研究团队还利用先进的推理模型生成多样化的合成推理数据。这就像请经验丰富的老师为学生创建定制习题集。他们选择标记为具有高推理深度的STEM内容,并提示模型基于这些材料进行深入思考和分析。同时,他们收集数学和代码问题,让推理模型来解决它们,并将解答过程纳入训练数据。
有趣的是,研究发现,与普通数据不同,合成推理数据可以被训练极高的轮次而不会出现过拟合风险 - 就像解决复杂问题的思路永远不会过时一样。
三阶段数据混合策略
最后,研究团队采用了三阶段数据混合策略来优化预训练数据分布:
在第一阶段,他们纳入了除推理任务合成响应外的所有数据源。他们减少了过度代表的内容(如广告、新闻、招聘信息)的比例,并增加了来自专业领域的高价值数据的比例。这就像为学生提供广泛但精心筛选的基础知识。
在第二阶段,他们显著增加了数学和代码相关数据的比例,使其占混合数据的约70%。这就像让学生开始专注于数学和编程的训练,但不放弃其他基础知识的学习。
在第三阶段,为了提升解决复杂任务的能力,他们进一步纳入了约10%的数学、代码和创意写作查询的合成响应。同时,他们将上下文长度从8,192个词元扩展到32,768个词元,这就像训练学生处理更长、更复杂的问题。
通过这个过程,研究团队构建了一个包含约25万亿词元的大规模高质量预训练数据集。这相当于让模型阅读了数以亿计的书籍和文档,吸收了海量的知识和推理模式。
模型架构与超参数
MiMo-7B遵循通用的仅解码器Transformer架构,包括分组查询注意力(GQA)、预RMS归一化、SwiGLU激活和旋转位置嵌入(RoPE)。
为了解决推理模型在长序列自回归生成过程中的推理速度瓶颈,研究团队引入了多词元预测(MTP)作为额外的训练目标。这就像训练一个学生不仅能一步步解题,还能预见几步后的解题路径,从而加速整个解题过程。
在预训练阶段,他们仅使用单个MTP层,因为初步研究表明多个MTP层不会带来进一步的改进。然而,在推理阶段,他们发现多个并行MTP层通过推测解码能显著加速推理。
这种设计使MiMo-7B能够在AIME24基准测试上实现约90%的第一MTP层接受率,即使是第三MTP层也能保持75%以上的接受率。这种高接受率使模型能够提供增强的解码速度,特别是在需要极长输出的推理场景中。
三、后训练:从潜力到卓越表现
预训练阶段打下了坚实的基础,就像为一个天才学生提供了广泛而深入的知识。但要使这位学生成为特定领域的专家,还需要更有针对性的训练。这就是后训练阶段的作用。
监督微调
研究团队首先对MiMo-7B-Base模型进行了监督微调(SFT)。他们使用了约50万个样本的SFT数据集,这些数据是开源和专有蒸馏数据的组合。为确保最佳质量和多样性,他们实施了三阶段预处理流程:
首先,他们消除了与评估基准有16-gram重叠的所有训练查询,防止数据泄漏。其次,他们排除了带有混合语言或不完整响应的样本。最后,他们将每个查询的响应数量上限设为八个,在保持多样性和防止冗余之间取得平衡。
强化学习数据策略
研究团队使用了两类可验证问题 - 数学和代码 - 来构建强化学习训练数据。他们的初步研究表明,高质量的问题集在稳定RL训练过程和进一步增强语言模型推理能力方面起着关键作用。
对于数学问题,他们从各种来源收集数据,包括开源数据集和专有收集的竞赛级收藏。为降低奖励黑客(reward hacking)风险,他们使用LLM过滤基于证明的问题和多选题。与最近的方法不同,他们保留了原始问题而非修改问题以确保整数答案,以最小化奖励黑客风险。
他们还使用基于模型的难度评估来进一步提高数据集质量。初始阶段,他们过滤掉无法被先进推理模型解决的问题,识别那些过于困难或包含错误答案的问题。对于剩余问题,他们让MiMo-7B的SFT版本进行16次尝试,淘汰通过率超过90%的简单问题。这个过程移除了原始问题集中约50%的简单问题。
对于代码问题,他们精心策划了一个包含开源数据集和新收集问题集的高质量训练集。他们移除了没有测试用例的问题。对于有黄金解决方案的问题,他们排除了黄金解决方案未能通过所有测试用例的问题。对于没有黄金解决方案的问题,他们丢弃了在先进推理模型16次尝试中没有测试用例可以被解决的问题。
在奖励函数设计上,他们仅使用基于规则的准确性奖励。对于数学数据,他们使用基于规则的Math-Verify库评估响应的正确性。对于代码问题,他们实现了一个基于测试难度的奖励机制,如后续章节所详细描述的。
强化学习算法改进
研究团队采用了修改版的组相对策略优化(GRPO)算法,并融入了研究社区最近提出的改进:
移除KL损失:简单地删除KL损失可以有效释放策略模型的全部潜力,而不会影响训练稳定性。
动态采样:在RL滚出阶段,他们过度采样并过滤掉通过率等于1和0的提示,在保持一致批量大小的同时保留所有具有有效梯度的提示。这种策略会自动校准整个策略训练过程中的问题难度。
更高上限裁剪:他们增加了方程中的上限裁剪边界ε_high,同时保持固定的下限裁剪边界ε_low。这可以缓解熵收敛问题并促进策略探索新解决方案。
针对代码任务的测试难度驱动奖励
对于算法代码生成任务,现有的RL工作通常采用基于规则的奖励策略,只有当生成的代码通过给定问题的所有测试用例时才给予奖励。然而,对于困难的算法问题,模型可能永远不会收到任何奖励,阻止它从这些具有挑战性的案例中学习,并降低动态采样的训练效率。
为解决这个限制,研究团队提出了一种新的奖励机制:测试难度驱动奖励。这个设计受到国际信息学奥林匹克(IOI)评分规则的启发。在IOI比赛中,每个完整问题被分为多个子任务,参与者将获得他们成功完成的每个子任务的分数。每个子任务都有不同难度的测试。
研究团队提出了一种基于难度对测试用例进行分组的技术。他们利用多个模型对每个问题进行多次尝试,并计算每个测试用例在所有模型生成解决方案中的通过率。然后,他们根据通过率将测试用例聚类为不同的难度级别,通过率越低表示难度越高。
基于这些难度级别,他们设计了两种奖励方案:严格方案和软方案。在严格奖励方案下,解决方案仅在通过该组中的所有测试以及所有低难度组中的测试时才能获得相应难度级别的奖励。相比之下,软奖励方案将每个组的总分平均分配给其测试。最终奖励是所有通过测试的分数总和。
简单数据过滤和重采样策略
在RL训练过程中,随着策略改进,越来越多的问题达到完美的通过率1。在动态采样机制下,这些问题随后从策略更新的批次中过滤掉。这种过滤导致采样效率急剧下降,因为需要更多的滚出来构建固定大小的批次。
为了在不冒策略崩溃风险的情况下提高采样效率,研究团队开发了一种简单数据重采样策略。在训练过程中,他们维护一个简单数据池,存储具有完美通过率的问题。在执行滚出时,有一定概率(在他们的实验中为10%)从这个简单数据池中采样数据。
RL基础设施优化
研究团队开发了无缝滚出引擎并增强了vLLM的鲁棒性,以实现高效的基于动态采样的RL训练。无缝滚出引擎通过高效的任务调度优化GPU利用率,减少连续操作过程中的空闲时间。
该引擎包括以下组件:连续滚出、异步奖励计算和早期终止。它在训练速度上实现了2.29倍的提升,在验证速度上实现了1.96倍的提升。
四、模型性能表现:小个子的大能量
所有的努力最终带来了什么成果?MiMo-7B系列模型在各种基准测试上都展现出了卓越的表现。
预训练模型的推理潜力
研究人员采用了pass@k指标来评估不同模型的推理能力边界。这个指标衡量的是,如果允许模型尝试k次解决一个问题,只要其中任何一次解答正确,就视为成功解决问题。
结果显示,MiMo-7B-Base在所有基准测试和评估的k值上都比同类规模的开源模型(包括Llama-3.1-8B、Gemma-2-9B和Qwen2.5-7B)取得了明显更高的pass@k分数。甚至与更大的32B基线模型相比,MiMo-7B-Base也表现出色。
特别值得注意的是,随着k值的增加,MiMo-7B-Base与其他基线模型之间的性能差距稳步扩大,尤其是在LiveCodeBench上。这些结果证明了MiMo-7B-Base具有卓越的推理潜力,为RL训练奠定了坚实的基础。
后训练模型的综合表现
在后训练阶段,研究团队评估了两个版本的模型:
MiMo-7B-RL-Zero:直接从MiMo-7B-Base进行RL训练得到
MiMo-7B-RL:从MiMo-7B-SFT版本进行RL训练得到结果显示,从基础模型进行RL训练的版本展现出更强的增长趋势,例如在AIME 2024上从32.9%提高到56.4%。然而,从SFT模型进行RL训练的版本达到了更高的性能上限,在所有评估的基准测试中都取得了最佳结果。
与其他代表性模型(包括GPT-4o-0513、Claude-Sonnet-3.5-1022、OpenAI-o1-mini、QwQ-32B-Preview、
DeepSeek-R1-Distill-Qwen-14B和
DeepSeek-R1-Distill-Qwen-7B)相比:在数学推理方面,MiMo-7B-RL在同等参数规模的模型中实现了顶级性能,在AIME 2024上仅略低于
DeepSeek-R1-Distill-Qwen-14B。在算法代码生成任务上,MiMo-7B-RL展示了极其出色的结果。在LiveCodeBench v5上,它显著优于OpenAI o1-mini,而在最新的LiveCodeBench v6上,该模型实现了49.3%的得分,超过QwQ-32B-Preview约10个百分点,展示了其稳健持久的能力。
值得注意的是,MiMo-7B-RL还保持了强大的通用性能,超过了QwQ-32B-Preview和
DeepSeek-R1-Distill-Qwen-7B,尽管RL训练仅包含数学和代码问题。五、研究洞见与挑战
研究团队还分享了一些有趣的观察和挑战:
格式对齐的作用
在从MiMo-7B-Base开始的初始RL训练步骤中,他们观察到模型主要学习适应答案提取函数,例如数学问题的"\boxed{}"格式。因此,他们探索了一种"轻量级"SFT来帮助基础模型与预期的答案格式对齐。
然而,结果显示,经过这种轻量级SFT的MiMo-7B-RL-LiteSFT模型在推理潜力和最终性能方面都表现不佳。虽然MiMo-7B-RL-LiteSFT一开始的性能高于MiMo-7B-RL-Zero,但在仅仅500步后就落后于基础模型的发展轨迹。此外,与经过"更重"SFT的MiMo-7B-RL相比,MiMo-7B-RL-LiteSFT展示了类似的增长趋势,但由于起点较低而导致最终结果明显较差。
不同领域之间的干扰
在从MiMo-7B-Base进行RL训练的后期阶段,维持数学和编码任务之间的性能平衡变得具有挑战性。在训练步骤2000到2500之间,模型在代码问题上表现持续改进,而在数学推理任务上的表现则波动并下降。
相比之下,对冷启动SFT模型的RL训练在两个领域都显示出一致的改进。对模型输出的分析揭示,基础模型凭借其强大的探索能力,倾向于黑入数学问题的奖励。然而,对于代码问题,基于测试用例的验证器使奖励利用变得困难得多。这凸显了高质量数学问题集对确保稳健RL训练的关键需求。
语言混合问题
像DeepSeek-R1-Zero一样,研究团队在对MiMo-7B-Base进行RL训练时也观察到语言混合问题。为缓解这个问题,他们在奖励函数中引入了语言混合惩罚。
然而,他们发现设计这样的惩罚函数具有挑战性。虽然在英语响应中检测中文字符相对简单,但反过来就困难得多,因为数学方程式和代码本身就包含英文单词。结果,惩罚不仅未能完全解决语言混合问题,还引入了奖励黑客的风险,例如无论问题语言如何,总是生成英语响应。
六、结语:为推理模型开辟新路径
小米团队的MiMo-7B系列模型展示了如何通过优化预训练和后训练过程来释放大语言模型的推理潜力。这项研究的意义在于,它证明了相对较小的模型(70亿参数)也能在复杂推理任务上表现出色,甚至超越参数量是它四倍的模型。
MiMo-7B-Base表现出的卓越推理潜力,以及MiMo-7B-RL在数学和代码任务上的优异表现,为构建更强大、更高效的推理模型开辟了新的道路。
让我们把这个故事拉回到现实世界的意义:在人工智能迅速发展的今天,计算资源的效率变得越来越重要。小米的这项研究表明,通过细致的数据处理、创新的训练方法和精心设计的模型架构,我们可以创造出既强大又高效的AI系统,这将使先进AI技术更容易被广泛应用,从而在更多领域发挥积极作用。
最终,MiMo-7B不仅是一个技术成就,也是朝着更可持续、更普及的AI未来迈出的重要一步。对于开发者、研究人员和AI爱好者来说,这项研究提供了宝贵的见解,启发我们思考如何构建下一代更聪明、更高效的AI系统。