KL Regularization：奖励塑形与显式损失#

1. 采样惩罚的数值与梯度#

假设两个动作的参考概率各为 $0.5$，采样策略给它们的概率是 $0.7$ 和 $0.3$。采样以后，策略概率改成 $0.8$ 和 $0.2$，但这批数据中各动作的样本计数没有改变。此时，直接对这批样本上的 KL 惩罚反传，是否等于对当前策略的完整 KL 求导？

PPO 的 KL 奖励与 GRPO 的采样 KL loss 可能都使用 KL 惩罚这一名称，实际更新方向却不同。需要区分动作概率的变化与采样表达式自身的变化。

2. 固定状态与两种求导路径#

一个期望是“概率乘样本值”的和。完整求导需要考虑两部分：（1）各动作的概率随参数变化；（2）各动作对应的样本值随参数变化。把 KL 放进停止梯度的奖励，和直接对采样 KL loss 反传，分别保留了不同的导数项。某条路径能否省略，要证明它的期望为零。方向选择见 Reference KL。

2.1 固定状态下的正则目标#

先固定状态 $s$，以下概率省略相同的条件 $s$。对 LLM，可以分析某个前缀下的 token 分布，也可以把完整回答视为 prompt 下的一次动作。$\pi_\theta$ 是当前策略，$\pi_{\text{ref}}$ 是固定参考策略。

任务奖励 $R^{\text{task}}(a)$、reference 和系数 $\beta$ 都不随 $\theta$ 改变。目标是 $J(\theta)=\mathbb E_{a\sim\pi_\theta}[R^{\text{task}}(a)]-\beta D_{\mathrm{KL}}(\pi_\theta\|\pi_{\text{ref}})$。

这里分析正的 KL 惩罚；放进最大化目标时乘 $-\beta$，作为最小化 loss 的惩罚时乘 $+\beta$。假定两种策略在相同且固定的支持上概率为正、均归一化，且允许交换求和与求导。

2.2 采样概率与样本表达式的变化#

需要区分两个操作：（1）对固定样本直接反传，只计算所采动作对应表达式的导数；（2）对参数相关的完整期望求导，还要考虑各动作在平均中出现的概率怎样变化。

假设采样时 A 占 70%、B 占 30%，一批 100 个样本恰好有 70 个 A、30 个 B。更新后模型的概率变成 80%/20%。

直接反传固定样本的平均惩罚时，虽然重算了各动作的惩罚值，加权比例仍由原有样本计数决定。对完整期望求导，则还要计算加权概率随参数变化的贡献。因此，需要分别判断：（1）样本值的平均是否等于 KL；（2）直接反传是否得到 KL 的梯度。

3. 奖励塑形与显式 loss 的关系#

3.1 参数相关期望的两条求导路径#

令 $\widehat D_\theta(a)$ 表示单个动作的采样 KL 表达式；最直接的选择是 $\log[\pi_\theta(a)/\pi_{\text{ref}}(a)]$，其在 current 下的期望就是 KL。令 score 向量 $z(a)=\nabla_\theta\log\pi_\theta(a)$，乘法法则给出：

$$\begin{aligned} \nabla_\theta\mathbb E_{a\sim\pi_\theta}[\widehat D_\theta(a)] &=\sum_a\nabla_\theta\pi_\theta(a)\,\widehat D_\theta(a) +\sum_a\pi_\theta(a)\nabla_\theta\widehat D_\theta(a)\\ &=\mathbb E[\widehat D_\theta(a)z(a)] +\mathbb E[\nabla_\theta\widehat D_\theta(a)]. \end{aligned}$$

• 概率变化项 $\mathbb E[\widehat D_\theta z]$：动作出现的比例变化，由 policy gradient 计算。
• 样本值变化项 $\mathbb E[\nabla_\theta\widehat D_\theta]$：每个动作对应的采样 KL 数值变化，由直接反传计算。

固定样本不会随参数变化而自动更换动作。对其中的 $\widehat D_\theta(a)$ 直接反传，只计算样本值变化项。

把采样策略记作 old。旧样本的平均值估计的是固定采样分布下的期望，即 $\sum_a\pi_{\text{old}}(a)\widehat D_\theta(a)$；完整期望则用随参数变化的 $\pi_\theta(a)$ 加权。即使在 current 等于 old 的起点，两者数值相同，导数也可能不同。增加样本量只能提高固定权重平均的估计精度，不能恢复遗漏的概率变化项。

“数值无偏”要求采样表达式的期望等于 KL；“梯度无偏”要求采样梯度的期望等于完整 KL 的梯度。 是否能只保留一条路径，要看另一项的期望是否为零，不能只由 $\mathbb E[\widehat D_\theta]=D_{\mathrm{KL}}$ 判断。[1]

3.2 log-ratio 奖励对应的完整梯度#

这里把 $D_{\mathrm{KL}}(\pi_\theta\|\pi_{\text{ref}})$ 称为 reverse KL，其期望按 current 加权。

取 $\widehat D_\theta(a)=\log[\pi_\theta(a)/\pi_{\text{ref}}(a)]$。它在当前策略下的期望就是 KL，而样本值自身的导数是 $z(a)$。这个导数的期望恰好为零：

$$\begin{aligned} \mathbb E_{a\sim\pi_\theta}[z(a)] &=\sum_a\pi_\theta(a)\nabla_\theta\log\pi_\theta(a)\\ &=\nabla_\theta\sum_a\pi_\theta(a) =\nabla_\theta 1=0. \end{aligned}$$

因此，reverse KL 使用 log-ratio 时，样本值变化项的期望为零，完整梯度只剩概率变化项：

$$\boxed{ \nabla_\theta D_{\mathrm{KL}}(\pi_\theta\|\pi_{\text{ref}}) =\mathbb E_{a\sim\pi_\theta}\left[ \log\frac{\pi_\theta(a)}{\pi_{\text{ref}}(a)}\,z(a) \right]. }$$

把上式中的 reverse KL 梯度与任务奖励梯度合并，得到：

$$\boxed{ \nabla_\theta J =\mathbb E_{a\sim\pi_\theta}\left[ \left(R^{\text{task}}(a) -\beta\log\frac{\pi_\theta(a)}{\pi_{\text{ref}}(a)}\right)z(a) \right]. }$$

括号内就是含 KL 的奖励。按当前策略采样，用当前概率计算这份奖励，把它作为停止梯度的权重，再乘 score，就能无偏估计同一个正则目标的梯度。 停止奖励的梯度，舍去的是期望为零的样本值变化项。[2]

PPO 也使用固定样本，但通过 importance ratio 的导数得到“奖励权重 × score”。这就在估计概率变化项，无需改变固定样本中的动作。

若动作代表完整回答，log-ratio 可分解为逐 token 之和，再分配到逐 token 奖励，并用完整 return 做信用分配。这里按完整终止、$\gamma=1$ 分析；GAE、价值函数误差和 advantage 归一化带来的近似另计。

3.3 k2 loss 与奖励塑形的梯度对应#

固定前缀并按 current 采样时，一个能无偏估计 reverse KL 梯度的样本 loss 是 $\tfrac12\big(\log[\pi_\theta(a)/\pi_{\text{ref}}(a)]\big)^2$。保持采样动作固定，对它求导：

$$\nabla_\theta\frac12\left( \log\frac{\pi_\theta(a)}{\pi_{\text{ref}}(a)}\right)^2 =\log\frac{\pi_\theta(a)}{\pi_{\text{ref}}(a)}\,z(a).$$

其在当前策略下的期望梯度，正好是完整 KL 的梯度。这两种正则方式分别称为“$k_1$ in reward”和“$k_2$ as loss”。平方项的数值期望未必等于 KL；这里要求样本梯度的期望等于目标 KL 的梯度。 [2]

4. k3 的期望数值与反传梯度#

4.1 采样惩罚 k3 的定义#

原始 GRPO 把 KL 单列为 loss，使用的采样表达式通常称为 $k_3$。它在普通 log-ratio 上多加了一项：

$$\widehat D_\theta(a) =\log\frac{\pi_\theta(a)}{\pi_{\text{ref}}(a)} +\frac{\pi_{\text{ref}}(a)}{\pi_\theta(a)}-1.$$

等价地，它也可写成 $\pi_{\text{ref}}/\pi_\theta-\log(\pi_{\text{ref}}/\pi_\theta)-1$。[3]

4.2 两动作分布的期望数值#

固定前缀，只有动作 A、B。Reference 各给 50%，采样时 current 给 A 70%、B 30%。设 100 条样本恰好有 70 个 A、30 个 B，排除采样频率误差，只比较求导路径。

现在让 current 变成 A 80%、B 20%。完整期望既要按新比例加权，也要用新概率重算 $k_3$。先计算每个动作的值：

起点的完整期望与这批样本的平均值相同，都是 $\text{起点 KL 值}=0.7\times0.050758+0.3\times0.155841\approx0.082283$。

参数变化后，完整期望既更新动作数值，也更新加权比例，得到 $\text{新 current 的 KL 值}=0.8\times0.095004+0.2\times0.583709\approx0.192745$。这正是新 current 与 reference 的 KL。

直接对原来那批样本反传，只更新动作数值，样本仍然是 70 个 A、30 个 B，对应的平均值为 $\text{固定旧样本的估计值}=0.7\times0.095004+0.3\times0.583709\approx0.241615$。参数变化不会改变这批已采样数据中 A、B 的数量。

4.3 k3 的期望与直接反传梯度#

“动作按 current 采样”意味着用 $\pi_\theta(a)$ 给每个动作加权。因此，概率比的期望是：

$$\mathbb E_{a\sim\pi_\theta}\left[\frac{\pi_{\text{ref}}(a)}{\pi_\theta(a)}\right] =\sum_a\pi_\theta(a)\frac{\pi_{\text{ref}}(a)}{\pi_\theta(a)} =\sum_a\pi_{\text{ref}}(a)=1.$$

采样概率与分母抵消，剩下 reference 的全部概率之和。单个动作的概率比不必为 1，按 current 加权后的平均才是 1。

多加的 $\pi_{\text{ref}}/\pi_\theta-1$ 平均为零，所以 $k_3$ 的完整期望仍然是同一个 KL。 对这个完整期望求导，结果也必然相同。但直接反传固定样本只计算样本值变化项，并未包含概率变化项。

GRPO 把 $k_3$ 作为每个采样 token 的独立惩罚，并保留对 current 的梯度。对固定动作直接反传，得到：

$$\nabla_\theta\widehat D_\theta(a) =\left(1-\frac{\pi_{\text{ref}}(a)}{\pi_\theta(a)}\right)z(a).$$

固定前缀、reference 固定、动作按 current 采样时，这条单样本梯度的期望正是 forward KL 的梯度。 因为 $\mathbb E_{\pi_\theta}[z]=0$，概率比把 current 加权的 score 转成 reference 加权的 score：

$$\mathbb E_{a\sim\pi_\theta}\left[\left(1-\frac{\pi_{\text{ref}}(a)}{\pi_\theta(a)}\right)z(a)\right] =-\sum_a\pi_{\text{ref}}(a)z(a) =\nabla_\theta D_{\mathrm{KL}}(\pi_{\text{ref}}\|\pi_\theta).$$

最后一步使用了 reference 固定时的 forward KL 梯度表达式。这说明同一个 $k_3$ 可以有 reverse KL 的期望数值，却在直接反传时给出 forward KL 的期望梯度。[1]

两种系数一般不同。对指定的 reverse KL 而言，$k_3$ 直接反传遗漏了通常非零的概率变化项，因而有偏；对这里的局部 forward KL 而言，它已是完整梯度。[2]

上述结论限于固定前缀和 current 动作采样。实际 GRPO 的前缀分布、任务项及旧数据复用还需另行分析。

4.4 在同一起点核对梯度#

在同一个 70%/30% 起点，完整期望与固定样本均值取值相同。比较两者对参数的导数：参数 $\theta$ 取两个动作的 logit 差，因此 $\mathrm d\pi_\theta(A)/\mathrm d\theta=0.7\times0.3=0.21$，另一动作的导数为 $-0.21$。

仍只比较正的 KL 惩罚，不乘 $\beta$：

最后一行的 $0.2$ 是 forward KL 梯度，也是 $k_3$ 完整期望求导中的样本值变化项。起点处 A、B 的 $k_3$ 值分别为 $0.050758$、$0.155841$，此时 $\text{概率变化项}=0.21\times(0.050758-0.155841)\approx-0.022067$。将两条求导路径相加，得到 $\text{完整 reverse KL 梯度}=0.2-0.022067=0.177933$，恢复完整 reverse KL 的梯度。

增加样本量会让直接反传的估计更接近 $0.2$，不能恢复遗漏的 $-0.022067$。以上是手算结果，保留六位小数。

5. 正则方式的等价条件与旧数据复用#

5.1 三种正则方式的适用条件#

差别来自表达式和求导路径。显式 loss 也能给出 reverse KL 的完整梯度，不必采用 $k_3$。

如果直接对完整的 $\sum_a\pi_\theta(a)\log[\pi_\theta(a)/\pi_{\text{ref}}(a)]$ 求导，两条求导路径均被保留，结果与 log-ratio 奖励所估计的 reverse KL 梯度一致。但 LLM 的完整回答无法逐个枚举，实际通常需要采样。

因此，“KL 放进 reward 或单列为 loss 梯度等价”，需要指明表达式与条件。这里比较的是三种具体的 KL 配置，PPO、GRPO 的其他实现可以选择不同配置。

5.2 复用 old 数据后的限制#

前面按 current 采样比较梯度；实际 PPO 会复用 rollout，因此还要区分采样起点与后续更新。

在 PPO 的 KL 奖励计算中，采样时 current 等于 old，奖励在更新前计算一次，再固定本批 return 和 advantage。所以下标 old 记录的是计算这批奖励时的策略，不是因为理论要求奖励必须使用一个旧模型。

后续优化步中 current 继续变化，KL 奖励却没有重算；这一冻结方式加上 clipping，一般不再与每步计算当前 KL 的 loss 严格等价。“Current 与 old 接近”是复用数据时的近似依据，不能代替这两个阶段的区分。

也可以设计每步用 current 重算 KL 奖励的更新，但作为 policy-gradient 权重时仍要停止其梯度，并相应重算 return、advantage 等目标；若继续使用旧数据，还需匹配采样校正。这已改变了固定本批学习目标的 PPO 流程。

参考资料#

1. A Comedy of Estimators（2025）. 原文：采样 KL 数值与直接反传梯度。
2. Rethinking KL Regularization in RLHF（2025）. 原文：采样梯度、奖励接法与显式 loss。
3. DeepSeekMath（2024，v3）. 原文：原始 GRPO 的目标与流程。