Mini-R1: GRPOとCountdown GameによるDeepSeek-R1の推論の再現

Hugging Faceは、Group Relative Policy Optimization (GRPO) と Countdown Game と呼ばれる数学パズルを使用して、DeepSeek-R1 で観察された「aha moment」（モデルが人間のフィードバックなしに、より多くの思考時間を割り当て、アプローチを再評価することを学習する現象）をどのように再現できるかを示すテクニカルチュートリアルを公開しました。

Group Relative Policy Optimization (GRPO) の解説

GRPOは、大規模言語モデル (LLM) の推論能力を強化するために設計された強化学習 (RL) アルゴリズムです。DeepSeekMathの論文で導入されたGRPOは、従来のProximal Policy Optimization (PPO) から価値関数モデルの必要性を排除することで修正されています。代わりに、グループのスコアからベースラインを推定するため、メモリ使用量と計算オーバーヘッドが大幅に削減されます。

GRPOのプロセスは、主に以下の4つのステップに従います：

  1. Sampling: 現在のポリシーが、単一のプロンプトに対して複数の出力を生成します。
  2. Reward Scoring: 各生成結果が、報酬関数（ルールベースまたは結果ベース）を通じてスコア付けされます。
  3. Advantage Calculation: グループの平均報酬がベースラインとして機能します。各個別の解のアドバンテージは、この正規化されたグループ平均に対して計算されます。
  4. Policy Optimization: 計算されたアドバンテージとKLダイバージェンス項を組み込み、GRPOの目的関数を最大化するようにポリシーが最適化されます。

技術的な実装とセットアップ

Mini-R1の実験では、以下の技術スタックと構成が使用されました：

  • Base Model: Qwen/Qwen2.5-3B-Instruct。3Bパラメータモデルの選択は、モデルが推論プロセスを効果的に学習するには一般的に1.5B以上のパラメータが必要であるという観察に基づいています。
  • Dataset: Jiayi-Pan/Countdown-Tasks-3to4。3〜4個の数字を用いたパズルで構成されています。
  • Hardware: NVIDIA H100 80GB GPU x 4枚。
  • Software Stack: Hugging Face trl (GRPOTrainer用), transformers, datasets, accelerate, deepspeed, および生成を加速させるための vLLM。

報酬関数

人間のラベルなしで正解性を検証するために、トレーニングでは2つのルールベースの報酬関数が採用されました：

  • Format Reward: モデルが <think> [thinking] </think><answer> [answer] </answer> という構造に従うことを保証します。
  • Accuracy Reward: <answer> タグから方程式を抽出し、それがターゲットの数値と一致するか、および提供された各数字を正確に一度ずつ使用しているかを確認します。

分散トレーニング構成

トレーニングは DeepSpeed Zero-3 と vLLM を使用して実行されました。4-GPU構成では、3つのGPUがトレーニング (num_processes 3) に割り当てられ、最後の1つのGPUは vLLM による生成用に予約されました。450ステップのフルトレーニングランには約6時間かかり、各ステップには45〜60秒を要しました。

トレーニング結果と観察

トレーニングのパフォーマンスは TensorBoard を通じて追跡され、25ステップごとにチェックポイントが保存されました。モデルの挙動は、3つの異なるフェーズを経て進化しました：

  • Step 50: モデルは必要な <think> と <answer> のフォーマットを正常に学習しました。
  • Step 100: 成功率は約25%に達しました。モデルは試行錯誤のプロセスを自然言語の記述を用いて「推論」し始めました。
  • Step 200: 成功率は約40%に達しました。モデルは自然言語による推論から「プログラム的な実行」スタイルへと移行し、複数の組み合わせをリストアップして結果を体系的に確認するようになりました。
  • Step 450: 成功率は50%に達しました。モデルはプログラム的な推論形式を維持し、パフォーマンスは緩やかに向上し続けました。

ハイパーパラメータのチューニング

DeepSeekMathのハイパーパラメータ（学習率 1e-6、beta 0.04）を使用した初期の試みでは、150ステップ後にトレーニングが不安定になりました。学習率を 5e-7 に、beta (KL係数) を 0.001 に下げることで安定性が達成されました。

推論のシフトに関する分析

言葉ベースの推論からプログラム的な実行への移行は、いくつかの潜在的な要因に起因しています：

  • モデルの容量: Qwen 2.5 3Bは、DeepSeekで使用されているような大規模モデルと比較すると、複雑な自然言語推論を維持するには小さすぎる可能性があります。
  • 報酬の仕様: 報酬関数が、モデルが自然言語よりも効率的な数学的短縮記法を見つけるという「報酬ハッキング」を意図せず促進した可能性があります。
  • タスクの特異性: Countdown Gameのみでトレーニングを行ったことが、その特定のパズルタイプに対して最も効果的な解決方法へとモデルを自然に押し上げた可能性があります。
  • トレーニング期間: モデルのトレーニング期間が不十分であった可能性があります。元のR1の論文では、8,000ステップ以上のトレーニングが記録されています。

Sources