f-VAE: 基于流的变分自编码器.

1. 生成模型概述

对于随机变量$x$,假设其固有分布为$p(x)$,由已有数据集可构造的证据分布为$q(x)$。生成模型的基本思路是希望用如下的分布形式来拟合给定数据集的分布$p(x)$:

\[p(x) = \int p(z)p(x|z)dz\]

其中先验分布$p(z)$一般取标准高斯分布,而生成分布$p(x|z)$一般取高斯分布(VAE)或者狄拉克分布(GAN和流模型)。理想情况下,优化方式是最大化似然函数$\Bbb{E}[\log p(x)]$,或者等价地最小化$KL(q(x)||p(x))$。

由于$p(x)$表达式中的积分可能难以显式计算,所以需要一些特殊的求解技巧,这导致了不同的生成模型。

⚪ VAE

VAE引入了后验分布$q(z|x)$,将优化目标更改为更容易计算的上界$KL(q(x)q(z|x)||p(z)p(x|z))$。VAE收敛快、训练稳定,但一般情况下生成图像存在模糊。

⚪ Flow-based Model

在流模型中生成模型设置为狄拉克分布$p(x|z) = \delta(x-G(z))$,通过精心设计$G(z)$直接把$p(x)$表达式中的积分计算出来,从而可以直接完成最大似然优化。其主要组件是耦合层,即将输入$x$划分为两部分$x_1,x_2$,然后构造可逆变换:

\[\begin{aligned} y_1&=x_1 \\ y_2&= s(x_1) \otimes x_2+ t(x_1) \end{aligned}\]

耦合层需要保证其雅可比行列式容易计算,因此单个耦合层的设计比较简单,只能带来非常弱的非线性能力;通过大量耦合层的组合可以得到复杂的非线性变换。这导致流模型通常相当庞大,训练时间长。

2. f-VAE

作者提出将流模型引入VAE中,即用流模型拟合更一般的后验分布$q(z|x)$,而不是简单地设为高斯分布。该模型称为基于流的变分自编码器(Flow-based Variational Autoencoders, f-VAEs)。

相比于标准的VAEsf-VAEs跳出了关于后验分布为高斯分布的局限,能够生成更清晰的图像;相比于标准的流模型,f-VAEs的编码器给模型带来了更强的非线性能力,从而可以减少对耦合层的依赖,通过更小的模型规模来达到同样的生成效果。

⚪ 目标函数

VAEs的目标函数出发:

\[KL(q(x)q(z|x))||p(z)p(x|z) = \int \int q(x) q(z|x)\log \frac{q(x)q(z|x)}{p(z)p(x|z)}dzdx\]

与标准VAEs不同的是,后验概率$q(z|x)$不再假设为高斯分布,而是通过流模型构建:

\[q(z|x) = \int \delta(z-F_x(u))q(u)du\]

其中$q(u)$是标准高斯分布,$F_x(u)$是关于$u$的流模型,其参数与$x$有关,称为条件流。将流模型带入目标函数有:

\[\int \int \int q(x) \delta(z-F_x(u))q(u)\log \frac{q(x)\int \delta(z-F_x(v))q(v)dv}{p(z)p(x|z)}dzdudx \\ = \int \int q(x) q(u)\log \frac{q(x)\int \delta(F_x(u)-F_x(v))q(v)dv}{p(F_x(u))p(x|F_x(u))}dudx\]

记$v’=F_x(v)$, $v=H_x(v’)$,则有如下关系:

\[\det[\frac{\partial v}{\partial v'}] = \frac{1}{\det[\frac{\partial v'}{\partial v}]} = \frac{1}{\det[\frac{\partial F_x(v)}{\partial v}]}\]

则目标函数进一步写作:

\[\begin{aligned} & \int \int q(x) q(u)\log \frac{q(x)\int \delta(F_x(u)-F_x(v))q(v)dv}{p(F_x(u))p(x|F_x(u))}dudx \\ &= \int \int q(x) q(u)\log \frac{q(x)\int \delta(F_x(u)-v')q(H_x(v'))dH_x(v')}{p(F_x(u))p(x|F_x(u))}dudx \\ &= \int \int q(x) q(u)\log \frac{q(x)\int \delta(F_x(u)-v')q(H_x(v'))|\det[\frac{\partial v}{\partial v'}]|dv'}{p(F_x(u))p(x|F_x(u))}dudx \\ &= \int \int q(x) q(u)\log \frac{q(x)\int \delta(F_x(u)-v')q(H_x(v'))|\frac{1}{\det[\frac{\partial F_x(v)}{\partial v}]}|dv'}{p(F_x(u))p(x|F_x(u))}dudx \\ &= \int \int q(x) q(u)\log \frac{q(x)q(H_x(F_x(u)))|\frac{1}{\det[\frac{\partial F_x(v)}{\partial v}]}|_{v'=F_x(u)}}{p(F_x(u))p(x|F_x(u))}dudx \\ &= \int \int q(x) q(u)\log \frac{q(x)q(u)}{p(F_x(u))p(x|F_x(u))|\det[\frac{\partial F_x(u)}{\partial u}]|}dudx \end{aligned}\]

上式给出了一般化的f-VAE框架,通过选择不同的$F_x(u)$对应于不同的生成模型。注意到$q(x)$是未知但固定的常数,因此上述目标函数等价于:

\[\begin{aligned} &\int \int q(x) q(u)\log \frac{q(x)q(u)}{p(F_x(u))p(x|F_x(u))|\det[\frac{\partial F_x(u)}{\partial u}]|}dudx \\ &\to \int \int q(x) q(u)\log \frac{q(u)}{p(F_x(u))p(x|F_x(u))|\det[\frac{\partial F_x(u)}{\partial u}]|}dudx \\ &= \int q(x) [ \int q(u)\log \frac{q(u)}{p(F_x(u))p(x|F_x(u))|\det[\frac{\partial F_x(u)}{\partial u}]|}du]dx \\ &= \int q(x) [ \int q(u)[\log \frac{q(u)}{p(F_x(u))}-\log |\det[\frac{\partial F_x(u)}{\partial u}]|-\log p(x|F_x(u))]du]dx \end{aligned}\]

⚪ 例子1:退化为VAE

如果设置:

\[F_x(u) = \sigma(x) \otimes u +\mu(x)\]

此时目标函数中存在:

\[\begin{aligned} &\int q(u)\log \frac{q(u)}{p(F_x(u))}du = \frac{1}{2} \sum_{i=1}^{d} (\mu_i^2(x)+\sigma_i^2(x)-1) \\ &\int - q(u)\log |\det[\frac{\partial F_x(u)}{\partial u}]|du = -\sum_{i} \log \sigma_i(x) \end{aligned}\]

上述两式的组合恰好是后验分布和先验分布的KL散度;代入到目标函数中恰好是标准的VAE损失。值得一提的是,该结果自动包含了重参数的过程。

⚪ 例子2:退化为流模型

如果设置:

\[F_x(u) = F(\sigma u + x), p(x|z) = \mathcal{N}(x;F^{-1}(z),\sigma)\]

此时目标函数中存在:

\[-\log p(x|F_x(u)) = -\log \mathcal{N}(x;F^{-1}(F(\sigma u + x)),\sigma) \\ = -\log \mathcal{N}(x;\sigma u + x,\sigma) = \frac{d}{2}\log 2\pi \sigma^2 + \frac{1}{2}||u||^2\]

此时该目标项未包含训练参数。则原目标函数进一步写作:

\[\begin{aligned} & \int q(x) [ \int q(u)[\log \frac{q(u)}{p(F_x(u))}-\log |\det[\frac{\partial F_x(u)}{\partial u}]|-\log p(x|F_x(u))]du]dx \\ &\to \int q(x) [ \int q(u)[\log \frac{q(u)}{p(F_x(u))}-\log |\det[\frac{\partial F_x(u)}{\partial u}]|]du]dx \\ &\to - \int q(x) [ \int q(u)[\log p(F_x(u))|\det[\frac{\partial F_x(u)}{\partial u}]|]du]dx \\ &= - \int q(x) [ \int q(u)[\log p(F(\sigma u + x))|\det[\frac{\partial F(\sigma u + x)}{\partial u}]|]du]dx \end{aligned}\]

上式等价于在标准的流模型输入中增加方差为$σ^2$的高斯噪声。在常用的流模型(如Glow模型)中确实都会在训练的时候给输入图像加入一定量的噪声。

⚪ 例子3:最终设置

作者设置:

\[F_x(u) = F(\sigma_1 u + E(x)), p(x|z) = \mathcal{N}(x;G(F^{-1}(z)),\sigma_2)\]

其中$\sigma_1,\sigma_2$是可训练参数,$E(\cdot),D(\cdot)$是编码器和解码器,$F(\cdot)$是参数与$x$无关的流模型。目标函数等价于:

\[\begin{aligned} & \int q(x) [ \int q(u)[\log \frac{q(u)}{p(F_x(u))}-\log |\det[\frac{\partial F_x(u)}{\partial u}]|-\log p(x|F_x(u))]du]dx \\ & \to \int q(x) [ \int q(u)[\frac{1}{2}F^2(\sigma_1u+E(x))-\log |\det[\frac{\partial F(\sigma_1 u + E(x))}{\partial u}]| \\ & \qquad \qquad \qquad \qquad + \frac{1}{2\sigma_2^2}||G(\sigma_1u+E(x))-x||^2]du]dx \end{aligned}\]

模型的生成采样过程为:

\[u \text{~} q(u), z = F^{-1}(u) , x=G(z)\]

3. 实验分析

作者对比了VAE、流模型以及f-VAE的实验结果,在同样的复杂度和同样的训练时间下,f-VAEs有着更好的表现。