SOTAのさらに先へ：厳しい推論制約下での高性能モデルのPost-Training

Reference:
https://www.kaggle.com/competitions/ai-mathematical-olympiad-progress-prize-3/writeups/pushing-the-limits-post-training-high-capability

Transcript

1. SOTAのさらに先へ 厳しい推論制約下での 高性能モデルの Post-Training Pushing the Limits: Post-Training High-Capability Models under Strict Inference 2026年5月22日 @ 関西Kaggler会 Sakana AI Hiroshi Yoshihara (RabotniKuma) 1 1
2. チーム & 結果 Team: FishMath Public LB 44 (27位) Private LB 41.5 (899位) 2 2
3. コンペ概要: AIMO-3 AI Mathematical Olympiad Prize 第 3 回 賞金総額 $10M 問題: 110 問 (Reference 10 + Public 50 + Private 50 ) 全問オリジナル（汚染リスクなし） 範囲: 代数・組合せ・幾何・数論 今回は難易度↑ (National Olympiad ～ IMO 級) 同じsubmissionを2回実行して平均点で評価 3 3
4. 背景・課題 gpt-oss-120b は AIME'25 で既に ~98% 既存ベンチは飽和、データ汚染も疑い Kaggle 制約: 5h / H100 x 1 Post-training と推論設計を同時に最適化する必要 4 4
5. アプローチ 1. 汚染なし評価ベンチを自作 2. SFT で正答性を底上げ 3. RL (GRPO) で分布を尖らせる 4. 推論最適化（温度・EAGLE-3）で 5h に収める シンプルさと実行安定性を最優先 5 5
6. ベンチマーク構成 Benchmark Source Number of Problems AIMO 30 AIME 2025 30 HMMT Feb 2025 13 HMMT Nov 2025 22 Brumo 2025 16 CMIMC 2025 15 APEX 2025 6 SMT 2025 25 AMO-Bench 20 Nemotron-Math-v2 31 AoPS 2025 18 JMO 15 IMO-Answer Bench 44 Total 285 285 problems Primary category composition Combinatorics: 91 (31.9%) Algebra: 65 (22.8%) Geometry: 48 (16.8%) Number Theory: 45 (15.8%) Unlabeled: 36 (12.6%) 6 6
7. 評価ベンチマーク 285 問 動機: gpt-oss-120b は AIME'25 で ~98% → 既存ベンチ飽和 + 汚染懸念 真の zero-shot 推論能力を測れる 汚染なし整数回答ベンチを自作 7 7
8. ベンチマークのソース 2025 AIME / HMMT (Feb, Nov) — 最新版で汚染リスク低 MathArena 由来の整数問題 JMO — 英語学習データの外側を狙う Nemotron-Math-v2 — gpt-oss-120b が low/med effort で全滅 AoPS 2025 — LLM ループで書き起こし & 検証 success rate ∈ (0, 1) のみ採用 IMO-AnswerBench — proof 問題を整数回答に再構成 8 8
9. 時間制約ベンチ (50 問) 目的: Kaggle 環境（5h / H100 x 1）を再現 ソース: Crystal-Math-Preview 提出 runtime に合うように問題を選定 温度・EAGLE-3・wall time の検証に使用 9 9
10. 評価指標 pass@1 — 平均正答率（1 問あたり 16 trace） Maj@8 — 多数決精度（2,000 bootstrap で安定化） Coverage = pass@16 — 16 試行で 1 回でも正解した割合 TIR あり / なし の両モードで測定 10 10
11. SFT データ構築方針 難問フィルタ + 強モデルで再生成 ソース: Nemotron-Math-v2, Crystal-Math-Preview （評価ベンチと重複する問題は除外） 各問題で 4 trace 生成 → 正常終了 & 正答のみ採用 検証: 整数は完全一致、非整数は gpt-oss-120b judge 11 11
12. SFT データ内訳 合計 22,287 traces / 669.7M tokens ソース 構成 traces tokens Nemotron-Math-v2 (No TIR) Kimi-K2.5, DS-V3.2 8,403 195.7M Nemotron-Math-v2 (TIR) Kimi-K2.5, GLM-4.7 1,296 32.8M Crystal-Math (No TIR) Kimi-K2.5, GLM-5 8,073 290.3M Crystal-Math (TIR) Kimi-K2.5 4,515 150.9M TIR を 3× upsampling → 1 epoch あたり 551M / 1.037B tokens を TIR が占有 12 12
13. モデル並列の用語 略称 名称 仕組み PP Pipeline Parallel レイヤーを縦に分割し GPU 間でステージとして流す TP Tensor Parallel 1 レイヤーの重み行列を横方向に分割し並列計算 EP Expert Parallel MoE の各 expert を別 GPU に配置しルーティング後に分散実行 13 13
14. SFT 学習レシピ Megatron-SWIFT, 4 node × 8 GPU PP=4, attention TP=4, MoE EP=8 Max seq 65,536 + sequence packing Adam (β=0.9/0.95), LR 8e-6 → 8e-7 cosine, warmup 5% Weight decay 0.1, grad clip 1.0, bf16 Global batch 256 × 10 epoch ≈ 10.37B tokens 14 14
15. SFT 結果（TIR あり） pass@1 Maj@8 Coverage Base 72.61 77.86 92.28 Epoch 2 68.60 75.88 87.02 Epoch 5 70.88 76.85 88.77 Epoch 10 73.38 78.62 90.88 TIR なしでも Base 64.65 → Ep10 65.79 (pass@1) 15 15
16. SFT 観察 2ep / 5ep は 一旦劣化、10ep で初めて Base 超え Coverage は Base が最大 — SFT は解空間を広げない 改善の正体は 「平均正答性と多数決の信頼性向上」 厳選した難問 trace への反復露出が効く 16 16
17. GRPO とは？ Group Relative Policy Optimization 1 prompt から K 個の trace をサンプル グループ内の平均報酬を基準に 相対 advantage を計算 critic（価値関数）不要 → メモリ・計算が軽い 数学・コード推論 RL のデファクトに Shao et al., 2024 (DeepSeek-Math) · arXiv:2402.03300 17 17
18. RL データ準備 allenai/Dolci-Think-RL-7B の数学サブセット 60,360 → 13,358 問 (≈22%) フィルタは gpt-oss-120b (high reasoning + TIR) で 16 試行して決定: 答えが 整数 [0, 99999] 正答数が 0–11 / 16（= 確実に解ける問題は除外） 「確率変動の大きい境界問題」だけを残し RL の学習効率を上 げる 18 18
19. RL アルゴリズムとレシピ GRPO + Clip-Higher + Dynamic Sampling (DAPO) + TIS + KL loss を除去 Max seq 65,536, samples/prompt = 16 Rollout batch 64, global batch 1024 AdamW (β=0.9/0.98), LR 1e-6 constant Clip range [0.2, 0.28], KL=0, entropy=0 Long-tail 対策に partial rollout checkpoint = step 70 を採用 DAPO: arXiv:2503.14476 · partial rollout: arXiv:2507.20534 19 19
20. Partial Rollout とは？ 通常の GRPO は K 本の rollout が全部終わるまで待つ → 1 本でも長いと律速 (long-tail problem) rollout を途中で打ち切り、次 iteration で続きを生成 待ち時間を削減 → throughput 改善 続きは最新ポリシーで生成 → 軽い off-policy 性 arXiv:2507.20534 20 20
21. RL 結果 TIR あり pass@1 Maj@8 Coverage Base 72.61 77.86 92.28 RL 73.61 78.78 90.17 TIR なし: pass@1 64.65 → 65.61（Maj@8/Coverage はや や低下） 21 21
22. SFT vs RL の挙動差 Public LB ピーク: SFT 43 / RL 44 RL は正解寄りに確率質量を集中 → pass@1 / Maj@8 ↑ / Coverage ↓ RL モデルは生成時間がやや長い → 5h+30min 余裕で動かすには T ≤ 0.8 に制約 SFT は T=1.0 でも完走 → 最終提出は SFT 22 22
23. 温度ごとの性能比較 67 68 69 70 71 72 73 74 75 Accuracy (%) baseline Mean accuracy 72 74 76 78 80 Accuracy (%) baseline Bootstrap majority 0.5 0.6 0.7 0.8 0.9 1.0 1.1 1.2 1.3 1.4 1.5 Temperature 82 84 86 88 90 92 94 Problems solved (%) baseline Pass coverage 0.5 0.6 0.7 0.8 0.9 1.0 1.1 1.2 1.3 1.4 1.5 Temperature 0 5 10 15 20 25 30 35 40 45 Mean tokens (k) ± 1 std baseline Token length 23 23
24. 5h 制約下での温度選択 Temperature sweep under a five-hour submission budget Pass@1 72 72.5 73 73.5 0.6 0.7 0.8 0.9 1.0 Score (%) Maj@8 76.5 77 77.5 78 78.5 0.6 0.7 0.8 0.9 1.0 Coverage 89 90 91 92 0.6 0.7 0.8 0.9 1.0 Decoding temperature T RL (T <= 0.8 feasible) SFT (T = 1.0 feasible) Base RL T > 0.8 Selected 24 24
25. 推論最適化: 温度の検証 T=0.5 → 1.0 で 精度↑、トークン長も↑ pass@1 / Maj@8 ともに T=1.0 がピーク T=1.5 では劣化 Coverage は T≈0.6 で頭打ち 25 25
26. EAGLE-3 とは？ 投機的デコーディング (speculative decoding) の一種 小さな draft model が複数トークンを先読み target model が 並列に検証 → 受理ぶんだけ採用 出力分布は元モデルと 等価（lossless） EAGLE-3 は draft に target の 複数層の隠れ状態を入力 → accept rate↑ EAGLE-3: arXiv:2503.01840 · speculative decoding: arXiv:2211.17192 26 26
27. 推論最適化: EAGLE-3 50 問ベンチでの wall time: speculative なし: 342 min T=0.5 + EAGLE-3: 226 min（accept 0.50） T=1.0 + EAGLE-3: 253 min（accept 0.37） 5h 制約内で T=1.0 が走り切る 27 27
28. SAR とは？ Self-Aggregation Refinement 複数の完全な reasoning trace を 1 つの prompt に連結 モデルに「これらを総合してより良い答えを出して」と再生 成 多数決とは違い、推論過程そのものを参照できる 期待 +16.9% だが context 長を強く消費 Recursive Self-Aggregation (RSA) ベース · arXiv:2509.26626 28 28
29. GenSelect とは？ AIMO-2 優勝チーム (Moshkov et al., 2025) 提案 複数の候補解を並べ、モデル自身に最良を選ばせる 候補数が多い時は トーナメント方式 多数決が割れた時のタイブレーカーとして使える match=4 で +8pt の改善ポテンシャル arXiv:2504.16891 29 29
30. 採用しなかった推論技法 SAR — 期待 +16.9% だが context 制約で実用困難 GenSelect — match=4 で +8pt の伸びしろ、ただし分散 大 DFlash — 低温度では速いが T=1.0 では EAGLE-3 優位 → シンプルな 温度スケーリング + EAGLE-3 が最安定 DFlash: Chen et al., 2026 · arXiv:2602.06036 30 30
31. 最終提出 採用: SFT モデル @ T=1.0 + EAGLE-3 RL モデルは推論オーバーヘッドで T ≤ 0.8 制約 → 採用見送り 30 分のセーフティマージンを確保 Public LB 43 Private LB 41.5 31 31
32. Key Takeaways Post-training と推論設計は結合して最適化 精度だけでなく遅延・実行安定性を見る コンペの評価は2回のみ実行の平均点 → 学習による改善幅 < ランダム性 32 32
33. Resources Kaggle Writeup (Pushing the Limits) SakanaAI/gpt-oss-120b-sft-aimo3-fishmath SakanaAI/FishMath-SFT-Data Thank you! 33 33