Xiaomi MiMo:為“推理而生”
在數學推理(AIME 24-25)和 代碼競賽(LiveCodeBench v5)公開測評集上,MiMo 僅用 7B 的參數規模,超越了 OpenAI 的閉源推理模型 o1-mini 和阿里 Qwen 更大規模的開源推理模型 QwQ-32B-Preview。
預訓練
預訓練階段對于MiMo-7B模型至關重要,它為模型后續展現強大的推理能力奠定了堅實基礎。這一階段主要涵蓋了數據處理、模型架構設計以及超參數設置等關鍵環節,每個環節都經過精心優化,旨在增強模型的推理潛能。
預訓練數據
預訓練數據是模型學習推理模式的重要信息來源,MiMo-7B的預訓練語料庫整合了網頁、學術論文、書籍、編程代碼以及合成數據等多種不同的來源。
優化推理數據提取:常用的文本提取工具在處理網頁中的數學公式和代碼片段時存在不足,小米為此開發了專門針對數學內容、代碼塊和論壇網站的HTML提取工具,同時對PDF解析工具包進行增強,以更好地處理STEM和代碼內容,成功保留了大量推理模式。
快速全局去重:為提高訓練效率、減少過擬合,采用URL去重和MinHash去重方法對網頁數據進行處理。并且,根據多維度質量分數調整最終數據分布,確保去重過程不會影響數據質量。
多維數據過濾:傳統的啟發式規則過濾器可能會錯誤地過濾掉包含大量數學和代碼內容的高質量網頁。因此,小米通過微調小型LLMs作為數據質量標記器,進行領域分類和多維質量評估,篩選出高質量的預訓練數據。
合成推理數據:利用先進推理模型生
成合成推理數據,采用多種策略確保數據的多樣性。例如,選擇具有高推理深度的STEM內容,促使模型進行深入分析;收集數學和代碼問題讓模型求解;引入通用領域查詢和創意寫作任務等。而且,合成推理數據在訓練時不易出現過擬合現象。
三階段數據混合:為優化數據分布,采用三階段數據混合策略。第一階段,整合除推理任務查詢的合成響應之外的所有數據源,對過度代表的內容進行下采樣,對高價值數據進行上采樣;第二階段,將數學和代碼相關數
據比例提升至約70%,在不影響通用語言能力的情況下增強專業技能;第三階段,加入約10%的數學、代碼和創意寫作查詢的合成響應,并將上下文長度從8192擴展到32768,從而構建出約25萬億token的高質量預訓練數據集。
模型架構
MiMo-7B 采用通用的解碼器架構,這種架構在語言模型中應用廣泛,具有良好的性能表現。它包含分組查詢注意力(GQA)、預 RMSNorm、SwiGLU 激活函數以及旋轉位置嵌入(RoPE)等組件,與 Llama 和 Qwen 等模型類似。
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 值增加,與其他基線模型的性能差距不斷擴大,這充分展示了其卓越的推理潛力。
評估結果:在通用推理方面,MiMo-7B-Base 在常識和推理任務中表現出色,超越了類似規模的開源模型。例如在評估語言推理能力的 BBH 基準測試中,得分達到 75.2,超過 Qwen2.5-7B 約 5 分;在 SuperGPQA 測試中,也展現出解決研究生水平問題的強大性能;在 DROP 閱讀理解基準測試中,同樣表現優異,體現了其先進的語言理解能力。在代碼和數學推理方面,MiMo-7B-Base 在編碼和數學任務中表現出很強的熟練度。
在LiveCodeBench v5 和 AIME 2024 測試中,得分均遠超其他類似規模的基礎模型,凸顯了其強大的問題解決能力和在復雜推理任務中的巨大潛力。在長上下文理解方面,對于長上下文檢索任務,MiMo-7B 在 32K 上下文窗口內的所有位置都能實現近乎完美的檢索性能。在需要長上下文推理的任務中,如常見詞提取(CWE)、頻繁詞提取(FWE)和變量跟蹤(VT),也表現出色,超越了Qwen2.5-7B,驗證了預訓練階段融入高質量推理模式數據策略的有效性。
后訓練
在完成預訓練后,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)評分規則的啟發,提出測試難度驅動獎勵機制。通過對每個問題進行多次模型滾動評估,計算每個測試用例的通過率,根據通過率將測試用例聚類為不同難度級別。基于難度級別設計了兩種獎勵方案:嚴格獎勵方案要求解決方案通過某難度級別及以下所有組的測試才能獲得相應獎勵;軟獎勵方案則將每組的總分平均分配到該組的測試中,最終獎勵為所有通過測試的得分總和。實驗結果表明,該獎勵機制能夠更好地利用困難示例進行訓練,提升模型性能。
簡單數據過濾和重采樣:在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 空閑時間和樣本浪費。
基于vLLM的推理引擎:RL 系統采用 vLLM 作為推理引擎,并對其進行擴展以適應模型的新特征。實現并開源了對 MTP 的支持,使模型能夠高效推理。同時,增強了引擎的魯棒性,在搶占時清除前綴緩存中的計算塊以保持KVCache 一致性,增加調度器步驟時禁用異步輸出處理以確保兼容性和優化性能。
后訓練評估
評估設置:在多種基準測試中對推理模型進行全面評估,涵蓋語言理解與推理、科學問答、指令跟隨、閱讀理解、數學推理和編碼等領域。評估時,對所有基準測試設置采樣溫度為 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 訓練則能達到更高的性能上限。
討論
SFT對格式對齊的影響:在從 MiMo-7B-Base 進行初始 RL 訓練時,發現模型主要學習適應答案提取功能,如數學問題的特定格式。為此研究了“輕量級”SFT 以幫助基礎模型對齊預期答案格式,但結果表明,MiMo-7B-RL-Lite SFT 模型在推理潛力和最終性能上均表現不佳,其起始性能雖高于 MiMo7B-RL-Zero,但在訓練 500 步后落后于基礎模型的訓練軌跡,且與經過“重量級”SFT 的 MiMo-7B-RL 相比,最終結果更差。
不同領域之間的干擾:在從 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。小米的這一研究成果為開發更強大的推理模型提供了寶貴的經驗和思路,有望推動整個大語言模型推理領域的進一步發展。
本文轉載自???AIGC前沿技術追蹤???,作者:AIGC前沿技術追蹤
