Mini-R1: 通过 GRPO 和倒计时游戏再现 DeepSeek-R1 推理

Hugging Face 已发布一份技术教程，演示了如何使用 Group Relative Policy Optimization (GRPO) 和一种称为倒计时游戏的数学谜题来再现 DeepSeek-R1 中观察到的 'aha moment'——在这种情境下，模型学会在没有人工反馈的情况下分配更多思考时间并重新评估其方法。

Group Relative Policy Optimization (GRPO) 解释

GRPO 是一种旨在增强大型语言模型 (LLM) 推理能力的强化学习 (RL) 算法。它在 DeepSeekMath 论文中被提出，通过去除价值函数模型的需求，GRPO 对传统的近端策略优化 (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 个数字的谜题。
  • 硬件：4x NVIDIA H100 80GB GPU。
  • 软件栈：Hugging Face trl（用于 GRPOTrainer）、transformers、datasets、accelerate、deepspeed 和 vLLM 以实现加速生成。

奖励函数

为了在没有人工标签的情况下验证正确性，训练采用了两种基于规则的奖励函数：

  • 格式奖励：确保模型遵循结构 $$ [thinking] $$ <answer> [answer] </answer>。
  • 准确率奖励：从 <answer> 标签中提取方程，并验证其是否等于目标数字，并且每个提供的数字恰好使用一次。

分布式训练配置

训练使用 DeepSpeed Zero-3 和 vLLM 执行。在 4 GPU 配置中，三个 GPU 用于训练（num_processes 3），而最后一个 GPU 保留用于 vLLM 生成。完整的 450 步训练大约需要 6 小时，每步耗时 45-60 秒。

训练结果与观察

训练性能通过 TensorBoard 追踪，每 25 步保存一次检查点。模型的行为经历了三个不同的阶段：

  • 步骤 50：模型成功学习了所需的 $$ 和 <answer> 格式。
  • 步骤 100：成功率达到约 25%。模型开始使用自然语言描述其试错过程进行 "reasoning"。
  • 步骤 200：成功率达到约 40%。模型从自然语言 "reasoning" 转向 "programmatic execution" 风格，其中它列出多种组合并系统地审查结果。
  • 步骤 450：成功率达到 50%。模型保持了 "programmatic execution" 推理格式，性能继续缓慢提升。

超参数调优

初次使用 DeepSeekMath 超参数（学习率 1e-6，beta 0.04）尝试后，在 150 步后训练变得不稳定。通过将学习率降低到 5e-7 并将 beta（KL 系数）降低到 0.001，实现了稳定。

推理转变分析

从基于词的推理转向程序化执行的转变可归因于以下几个潜在因素：

  • 模型容量：与 DeepSeek 使用的较大模型相比，Qwen 2.5 3B 可能太小，无法维持复杂的自然语言推理。
  • 奖励规范：奖励函数可能无意中鼓励了 "reward hacking"，即模型发现了一种比自然语言更高效的数学速记法。
  • 任务特异性：仅在 Countdown Game 上训练可能自然地将模型推向针对该特定谜题类型的最有效求解方法。
  • 训练时长：模型可能训练时间不足；原始 R1 论文记载了超过 8,000 步的训练。

Sources