现在各个训练框架中的KL散度计算都参考了John Schulman的博客。

如果要准确的计算两个分布之间的KL散度，离散分布需要遍历所有的 $x$，而连续分布需要求积分，在复杂的深度网络中几乎是无法计算的。因此在训练过程中，更多的是通过蒙特卡洛，对训练数据进行采样。

$$D_{KL}(P||Q)=\sum P(x)log(\frac{P(x)}{Q(x)}) = \mathbb{E}_{x\sim P}\left[log(\frac{P(x)}{Q(x)}) \right]$$

$k_1$#

最简单的估计 $k_1$ 为 $log(\frac{P(x)}{Q(x)})=-log(r)$， 其中$r=\frac{Q(x)}{P(x)}$ ，它是无偏的，有正确的平均值，但方差极高，尤其是 $-log(r)$ 的方向有正有负，而KL散度是不会小于0的，因此这个简单的估计在工程中会引入极大的不稳定。

这里提前算一下$r$的期望，由于我们是在$P(x)$的视角下求期望，因此：

$$\mathbb{E}_{x\sim P}(r)=\mathbb{E}_{x\sim P}\left[ \frac{Q(x)}{P(x)} \right]=\sum_{x} \left[ P(x) \frac{Q(x)}{P(x)}\right] =\sum_{x} Q(x)=1$$

$k_2$#

估计量 $k_2$：$\frac{1}{2}(log(r))^2$ 是从f-散度（KL散度是它的一个特例）推导出来的，它始终>0，这让它的方差很低，不会像$k_1$那样在正负之间来回横跳。根据泰勒展开，$k_2$和$k_1$的期望值在二阶近似上是完全相等的，因此当P和Q两个分布接近时，$k_2$可以用来估计KL散度，但由于在三阶近似上存在差异，因此$k_2$是有偏估计，但它低方差、非负。

$k_3$#

那么怎么能得到一个无偏且低方差，又始终为正的估计呢？

给 $k_1$打个补丁，这个补丁有两个要求：

1. 保持$k_1$的无偏：补丁的期望为0，在多次采样中，这个补丁的平均值为0
2. 消除负数：$k_1$加上这个补丁后，必须始终为正

John Schulman找到了一个补丁：$r-1$，$-log(r)$是一个开口向上的凸函数，在几何学中它的曲线始终在切线的上方，而 $1-r$ 正是 $-log(r)$ 在 $r=1$ 处的切线(注意这里的$log$是$ln$)。因此：

$$-log(r) \ge 1-r$$ $$-log(r) + r - 1 \ge 0$$

而 $r$ 的期望是1，因此 $r-1$ 的期望是0，完美的满足上述两个条件，因此估计量 $k_3$ 是一个无偏、低方差、非负 的KL散度估计量。

John Schulman的博客中举了两个例子，在 $P=N(0, 1)$ 和 $Q=N(0.1, 1)$ 两个分布相差不大时（此时的KL散度是0.005），可以看出$k_1$的方差非常大，$k_2$的偏差较小：

	bias	stdev
$k_1$	0	20
$k_2$	0.002	1.42
$k_3$	0	1.42

而在 $P=N(0, 1)$ 和 $Q=N(0.5, 1)$ 两个分布相差较大时（此时KL散度是0.5），可以看出 $k_2$此时的偏差就较大了。

	bias	stdev
$k_1$	0	2
$k_2$	0.25	1.73
$k_3$	0	1.7

尽管 $k_3$ 这个估计量有这么多优秀的特征，但 $k_3$ 存在一个工程上的计算问题，训练框架在实际算logit的概率时，为了防止极小值的下溢出，训练框架中通常只会计算 $log(p(x))$，而不是 $p(x)$，所以在计算 $k_3$ 时，我们只有 $log(r)$，为了还原$k_3 = r - 1 - \log r$ 中的$r$，必须要加一个 $e$ 的底数 $r=e^{log (r)}$ ，这导致当 $P$ 和$Q$相差较远时，例如当 $log(r)=20$，$r=e^{log r}=e^{20}$，是一个非常夸张的数字，因此使用 $k_3$时，则通常会加上一个clamp避免数值问题（GPRO和PPO的常见做法）。

$k_1$ 和 $k_3$ 的梯度偏差#

在使用 $k_1$ 和 $k_3$ 直接作为loss时，会在求梯度时遇上严重的梯度偏差。

以 $k_1$ 为例，正确的优化目标应该是：

$$\mathcal{J}_{k_1}(\theta) = \mathbb{E}_{y\sim\pi_{\theta}}[k_1] = \mathbb{E}_{y\sim\pi_{\theta}}[\log \pi_{\theta}(y)-\log \pi _{ref}(y)]$$

展开为概率求和形式：

$$\mathcal{J}_{k_1}(\theta) = \sum_{y}\pi_{\theta}(y)\cdot(\log \pi_{\theta}(y)-\log \pi _{ref}(y))$$

现在对参数 $\theta$ 求导。根据微积分的乘积求导法则 $(uv)' = u'v + uv'$，得到：

$$\nabla_\theta J_{k_1}(\theta) = \underbrace{\sum_y \nabla_\theta \pi_\theta(y) \cdot (\log \pi_\theta(y) - \log \pi_\text{ref}(y))}_{\text{项 A：分布漂移拉力}} + \underbrace{\sum_y \pi_\theta(y) \cdot \nabla_\theta (\log \pi_\theta(y) - \log \pi_\text{ref}(y))}_{\text{项 B：内部直接导数}}$$

项A

由于pytorch这种自动求导的机制是从采样出的token进行求导，而采样这个动作本身是不可导的，因此采样出的token已经是固定的数据，项A中的 $\nabla_\theta\pi_\theta(y)$ 不会出现在计算图中。换句话说，直接把 $k_1$ 作为loss时，pytorch只能计算项B。

但项B的期望实际上为0：

$$\sum_y \pi_\theta(y)\nabla_\theta\log\pi_\theta(y) =\sum_y\nabla_\theta\pi_\theta(y) =\nabla_\theta\sum_y\pi_\theta(y) =\nabla_\theta 1=0$$

所以直接把 $k_1$ 当作loss，看起来每个token都有梯度，但这些梯度在期望上会完全抵消。真正让分布向Reference Model靠近的项A，反而因为采样不可导而丢失了。也就是说，$k_1$ 的值是KL散度的无偏估计，但通过它直接反向传播得到的梯度，却不是KL散度的无偏梯度估计。

$k_3$ 也存在同样的问题。此时 $r=\frac{\pi_{ref}(y)}{\pi_\theta(y)}$，直接对 $k_3=r-1-\log r$ 求导，得到：

$$\nabla_\theta k_3=(1-r)\nabla_\theta\log\pi_\theta(y)$$

再对采样求期望：

$$\begin{aligned} \mathbb{E}_{y\sim\pi_\theta}[\nabla_\theta k_3] &=\mathbb{E}_{y\sim\pi_\theta}[(1-r)\nabla_\theta\log\pi_\theta(y)]\\ &=-\mathbb{E}_{y\sim\pi_{ref}}[\nabla_\theta\log\pi_\theta(y)]\\ &=\nabla_\theta D_{KL}(\pi_{ref}||\pi_\theta) \end{aligned}$$

可以看到，$k_3$ 的值估计的是 $D_{KL}(\pi_\theta||\pi_{ref})$，直接反向传播得到的却是反方向 $D_{KL}(\pi_{ref}||\pi_\theta)$ 的梯度。因此 $k_1$ 和 $k_3$ 都可以用来估计KL散度的值，但不能因为它们是好的估计量，就直接把它们当成loss。KL散度的值是否估计准确，和通过这个值反向传播得到的梯度是否准确，是两个不同的问题。