1. Nouveau VAE

变分自编码器(VAE)优化对数似然的变分下界:

\[\mathcal{L} = KL[q(z|x)||p(z)]+\mathbb{E}_{z \text{~} q(z|x)} [-\log p(x | z)]\]

VAE需要建模$q(z|x),p(z),p(x|z)$。建模要求是概率能写出解析表达式，并且方便采样。最常用的建模方法是各分量独立的高斯分布：

\[q(z|x) = \mathcal{N}(\mu_1,\sigma_1^{2}),p(z) = \mathcal{N}(0,I),p(x|z) = \mathcal{N}(\mu_2,\sigma_2^{2})\]

然而各分量独立的高斯分布不能拟合任意复杂的分布，VAE的优化目标无法达到$KL(q(x,z)||p(x,z))=0$。这意味从理论来讲让$p(x,z),q(x,z)$相互逼近只能得到一个平均的结果，因此VAE生成的图像通常比较模糊。

Nouveau VAE设计的出发点便是增强$q(z|x),p(z),p(x|z)$这几个分布的建模。为了保证生成过程的并行性，$p(x|z)$仍然采用高斯分布$\mathcal{N}(\mu_2,\sigma_2^{2})$。

对于后验分布$q(z|x)$与先验分布$p(z)$，将其建模为自回归模型。具体地，对隐变量进行分组$z=\{z_1,z_2,\cdots z_L\}$，然后构造:

\[\begin{aligned} p(z)& = \prod_{l=1}^{L} p(z_l|z_{\lt l}) = p(z_1)p(z_2|z_1)p(z_3|z_1z_2)\cdots p(z_L|z_1z_2\cdots z_{L-1}) \\ q(z|x) &= \prod_{l=1}^{L} q(z_l|z_{\lt l},x) = q(z_1|x)q(z_2|z_1,x)q(z_3|z_1z_2,x)\cdots q(z_L|z_1z_2\cdots z_{L-1},x) \end{aligned}\]

注意到$z_l$仍然是向量。对$p(z_l|z_{\lt l}),q(z_l|z_{\lt l},x)$仍然建模为高斯分布，则$q(z|x),p(z)$建模为自回归高斯模型。

Nouveau VAE给出了一种对$p(z_l|z_{\lt l}),q(z_l|z_{\lt l},x)$建模的相对设计。即建模先验分布$p(z)$的均值和方差，建模后验分布$q(z|x)$与先验分布$p(z)$均值方差的相对值。

\[\begin{aligned} p(z_l|z_{\lt l}) &= \mathcal{N}(z_l;\mu(z_{\lt l}),\sigma^{2}(z_{\lt l})) \\ q(z_l|z_{\lt l},x) &= \mathcal{N}(z_l;\mu(z_{\lt l})+\Delta \mu(z_{\lt l},x),\sigma^{2}(z_{\lt l})\otimes \Delta \sigma^{2}(z_{\lt l},x)) \end{aligned}\]

这种设计能使得训练更加稳定。此时后验分布$q(z|x)$与先验分布$p(z)$的KL散度项为：

\[\begin{aligned} KL(q(z|x)||p(z))& = KL(\prod_{l=1}^{L} q(z_l|z_{\lt l},x)||\prod_{l=1}^{L} p(z_l|z_{\lt l})) \\ &= \Bbb{E}_{\prod_{l=1}^{L} q(z_l|z_{\lt l},x)}[\log \frac{\prod_{l=1}^{L} q(z_l|z_{\lt l},x)}{\prod_{l=1}^{L} p(z_l|z_{\lt l})}] = \Bbb{E}_{\prod_{l=1}^{L} q(z_l|z_{\lt l},x)}[\sum_{l=1}^{L}\log \frac{ q(z_l|z_{\lt l},x)}{p(z_l|z_{\lt l})}] \\ &= \Bbb{E}_{\prod_{l=1}^{L} q(z_l|z_{\lt l},x)}[\log \frac{ q(z_1|x)}{p(z_1)}]+\Bbb{E}_{\prod_{l=1}^{L} q(z_l|z_{\lt l},x)}[\sum_{l=2}^{L}\log \frac{ q(z_l|z_{\lt l},x)}{p(z_l|z_{\lt l})}] \\ &= \Bbb{E}_{ q(z_1|x)}[\log \frac{ q(z_1|x)}{p(z_1)}]+\sum_{l=2}^{L} \Bbb{E}_{\prod_{l=1}^{l} q(z_l|z_{\lt l},x)}[\log \frac{ q(z_l|z_{\lt l},x)}{p(z_l|z_{\lt l})}] \\ &= KL(q(z_1|x)||p(z_1))+\sum_{l=2}^{L} \Bbb{E}_{ q(z_{\lt l}|x)}[KL( q(z_l|z_{\lt l},x)||p(z_l|z_{\lt l}))] \end{aligned}\]

其中$p(z_l|z_{\lt l}),q(z_l|z_{\lt l},x)$的KL散度项为：

\[KL(q(z_l|z_{\lt l},x)||p(z_l|z_{\lt l})) = \frac{1}{2}\sum_{i=1}^{|z_l|}(-\log \Delta \sigma_{(i)}^2 + \frac{\Delta \mu_{(i)}^2}{\sigma_{(i)}^2}+\Delta \sigma_{(i)}^2-1)\]

2. 实现技巧

⚪ 层次设计

Nouveau VAE对隐变量进行分组$z=\{z_1,z_2,\cdots z_L\}$，其结构设计为多层次的编码器-解码器结构。

编码器经过层次编码，得到最顶层的编码向量$z_1$，然后再逐渐从顶层往回走，逐步得到底层特征$z_2,\cdots,z_L$；解码器也是自上而下地生成$z_1,z_2,\cdots,z_L$。其中$r$代表残差模块，$h$代表可训练参数，蓝色部分是参数共享的。

⚪ 残差模块

作者设计了残差模块，用到了深度可分离卷积、BatchNorm、Swish激活函数等结构。

值得一提的是很多生成模型已经弃用BatchNorm，改用LayerNorm或WeightNorm，因为实验发现用BatchNorm会损失性能。作者认为BatchNorm对训练有帮助，但对预测有害，因为预测阶段所使用的滑动平均均值方差不够好。作者提出在模型训练完后，通过多次采样同样batch_size的样本来重新估算均值方差，能够保证BatchNorm的预测性能。此外，为了保证训练的稳定性，为BatchNorm中$γ$的模长增加一个正则项。

作者为每一个卷积层增加了谱正则化，使得模型的Lipschitz常数变小，进而使得模型的损失Landscape更为光滑，有利于模型稳定训练。

3. 实验分析

作者给出了Nouveau VAE在CelebA HQ和FFHQ上的生成效结果。值得一提的是Nouveau VAE是第一个FFHQ数据集($256\times 256$)上进行图像生成的VAE类模型。

下表给出了Nouveau VAE的训练成本。以FFHQ数据集上训练的模型为例，隐变量共有$4+4+4+8+16=36$组，每组隐变量大小分别是{$82,162,322,642,1282$}$×20$。则生成一张$256×256×3$(约$20$万个元素)的FFHQ图像，需要总大小为$(4×82+4×162+4×322+8×642+16×1282)×20=6005760$维的随机变量。