NuminaMath 如何贏得第一屆 AIMO 進步獎

今年，Numina 和 Hugging Face 合作，參加了首屆 AI 數學奧林匹克 (AIMO) 進步獎的角逐。這項比賽涉及對開放式 LLM 進行微調，以解決高中生為國際數學奧林匹克訓練而使用的難題。我們很高興地宣佈，我們的模型——NuminaMath 7B TIR——榮獲冠軍，並在私有測試集中成功解決了 50 道問題中的 29 道 🥳！

在這篇博文中，我們將介紹 Numina 倡議以及我們獲獎解決方案背後的技術細節。如果您想直接測試模型以解決最難的數學問題，請檢視我們的演示。

讓我們深入瞭解！

1. Numina 介紹——一個開放的 AI4Maths 倡議
2. AI 數學奧林匹克 (AIMO) 獎
3. 我們贏得第一個進步獎的解決方案
4. 訓練方案
5. 好資料是你的全部所需
6. 透過自我一致性和工具整合推理（SC-TIR）馴服方差
7. 避免過擬合的詛咒
8. 我們嘗試過的其他想法
9. Numina 的未來——尋找貢獻者和合作夥伴！
10. 致謝

Numina 介紹——一個開放的 AI4Maths 倡議

數學有一些非常特別的地方。

數學是一個人人都能接觸的領域，甚至在兒童學會閱讀之前就能接觸。歷史上最偉大的數學家之一 斯里尼瓦瑟·拉馬努金 是自學成才的，他於 1887 年出生在印度的一個普通家庭。而對其他人來說，數學從休閒到職業，無所不包。

數學對人類至關重要，它是我們構建商業、iPhone 和核電站等一切事物的支柱。然而，即使是解決關鍵應用的數學問題，也可能是一種愉快的體驗。

純粹的數學超越了智力，如同浩瀚的海洋，只有心靈才能航行。

這就是為什麼當我們啟動 Numina 時，選擇開源和開放資料集是自然而然的選擇。就像人類智慧一樣，我們相信人工智慧在數學領域的進步應該普世。如果計算機是心靈的腳踏車，那麼人工智慧就是它的引擎——為我們這個時代的拉馬努金們開啟了新的視野。

在 Mistral AI 的初步支援下，Numina 於 2023 年末由一群對 AI 和數學充滿熱情的人士（李佳、Yann Fleureau、Guillaume Lample、Stan Polu 和 Hélène Evain）共同創立，靈感來自於 Alex Gerko 和 XTX Markets 發起的 AI 數學奧林匹克（AIMO）競賽。

2024年初，Hugging Face 的兩位 LLM 微調專家（👋 Lewis Tunstall 和 Ed Beeching）加入 Numina 團隊，以應對 2024 年 AIMO 進步獎的挑戰。我們還獲得了 General Catalyst 和 Answer.ai 的額外支援，截至 2024 年 3 月，Numina 已匯聚了來自 世界各地的頂尖人才。

團隊組建完畢，是時候應對 AIMO 挑戰了！

AI 數學奧林匹克 (AIMO) 獎

每年，來自世界各地的高中生都會參加國際數學奧林匹克競賽——這項競賽旨在解決代數、幾何和數論等領域的六個具有挑戰性的問題。為了讓您瞭解其難度，以下是去年的問題之一

2023 年 11 月，AIMO 獎啟動，旨在推動在數學推理方面表現出色的人工智慧模型的開放開發。凡能創造出在 IMO 中獲得金牌的人工智慧模型者，將獲得 500 萬美元的鉅額獎金。除了大獎，AIMO 還引入了一系列“進步獎”來標記實現這一最終目標的里程碑。第一個進步獎以Kaggle 競賽的形式舉行，其問題比 IMO 中的問題“不那麼具有挑戰性”，但達到了 IMO 預選賽的水平。以下是一個示例問題，它比上面的 IMO 示例更容易解決，但對 LLM 來說仍然很棘手

設 $k, l > 0$ 為引數。拋物線 $y = kx^2 - 2kx + l$ 與直線 $y = 4$ 在 $A$ 和 $B$ 兩點相交，這兩點之間距離為 6。那麼 $A$ 和 $B$ 到原點距離的平方和是多少？

本次競賽分為兩組各 50 道題目，分別用於公共和私人排行榜，題目對參賽者隱藏。這些問題的難度與 AMC12 和 AIME 考試相當，需要整數輸出進行驗證。最終排名由私人排行榜決定。參賽者每天可以提交兩次解決方案，並且只能使用在 2 月 23 日之前釋出的開放權重模型。每次提交分配一塊 P100 GPU 或兩塊 T4 GPU，最長 9 小時用於解決 50 道問題。

考慮到這些限制和規則，戰略性選擇對於開發我們的獲勝解決方案至關重要。

我們贏得第一個進步獎的解決方案

經過競賽中的多次迭代，我們針對第一個進步獎的解決方案包含三個主要部分：

• 一個微調 DeepSeekMath-Base 7B 的方案，使其能夠作為一個“推理代理”，透過自然語言推理和使用 Python REPL 計算中間結果來解決數學問題。
• 一種新穎的工具整合推理（TIR）解碼演算法，帶有程式碼執行反饋，用於在推理過程中生成候選解決方案。
• 我們使用各種內部驗證集來指導模型選擇，並避免對公共排行榜過擬合。

我們使用了多種開源庫來訓練模型，特別是 TRL、PyTorch、vLLM 和 DeepSpeed。在 8 塊 H100 GPU 的一個節點上，我們的模型訓練耗時 10 小時。

訓練方案

我們的微調方案主要基於 MuMath-Code 論文，其中包括兩個階段的模型訓練：

MuMath-Code 論文中的兩階段訓練方法

• 第一階段：在一個大型、多樣化的自然語言數學問題和解決方案資料集上對基礎模型進行微調，其中每個解決方案都使用思維鏈 (CoT) 進行模板化，以方便推理。

• 第二階段：在第一階段模型的基礎上，使用合成的工具整合推理資料集進行微調。在該資料集中，每個數學問題都被分解為一系列的推理、Python 程式及其輸出。在這裡，我們遵循微軟的 ToRA 論文，並提示 GPT-4 以 ToRA 格式生成帶有程式碼執行反饋的解決方案。對這些資料進行微調，可以生成一個推理代理，該代理可以透過自然語言推理和使用 Python REPL 計算中間結果的組合來解決數學問題（參見下面的截圖）。

  

  ToRA 論文中關於我們訓練模型所使用的工具整合推理格式的圖示。

我們在兩個階段都進行了“完全微調”，即在反向傳播過程中更新了所有模型權重。換句話說，我們沒有使用 LoRA 或 DoRA 等引數高效技術，因為我們不確定它們能否在不進行大量實驗的情況下達到完全微調的效能。我們使用了 TRL 的 SFTTrainer 的“打包”功能，將多個樣本連線成一個 2048 個令牌的單一塊。所有模型都使用梯度檢查點進行訓練，並使用 DeepSpeed ZeRO-3 協議進行分片，以確保權重、梯度和最佳化器狀態能夠適應可用的 VRAM。下面是我們每個階段使用的主要超引數：

我們最初的提交只使用了在第一階段微調過的 DeepSeek 7B 模型，但我們發現其效能相當有限，在公共排行榜上使用 maj@32 的最佳得分僅為 8/50。是 Abdur Rafae 的 公共獎項筆記本 促使我們考慮將程式碼執行整合到訓練方案中。最初，我們專注於 最小最優集混合 (MMOS) 資料集，正如筆記本標題所述。我們發現使用 MMOS 確實提高了效能，但在公共排行榜上使用 maj@32 的得分仍然停留在 16/50，這可能是因為 MMOS 只包含單步解決方案（即模型只生成一個 Python 程式，這對於難題來說是不夠的）。我們後來意識到 MMOS 是一個誤稱，Kaggle 筆記本實際上執行的是 DeepSeekMath 7B RL 模型，該模型能夠進行多步推理和程式碼執行。

此時，我們專注於生成一個類似於 DeepSeekMath Instruct / RL 模型所使用的資料集，再加上 MuMath-Code 方案，這帶來了顯著的改進。

讓我們來看看我們是如何構建這些資料集的。

好資料是你的全部所需

在資料集方面，我們廣泛參考了 DeepSeek Math 和其他學者的研究方法，並對其進行了顯著擴充套件。這形成了一個包含數十萬個問題-解決方案對的微調資料集，涵蓋了從高中數學到競賽級數學的各種主題。該資料集將在未來幾周內完全開源，並可能包含更大的模型，以觀察我們的方案如何擴充套件。有關資料集構建的詳細資訊，請參閱我們即將釋出的資料集技術報告。在本次進步獎競賽中，我們構建了兩個資料集來微調我們的模型。

思維鏈

該資料集包含數十萬個問題，每個問題都以思維鏈的方式給出解決方案。資料集的來源涵蓋了從中國高中數學練習題到美國和國際數學奧林匹克競賽題。這些資料主要從線上試卷 PDF 和數學討論論壇收集而來。

處理步驟包括：

1. 從原始 PDF 進行 OCR。
2. 將問題-解決方案對進行分割。
3. 翻譯成英文。
4. 重新對齊以生成思維鏈推理格式。
5. 最終答案格式化。

工具整合推理

工具整合推理（TIR）在本次競賽中扮演著至關重要的角色。然而，收集和標註此類資料既耗時又昂貴。為了解決這個問題，我們從 Numina 資料集中選擇了大約 60,000 個問題，重點關注那些具有數值輸出（其中大部分是整數）的問題。

然後，我們利用 GPT-4 生成 TORA 樣式的推理路徑，執行程式碼並生成結果，直到解決方案完成。我們過濾掉最終答案與參考不匹配的解決方案，並重復此過程三次，以確保準確性和一致性。這種迭代方法使我們能夠高效地生成高質量的 TORA 資料。

作為參考，以下是我們的第一階段模型 NuminaMath-7B-CoT 和最終第二階段模型 NuminaMath-7B-TIR 在 MATH 基準測試上的表現，與其他開放和專有模型進行對比：

MATH 基準測試上的表現。除非另有明確說明，所有資料均透過零樣本貪婪解碼獲得。

透過自我一致性和工具整合推理（SC-TIR）馴服方差

正如其他參賽者所指出的，本次比賽在模型提交和評估方面存在若干挑戰：

• 評估 API 以隨機順序提供問題，因此早期停止等策略會產生高方差，因為一次執行可能在一開始就有更多難題，從而導致剩餘時間減少（反之亦然）
• LLM 推理中的大多數創新都需要訪問現代 GPU，因此 Flash Attention 2 或 torch.compile 等標準方法無法在 T4 GPU 上執行。同樣，bfloat16 等現代資料型別也不受支援，這促使我們探索 AWQ 和 GPTQ 等後訓練量化方法。

最初，我們使用 Abdur Rafae 的 公共筆記本 進行提交，但發現高方差是個問題。為了解決這個問題，我們採用了基於工具整合推理的不同方法：

1. 對於每個問題，將輸入複製 N 次，以定義要饋送 vLLM 的初始提示批次。這有效地定義了用於多數投票的候選數量。
2. 取樣 N 個不同的補全，直到生成完整的 Python 程式碼塊。
3. 執行每個 Python 程式碼塊並連線輸出，包括追溯（如果出現）。
4. 重複 M 次，生成大小為 N 且深度為 M 的批次，允許模型使用追溯自糾程式碼錯誤。如果樣本未能產生合理輸出（例如，不完整的程式碼塊），則剪枝該結果。
5. 對候選解決方案進行後處理，然後應用多數投票選擇最終答案。

對於我們的獲勝提交，我們生成了 N=48 個候選，深度為 M=4。增加任何一個引數都不會提高效能，因此我們採取了保守的方法，以確保在時間限制內完成。實際上，此演算法透過工具整合推理增強了與 CoT 的自我一致性（如下所示）。

我們發現我們的 SC-TIR 演算法在我們的內部評估和公共排行榜上都產生了更穩健的結果，方差顯著降低。

值得一提的一個技術細節是，我們發現將模型量化為 8 位精度很有幫助。原因有三：

• 將模型上傳到 Kaggle Hub 非常慢，壓縮模型使這一步的速度提高了一倍。
• T4 GPU 不支援 bfloat16，轉換為 float16 會導致模型效能下降。轉換為 float32 不可能，因為那會超出可用的 GPU 記憶體。
• 此外，一個 16 位模型僅載入權重就需要消耗大約 32GB 的視訊記憶體。對於 2xT4s，這將需要操作 KV 快取才能快速執行，我們發現權衡模型精度以換取速度是有益的。

我們使用 AutoGPTQ 和訓練資料集進行校準，對模型進行了量化。實際上，這導致了準確性的小幅下降，但為適應 Kaggle 平臺上的評估所施加的約束提供了最佳折衷方案。

避免過擬合的詛咒

過擬合公共排行榜是 Kaggle 競賽中的常見風險，尤其是在測試集只有 50 個問題時更是如此。此外，規則規定每天最多提交兩次，這使得魯棒的內部驗證資料集對於我們控制開發進度至關重要。根據 AIMO 團隊的規定，測試題難度中等，介於 AMC12 和 AIME 之間，並且需要整數輸出。

為了指導模型選擇，我們使用了四個內部驗證集來評估模型在不同難度數學問題上的效能。為了避免基礎模型中可能存在的汙染，我們從 AMC12（2022、2023）和 AIME（2022、2023、2024）中選擇了問題，建立了兩個內部驗證資料集：

• AMC（83 個問題）：我們選擇了 AMC12 22、AMC12 23 中的所有問題，並保留了那些可以轉換為整數輸出的問題。這形成了一個包含 83 個問題的資料集。這個驗證集旨在代表 Kaggle 上的私有測試集，因為我們從比賽描述中得知這些問題的難度與此級別或更高。我們發現我們的模型可以解決這些問題的大約 60-65%。為了衡量方差，我們每次評估都使用 5-10 個不同的種子進行，通常使用我們的 SC-TIR 演算法可以看到大約 1-3% 的變化。
• AIME（90 個問題）：我們選擇了 AIME 22、AIME 23 和 AIME 24 中的所有問題，以衡量我們的模型在難題上的表現，並評估最常見的失敗模式。與上文一樣，我們每次評估都使用 5-10 個種子來衡量變化。

由於 AMC/AIME 驗證集規模較小，模型在這些資料集上的效能容易受到噪聲影響，這與公共排行榜類似。為了更好地評估我們模型的效能，我們還使用 MATH 測試集的一個子集對其進行了評估，該測試集包含 5,000 個問題。我們只保留了具有整數輸出的問題，以簡化多數投票並模仿競賽評估。這產生了兩個額外的驗證集：

• MATH 等級 4（754 個問題）
• MATH 等級 5（721 個問題）

透過使用這四個驗證集，我們能夠從不同的訓練階段中挑選出最有希望的模型，並縮小超引數的選擇範圍。我們發現，在這種特殊的比賽中，結合小型但具有代表性的驗證集和大型驗證集很有用，因為每次提交都會受到抽樣帶來的一些隨機性影響。

我們嘗試過的其他想法

如上所述，我們嘗試了幾種方法，但最終都放棄了，轉而使用 MuMath-Code 配方：

• 訓練純 CoT 模型並使用多數投票進行評估。
• 訓練 MMOS 模型以一步解決 Python 問題。

我們嘗試的另一種技術是對 SFT 模型取樣的新補全應用 Kahneman-Tversky 最佳化 (KTO)。這裡的方法與 OrcaMath 類似，即：

• 使用 SFT 模型為每個問題取樣 4 個補全，交錯推理和程式碼執行。我們使用第二階段的 SFT 資料集作為提示源。
• 提取答案並與真實值進行比較。如果正確，將樣本標記為正，否則標記為負。
• 將 KTO 應用於該資料集上的 SFT 模型。

我們發現這種形式的線上 KTO 產生了一個比 SFT 模型略好的模型（在我們的內部評估中提高了幾個百分點），並在公共排行榜上獲得了 27/50 的分數。

KTO 的一個優點是可以在訓練期間跟蹤隱式獎勵，這確實有助於除錯執行——例如，這是我們成功的訓練日誌之一，可以看到所選（即正確解決方案）獎勵在訓練過程中增加，而被拒絕的獎勵則受到抑制。

不幸的是，我們沒有足夠的時間將這種方法應用於我們的最終 SFT 模型，所以我們可能能夠多解決 1-2 個問題！

我們還嘗試將 SFT 方案應用於更大的模型，如 InternLM-20B、CodeLlama-33B 和 Mixtral-8x7B，但發現 (a) DeepSeek 7B 模型由於在數學上的持續預訓練而難以超越，(b) 在 2xT4 GPU 上的推理速度非常慢，並且我們遇到了許多無法追蹤根源的神秘超時問題。

另一個失敗的實驗包括嘗試使用強化學習（特別是近端策略最佳化演算法和 REINFORCE-leave-one-out (RLOO) 演算法），透過程式碼執行反饋和針對編寫程式碼以及獲得正確/不正確解決方案的形狀獎勵。我們將其應用於 DeepSeekMath 7B RL 模型。雖然我們看到了一些有希望的獎勵曲線，但我們沒有看到效能的顯著提升。鑑於 RLOO 等線上方法受到文字生成和迭代緩慢的瓶頸，我們放棄了強化學習，轉而嘗試 KTO。

在推理方面，我們還嘗試了：

• 使用靜態 KV 快取和 torch 編譯。我們發現在 H100 上可以將原生 transformers 程式碼的生成速度提高 2-3 倍，但在 Kaggle T4s 上遇到了各種神秘錯誤，主要是由於 Accelerate 中 torch 編譯不支援模型分片。

各種模型合併技術，如 DARE、TIES 和 WARP。在這裡，我們使用 mergekit 來合併 SFT 和 KTO 模型，或者 SFT 模型與公共 DeepSeekMath 模型。總的來說，我們發現這些合併導致了我們在內部評估中出現了顯著的效能下降，並且我們沒有足夠的時間更深入地探索這一點。

Numina 的未來——尋找貢獻者和合作夥伴！

在 Numina 成功贏得 AIMO 2024 進步獎之後，我們現在旨在繼續我們的使命，即促進人工智慧和人類智慧在數學領域的發展。您可以訪問我們的網站了解更多專案資訊，並隨時透過 contact@projectnumina.ai 與我們聯絡。

Numina，如同數學本身，旨在向所有願意透過人工智慧推動數學進步的全球人才和支持者開放！

致謝

我們感謝 Thomas Wolf 和 Leandro von Werra 促成了 Numina 和 Hugging Face 的合作。我們還要感謝 Hugo Larcher 幫助 Hugging Face 叢集上的 GPU 全速執行，Colin Raffel 在模型合併方法方面的建議，以及 Omar Sanseviero 對博文的反饋。

我們還要向 Mistral.ai、General Catalyst、Answer.AI 以及北京大學北京國際數學研究中心致以誠摯的謝意，他們從專案伊始就給予了支援。

最後，我們感謝 AIMO 獎團隊舉辦瞭如此激動人心和富有啟發性的競賽！