Xiaomi MiMo:为“推理而生”

发布于 2025-5-9 06:13
浏览
0收藏

在数学推理(AIME 24-25)和 代码竞赛(LiveCodeBench v5)公开测评集上,MiMo 仅用 7B 的参数规模,超越了 OpenAI 的闭源推理模型 o1-mini 和阿里 Qwen 更大规模的开源推理模型 QwQ-32B-Preview。


Xiaomi MiMo:为“推理而生”-AI.x社区


Xiaomi MiMo:为“推理而生”-AI.x社区



预训练

预训练阶段对于MiMo-7B模型至关重要,它为模型后续展现强大的推理能力奠定了坚实基础。这一阶段主要涵盖了数据处理、模型架构设计以及超参数设置等关键环节,每个环节都经过精心优化,旨在增强模型的推理潜能。


Xiaomi MiMo:为“推理而生”-AI.x社区

预训练数据

预训练数据是模型学习推理模式的重要信息来源,MiMo-7B的预训练语料库整合了网页、学术论文、书籍、编程代码以及合成数据等多种不同的来源。

优化推理数据提取:常用的文本提取工具在处理网页中的数学公式和代码片段时存在不足,小米为此开发了专门针对数学内容、代码块和论坛网站的HTML提取工具,同时对PDF解析工具包进行增强,以更好地处理STEM和代码内容,成功保留了大量推理模式。

快速全局去重:为提高训练效率、减少过拟合,采用URL去重和MinHash去重方法对网页数据进行处理。并且,根据多维度质量分数调整最终数据分布,确保去重过程不会影响数据质量。

多维数据过滤:传统的启发式规则过滤器可能会错误地过滤掉包含大量数学和代码内容的高质量网页。因此,小米通过微调小型LLMs作为数据质量标记器,进行领域分类和多维质量评估,筛选出高质量的预训练数据。

合成推理数据:利用先进推理模型生

成合成推理数据,采用多种策略确保数据的多样性。例如,选择具有高推理深度的STEM内容,促使模型进行深入分析;收集数学和代码问题让模型求解;引入通用领域查询和创意写作任务等。而且,合成推理数据在训练时不易出现过拟合现象。

三阶段数据混合:为优化数据分布,采用三阶段数据混合策略。第一阶段,整合除推理任务查询的合成响应之外的所有数据源,对过度代表的内容进行下采样,对高价值数据进行上采样;第二阶段,将数学和代码相关数

据比例提升至约70%,在不影响通用语言能力的情况下增强专业技能;第三阶段,加入约10%的数学、代码和创意写作查询的合成响应,并将上下文长度从8192扩展到32768,从而构建出约25万亿token的高质量预训练数据集。

模型架构

MiMo-7B 采用通用的解码器架构,这种架构在语言模型中应用广泛,具有良好的性能表现。它包含分组查询注意力(GQA)、预 RMSNorm、SwiGLU 激活函数以及旋转位置嵌入(RoPE)等组件,与 Llama 和 Qwen 等模型类似。


Xiaomi MiMo:为“推理而生”-AI.x社区


MTP模块:受 DeepSeek-V3 启发,引入多令牌预测(MTP)模块作为额外训练目标。预训练时,使用单 MTP 层,因为实验发现多层 MTP 层不会带来进一步提升;

推理时,将预训练的单 MTP 层复制为两个相同副本,冻结主模型和第一个 MTP 层,微调两个新的 MTP 层,通过推测解码显著加速推理。在 AIME24 基准测试中,MTP 层表现出较高的接受率,能够有效提升解码速度,尤其在需要长输出的推理场景中优势明显。

超参数

模型超参数:设置了一系列关键的模型超参数,如 Transformer 层数为 36,隐藏维度为 4096,前馈网络(FFN)的中间隐藏维度为 11008,注意力头数为 32,键值组为 8。这些参数的设置影响着模型的性能和推理能力。

训练超参数:优化器选择 AdamW,设置 β1=0.9,β2=0.95,权重衰减为 0.1,并应用梯度裁剪,最大范数为1.0。在不同的预训练阶段,对最大序列长度、RoPE 基数、学习率和批量大小等参数进行调整。例如,前两个阶段最大序列长度为 8192,第三阶段扩展到 32768;学习率在不同阶段也按照特定的线性预热、恒定和余弦衰减策略进行变化;批量大小在不同阶段也有相应的调整。此外,MTP 损失权重在预训练过程中也会进行调整。

预训练评估

评估设置:在一系列基准测试中对 MiMo-7B-Base 进行评估,涵盖自然语言理解与推理、科学问答、阅读理解、数学推理、编码、中文理解以及长上下文理解等多个方面。这些基准测试能够全面评估模型的能力。同时,将MiMo-7B-Base 与其他类似规模的开源基础模型进行比较,保证评估设置的一致性。

推理能力上限评估:传统评估方法可能会低估模型的真实推理潜力,因此采用 pass@k 指标进行评估。该指标认为,若 k次 采样的解决方案中有任何一个正确,则视为问题解决。结果显示,MiMo-7B-Base 在所有基准测试和评估的 k值 下,pass@k 得分均显著高于其他对比模型,包括 320亿 参数的基线模型,且随着 k 值增加,与其他基线模型的性能差距不断扩大,这充分展示了其卓越的推理潜力。


Xiaomi MiMo:为“推理而生”-AI.x社区


评估结果:在通用推理方面,MiMo-7B-Base 在常识和推理任务中表现出色,超越了类似规模的开源模型。例如在评估语言推理能力的 BBH 基准测试中,得分达到 75.2,超过 Qwen2.5-7B 约 5 分;在 SuperGPQA 测试中,也展现出解决研究生水平问题的强大性能;在 DROP 阅读理解基准测试中,同样表现优异,体现了其先进的语言理解能力。在代码和数学推理方面,MiMo-7B-Base 在编码和数学任务中表现出很强的熟练度。


Xiaomi MiMo:为“推理而生”-AI.x社区


在LiveCodeBench v5 和 AIME 2024 测试中,得分均远超其他类似规模的基础模型,凸显了其强大的问题解决能力和在复杂推理任务中的巨大潜力。在长上下文理解方面,对于长上下文检索任务,MiMo-7B 在 32K 上下文窗口内的所有位置都能实现近乎完美的检索性能。在需要长上下文推理的任务中,如常见词提取(CWE)、频繁词提取(FWE)和变量跟踪(VT),也表现出色,超越了Qwen2.5-7B,验证了预训练阶段融入高质量推理模式数据策略的有效性。


Xiaomi MiMo:为“推理而生”-AI.x社区


后训练

在完成预训练后,MiMo-7B 进入后训练阶段。这一阶段基于预训练的基础,进一步挖掘模型的推理潜力,通过监督微调、RL 数据整理、RL 训练方法改进以及 RL 基础设施优化等一系列操作,提升模型在推理任务上的性能。

监督微调

SFT数据:SFT 数据由开源和专有蒸馏数据组合而成。为确保数据质量和多样性,进行了三阶段预处理。首先,去除与评估基准有 16-gram 重叠的训练查询,防止数据泄露;其次,排除语言混合或响应不完整的样本;最后,将每个查询的响应数量限制为 8 个,平衡数据多样性和冗余度。经过预处理,最终的 SFT 数据集包含约 500K 个样本。

SFT超参数:使用恒定学习率 3×10⁻⁵ 和批量大小 128 对 MiMo-7B-Base 模型进行微调。在训练过程中,将样本打包至最大长度 32768 token。

RL 数据整理

数学数据:数学问题集来源广泛,包括开源数据集和专有收集的竞赛级集合。为降低奖励作弊风险,利用 LLM 过滤基于证明和多项选择的问题,保留原始问题,并进行全局 n-gram 去重和与评估基准的去污染处理。通过基于模型的难度评估,进一步提高数据集质量。先过滤掉先进推理模型无法解决的问题,再对剩余问题 用MiMo-7B 的 SFT 版本进行 16 次滚动评估,去除通过率超过 90% 的简单问题,最终得到 100K 个数学训练问题。

代码数据:代码训练集由开源数据集和新收集的问题集组成。去除无测试用例的问题,对于有黄金解决方案的问题,排除黄金解决方案无法通过所有测试用例的情况;对于无黄金解决方案的问题,丢弃先进推理模型在 16 次滚动中无测试用例能解决的问题。同样利用 MiMo-7B 的 SFT 版本过滤掉容易的问题,最终得到 30K 个代码问题。

奖励函数:训练过程中仅采用基于规则的准确率奖励。对于数学数据,使用基于规则的 Math-Verify 库评估响应正确性;对于代码问题,采用测试难度驱动奖励(在后续章节详细介绍)。不引入额外奖励,如格式奖励和长度惩罚奖励。

RL 训练方法

采用改进版的分组相对策略优化(GRPO)算法,并结合了研究社区最近提出的改进方法。在训练过程中,针对遇到的两个关键挑战,即代码问题的稀疏奖励和动态采样的采样效率降低问题,分别提出了测试复杂度驱动奖励函数和简单数据重采样方法。

测试难度驱动奖励:在算法代码生成任务中,传统的基于规则的奖励策略可能导致模型在面对困难算法问题时难以获得奖励,影响训练效率。受国际信息学奥林匹克竞赛(IOI)评分规则的启发,提出测试难度驱动奖励机制。通过对每个问题进行多次模型滚动评估,计算每个测试用例的通过率,根据通过率将测试用例聚类为不同难度级别。基于难度级别设计了两种奖励方案:严格奖励方案要求解决方案通过某难度级别及以下所有组的测试才能获得相应奖励;软奖励方案则将每组的总分平均分配到该组的测试中,最终奖励为所有通过测试的得分总和。实验结果表明,该奖励机制能够更好地利用困难示例进行训练,提升模型性能。


Xiaomi MiMo:为“推理而生”-AI.x社区


简单数据过滤和重采样:在RL训练中,随着策略的改进,越来越多的问题通过率达到 1,这些问题在动态采样机制下会被过滤掉,导致采样效率下降。若直接从训练数据中完全移除通过率为1的问题,会使策略更新不稳定。为此,开发了简单数据重采样策略,在训练过程中维护一个简单数据池,存放通过率为 1 的问题。在进行滚动评估时,以一定概率(实验中为10%)从简单数据池中采样数据,有效稳定了策略更新,提高了采样效率,尤其在RL训练后期效果显著。

超参数:实验中设置训练批量大小为 512,演员小批量大小为 32,每次训练迭代执行 16 次梯度更新,学习率为1e-6。最大序列长度设置为 32768 token,以支持复杂推理任务。训练阶段,将温度和 top-p 参数都配置为1.0,以促进输出多样性。

RL 基础设施

为实现高效的基于动态采样的 RL 训练,开发了无缝滚动引擎并增强了 vLLM 的鲁棒性。整个 RL 系统基于开源 RL 训练库 ver l构建,verl 使用 Ray 管理计算和通信,但存在 GPU 空闲时间长等问题。

无缝滚动引擎:无缝滚动引擎通过高效的任务调度优化 GPU 在滚动工作中的利用率,减少空闲时间。它包含连续滚动、异步奖励计算和早期终止三个组件。连续滚动主动处理完成的滚动任务并启动新任务,消除生成和奖励阶段之间的同步障碍;异步奖励计算利用 Ray 启动异步奖励计算,解决代码相关数据奖励计算耗时的问题,防止滚动管道出现瓶颈;早期终止在有效样本数量达到训练批量大小时,采用先进先出的选择策略终止正在进行的任务,确保数据分布的完整性。实验表明,无缝滚动引擎在训练和验证过程中分别实现了 2.29 倍和 1.96 倍的加速,有效减少了 GPU 空闲时间和样本浪费。


Xiaomi MiMo:为“推理而生”-AI.x社区


Xiaomi MiMo:为“推理而生”-AI.x社区


基于vLLM的推理引擎:RL 系统采用 vLLM 作为推理引擎,并对其进行扩展以适应模型的新特征。实现并开源了对 MTP 的支持,使模型能够高效推理。同时,增强了引擎的鲁棒性,在抢占时清除前缀缓存中的计算块以保持KVCache 一致性,增加调度器步骤时禁用异步输出处理以确保兼容性和优化性能。

后训练评估

Xiaomi MiMo:为“推理而生”-AI.x社区


评估设置:在多种基准测试中对推理模型进行全面评估,涵盖语言理解与推理、科学问答、指令跟随、阅读理解、数学推理和编码等领域。评估时,对所有基准测试设置采样温度为 0.6,top-p 为 0.95。对于数学推理、编码和科学问答基准测试,将最大生成长度设置为 32768 token,其他基准测试设置为 8192 token。将 MiMo-7B-RL 与多个强大的基线模型进行比较,包括非推理模型 GPT-4o-0513、Claude-Sonnet-3.5-1022 和推理模型OpenAI-o1-mini、QwQ-32BPreview、DeepSeek-R1-Distill-Qwen-14B、DeepSeek-R1-Distill-Qwen-7B。

评估结果:在数学推理方面,MiMo-7B-RL 在可比参数规模的模型中表现卓越,仅略落后于 DeepSeek-R1-Distill-Qwen-14B。在算法代码生成任务中,MiMo-7B-RL 表现极为出色,在 LiveCodeBench v5 和 v6 测试中,均显著超越 OpenAI o1-mini 和 QwQ-32B-Preview 等模型,展示了其强大且稳定的能力。在通用性能方面,MiMo-7B-RL 也表现出色,超过了 QwQ-32B-Preview 和 DeepSeek-R1-Distill-Qwen-7B。此外,对不同版本的 MiMo-7B 模型进行评估发现,从基础

模型进行 RL 训练呈现出较强的增长趋势,而从 SFT 模型进行 RL 训练则能达到更高的性能上限。


Xiaomi MiMo:为“推理而生”-AI.x社区

讨论

SFT对格式对齐的影响:在从 MiMo-7B-Base 进行初始 RL 训练时,发现模型主要学习适应答案提取功能,如数学问题的特定格式。为此研究了“轻量级”SFT 以帮助基础模型对齐预期答案格式,但结果表明,MiMo-7B-RL-Lite SFT 模型在推理潜力和最终性能上均表现不佳,其起始性能虽高于 MiMo7B-RL-Zero,但在训练 500 步后落后于基础模型的训练轨迹,且与经过“重量级”SFT 的 MiMo-7B-RL 相比,最终结果更差。


Xiaomi MiMo:为“推理而生”-AI.x社区


不同领域之间的干扰:在从 MiMo7B-Base 进行 RL 训练的后期,难以维持数学和编码任务之间的性能平衡。在训练步骤 2000 到 2500之 间,模型在代码问题上持续改进,而在数学推理任务上性能波动下降。而从冷启动的SFT 模型进行 RL 训练在两个领域均有持续改进。分析发现,基础模型的强探索能力可能导致其在数学问题上进行奖励作弊,而代码问题基于测试用例的验证机制使奖励利用更困难,这凸显了高质量数学问题集对稳健RL训练的重要性。

语言混合惩罚:与 DeepSeek-R1-Zero 类似,在 MiMo-7B-Base 的 RL 训练中也观察到语言混合问题。为缓解该问题,在奖励函数中引入语言混合惩罚,但设计这样的惩罚函数具有挑战性。检测英文响应中的中文字符相对容易,但检测中文响应中的英文单词因数学方程和代码中本身包含英文单词而变得困难,导致惩罚

不仅无法完全解决语言混合问题,还可能引入奖励作弊风险,如无论问题语言如何都生成英文响应。

结论

MiMo-7B系列大语言模型通过优化预训练和后训练过程,成功解锁了先进的推理能力。在预训练阶段,通过优化数据处理、设计合理的模型架构和设置超参数,使 MiMo-7B-Base 接触到丰富多样的推理模式,具备了卓越的推理潜力,超越了规模大得多的模型。在后训练阶段,借助强大高效的RL框架,训练出的 MiMo-7B-RL-Zero 和MiMo-7B-RL 在数学、代码和通用推理任务中展现出优异的推理能力。特别是 MiMo-7B-RL 在 LiveCodeBench v6 和 AIME 2025 测试中分别取得了 49.3% 和 55.4% 的成绩,超越了 OpenAI的o1-mini。小米的这一研究成果为开发更强大的推理模型提供了宝贵的经验和思路,有望推动整个大语言模型推理领域的进一步发展。


Xiaomi MiMo:为“推理而生”-AI.x社区

本文转载自​​AIGC前沿技术追踪​​,作者:AIGC前沿技术追踪

收藏
回复
举报
回复
相关推荐