理解 KL 散度的近似

作者：John Schulman（OpenAI）

译者：朱小虎 Xiaohu (Neil) Zhu（CSAGI / University AI）

原文链接：http://joschu.net/blog/kl-approx.html

术语： $\text{KL}$ 散度（ $\text{KL}$ Divergence）; 近似（Approximation）; Monte-Carlo 估计（Monte-Carlo estimator）

本文讨论 $\text{KL}$ 散度的 Monte-Carlo 近似：

$\text{KL}[q,p] = \sum_{x}q(x)\log \frac{q(x)}{p(x)} = E_{x\sim q}[\log \frac{q(x)}{p(x)}]$

这解释了之前使用了一个技巧，针对来自 $q$ 中的样本 $x$ 以 $\frac{1}{2}(\log p(x)-\log q(x))^2$ 样本平均来近似 $\text{KL}[q,p]$ ，而不是更加标准的 $\log \frac{q(x)}{p(x)}$ . 本文谈谈为何该表达是一个 KL 散度的好的估计（尽管有偏 biased），以及如何让其变得无偏（unbiased）保证其低方差。

我们计算 KL 的选择取决于对 $p$ 和 $q$ 的访问方式。这里，我们假设能够对任意 $x$ 计算概率（或者概率密度） $p(x)$ 和 $q(x)$ ，但是我们不能解析地跑遍 $x$ 求和。为何我们不能解析地计算呢？

准确地计算此和需要太多计算或者内存
没有闭式形式
我们可以通过仅仅存储 $\text{log-prob}$ 而非整个分布来简化代码。只要 $\text{KL}$ 仅仅用来作为诊断工具这是一个合理的选择，这也是强化学习中常见情况

最常用的估计求和或者积分的策略是使用 Monte-Carlo 估计。给定样本 $x_1, x_2, \cdots \sim q$ ，我们如何构造好的估计？

一个好的估计是无偏的（即有正确的均值）并且低方差。我们知道一个无偏估计（在从 $q$ 中采样的样本下）是 $\log \frac{q(x)}{p(x)}$ 。但是，它有高方差，因为它对样本的一半是负，而 KL 总是为正。让我们称此简易估计 $k_1 = \log \frac{q(x)}{p(x)} = - \log r$ ，其中我们已经定义了比例 $r = \frac{p(x)}{q(x)}$ 后面也会多次出现此值。

另一个替代估计有低的方差不过是有偏的，即 $\frac{1}{2} (\log \frac{p(x)}{q(x)})^2 = \frac{1}{2} (\log r)^2$ 。我们不妨称此为 $k_2$ 。直觉上看， $k_2$ 看起来更加好因为每个样本告诉了我们 $p$ 和 $q$ 之间相距多远，并且总为正。实验上看， $k_2$ 实际有比 $k_1$ 更低的方差，并也有相当低的偏差（bias）。（下面在实验中给出此点）。

关于估计 $k_2$ 为何有低偏差有一个很好的原因：其期望是一个 f-散度（divergence）。一个 f-散度 被定义为关于一个凸函数 $f$ ， $D_f(p,q) = \text{E}_{x\sim q} [f(\frac{p(x)}{q(x)})]$ 。KL 散度和其他有名的概率距离均是 $f$ -散度。现在这是关键的难以被发现的事实：所有具有可微函数 $f$ 的 $f$ -散度与 $\text{KL}$ 散度当 $q$ 接近 $p$ 时的二阶。也就是说，对一个参数化分布 $p_{\theta}$ ，

$D_f(p_0,p_\theta) = \frac{f''(1)}{2} \theta^\intercal F \theta + O(\theta^3)$

其中 $F$ 是关于 $p_\theta$ 的 Fisher 信息矩阵在 $p_\theta = p_0$ 的值。

期望 $E_q[k_2] = E_q[\frac{1}{2} (\log r)^2]$ 是 f-散度，其中 $f(x) = \frac{1}{2}(\log x)^2$ ，而 $\text{KL}[q,p]$ 对应于 $f(x) = - \log x$ 。易见，两者均有 $f''(1) = 1$ ，所以两者看起来对 $p\approx q$ 有相同的二阶距离函数。

是否可以写出一个 $\text{KL}$ 散度估计既是无偏又是低方差的呢？一般达成低方差的方法是通过一个控制变量。就是说，取 $k_1$ 并加上某个期望为零但是与 $k_1$ 负相关的量。保证期望为零唯一有趣的量是 $\frac{p(x)}{q(x)} - 1 = r - 1$ 。所以，对任意的 $\lambda$ ，表达式 $-\log r + \lambda (r - 1)$ 是 $KL[q,p]$ 的无偏估计。我们可以做一些计算来最小化这个估计的方差，对 $\lambda$ 求解。但不幸的是，我们获得一个表达式，它依赖于 $p$ 和 $q$ 并难以解析地计算。

但是，我们可以使用一个更为简单的策略来选择一个好的 $\lambda$ 。注意因为 $\log$ 是凹函数， $\log(x) \leq x - 1$ 。因此，如果我们令 $\lambda = 1$ ，上面的表达式会确保为正。它度量了 $\log(x)$ 和它的切线的竖直距离。这让我们有了估计 $k_3 = (r - 1) - \log r$ 。

通过看凸函数和它切平面的差距来度量距离的想法出现在很多领域。这杯称为 Bregman 散度并有很多优美性质。

我们可以推广上面想法来获得一个好的，总是为正的对任何 f-散度的估计，大多数明显是另一个 $\text{KL}$ 散度即 $\text{KL}[p,q]$ （注意这里的 $p$ 和 $q$ 调换了次序）。因为 $f$ 是凸函数，并且 $E_q[r] = 1$ ，下面是 f-散度的一个估计： $f(r) - f'(1)(r-1)$ 。这总是为正因为它是 f 和其在 $r=1$ 处的距离，并且凸函数在它们的切线上方。现在 $\text{KL}[p,q]$ 对应于 $f(x) = x \log x$ ，其有 $f'(1) = 1$ ，使得我们有了估计 $r \log r - (r - 1)$ 。

总结一下，我们有下列估计（对样本 $x\sim q$ 和 $r = \frac{p(x)}{q(x)}$ ）：

$\text{KL}[p,q] : r \log r - (r - 1)$
$\text{KL}[q,p] : (r - 1) - \log r$

现在我们比对这三个对 $\text{KL}[p,q]$ 估计的偏差和方差。假设 $q=N(0,1), p=N(0.1,1)$ 。这里正确的 KL 散度为 $0.005$ 。

k	bias/true	stdev/true
k1	0	20
k2	0.002	1.42
k3	0	1.42

注意 k2 的偏差非常低：为 0.2%。

现在我们尝试对大一些的 $\text{KL}$ 散度近似。 $p=N(1,1)$ 给我们一个真实 $\text{KL}$ 散度为 $0.5$ 。

k	bias/true	stdev/true
k1	0	2
k2	0.25	1.73
k3	0	1.7

这里，k2 的偏差更大一些。k3 甚至有比 k2 更低的标准差同时还是无偏的，所以它看起来也是在一个严格意义上更好的估计。

这里是我用来产生这些结果的代码：

import torch.distributions as dis
p = dis.Normal(loc=0, scale=1)
q = dis.Normal(loc=0.1, scale=1)
x = q.sample(sample_shape=(10_000_000,))
truekl = dis.kl_divergence(p, q)
print("true", truekl)
logr = p.log_prob(x) - q.log_prob(x)
k1 = -logr
k2 = logr ** 2 / 2
k3 = (logr.exp() - 1) - logr
for k in (k1, k2, k3):
    print((k.mean() - truekl) / truekl, k.std() / truekl)