理解 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 近似：





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



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



1. 准确地计算此和需要太多计算或者内存

2. 没有闭式形式

3. 我们可以通过仅仅存储 $\text{log-prob}$ 而非整个分布来简化代码。只要 $\text{KL}$ 仅仅用来作为诊断工具这是一个合理的选择，这也是强化学习中常见情况



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



一个好的估计是无偏的（即有正确的均值）并且低方差。我们知道一个无偏估计（在从 $q$ 中采样的样本下）是 $\log \frac{q(x)}{p(x)}$。但是，它有高方差，因为它对样本的一半是负，而 $\text{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）](https://en.wikipedia.org/wiki/F-divergence)。一个 $f$，$Df(p,q)=\text{E}x\sim q[f(\frac{p(x)}{q(x)})]$。$\text{KL}$ 散度和其他有名的概率距离均是$f$-散度。现在这是关键的难以被发现的事实：所有具有可微函数 $f$ 的 $f$-散度与 $\text{KL}$ 散度当 $q$ 接近 $p$ 时的二阶。也就是说，对一个参数化分布 $p_{\theta }$，

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



期望 $E_q[k2]=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^{\prime \prime }(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)$ 是 $\text{KL}[q,p]$ 的无偏估计。我们可以做一些计算来最小化这个估计的方差，对 $\lambda$ 求解。但不幸的是，我们获得一个表达式，它依赖于 $p$ 和 $q$ 并难以解析地计算。



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



通过看凸函数和它切平面的差距来度量距离的想法出现在很多领域。这杯称为 [Bregman 散度](https://en.wikipedia.org/wiki/Bregman_divergence)并有很多优美性质。



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



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



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





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



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



import torch.distributions as dis
p = dis.Normal(loc=0, scale=1)
q = dis.Normal(loc=0.1, scale=1)
x = q.sample(sampleshape=(10000000,))
truekl = dis.kldivergence(p, q)
print("true", truekl)
logr = p.logprob(x) - q.logprob(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)