近似KL散度

作者:John Schulman(OpenAI)

译者:Xiaohu Zhu(CSAGI)

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

本文讨论 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 的选择取决于对 pq 的访问方式。这里,我们假设能够对任意 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)}。但是,它有高方差,因为它对样本的一半是负,而 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 看起来更加好因为每个样本告诉了我们 pq 之间相距多远,并且总为正。实验上看,k_2 实际有比 k_1 更低的方差,并也有相当低的偏差(bias)。(下面在实验中给出此点)。

关于估计 k_2 为何有低偏差有一个很好的原因:其期望是一个 f-散度(divergence)。一个 f-散度 被定义为关于一个凸函数 fD_f(p,q) = \text{E}_{x\sim q} [f(\frac{p(x)}{q(x)})]。KL 散度和其他有名的概率距离均是 f-散度。现在这是关键的难以被发现的事实:所有具有可微函数 f 的 f-散度与 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 有相同的二阶距离函数。

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

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

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

我们可以推广上面想法来获得一个好的,总是为正的对任何 f-散度的估计,大多数明显是另一个 KL 散度即 \text{KL}[p,q] (注意这里的 pq 调换了次序)。因为 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 qr = \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。

bias/truestdev/true
k1020
k20.0021.42
k301.42

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

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

bias/truestdev/true
k102
k20.251.73
k301.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)

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out /  Change )

Google photo

You are commenting using your Google account. Log Out /  Change )

Twitter picture

You are commenting using your Twitter account. Log Out /  Change )

Facebook photo

You are commenting using your Facebook account. Log Out /  Change )

Connecting to %s