Mini-R1: 透過 GRPO 和倒數遊戲重現 DeepSeek-R1 推理

Hugging Face 已發布一個技術教程，示範如何使用 Group Relative Policy Optimization (GRPO) 和一個稱為倒數遊戲的數學謎題來重現 DeepSeek-R1 觀察到的 'aha moment'——在此情境下，模型學會在無人類回饋的情況下分配更多思考時間並重新評估其方法。

Group Relative Policy Optimization (GRPO) 說明

GRPO 是一種強化學習（RL）演算法，旨在提升大型語言模型（LLM）的推理能力。在 DeepSeekMath 論文中首次提出，GRPO 通過移除對價值函數模型的需求來修改傳統的 Proximal Policy Optimization（PPO），而是透過群組分數估計基線，這樣顯著降低了記憶體使用量和計算開銷。

GRPO 流程包含四個主要步驟：

  1. 取樣：當前策略為單個提示生成多個輸出。
  2. 獎勵評分：每個生成透過獎勵函數進行評分（可以是基於規則或基於結果的）。
  3. 優勢計算：群組的平均獎勵作為基線；每個個別解決方案的優勢相對於此標準化群組平均值進行計算。
  4. 策略優化：策略透過併入計算出的優勢和 KL 散度項來優化，以最大化 GRPO 目標。

技術實作與設置

Mini-R1 實驗使用了以下技術棧與配置：

  • 基礎模型：Qwen/Qwen2.5-3B-Instruct。選擇 3B 參數模型是基於觀察：模型通常需要超過 1.5B 參數才能有效學習推理過程。
  • 資料集：Jiayi-Pan/Countdown-Tasks-3to4，包含 3 到 4 個數字的謎題。
  • 硬體：4 顆 NVIDIA H100 80GB GPU。
  • 軟體棧：Hugging Face trl（用於 GRPOTrainer）、transformers、datasets、accelerate、deepspeed 和 vLLM 以加速生成。

獎勵函數

為了在無人工標籤的情況下驗證正確性，訓練採用了兩個基於規則的獎勵函數：

  • 格式獎勵：確保模型遵循結構 `\ [thinking]

[answer] `。

  • 準確率獎勵：從 <answer> 標籤中提取方程式，並驗證其是否等於目標數字，且每個提供的數字恰好使用一次。

分散式訓練配置

訓練使用 DeepSpeed Zero-3 和 vLLM 執行。在 4 GPU 配置中，三個 GPU 專門用於訓練（num_processes 3），而最後一個 GPU 則保留給 vLLM 用於生成。完整的 450 步訓練大約需要 6 小時，每步耗時 45-60 秒。

訓練結果與觀察

訓練透過 TensorBoard 追蹤效能，每 25 步儲存一次檢查點。模型的行為經歷了三個不同階段：

  • 步驟 50：模型成功學會了所需的 \[thinking\] 和 <answer> 格式。
  • 步驟 100：成功率達到約 25%。模型開始使用自然語言描述其試誤過程來進行「推理」。
  • 步驟 200：成功率達到約 40%。模型從自然語言推理轉向「程式化執行」風格，即列出多種組合並系統地檢查結果。
  • 步驟 450：成功率達到 50%。模型保持程式化推理格式，效能繼續緩慢提升。

超參數調整

最初使用 DeepSeekMath 超參數（學習率 1e-6，beta 0.04）進行嘗試時，在 150 步後訓練變得不穩定。透過將學習率降低至 5e-7 以及將 beta（KL 係數）降至 0.001，達成了穩定。

推理轉變分析

從基於文字的推理轉向程式化執行的過渡歸因於以下幾個潛在因素：

  • 模型容量：與 DeepSeek 使用的較大模型相比，Qwen 2.5 3B 可能太小，無法維持複雜的自然語言推理。
  • 獎勵規範：獎勵函數可能無意間鼓勵了「獎勵 hacking」，模型在此找到了比自然語言更有效率的數學速記。
  • 任務特異性：僅在 Countdown Game 上進行訓練可能自然地將模型推向針對該特定謎題類型最有效的求解方法。
  • 訓練時長：模型可能訓練時間不足；原始 R1 論文記錄了超過 8,000 步的訓練。

Sources