今年，Numina 和 Hugging Face 合作角逐 AI 數學奧林匹克 (AI Math Olympiad，AIMO) 的首屆進步獎。此次比賽旨在對開放 LLM 進行微調，以使其能解決高中難度的國際數學奧林匹克訓練題。我們很高興向大家報告: 我們的模型 - NuminaMath 7B TIR - 在比賽中脫穎而出，成功解決了私有測試集 50 道題中的 29 道🥳！
本文介紹了 Numina 計劃以及我們獲勝方案背後的技術細節。如果你想先直接用你的數學難題測一下模型，可跳到這個 演示 玩一玩。
我們開始吧！
Numina 簡介 - 開放的 AI For Math 計劃
數學總有點與眾不同在身上！
人人在日常中都會接觸數學，孩子們甚至還未識字就先接觸了數學。有史以來最偉大的數學家之一 拉馬努金 於 1887 年出生在印度的一個普通家庭，靠自學成為一代大師。每個人與數學都有或大或小的交集，從用它消遣到靠它吃飯各有不同。
無法否認的是，數學對人類至關重要，商業社會的一切事物，從 iPhone 到核電站等等，都根植於數學之上。但，就算是純面向應用的數學問題，也自有其趣味之處。
純粹數學超越了智力，就如無邊無際的海洋，唯有心靈才可徜徉其中。
這就是為什麼當我們啟動 Numina 時，開源和開放資料整合了自然之選。相對於人類智慧，我們認為人工智慧對數學的進步也理應起到應有的廣泛作用。如果計算機是思維的腳踏車，那麼人工智慧就是它的引擎 —— 它為我們這個時代的拉馬努金開啟新的視野。
肇始，在 Mistral AI 的支援下，一群對人工智慧和數學充滿熱情的人於 2023 年底集體創立 ( Jia Li、Yann Fleureau、Guillaume Lample、Stan Polu 以及 Hélène Evain) 了 Numina，其靈感來自於由 Alex Gerko 和 XTX Markets 發起的人工智慧數學奧林匹克 (AI Math Olympiad，AIMO) 競賽。
2024 年初，Numina 團隊獲得了兩位來自 Hugging Face 的 LLM 微調專家的支援 (👋 Lewis Tunstall 和 Ed Beeching) 從而開始競逐 2024 AIMO 進步獎。隨後，我們又獲得了 General Catalyst 和 Answer.ai 的支援。到 2024 年 3 月，Numina 已聚集了一支 來自世界各地的頂尖人才 團隊。
團隊就位後，是時候對 AIMO 發起挑戰了！
AIMO 獎
每年，來自世界各地的高中生都會參加 國際數學奧林匹克競賽 - 一項包含六道富有挑戰性的題目，橫跨代數、幾何、數論等領域的競賽。為了讓大家瞭解競賽的難度，下面給出了 去年的一道題:
2023 年 11 月，AIMO 獎 啟動，旨在推動擅長數學推理的人工智慧模型的開放式開發。誰能訓出能夠贏得 IMO 金牌的 AI 模型，誰就會獲得 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 或 2xT4 GPU，最多給 9 個小時來解決 50 道題。
考慮到上述規則和限制，策略選擇對於我們開發制勝方案至關重要。
我們的首個進步獎獲獎解決方案
經過整個比賽的多輪迭代，我們的首個進步獎解決方案主要由三個部分組成:
- 微調 DeepSeekMath-Base 7B 的攻略。透過該攻略，我們將模型打造成可以解決數學題的“推理代理”，其可以透過把自然語言推理和使用 Python REPL 計算中間結果相結合以最終解決問題。
- 一種帶程式碼執行反饋的、為工具整合推理 (tool-integrated reasonin，TIR) 設計的新解碼演算法，以在推理過程中生成候選解答。
- 用來指導模型選擇並避免過擬合公開排行榜的各種內部驗證集。
我們使用了多個開源庫來訓練我們的模型，主要有 TRL、PyTorch、vLLM 以及 DeepSpeed。在一個 8xH100 GPU 節點上，我們花了 10 個小時訓成了模型。
訓練攻略
我們採用的微調方法主要基於 MuMath-Code 論文，其模型訓練過程分為兩個階段:
MuMath-Code 論文中的兩階段訓練方法
- 第 1 階段: 在自然語言“數學題 + 解答”的大規模、多樣化資料集上微調基礎模型，其中每個解答都需套用思維鏈 (CoT) 模板以促使 LLM 進行推理。
- 第 2 階段: 在工具整合推理的合成資料集上微調第 1 階段得到的模型，其中每個數學題都分解為一系列推理、Python 程式及其輸出。此時，我們遵循微軟的 ToRA 論文 的做法，提示 GPT-4 以 ToRA 格式生成帶有程式碼執行反饋的解答。對這些資料進行微調會產生一個推理代理，它可以透過將自然語言推理和使用 Python REPL 來計算中間結果結合起來以解決數學問題 (請參見下圖)。
來自 ToRA 論文的圖，該論文介紹了我們用來訓練模型的工具整合推理格式。
這兩個階段，我們都用了“全模型微調”，所有模型權重在反向傳播期間都得到了更新。換句話說，我們沒有使用像 LoRA 或 DoRA 這樣的引數高效技術，因為沒有大量實驗表明它們能夠媲美全模型微調的效能。我們使用 TRL 的 SFTTrainer 中的“填充”功能將多個樣本串接到一個 2048 個詞元的塊中。所有模型模型都使能了梯度 checkpointing 訓練，並使用 DeepSpeed ZeRO-3 進行分片，以確保權重、梯度和最佳化器狀態能夠放進 VRAM。兩個階段使用的主要超參如下:
| 1 階段 | 2 階段 | |
|---|---|---|
| 學習率 | 2.0 E-5 | 2.0 E-5 |
| 總 batch size | 32 | 32 |
| 塊大小 | 2048 | 1024 |
| epoch 數 | 3 | 4 |
| 學習率排程器 | cosine | cosine |
| 預熱率 | 0.1 | 0.1 |
首次提交時，我們使用了 DeepSeek 7B 模型，我們僅對它進行了第 1 階段微調，但我們發現效能相當有限，其在公開排行榜上的最佳 maj@32 成績僅為 8/50。Abdur Rafae 的 公開筆記本 促使我們考慮在訓練方案中加入程式碼執行。最初，我們專注於 MMOS (Mix of Minimal Optimal Sets) 資料集。我們發現使用 MMOS 雖然提高了效能，但在公開排行榜上的 maj@32 最高分仍只有 16/50，我們當時望文生義地猜測其原因是 MMOS 僅包含單輪解 (即該模型僅生成單個 Python 程式，這不足以解決難題)。後來，我們意識到 MMOS 是一個誤稱，該 Kaggle 筆記本實際上使用的是 DeepSeekMath 7B RL 模型，也就是說它能夠進行多步推理及程式碼執行。
經此一役，我們想集中精力生成一個與 DeepSeekMath Instruct/RL 模型使用的資料集類似的資料集，這一做法與 MuMath-Code 攻略結合後，帶來了顯著的改進。
下面，一起來看看我們是如何構建這些資料集的吧。
所需惟資料
在構建資料集時，我們廣泛參考了 DeepSeek Math 和其他學者的方法，並對它們進行了大幅擴充套件。我們生成了含數十萬 數學題 - 解答 對的微調資料集，涵蓋從高中數學到競賽級數學的各種知識點。接下來的幾周，我們會將該資料集完全開源。同時。我們還可能會用更大的模型來檢查我們攻略的可擴充套件性。有關資料集構建的詳細資訊，請參閱我們即將釋出的資料集技術報告。
具體到這次進步獎，我們為此構建了兩個資料集以微調模型。
思維鏈
該資料集由數十萬個題目組成，每題都有以思維鏈的方式編寫的解答。資料集的來源範圍有中國高中數學練習以及美國及國際數學奧林匹克競賽題目。資料主要來自線上試卷 PDF 和數學論壇。
處理步驟如下:
- 對原始 PDF 進行 OCR。
- 分割為“題目 - 解答”對。
- 翻譯成英文。
- 重新調整以變成思維鏈推理格式。
- 格式化為最終答案。
工具整合推理
工具整合推理 (TIR) 在本次比賽中發揮了至關重要的作用。然而，收集和標註此類資料既昂貴又耗時。為了解決這個問題，我們從 Numina 資料集中選擇了大約 6 萬道題，重點關注那些答案為數字的題，其中大多數答案是整數。
然後，我們利用 GPT-4 的流水線生成類似 TORA 的推理路徑，執行程式碼並生成結果，直到生成完整解答。我們篩選出最終答案與參考答案不匹配的解答，並重復此過程三次，以確保準確性和一致性。這種迭代方法使我們能夠高效地生成高質量的 TORA 資料。
作為參考，以下是我們訓得的第 1 階段模型 NuminaMath-7B-CoT 和第 2 階段模型 NuminaMath-7B-TIR 在 MATH 基準 上與其他開放及私有模型的跑分對比:
| 模型 | MATH (%) |
|---|---|
| 思維鏈推理 | |
| GPT-4 (2023) | 42.5 |
| GPT-4o | 76.6 |
| Claude 3.5 Sonnet | 71.1 |
| DeepSeekMath-7B-Instruct | 46.8 |
| DeepSeekMath-7B-RL | 51.7 |
| NuminaMath-7B-CoT | 56.3 |
| 工具整合推理 | |
| DeepSeekMath-7B-Instruct | 57.4 |
| DeepSeekMath-7B-RL | 58.8 |
| NuminaMath-7B-TIR | 68.2 |
各模型在 MATH 基準上的表現。除非明確說明，所有跑分均由零樣本貪心解碼獲得。
透過自一致性工具整合推理 (SC-TIR) 來抑制高波動
正如其他參賽者指出的那樣，本次比賽在模型提交和評估方面都帶來了一些挑戰:
- 評估 API 以隨機順序出題，因此提前停止等策略會產生較高的波動，因為可能一開始就會遇到很多難題，這就導致留給剩餘部分的時間就不多了 (反之亦然)。
- LLM 推理中的大多數創新都是基於最新的 GPU 的，因此
Flash Attention 2或torch.compile等標準方法不適用於 T4 GPU。同樣，老 GPU 並不支援 bfloat16 等新資料型別，這促使我們探索 AWQ 和 GPTQ 等訓後量化方法。
最初，我們使用 Abdur Rafae 的 公開筆記本 來提交，但發現高波動是個大問題。為了解決這個問題，我們採取了一種基於工具整合推理的新方法:
- 將每道題複製 N 次以生成 vLLM 的一個 batch。N 可以看成多數投票時的候選數量。
- 對這 N 個輸入進行取樣解碼，直至生成完整的 Python 程式碼塊。
- 執行每個 Python 程式碼塊並將其輸出串接在程式碼後面，包括棧回溯 (如有)。
- 重複 M 次以生成 N 個、深度為 M 的生成，允許模型使用棧回溯自糾正程式碼錯誤。如果某個樣本無法生成合理的輸出 (如，生成了不完整的程式碼塊)，就刪除之。
- 對候選解答進行後處理，並使用多數投票來選擇最終答案。
我們獲勝的提交使用的 N=48，M=4 。因為增加任一引數的數值並不會提高效能，所以我們就選擇了這兩個最小值以保證滿足時間限制。實際上，該演算法透過工具整合推理增強了 CoT 的自一致性 (如下所示)。
我們發現，我們的 SC-TIR 演算法產生了更穩健的結果，並且顯著降低了在內部評估集和公開排行榜上的波動。
值得一提的一個技術細節是，我們發現以 8 位元精度量化模型很有用。原因有三:
- 將模型上傳到 Kaggle Hub 非常慢，壓縮模型使上傳速度提高了一倍。
- T4 GPU 不支援 bfloat16，轉換為 float16 會導致模型效能下降。又沒法轉換為 float32，因為超出了 GPU 可用記憶體。
- 此外，16 位模型僅用於載入權重就需消耗約 32GB VRAM。對於 2xT4，需要使能 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 道題)
透過使用這四個驗證集，我們能夠在不同的訓練階段選擇最有潛力的模型，並縮小超參的選擇範圍。我們發現，對本 AIMO 賽程而言，將小型但具代表性的驗證集與較大的驗證集相結合是有用的，因為每個提交都受到抽樣隨機性的影響。
我們嘗試過的其他點子
上文有提及，我們在過程中還嘗試了一些其他方法，但最終放棄，轉而採用 MuMath-Code 的方法。我們嘗試過的方法有:
- 訓練純 CoT 模型並使用多數投票進行評估
- 訓練 MMOS 模型以透過 Python 一步解決問題
我們還試過對 SFT 模型生成的補全應用 Kahneman-Tversky Optimization (KTO)，具體想法有點類似於 OrcaMath，即:
- 交織使用推理和程式碼執行，每道題用 SFT 模型取樣出 4 個補全。我們使用第 2 階段的 SFT 資料集作為提示。
- 提取答案並將其與標註答案進行比較。如果正確，則將樣本標記為正，否則標記為負。
- 在此資料集上對 SFT 模型應用 KTO。
我們發現這種形式的同策 KTO 生成的模型比 SFT 模型稍好 (內部評估好幾個百分點)，在公開排行榜上得分為 27/50。
KTO 的一個很好的功能是，你可以在訓練期間跟蹤隱式獎勵，這確實有助於除錯 - 如，下圖展示了我們成功的訓練日誌之一，其中人們可以看到正確答案的獎勵隨著訓練而增加，而錯誤答案的獎勵則被抑制。
但，由於時間關係，我們最終沒有將此方法應用於最終的 SFT 模型。如果我們做了的話，可能還能多做對 1-2 道題！
我們還嘗試將我們的 SFT 攻略應用於 InternLM-20B、CodeLama-33B 和 Mixtral-8x7B 等更大的模型，但發現 (a) DeepSeek 7B 模型由於已在數學上進行過增量預訓練而很難被擊敗，且 (b) 在 2xT4 GPU 上推理速度非常慢，並且我們遇到了許多神秘的超時，但我們無法分析到其根因。
還有一個失敗的實驗是嘗試將強化學習 (特別是 PPO 演算法及 Reinforce-Leave-One-Out (RLOO) 演算法) 和程式碼執行反饋結合起來以生成對編寫程式碼及獲得正確/錯誤解答的獎勵。我們將其應用於 DeepSeekMath 7B RL 模型。雖然我們看到了一些很不錯的獎勵曲線，但我們沒有看到效能有任何顯著的提升。鑑於像 RLOO 這樣的線上方法受限於文字生成的效能並且迭代緩慢，我們放棄了強化學習，轉而嘗試 KTO。
在推理方面，我們也進行了如下實驗:
- 使用靜態 KV 快取和 torch 編譯。我們發現我們能夠在 H100 上將原生
transformers程式碼的生成速度加快 2-3 倍，但在 Kaggle T4 上會遇到各種神秘錯誤，主要是由於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 團隊發起瞭如此令人激動、鼓舞人心的比賽！
英文原文: https://hf.co/blog/winning-aimo-progress-prize
原文作者: Yann Fleureau，Li Jia，Edward Beeching，Lewis Tunstall，Ben Lipkin，Roman Soletskyi，Shengyi Costa Huang，Kashif Rasul
譯者: Matrix Yao (姚偉峰)，英特爾深度學習工程師，工作方向為 transformer-family 模型在各模態資料上的應用及大規模模型的訓練推理。