ELBO 梯度估计

作者: WilliamY | 来源:发表于2019-11-22 11:32 被阅读0次

问题

对于观察（observation） $x$ 和隐变量 $z$ ，其联合概率密度为
$p_\theta(x, z) = p_\theta(x|z)p_\theta(z)$
变分分布为 $q_\phi(z)$ 。这里的 $\phi$ 和 $\theta$ 为模型（model）和变分分布（guide）的参数。【注：所谓变分就是将原始函数换作另一（易处理的）函数的数学技巧】
目标为最大化证据（evidence）的对数形式 $\log p_\theta(x)$ 。而通常做法是最大化“（对数的）证据下限”ELBO（evidence lower bound），其形式如下：
$ELBO \equiv \mathbb{E}_{q_\phi(z)}[\log p_\theta(x, z) - \log q_\phi(z)]$
ELBO和证据的对数，二者之差为：
$\log p_\theta(x)-ELBO=KL[q_\phi(z)||p_\theta(z|x)]$
ELBO的无偏梯度估计为：
$\nabla_{\theta, \phi} ELBO=\nabla_{\theta, \phi}\mathbb{E}_{q_\phi(z)}[\log p_\theta(x, z) - \log q_\phi(z)]$
我们考虑一个更一般的形式：
$\nabla_{\phi} \mathbb{E}_{q_\phi(z)}[f_\phi(z)]$
【注：这里的 $\phi$ 指代更一般化的参数，和ELBO中狭义的 $\phi$ 不同。】

题外话： $q_\phi(z)$ 的最优值等于 $p_\theta(z|x)$

我们借用结论 $\int q_\phi(z)dz = 1$ ，我们使用拉格朗日法：
$\frac{\delta}{\delta q_\phi(z)}[-ELBO-\lambda(x)\int q_\phi(z)dz]$
$=-\log p_\theta(x, z) + \log q_\phi(z)-\lambda(x)=0$
$q^*_\phi(z)=p_\theta(x, z)e^{\lambda(x)}$
当我们取 $\lambda(x)=-\log p_\theta(x)$ 时，
$q^*_\phi(z)=p_\theta(z|x)$
这时的后验概率恰好有 $ELBO=-\lambda(x)=\log p_\theta(x)$

易处理的情况：可重参数化的随机变量

假如我们能够对随机变量 $z$ 重参数化为
$\mathbb{E}_{q_\phi(z)}[f_\phi(z)]=\mathbb{E}_{q(\epsilon)}[f_\phi(g_\phi(\epsilon))]$
也就是说，我们把对 $\phi$ 依赖的项，全部放在求期望的范围里（即用E给“框”在里面），这时候 $q(\epsilon)$ 就不再依赖于 $\phi$ 了。
这样的重参数化操作，可以对许多分布使用（比如高斯分布）。这样我们就得到梯度估计：
$\nabla_\phi \mathbb{E}_{q(\epsilon)}[f_\phi(g_\phi(\epsilon))]=\mathbb{E}_{q(\epsilon)}[\nabla_\phi f_\phi(g_\phi(\epsilon))]$
这里假定 $f$ 和 $g$ 都是光滑的（即可导的），我们就可以用蒙特卡洛法（将多次观察求平均）求解上述无偏的梯度估计了。

取巧的情况：非重参数化的随机变量

如果不能使用重参数化，例如分布是离散的，这时上面的技巧就不管用了。
我们将梯度估计量展开：
$\nabla_\phi \mathbb{E}_{q_\phi(z)}[f_\phi(z)]=\nabla_\phi \int q_\phi(z)f_\phi(z)dz$
由链式法则，我们继续展开：
$\int \{(\nabla_\phi q_\phi(z)) f_\phi(z) + q_\phi(z) \nabla_\phi f_\phi(z)\}dz$
对于 $\nabla_\phi q_\phi(z)$ 存在恒等式
$\nabla_\phi q_\phi(z) = q_\phi(z)\nabla_\phi \log q_\phi(z)$
代入上式得：
$\mathbb{E}_{q(\epsilon)}[(\nabla_\phi \log q_\phi(z)) f_\phi(z) + \nabla_\phi f_\phi(z)]$
我们把求期望和梯度的项写在一起，称为“代理目标”（surrogate objective）：
$surrogate \ objective \equiv \log q_\phi(z) \overline{f_\phi(z)} + f_\phi(z)$
于是ELBO的梯度无偏估计为
$\nabla_\phi ELBO = \mathbb{E}_{q_\phi(z)}[\nabla_\phi (surrogate \ objective)]$ 。
$\overline{f_\phi(z)}$ 的横线表示该项对 $\phi$ 来说是常数，不对 $\phi$ 求导数。

减少梯度估计的方差

考虑下面的等式:
$\mathbb{E}_{q_{\phi}({\bf z})} \left [\nabla_{\phi} (\log q_{\phi}({\bf z}) \cdot b) \right]=0$
其中 $b$ 为任意的常数。这是因为：
$\mathbb{E}_{q_{\phi}({\bf z})} \left [\nabla_{\phi} \log q_{\phi}({\bf z}) \right]= \int \!d{\bf z} \; q_{\phi}({\bf z}) \nabla_{\phi} \log q_{\phi}({\bf z})= \int \! d{\bf z} \; \nabla_{\phi} q_{\phi}({\bf z})= \nabla_{\phi} \int \! d{\bf z} \; q_{\phi}({\bf z})=\nabla_{\phi} 1 = 0$
于是，对于 $\log q_{\phi}({\bf z}_i) \overline{f_{\phi}({\bf z})}$ ，我们利用上述等式，用下面的项代替：
$\log q_{\phi}({\bf z}_i) (\overline{f_{\phi}({\bf z})}-b)$
二者的梯度的期望是相同的。更妙的是， $b$ 不必是常数，只要对下游任务没有影响即可。所以 $b$ 可设为上游任务，自变量为 $z_i$ 的函数。
参考文献：http://pyro.ai/examples/svi_part_iii.html

ELBO 梯度估计