Open R1 アップデート #2: OpenR1-Math-220k データセットと推論の洞察

Hugging Face は OpenR1-Math-220k をリリースしました。これは、DeepSeek R1 の蒸留能力を再現するように設計された大規模な数学的推論データセットです。このリリースは、高性能な推論モデルを作成するために使用されたトレーニングパイプラインと合成データ生成方法をオープンソース化することを目的とする広範な Open R1 プロジェクトの一部です。

OpenR1-Math-220k データセット

OpenR1-Math-220k データセットは、蒸留を通じて小さなモデルに高度な推論能力を移転する基盤を提供し、元の DeepSeek R1 蒸留で使用された非公開の推論トレースによるギャップを埋めます。

データ生成パイプライン

Hugging Face は 512 個の H100 GPU を使ってローカルでデータセットを生成し、vLLM と SGLang を活用して推論トレースを生成しました。SGLang への移行により、ほぼ 2 倍のスピードアップが実現され、H100 あたりのスループットが 1 時間あたり 15 から 25 ソリューションに増加し、1 日あたり 180k トレースの生成が可能になりました。

生成プロセスの主要な技術仕様は以下の通りです：

  • ソース素材: NuminaMath 1.5 に基づく、NuminaMath-CoT の改良版。
  • スケール: 400k 問題に対して 800k の R1 推論トレースを生成（1 問あたり 2〜4 つの解答）。
  • プロンプト: DeepSeek R1 に「ステップバイステップで理由付けを行い、最終答案を \boxed{} 内に記述してください」と指示。
  • トークン制限: 分析により 25% の問題が 8k トークン以上を必要とすることが示されたため、1 生成あたり 16k トークンの制限を設定。

自動フィルタリングと品質管理

高品質な推論トレースを確保するため、Hugging Face はハイブリッド検証システムを採用しました：

  1. ルールベース検証: Math Verify を使用して最終答案を抽出し、グランドトゥルースデータと比較。55% の問題に少なくとも 1 つの正解が存在しました。
  2. LLMベース評価: グランドトゥルースの答案が不正または空である場合、Llama-3.3-70B-Instruct が判定者として機能し、モデルの解答が参照と等価かどうかを判定。これにより、以前に却下された 28,000 問題を回復しました。

最終データセットは、検証済み推論トレースを持つ 220k 問題から構成されます。2 つのスプリットで利用可能です：default スプリット（94k 問題）と extended スプリット（131k 問題）。default スプリットは、スーパーバイズドファインチューニング（SFT）後に最高のパフォーマンスを示しました。

パフォーマンスベンチマーク

Qwen2.5-Math-Instruct を default スプリットで 3 エポック（学習率 5e-5、32k コンテキスト長に対して RoPE 周波数を 300k に増加）ファインチューニングした結果、得られた OpenR1-Qwen-7B モデルは DeepSeek-Distill-Qwen-7B と競合するパフォーマンスを示しました：

モデル MATH-500 AIME24 AIME25
DeepSeek-Distill-Qwen-7B 91.6 43.3 40
OpenR1-Qwen-7B 90.6 36.7 40
OpenThinker-7B 89.6 30.0 33.3

Math-Verify の改善

Hugging Face は Math-Verify (v0.5.2) を更新し、複雑な数学的表現をより堅牢に扱えるようにしました。主な改善点は以下の通りです：

  • テキストのみの答案および答案のリストのパースを強化。
  • 単一の LaTeX 環境内での複数の箱入り答案をサポート。
  • ゴールド答案に基づいてセットとタプルを区別するための順序付きタプルを導入。
  • 関係式（例：$1 < x < 2$）および区間（例：$(1,2)$）のサポートを追加。

推論と GRPO に関するコミュニティの洞察

最近のコミュニティの動向は、オープンモデルにおける推論能力の引き出し方に変化が見られることを示しています。

GRPO とベースモデルの能力

  • ベースモデルの推論: Qwen2.5-0.5B ベースモデルに GRPO を適用した実験では、Instruct バージョンよりも GSM8k で 10 ポイントの改善が見られました。一部の研究者は、DeepSeek-R1 で述べられた「あ-ha」瞬間は、RL プロセス単独ではなく、ベースモデルの固有の能力の症状である可能性があると示唆しています。
  • 効率性: Unsloth は GRPO を最適化し、15GB の VRAM しかなくても 15B パラメータまでのモデルをトレーニングできるようにしました。
  • 検証不可能なドメイン: GRPO は詩にも正常に適用され、数学やコードのように従来検証可能なタスク以外にもその有用性を示しています。

データ効率と潜在空間推論

  • 小規模高品質データ: s1K（1,000 サンプル）や LIMO（817 サンプル）などのデータセットに関する研究は、モデルに十分な事前学習知識がある場合、非常に少ない数の高品質でよく構造化された例だけで高度な推論を引き出せることを示唆しています。
  • 潜在空間推論: 最近の論文では、再帰言語モデルを使って潜在空間で暗黙的に推論を行うことにより、テストタイムの計算をスケーリングし、広範な自然言語の「思考」トークンを生成するよりも計算効率が高い手法が提案されています。

思考の連鎖（CoT）長さの制御

  • バジェットフォーシング: 「Wait」または思考終了トークンを追加することで推論を延長または切り詰めるこの技術は、思考時間の増加が精度の向上に相関するテストタイムスケーリングを可能にします。
  • コサインリワード: 正解に対して短い CoT を、不正解に対して長い CoT をインセンティブとする新しい報酬関数を導入。これにより RL トレーニングが安定し、報酬ハッキング（モデルが繰り返しによって CoT 長を増やして報酬を得ること）を防止します。

Sources