AIMO-2 Winning Solution: Building State-of-the-Art Mathematical Reasoning Models with OpenMathReasoning dataset
Paper • 2504.16891 • Published • 27
日本語数学推論のためのTool-Integrated Reasoning (TIR) モデルです。 CoTモデル(llm-jp-4-8bをCoTデータでSFTしたモデル)をベースに、134,834件のTIRデータセットでさらにSFTを行った最終モデルです。
自然言語による推論とPythonコード実行を組み合わせて数学問題を解くことができます。 開発用評価データセット(100問)において**正解率88.0%**を達成しました。
本モデルはFT-LLM2026コンペティションにおける我々の提出モデルです。
| 項目 | 値 |
|---|---|
| ベースモデル | llm-jp-4-8b (80億パラメータ) |
| 学習手法 | 2段階SFT(CoT SFT → TIR SFT)/ フルパラメータ |
| 学習データ | CoTデータ(306,366件)→ TIRデータ(134,834件) |
| フレームワーク | NeMo + Megatron-LM |
| 精度 | bf16-mixed |
OpenMathReasoning(NVIDIAのAIMO-2優勝手法)に倣い、2段階のSFTを実施しました。
llm-jp-4-8b (ベースモデル)
↓ CoT SFT(306,366件)
CoTモデル
↓ TIR SFT(134,834件)
TIRモデル(本モデル)★
| パラメータ | CoT SFT(第1段階) | TIR SFT(第2段階) |
|---|---|---|
| ベースモデル | llm-jp-4-8b | CoTモデル |
| 学習率 | 2e-5 | 2e-5 |
| エポック数 | 2 | 2 |
| バッチサイズ | 96 | 64 |
| 最大系列長 | 4,096 | 4,096 |
| オプティマイザ | AdamW (β₁=0.9, β₂=0.98) | AdamW (β₁=0.9, β₂=0.98) |
| 重み減衰 | 0.1 | 0.1 |
| ウォームアップ | 20 steps | 20 steps |
| スケジューラ | CosineAnnealing | CosineAnnealing |
| 精度 | bf16-mixed | bf16-mixed |
| PEFT | なし(フルパラメータ) | なし(フルパラメータ) |
開発用評価データセット(100問)における正解率です。推論時には1問あたり15個の回答候補を生成し、多数決(majority voting)により最終回答を決定しています。
| カテゴリ | 正解率 |
|---|---|
| 中1 | 1.000 |
| 中2 | 1.000 |
| 中3 | 0.938 |
| 数学I・A | 0.867 |
| 数学II・B | 0.731 |
| 数学III・C | 0.800 |
| 全体 | 0.880 |
| パラメータ | 値 |
|---|---|
| サンプル数(多数決用) | 15 |
| 最大コード実行回数 | 5回 |
| 最大生成トークン数 | 4,096 |
from transformers import AutoModelForCausalLM, AutoTokenizer
model_name = "AcroYAMALEX/acro-yamalex-llmjp-4-math-tir"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype="bfloat16", device_map="auto")
prompt = """あなたは「自然言語の推論」と「Pythonコードの実行」を組み合わせて数学問題を解くアシスタントです。
### 指示:
1から100までの自然数の和を求めてください。
### 応答:
"""
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
outputs = model.generate(**inputs, max_new_tokens=2048, temperature=0.7)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
本モデルの性能を最大限に引き出すには、生成されたPythonコードブロックを検出・実行し、結果をモデルにフィードバックするTIR推論パイプラインが必要です。
```python ブロックを検出し、サンドボックスで実行```output ブロックとしてモデルにフィードバック\boxed{}形式)が得られるまで繰り返し(最大5回)システムプロンプトでTIR推論の役割を指定し、ユーザーメッセージとして数学の問題を与えます。
<think>タグで囲まれた思考過程、Pythonコードブロック、コード実行結果を含むマルチターン対話を生成し、最終的に\boxed{}形式で回答を出力します。
<think>
Python を使ってシミュレーションしよう。
</think>
```python
total = sum(range(1, 101))
print(total)
5050
計算結果は5050です。
1から100までの自然数の和は:
## 関連リソース
| リソース | リンク |
|---|---|
| CoTデータセット | [AcroYAMALEX/acro-yamalex-llmjp-4-math-cot](https://huggingface.co/datasets/AcroYAMALEX/acro-yamalex-llmjp-4-math-cot) |
| TIRデータセット | [AcroYAMALEX/acro-yamalex-llmjp-4-math-tir](https://huggingface.co/datasets/AcroYAMALEX/acro-yamalex-llmjp-4-math-tir) |
| CoTモデル(中間モデル) | [AcroYAMALEX/acro-yamalex-llmjp-4-math-cot](https://huggingface.co/AcroYAMALEX/acro-yamalex-llmjp-4-math-cot) |
| 論文 | NLP2026にて発表予定 |
## 参考文献
- Ivan Moshkov et al. "AIMO-2 Winning Solution: Building State-of-the-Art Mathematical Reasoning Models with OpenMathReasoning dataset." arXiv:2504.16891, 2025.
- Ting Zhang et al. "StackMathQA: A Curated Collection of 2 Million Mathematical Questions and Answers Sourced from Stack Exchange." 2024.
- Keiran Paster et al. "OpenWebMath: An Open Dataset of High-Quality Mathematical Web Text." arXiv:2310.06786, 2024.
- DeepSeek-AI. "DeepSeek-V3 Technical Report." arXiv:2412.19437, 2024.
## 著者
佐々木峻・山本大輝・樋口慎・吉岡駿(アクロクエストテクノロジー株式会社)
## ライセンス
MIT License