λLLM Wiki
3 min read#grpo#kl-divergence#unbiased-estimator#derivation

Derivation: GRPO's Unbiased Non-Negative KL Estimator

M
memset0
Apr 10, 2026

本页推导 GRPO 使用的 k3 KL 估计器——一个无偏且恒非负的 KL 散度蒙特卡洛估计量。原始讨论出自 John Schulman 的博客 “Approximating KL Divergence” (2020)，DeepSeekMath 把它搬进了 GRPO。阅读前建议先熟悉 kl-divergence。

问题设定

两个分布 p,qp, q，我们只有从 qq 采样的样本 x∼qx \sim q，想估计：

DKL(q∥p)=Ex∼q[log⁡q(x)p(x)]=−Ex∼q[log⁡r(x)]D_\text{KL}(q\|p) = \mathbb{E}_{x \sim q}\Big[\log \frac{q(x)}{p(x)}\Big] = -\mathbb{E}_{x \sim q}[\log r(x)]

其中 r(x)=p(x)/q(x)r(x) = p(x)/q(x) 是似然比。在 RL/RLHF 语境下：

  • q=πθq = \pi_\theta：当前策略
  • p=πrefp = \pi_\text{ref}：参考策略
  • 每步采样得到 (x,log⁡πθ(x),log⁡πref(x))(x, \log\pi_\theta(x), \log\pi_\text{ref}(x))

GRPO 需要在 loss 里加 β DKL(πθ∥πref)\beta\, \mathbb{D}_\text{KL}(\pi_\theta\|\pi_\text{ref}) 作为惩罚项，因此需要一个”在只有单样本的情况下、数值稳定、方差低、恒非负”的估计器。

三个候选估计器

Schulman 定义了三个估计量，分别称为 k1、k2、k3。

k1：朴素 MC 估计

k1(x)=−log⁡r(x)=log⁡q(x)p(x)k_1(x) = -\log r(x) = \log\frac{q(x)}{p(x)}

这是最直接的 MC 估计：Ex∼q[k1(x)]=DKL(q∥p)\mathbb{E}_{x\sim q}[k_1(x)] = D_\text{KL}(q\|p)，无偏。

问题：

  • 方差高，尤其当 r(x)r(x) 可能很大或很小时。
  • 可以为负：对单个样本，如果 q(x)<p(x)q(x) < p(x) 则 k1(x)<0k_1(x) < 0。但真实 KL 散度非负，所以单个负值的估计没有意义——在 loss 里甚至会鼓励策略”朝负 KL 方向”走，造成训练崩溃。

k2：平方差估计

k2(x)=12(log⁡r(x))2k_2(x) = \frac{1}{2}(\log r(x))^2

恒非负，但有偏——期望不是 DKLD_\text{KL}，而是 DKLD_\text{KL} 在 r→1r \to 1 附近的二阶泰勒展开的近似值。对于 rr 在 1 附近的情形，k2k_2 是一个合理的低方差近似；但当 rr 偏离 1 较远时 bias 变大。

k3：无偏且非负

k3(x)=(r(x)−1)−log⁡r(x)\boxed{k_3(x) = (r(x) - 1) - \log r(x)}

写成更对称的形式：令 y=log⁡r(x)=log⁡πref(x)−log⁡πθ(x)y = \log r(x) = \log\pi_\text{ref}(x) - \log\pi_\theta(x)，则

k3(x)=ey−y−1=(r(x)−1)−log⁡r(x)k_3(x) = e^y - y - 1 = (r(x) - 1) - \log r(x)

这就是 GRPO / DeepSeekMath 里看到的：

DKL[πθ∥πref]=πref(ot∣o<t)πθ(ot∣o<t)−log⁡πref(ot∣o<t)πθ(ot∣o<t)−1\mathbb{D}_\text{KL}[\pi_\theta\|\pi_\text{ref}] = \frac{\pi_\text{ref}(o_t|o_{<t})}{\pi_\theta(o_t|o_{<t})} - \log\frac{\pi_\text{ref}(o_t|o_{<t})}{\pi_\theta(o_t|o_{<t})} - 1

k3 的关键性质

性质 1：恒非负

k3(x)=ey−y−1,y=log⁡r(x)∈Rk_3(x) = e^y - y - 1,\quad y = \log r(x) \in \mathbb{R}

对所有 y∈Ry \in \mathbb{R}，ey≥1+ye^y \ge 1 + y（这是指数函数的凸性，或者等价地由 Taylor 展开 ey=1+y+y2/2+⋯e^y = 1 + y + y^2/2 + \cdots 得到，因为从二阶起所有项都 ≥0\ge 0）。所以：

k3(x)=ey−y−1≥0k_3(x) = e^y - y - 1 \ge 0

等号只在 y=0y = 0（即 p(x)=q(x)p(x) = q(x)）时取到。

在 loss 里使用时，这保证了 KL 惩罚项始终是非负的——不会出现”鼓励负 KL”的诡异梯度。

性质 2：无偏

要证明 Ex∼q[k3(x)]=DKL(q∥p)\mathbb{E}_{x\sim q}[k_3(x)] = D_\text{KL}(q\|p)。

Ex∼q[k3(x)]=Ex∼q[r(x)−1−log⁡r(x)]\mathbb{E}_{x\sim q}[k_3(x)] = \mathbb{E}_{x\sim q}[r(x) - 1 - \log r(x)]

分开三项：

=Ex∼q[r(x)]−Ex∼q[1]−Ex∼q[log⁡r(x)]= \mathbb{E}_{x\sim q}[r(x)] - \mathbb{E}_{x\sim q}[1] - \mathbb{E}_{x\sim q}[\log r(x)]

逐项：

  • Ex∼q[r(x)]=∫q(x)⋅p(x)q(x) dx=∫p(x) dx=1\mathbb{E}_{x\sim q}[r(x)] = \int q(x)\cdot \frac{p(x)}{q(x)}\, dx = \int p(x)\, dx = 1
  • Ex∼q[1]=1\mathbb{E}_{x\sim q}[1] = 1
  • Ex∼q[log⁡r(x)]=Ex∼q[log⁡p(x)−log⁡q(x)]=−DKL(q∥p)\mathbb{E}_{x\sim q}[\log r(x)] = \mathbb{E}_{x\sim q}[\log p(x) - \log q(x)] = -D_\text{KL}(q\|p)

合起来：

Ex∼q[k3(x)]=1−1−(−DKL(q∥p))=DKL(q∥p)\mathbb{E}_{x\sim q}[k_3(x)] = 1 - 1 - (-D_\text{KL}(q\|p)) = D_\text{KL}(q\|p)

所以 k3 是 DKL(q∥p)D_\text{KL}(q\|p) 的无偏估计量。

性质 3：方差较低

直觉：k3 可以被看作在 k1 上加了一个”控制变量(control variate)“来降方差。具体地：

k3(x)=k1(x)⋅(−1)+(r(x)−1)k_3(x) = k_1(x)\cdot (-1) + (r(x) - 1)

上面已经证明 E[r(x)−1]=0\mathbb{E}[r(x) - 1] = 0，所以这一项不改变期望；但它与 k1k_1 相关性为正（rr 大时 k1k_1 小，两者反号，加起来相互抵消），整体方差比朴素的 −k1-k_1 更低。

Schulman 的博客还给出了几个数值实验，在 q,pq, p 为高斯分布的情形下，k3 的 RMSE 比 k1 小一到两个数量级。

为什么不用 k2

k2 也是非负的，而且 E[k2]≈DKL\mathbb{E}[k_2] \approx D_\text{KL} 当 p,qp, q 接近时。但它有 bias——并不等于 DKLD_\text{KL}。如果希望 KL 惩罚能准确反映两个分布的真实距离（而不是一个二次近似），就应该用 k3。

更进一步，在 RLHF 训练中策略可能偏离 reference 较远（rr 远离 1），此时 k2 的 bias 会变大，k3 仍然无偏。

代码实现

# 给定 per-token log-probs:
# logp_ref = log π_ref(o_t | o_<t)
# logp = log π_θ(o_t | o_<t)
log_ratio = logp_ref - logp  # y = log r
kl_per_token = log_ratio.exp() - log_ratio - 1  # k3 estimator

乘以 mask（屏蔽 padding token）后在 loss 里加 β⋅kl_per_token\beta \cdot \text{kl\_per\_token} 即可。注意：

  • 符号方向：k3 估计的是 DKL(πθ∥πref)D_\text{KL}(\pi_\theta \| \pi_\text{ref})。如果想反过来估 DKL(πref∥πθ)D_\text{KL}(\pi_\text{ref}\|\pi_\theta)，要把 log⁡r\log r 取反。两者在 RLHF 里物理含义不同（见 kl-divergence 里的正向/反向 KL 讨论），不能混用。
  • log_ratio 的数值范围：如果 policy 与 reference 差异极大，exp⁡(log⁡r)\exp(\log r) 可能数值溢出；实现时通常 clip 一下 log⁡r\log r 或者用 more stable 的方式组合。
  • token-level vs sequence-level：上面给出的是 per-token KL。GRPO 沿着 sequence 平均后再乘 β\beta 加到 loss 里。

与 PPO 的 KL 写法的差异

PPO 在 RLHF 里通常把 KL “减在 reward” 里：

rt′=rt−βlog⁡πθ(ot∣o<t)πref(ot∣o<t)r_t' = r_t - \beta\log\frac{\pi_\theta(o_t|o_{<t})}{\pi_\text{ref}(o_t|o_{<t})}

这里用的是 k1（朴素对数比），会引入高方差和可能为负的 KL 估计——但因为 PPO 用 GAE 把 reward 展成 advantage 后又做了标准化，原始 reward 里的噪声被大幅滤除，问题没那么严重。

GRPO 直接把 KL 项写进 loss，而且 reward 进入 advantage 计算时不带 KL——此时 KL 项是一个独立的梯度分量，所以它必须既无偏又非负，否则会直接污染 loss。这是 GRPO 选择 k3 而非 k1 的核心原因。

参考

Comments