Nemotron-Cascade的简化GRPO目标函数

引言

在大语言模型的强化学习训练中，如何设计一个既稳定又高效的目标函数一直是核心挑战。NVIDIA最新发布的Nemotron-Cascade论文提出了一个出人意料的简化方案：完全移除KL散度约束，将GRPO简化为最基础的REINFORCE算法。这个看似"倒退"的设计，却在数学推理任务上取得了显著效果。

本文将深入剖析这个目标函数的设计哲学、数学细节和实践意义。

核心目标函数

Nemotron-Cascade采用的GRPO目标函数形式如下：

$\mathcal{J}_{\text{GRPO}}(\theta) = \mathbb{E}_{(q,a)\sim\mathcal{D}, \{o_i\}_{i=1}^G\sim\pi_\theta(\cdot|q)} \left[\frac{1}{\sum_{i=1}^G |o_i|} \sum_{i=1}^G \sum_{t=1}^{|o_i|} \hat{A}_{i,t}\right]$

其中优势函数定义为：

$\hat{A}_{i,t} = \frac{r_i - \text{mean}(\{r_i\}_{i=1}^G)}{\text{std}(\{r_i\}_{i=1}^G)} \quad \text{对所有} \; t$

让我们逐个拆解这个公式的每个部分。

公式拆解

1. 采样过程


展开代码
(q, a) ∼ D

从数据集中采样一个问题-答案对
q: 问题（如数学题）
a: 标准答案


展开代码
{o_i}^G_{i=1} ∼ π_θ(·|q)

对同一个问题q，用当前策略生成G个不同的响应
这是GRPO的核心：群组采样（Group sampling）
典型的G值为4-8个

2. 奖励计算


展开代码
r_i = Reward(o_i, a)

对于数学推理任务（RLVR），奖励通常是二元的：

r_i = 1: 答案正确
r_i = 0: 答案错误

这个奖励是序列级（sequence-level）的，整个响应只有一个分数。

3. 组内标准化

这是GRPO的关键创新：

$\hat{A}_{i,t} = \frac{r_i - \mu_G}{\sigma_G}$

其中：

μ_G = mean({r₁, r₂, ..., r_G})：这一组响应的平均奖励
σ_G = std({r₁, r₂, ..., r_G})：这一组响应的奖励标准差

为什么这样做？

假设对于一道难题，模型生成了8个响应：

6个错误（r=0）
2个正确（r=1）

标准化后：

μ_G = 2/8 = 0.25
σ_G ≈ 0.43
错误响应的优势：Â = (0 - 0.25) / 0.43 ≈ -0.58（惩罚）
正确响应的优势：Â = (1 - 0.25) / 0.43 ≈ +1.74（奖励）

这种相对比较使得不同难度的问题产生的梯度尺度保持一致。

4. Token级传播

注意到优势函数 $\hat{A}_{i,t}$ 对所有token位置 t 都相同。这意味着：


展开代码
response: "Let's solve this step by step..."
          ↓    ↓     ↓    ↓    ↓   ↓
reward:   +1.74 +1.74 +1.74 +1.74 +1.74 +1.74 (如果答案正确)

虽然奖励是序列级的，但每个token都会被同等地奖励或惩罚。

5. 归一化因子

$\frac{1}{\sum_{i=1}^G |o_i|}$

这个因子对所有生成的token总数进行归一化：

如果8个响应平均长度200 tokens
总token数 = 8 × 200 = 1600
每个token的贡献被缩放为 1/1600

这确保了不同batch size和响应长度下梯度的可比性。

三大核心设计决策

决策1：移除KL散度项

传统RLHF/GRPO 通常包含KL约束：

$\mathcal{J} = \mathbb{E}[\text{reward}] - \beta \cdot \text{KL}(\pi_\theta || \pi_{\text{ref}})$

这个KL项的作用：

防止策略偏离初始模型太远
保持输出的多样性
提供正则化效果

Nemotron-Cascade的选择：完全移除KL项

为什么敢这么做？

严格on-policy训练：每次迭代都用最新策略采样，然后立即更新，importance ratio = 1
清晰的奖励信号：数学题的对错是明确的，不需要额外约束
熵正则化依然存在：组内采样本身就鼓励多样性

决策2：严格on-policy


展开代码
for iteration in training:
    # 1. 用当前策略采样
    responses = current_policy.generate(questions)
    
    # 2. 立即计算奖励
    rewards = evaluate(responses)
    
    # 3. 立即更新策略
    current_policy.update(responses, rewards)
    
    # 4. 丢弃这批数据，下次迭代重新采样

优点：

Importance sampling ratio = 1，无需复杂的重要性权重计算
避免off-policy带来的分布偏移
训练更稳定

代价：

数据效率较低（不能复用旧数据）
需要更多计算资源

决策3：群组标准化

对比三种优势估计方法：

方法A：全局baseline


展开代码
Â_i = r_i - baseline_all_data
问题：不同难度问题的梯度尺度差异巨大

方法B：价值函数baseline


展开代码
Â_i = r_i - V(state)
优点：方差更小
缺点：需要训练额外的价值网络

方法C：组内标准化（Nemotron采用）


展开代码
Â_i = (r_i - μ_group) / σ_group
优点：自动适应问题难度，无需额外网络
缺点：需要组内有足够的方差

梯度计算详解

让我们看一个具体例子，理解梯度如何流动。

场景设置：

问题：计算 23 × 47
生成4个响应（G=4）

响应	内容	长度	正确性	奖励 r_i
o₁	"23×47=1081"	50 tokens	✓	1.0
o₂	"Let me think... 1081"	80 tokens	✓	1.0
o₃	"23×47=1071"	45 tokens	✗	0.0
o₄	"I'll calculate... 1071"	60 tokens	✗	0.0

计算优势：


展开代码
μ = (1.0 + 1.0 + 0.0 + 0.0) / 4 = 0.5
σ = std([1.0, 1.0, 0.0, 0.0]) = 0.577

Â₁ = (1.0 - 0.5) / 0.577 = +0.866
Â₂ = (1.0 - 0.5) / 0.577 = +0.866
Â₃ = (0.0 - 0.5) / 0.577 = -0.866
Â₄ = (0.0 - 0.5) / 0.577 = -0.866

策略梯度（简化形式）：

$\nabla_\theta \mathcal{J} = \frac{1}{235} \left[\sum_{t=1}^{50} (+0.866) \nabla_\theta \log \pi_\theta(o_1^t) + \sum_{t=1}^{80} (+0.866) \nabla_\theta \log \pi_\theta(o_2^t) + \ldots \right]$

其中235是总token数（50+80+45+60）。

解读：

正确响应的每个token都获得正梯度（增加其概率）
错误响应的每个token都获得负梯度（降低其概率）
即使同样正确，较短的响应（o₁）每个token的贡献相对更大

与其他算法的对比

vs. 标准PPO

特性	Nemotron-GRPO	PPO
策略更新	On-policy	Near on-policy
KL约束	❌ 无	✓ Clipping或KL惩罚
价值网络	❌ 不需要	✓ 需要
优势估计	组内标准化	GAE（λ-return）
实现复杂度	⭐⭐	⭐⭐⭐⭐

vs. DPO

特性	Nemotron-GRPO	DPO
训练数据	生成的响应	预先标注的偏好对
在线/离线	在线（生成新数据）	离线（固定数据集）
奖励模型	可选	❌ 隐式
适用场景	明确奖励（数学）	偏好建模

vs. ReSTEM

ReSTEM（Reinforced Self-Training on Error-corrected Multi-step reasoning）的主要区别：

特性	Nemotron-GRPO	ReSTEM
训练方式	纯RL（策略梯度）	蒸馏+RL混合
数据利用	丢弃旧数据	保留高质量轨迹
迭代方式	单步更新	多轮蒸馏

理论洞察

为什么去掉KL反而更稳定？

这看似违反直觉，但有几个解释：

任务特性：数学推理的奖励信号足够dense和informative
on-policy机制：本身就提供了隐式的正则化
组内竞争：标准化机制让模型学习相对好坏，而非绝对分数
避免复杂性：移除KL项减少了超参数调优的困难

数学上的等价性

去掉KL项后，目标函数退化为经典REINFORCE：

$\nabla_\theta \mathcal{J} = \mathbb{E}_{\pi_\theta}\left[\hat{A} \cdot \nabla_\theta \log \pi_\theta(a|s)\right]$

这其实是最原始的策略梯度形式！但通过：

组内标准化（降低方差）
On-policy采样（减少偏差）
Token级归一化（统一尺度）

让这个简单的算法在现代LLM上依然有效。

方差-偏差权衡

组内标准化的巧妙之处：

方差：


展开代码
Var[Â] = Var[(r - μ_G) / σ_G] 
       = 1 / σ_G² · Var[r - μ_G]
       ≈ 1  (标准化后方差为1)

偏差：


展开代码
E[Â] = E[(r - μ_G) / σ_G] = 0  (期望为0)

这是一个无偏估计器，且方差被归一化，在不同问题间保持一致。

目录

引言