Semi-Amortized Variational Autoencoders阅读笔记

假设变分后验为高斯分布，变分参数是$\lambda$，也就是$\lambda =[\mu ,\Sigma ]$，输入为$x$，这里的$x$表示一个样本。注意，每个样本都有与之对应的变分参数，也就是样本集$X=[x^{(1)},x^{(2)},...,x^{(N)}]$，变分参数集$\Lambda =[{\lambda }^{(1)},{\lambda }^{(2)},...,{\lambda }^{(N)}]$。

变分推理就是求解变分参数$\lambda$的过程。有两种技术，第一种是Hoffman在2013年提出的stochastic variational inference（SVI），随机梯度变分推理。第二种是Kingma在2014年提出的Amortized Variational inference（AVI） ，叫做平摊变分推理。

VAE的变分推理是属于Amortized Variational inference，每个样本$x$的变分参数都是通过同一个网络推理出的，这个网络是共享的。神经网络就相当于一个函数逼近器，所以VAE有一个约束就是认为变分参数$\lambda$是输入$x$的函数。但是这个约束太强了，可能会使得VAE的ELBO与loglikelihood之间存在一个gap，在复杂的数据集上，这个gap会显著影响VAE的性能。

SVI是另一种推理方法，SVI是特定于数据的，SVI会随机初始化$\Lambda$，然后通过梯度的方法来最大化ELBO，然后迭代求解变分参数。也就是说，变分参数不是由网络推理出的，而是通过梯度来优化的。

可见VAE的优点就是快速，缺点是由gap太大，得到的变分参数通常是次优的。SVI优点是对每个样本都迭代执行变分推断，缺点是比VAE耗时。本文作者将二者结合起来，用VAE推理出的变分参数初始化SVI的变分参数，然后再用SVI 特定于样本微调变分参数。

Variational Inference

$$\begin{array}{rl}logp(x;\theta )\ge & E_{q(z;\lambda )}[logp(x|z;\theta )]-KL[q(z;\lambda )||p(z)]\\ =& ELBO(\lambda ,\theta ,x)\end{array}\text{\hspace{1em}(1)}$$
(1)式是变分推断的ELBO，$\lambda$是变分参数，$\theta$是生成部分$p(x|z)$的参数。通常情况下，给定样本集$X$，变分推断的任务是寻找最优的$\Lambda$和生成模型参数$\theta$，使得${\sum }_{i=1}^{N}ELBO({\lambda }^{(i)},\theta ,x^{(i)})$最大化。

Stochastic Variational Inference

SVI的推理过程如下：

$K$是迭代的次数，$\alpha$是学习速率。可以看到，变分参数$\lambda$的优化和生成模型的参数$\theta$的优化是分开的，并不是由同一个ELBO算出的，这样$\theta$与$\lambda$的适应性可能不是很好。

Amortized Variational Inference

VAE使用一个全局共享的推理网络计算每个样本的变分参数，推理网络表示为$enc(x;\varphi )$，表示输入为$x$，网络参数是$\varphi$

与SVI相比，AVI要更新推理网络的参数$\varphi$，$\lambda$由推理网络计算得到，并且$\varphi$与$\theta$是由相同的ELBO算出的，这样可以$\lambda 与\theta$二者可以更好的适应。另外，AVI是在输入上进行推理，而SVI是在ELBO的迭代优化上进行推理，所以AVI更快。

Semi-Amortized Variational Autoencoders

SA-VAE在输入上利用推理网络提供初始化变分参数，然后运行SVI对其进行优化。

这里为了更新编码器（推理网络）和解码器（生成网络）的参数（分别是$\varphi$和$\theta$），需要第$K$次迭代的$ELBO({\lambda }_K,\theta ,x)$的梯度能够传到$\theta$和$\varphi$。

这里主要分析一下$ELBO({\lambda }_K,\theta ,x)$的梯度如何传到$\varphi$，由于${\lambda }_0$是由参数为$\varphi$的推理网络算出的，所以根据链式法则，
$$\frac{dELBO({\lambda }_K,\theta ,x)}{d\varphi }=\frac{dELBO({\lambda }_K,\theta ,x)}{d{\lambda }_0}\ast \frac{d{\lambda }_0}{d\varphi }\text{\hspace{1em}(1)}$$

那么（1）式右边的两项如何求解，先看$\frac{d{\lambda }_0}{d\varphi }$，我们前面已经假设后验分布是高斯分布 了，那么${\lambda }_0=[{\mu }_0,{\Sigma }_0]$，注意，这里的下标不是指样本的下标，而是指第几轮SVI推理。

所以${\lambda }_0$应该是两个向量，也就是向量对网络参数$\varphi$求导，在pytorch中，只需要在backward()中传入一个和${\lambda }_0$维度一样的向量即可。

关键是$\frac{dELBO({\lambda }_K,\theta ,x)}{d{\lambda }_0}$如何求解，如果用$ELBO({\lambda }_K,\theta ,x).$backward方式，梯度是传不到${\lambda }_0$的，因为中间还有$ELBO({\lambda }_{K-1},\theta ,x)，ELBO({\lambda }_{K-2},\theta ,x)，...，ELBO({\lambda }_1,\theta ,x)$，先考虑$ELBO({\lambda }_1,\theta ,x)$的梯度如何传到${\lambda }_0$

用SVI得到新的变分参数${\lambda }_1={\lambda }_0+\alpha {\nabla }_{\lambda }ELBO({\lambda }_0,\theta ,x)$，令
$L=ELBO({\lambda }_1,\theta ,x)$，根据链式法则，有下式成立

上式中，$H_{\lambda ,\lambda }$表示海森矩阵，是标量对向量的二阶导数组成的矩阵，$\frac{dL}{d_{{\lambda }_1}}$是标量对向量的导数，是一个向量。所以上式中涉及到海森矩阵与向量的矩阵乘法，后面会提到用一种专门的算法来实现二者的乘积。

多次使用链式法则，便可求得$\frac{dELBO({\lambda }_K,\theta ,x)}{d{\lambda }_0}$。这样$ELBO({\lambda }_K,\theta ,x)$的梯度就可传到$\varphi$了。

所以上面的分析可以分为两个过程：
第一个过程是前向迭代$K$次得到最优的变分参数${\lambda }_K$：

• ${\lambda }_0=enc(x;\varphi )$
• ${\lambda }_{k+1}={\lambda }_k+\alpha {\nabla }_{\lambda }ELBO({\lambda }_k,\theta ,x)$

第二个过程是后向迭代$K$次，将$L_K=ELBO({\lambda }_K,\theta ,x)$的梯度传到${\lambda }_0$：

• 通过$ELBO({\lambda }_K,\theta ,x).backward()$得到$\frac{d{L}_K}{d{\lambda }_K}$
• $\frac{d{L}_K}{d{\lambda }_0}=\frac{d{\lambda }_K}{d{\lambda }_{K-1}}\ast ...\ast \frac{d{\lambda }_1}{-{\lambda }_0}\frac{d{L}_K}{d{\lambda }_K}$
• 通过（1）式得到$ELBO({\lambda }_K,\theta ,x)$对$\varphi$的梯度。

对$\theta$也是同理。

最后算法框架如下图：

关于海森矩阵与向量的乘积，使用了 Hessian-vector products方法求解，并没有算二阶导。