NuminaMath 7B TIR が AIMO Progress Prize を受賞 – 技術的まとめ

TL;DR

NuminaMath 7B TIR は、隠された 50 問の競技問題のうち 29 問を正しく解くことで、初の AI Math Olympiad (AIMO) Progress Prize を受賞しました。これにより、2 段階のファインチューニングレシピ、大規模で高品質な数学データ、そしてツール統合推論を伴う自己整合性（SC‑TIR）推論戦略が、オリンピックレベルの数学におけるオープンソース LLM の性能を劇的に向上させることが示されました。

Numina – オープン AI4Maths イニシアチブの紹介

Numina は、2023 年後半に開始されたオープンソースプロジェクトで、AI 主導の数式推論を加速することを目的としています。Jia Li、Yann Fleureau、Guillaume Lample、Stan Polu、Hélène Evain が創設し、Mistral AI の初期支援を受けました。2024 年初頭には、Hugging Face のファインチューニング専門家 Lewis Tunstall と Ed Beeching がチームに加わり、General Catalyst と Answer.ai からも追加支援がありました。コラボレーションの最初の具体的な目標は、2024 年 AIMO Progress Prize という Kaggle コンペで、50 問の事前選考レベルの数学問題（AMC 12/AIME に相当）に対し、2024 年 2 月 23 日以前にリリースされたオープンウェイトモデルのみを使用して LLM をテストすることでした。

AI Math Olympiad (AIMO) Progress Prize

AIMO Prize は、国際数学オリンピックで金メダルを獲得できる AI を創出することを目指し、500 万ドルのグランドプライズを提供します。最初のプログレス賞では、参加者は Kaggle プラットフォーム上で 1 日 2 回、解答を提出する必要があり、各提出は P100 または 2 台の T4 GPU 上で最大 9 時間実行されました。問題は整数出力の質問で、高校レベルの競技問題から抽出され、公開リーダーボード（50 問）と、最終順位を決定する非公開プライベートリーダーボードがあります。

勝利したソリューションのアーキテクチャ

勝利システムは、以下の 3 つのコアコンポーネントで構成されました。

  1. ファインチューニングされた推論エージェント – DeepSeekMath‑Base 7B を 2 段階でファインチューニングし、自然言語の推論と Python REPL 呼び出しを交互に行うモデルにしました。
  2. SC‑TIR デコーディングアルゴリズム – 多数の候補解を生成し、埋め込まれた Python コードを実行、自己整合性チェック後に多数決を取る新規推論手法です。
  3. 堅牢な内部検証 – AMC、AIME、そして MATH ベンチマークの 2 つのサブセットからなる 4 つのキュレーション済み検証セットを用いてモデル選択を行い、公開リーダーボードへの過学習を防止しました。

トレーニングはオープンソーススタック（TRL、PyTorch、vLLM、DeepSpeed）を使用し、8 台の H100 GPU を搭載した単一ノードで約 10 時間で完了しました。

2 段階トレーニングレシピ（MuMath‑Code）

このレシピは MuMath‑Code 論文に従い、以下の構成です。

  • Stage 1 – Chain‑of‑Thought (CoT) ファインチューニング – 基礎モデルは数十万件の自然言語数学問題と CoT 形式の解答から学習し、ステップバイステップの推論を促します。
  • Stage 2 – Tool‑Integrated Reasoning (TIR) ファインチューニング – ToRA 形式で GPT‑4 が生成した合成データセットを使用し、各問題に対して根拠、Python プログラム、実行結果をペアにします。これにより、モデルは中間計算のために Python REPL を呼び出す方法を学習します。

両ステージとも LoRA/DoRA を使用せずフルパラメータでファインチューニングし、2048 トークンのパッキング戦略、勾配チェックポイント、DeepSpeed ZeRO‑3 シャーディングを採用しました。主要ハイパーパラメータはステージ間で同一で、学習率 2e‑5、バッチサイズ 32、コサインスケジューラ、ウォームアップ比率は Stage 1 が 0.0、Stage 2 が 0.1 です。

データが全て

Numina は 2 つの大規模データセットを構築しました。

  • Chain‑of‑Thought データセット – 中国の高校ワークシート、米国の試験 PDF、国際オリンピックアーカイブから取得した数十万件の問題‑解答ペアです。パイプラインは OCR、セグメンテーション、英訳、CoT 形式への変換を行いました。
  • Tool‑Integrated Reasoning データセット – 約 60 000 件（主に整数出力）の問題を GPT‑4 パイプラインで処理し、ToRA 形式の推論を生成、コードを実行し、不一致を除外しました。各問題は正確性を保証するために最大 3 回再生成されました。

著者らは数週間以内にこの完全データセットをオープンソース化する予定です。

SC‑TIR: Self‑Consistency + Tool‑Integrated Reasoning

推論時、Kaggle が問題をランダム順に提供し GPU リソースが限られるため、モデルは高い分散を示します。SC‑TIR は以下の手順でこれを緩和します。

  1. 各問題を N 回複製（勝利ランでは N = 48）し、多様なプロンプトのバッチを作成。
  2. N 個の完了をサンプリングし、各々が Python コードブロックを生成。
  3. すべてのコードブロックを実行し、出力またはトレースバックを取得。
  4. 最大 M = 4 の深さまで再サンプリングし、過去のエラーに基づく自己修正を許可。
  5. 無効サンプルを除外し、最終的な数値解に多数決を適用。

8 ビット量子化（AutoGPTQ）により T4 GPU での実行が可能になり、アップロード時間が半減し、bfloat16 の非互換性を回避、VRAM 使用量も削減されましたが、精度への影響は僅かでした。

公開リーダーボードへの過学習防止

公開テストセットは 50 問のみであるため、チームは 4 つの内部検証スイートを構築しました。

  • AMC（83 件の整数出力問題） – 非公開テストセットを代表し、モデルは約 60‑65 % の正解率を示しました。
  • AIME（90 件） – 難易度が高く、失敗モードを露呈します。
  • MATH level 4（754 件） と MATH level 5（721 件） – 5 000 件の MATH ベンチマークから整数解のみを抽出したサブセットです。

5‑10 のランダムシードで繰り返し評価し、分散（SC‑TIR では通常 1‑3 %）を測定、信頼できるハイパーパラメータ調整を実現しました。

最終選考に至らなかった実験

MuMath‑Code パイプラインに落ち着くまで、チームは以下の代替案を検討しました。

  • 多数決付き純粋 CoT モデル。
  • MMOS（単ステップ Python）モデル – 最高 16/50。
  • Kahneman‑Tversky Optimisation（KTO）オンポリシーファインチューニング – 公開スコアは 27/50 に上昇したが、最終モデルの時間制限に間に合わなかった。
  • 強化学習（PPO と REINFORCE‑LOO）＋コード実行報酬 – 報酬曲線は有望だったが、精度向上は測定できず。
  • 大規模バックボーンへのスケーリング（InternLM‑20B、CodeLlama‑33B、Mixtral‑8x7B） – T4 GPU では推論が遅すぎた。
  • モデルマージング手法（DARE、TIES、WARP） – 内部指標での退化を引き起こした。

意義と今後の方向性

NuminaMath 7B TIR の成功は、高品質な数学データ、2 段階ファインチューニング体制、そして SC‑TIR による堅牢な推論を備えたオープンソース LLM が、オリンピックレベルの推論で競争力を持てることを示しています。この手法はより大規模なモデルにも拡張可能で（チームはより大きなバックボーン向けにデータセットを公開予定）、ツール利用が必要な他領域（例：記号計算、科学プログラミング）にも応用できます。

プロジェクトはオープンであるため、データセットの拡充、デコーディングアルゴリズムの改良、あるいは新たなモデルアーキテクチャの探索に貢献できる協力者やパートナーを歓迎します。コミュニティの継続的な取り組みにより、最終的な AIMO 目標である「金メダルを獲得できる AI」への道が一層近づくでしょう。

謝辞

著者らは、Thomas Wolf と Leandro von Werra が Numina–Hugging Face の協業を円滑にしてくれたこと、Hugo Larcher が GPU の提供を行ってくれたこと、Colin Raffel がモデルマージングに関する助言をくれたこと、そして Omar Sanseviero がブログ記事へのフィードバックをくれたことに感謝します。追加の支援は Mistral.ai、General Catalyst、Answer.ai、そして北京大学国際数学研究センターから提供されました。

モデルとデモへのアクセス

Sources