KL散度（KL Divergence）学习笔记

参考

Fantastic KL Divergence and How to (Actually) Compute It

笔记

我们用欧几里得距离衡量两个点之间的远近，那如何衡量两个概率分布的相似程度呢？

概率分布（Probability Distribution） 是用于描述随机变量 在各个可能取值上出现概率的数学函数或规则。对于离散随机变量，概率分布由概率质量函数（PMF）给出；对于连续随机变量，概率分布由概率密度函数（PDF）给出，其在全体定义域上的积分等于 1。

这两种分布都可以用 KL 散度来衡量：

$$D_{\mathrm{KL}}(P \parallel Q) = \sum_{i=1}^N P(x_i) \log \frac{P(x_i)}{Q(x_i)}$$ $$D_{\mathrm{KL}}(P \parallel Q) = \int_{-\infty}^{\infty} p(x) \log \frac{p(x)}{q(x)} \, dx$$

一种简要的计算方法是：

$$D_{\mathrm{KL}}(P \parallel Q) \approx \frac{Q(x)}{P(x)} - 1 - \log\left( \frac{Q(x)}{P(x)} \right) \quad \text{for } x \sim P$$

我们接下来来解释这些概念。

---

在信息论中，熵的定义是

$$H(X) = - \sum_{x} p(x) \log p(x)$$

对数通常以2为底。熵事实上是随机变量 X 的自信息的期望值，自信息的定义是

$$I(x) = \log \frac{1}{p(x)}$$

自信息可以看为某件事发生的“惊讶”程度。事件发生的概率越小，它的自信息越大，因为“罕见事件”带来的信息越多，越令人惊讶。反过来，概率越大越确定的事件自信息越小，几乎没什么新信息。

现在解释熵的单位。当对数的底数为 2 时，信息的单位是比特（bit）。简单地说，1 比特的信息量，等价于告诉你一个“是”或“否”的问题的答案。如果知道了 1 比特的信息，就是解决了一个有两个可能结果的不确定性。

来一些直观的例子：

如果明天有 50% 的可能是晴天，有 50% 的可能是雨天，那么熵就是 1 bits：

$$H(P)=-1/2\log{1/2}-1/2\log{1/2}=1 \text{bits}$$

你只需要问一个“是”或“否”的问题（e.g. 明天会下雨吗？），就能知道明天天气的结果。

事实上，“问题”代表“知道其中一个事件的答案”。

如果明天 100% 是晴天，那么熵是 0。不需要一个“是”或“否”的问题能知道结果。

此外，熵的比特数还表示可以用平均几位来对结果进行编码，如图：

现在来看交叉熵，它可以告诉你当你用 Q 进行预测,而真实分布是 P 时，你的平均“惊讶”程度。

事实上，P 是真实概率分布，Q 是预测概率分布/编码方案。

交叉熵 (H(P,Q)) 衡量的是，当使用编码方案（或预测分布）Q 来编码来自真实分布 P 的事件时，平均所需的比特数。简单来说，它告诉你你的预测分布（或编码）与真实分布匹配得有多好。交叉熵越低，意味着你的预测/编码越好。

当我们用错误的编码方案来预测真实概率分布时，交叉熵会大于真实分布的熵，如图，而它们的差就是 KL 散度：

因此，$D_{KL}$ 可以告诉我们分布 q 与 p 的偏差大小。

---

正向KL（Forward KL） 和 反向KL（Reverse KL） 是两个不同的计算方式。

正向KL（Forward KL）：

$$D_{\text{KL}}(P \| Q) = \sum_x P(x) \log \frac{P(x)}{Q(x)}$$

用真实分布 P 去衡量近似分布 Q 的“错位程度”。

反向KL（Reverse KL）：

$$D_{\text{KL}}(Q \| P) = \sum_x Q(x) \log \frac{Q(x)}{P(x)}$$

用近似分布 Q 去衡量它与真实分布 P 的差异。

当我们去 minimize 正/反向 KL 的时候，会产生不同的行为与结果，现在来解释。假设真实分布 P 是一个有两个峰值的概率分布，而我们用一个高斯分布 Q 来近似它。

当我们尝试 minimize 正向 KL 的时候，当 P 在某处有概率质量时，而Q较小，就会发生较大的惩罚。因此，结果就会鼓励 Q 要覆盖所有 P 的高概率区域，容易产生模式覆盖（mode covering）。

当我们尝试 minimize 反向 KL 的时候，当 P(x)=0 或接近 0 时，而 Q(x)>0，就会发生较大的惩罚。因此，结果就会鼓励 Q 只去匹配P 中概率最大的部分，容易产生模式崩塌（mode collapse）。

以上惩罚大的原因都是分母小、分子大。

场景	使用的KL	效果/倾向
变分推断（Variational Inference）	正向KL	模型试图覆盖更多可能
GAN的判别器训练	反向KL	容易聚焦于某几个模式
Reinforcement Learning 中的 Policy Optimization（如 TRPO、PPO）	正向KL	保证新策略不偏离旧策略太远

---

然而，计算$D_{KL}$的开销不小。例如，大语言模型一般有 200k+ 的 tokens，我们要计算的就是基于它们的离散分布。我们有方法来估计$D_{KL}$。

$$D_{\mathrm{KL}}(P \parallel Q) = \int_{-\infty}^{\infty} p(x) \log \frac{p(x)}{q(x)} \, dx$$ $$=\mathbb{E}_{x \sim P} \left[ \log \frac{P(x)}{Q(x)} \right]$$

接下来用 Monte Carlo 估计（蒙特卡洛估计）。我们不直接计算期望，而是从 P 中 sample 一些样本，如图：

这个估计是 unbiased 的，随着样本的增多，估计值会收敛到真实值。可以通过大数定理证明。

然而，这样处理的数据方差较大，并且当$Q(x_i) > P(x_i)$时，可能会出现负数的情况，而 KL 散度不会是负数。

有一些技巧可以减少方差：

但难以收敛：

总结：

如何兼得两种优点？可以为第一种方法添加一个有着 0 期望值的校正项，例如$\lambda(r(x_k)-\mathbb{E}_{x \sim P} \left[r(x)\right])$，这一项的期望是 0。

定义$r(x)=Q(x)/P(x)$，那么

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

现在，

$$\hat{D}_{\text{KL}}(P || Q) = \frac{1}{N} \sum_{k=1}^{N} \left[ \log \frac{P(x_k)}{Q(x_k)} + \lambda (r(x_k) - 1) \right]$$ $$=\frac{1}{N} \sum_{k=1}^{N} \left[ -\log r(x_k) + \lambda (r(x_k) - 1) \right]$$

如何决定$\lambda$？当它为 1 时，方差较小，大部分是非负的。这就是：

$$\hat{D}_{\text{KL}}(P || Q) =\frac{1}{N} \sum_{k=1}^{N} \left[ -\log r(x_k) + r(x_k) - 1 \right]$$

它的方差很小，并且可以收敛到真实值。