Deep Generative Models

Catalog:

VAEs

以图像生成为例,想象一下这个场景,我们现在有一堆马🐎的图片,我们现在希望学习 $p(x)$,从而生成新的图片。那么我们可以问下自己,我们如何去画出马的图片呢?换句话说,我们把自己当作一个生成模型,我们如何做这件事呢?也许我们会先勾勒出一匹马的大致轮廓,它的大小和形状,然后添加马蹄,填充头部的细节,给它上色等等。最后,我们可以考虑背景。一般来说,我们可以说数据中有一些因素(例如轮廓、颜色、背景)对于生成对象(这里是马)至关重要。 一旦我们决定了这些因素,我们就可以通过添加细节来生成它们。 当我们画某物时,这或多或少是我们生成一幅画的过程。

我们现在用数学来表达这个生成过程。 也就是说,我们有我们感兴趣的高维对象 $x\in \mathcal{X}^D$(例如,对于图像,$\mathcal{X}\in {0,1,2,…,255}$)和一个低维潜在变量 $z\in\mathcal{Z}^M$(例如,$\mathcal{Z}=\mathbb{R}$),我们可以称之为数据中的隐藏因素。 在数学上,我们可以将$\mathcal{Z}^M$称为低维流形。 那么,生成过程可以表示为:

  1. $z\sim p(z)$ (Figure1, In red)
  2. $x\sim p(x|z)$(Figure1, In Blue)
img

Figure 1. 潜在变量模型和生成过程的简图. 注意嵌入在高维空间(此处为3D)中的低维流形(此处为2D)

简单来说,我们首先采样𝐳(例如,我们想象我的马的大小、形状和颜色),然后创建具有所有必要细节的图像,即,我们从条件分布 $p(x|z)$中采样得到x。 由于许多各种外部因素,创建两次完全相同的图像几乎是不可能的。

潜在变量模型背后的idea是我们引入了潜在变量 $z$ ,并且联合分布被分解如下:$ 𝑝(𝐱,𝐳)=𝑝(𝐱|𝐳)𝑝(𝐳)$. 这个公式表达的就是上述的生成过程。 但是在训练时我们只能访问样本 $x$。 因此,根据概率推理,我们应该对未知的潜在变量,也就是𝐳 ,进行求和(sum out, 也叫边缘化 marginalize out)。最后,其边缘似然函数如下:

$$p({x}) = \int p({x} | {z}) p({z})\ \mathrm{d} {z} $$

那么问题来了,如何计算这个积分呢?总的来说,这是个非常复杂的任务。有两个可能的解决方法:1. 直接处理这个积分 (Probabilistic PCA );2. 利用近似的方法也来解决,也就是变分推断 (variational inference).

Probabilistic PCA

对于 $p(x|z)$是线性模型时,可以直接求解。但推导过程没看懂,先不写了。

Variational Inference for Non-linear Latent Variable Models

让我们再看一次积分,显然我们无法准确计算积分。 最简单的方法是使用蒙特卡罗近似:
$$
\begin{aligned}
p(\mathbf{x}) &= \int p(\mathbf{x} | \mathbf{z})\ p(\mathbf{z})\ \mathrm{d} \mathbf{z} \\
&= E_{\mathbf{z}\sim p(\mathbf{z})} \left[ p(\mathbf{x} | \mathbf{z}) \right] \\
&\approx \frac{1}{K} \sum_{k} p(\mathbf{x} | \mathbf{z}_{k}) \end{aligned}
$$
其中,在最后一行我们采用近似的方法来模拟这个期望/积分函数。我们基于潜在变量的的先验概率来得到samples,$z_k\sim p(z)$. 在当前计算机越来越快的情况下,这个方法是相对简单的,我们可以在短时间内采样出无数的点。然而,我们学过的统计学知识告诉我们,当潜在空间是$z\in\mathcal{Z}^M$多维,且$M$很大的时候,我们会陷入维度灾难。为了cover住样本空间,我们所需采样的样本是M的指数形式。如果我们采样的样本数不够,那么这个近似效果就不好。

我们当然可以采用一些更先进的蒙特卡洛方法,然而,它们始终会受到维度灾难的影响。另一个可选择的近似方法是变分推断variational inference (Jordan et al., 1999). 我们考虑一组由 $\phi$ 参数化的变分分布,${q_{\phi}(z)}$. 比如,我们假设$\phi$为高斯分布, $\phi={\mu, \sigma^2}$. 我们知道这些分布的形式,并假设它们将非零的概率分配给所有的潜在变量 $z\in \mathcal{Z}^{M}$. 然后,这个边缘概率分布可以近似推导如下:
$$
\begin{aligned}
\ln p(x) &= ln\int p(x|z)p(z)dz \\
&= ln\int \dfrac{q_{\phi}(z)}{q_{\phi}(z)}p(x|z)p(z)dz \\
&= lnE_{z\sim q_{\phi}(z)}[\dfrac{p(x|z)p(z)}{q_{\phi}(z)}] \\
&\ge E_{z\sim q_{\phi}(z)}ln[\dfrac{p(x|z)p(z)}{q_{\phi}(z)}] \\
&= E_{z\sim q_{\phi}(z)} [lnp(x|z)+lnp(z)-lnq_{\phi}(z)] \\
&= E_{z\sim q_{\phi}(z)} [lnp(x|z)] - E_{z\sim q_{\phi}(z)}[lnq_{\phi}(z) - lnp(z)]
\end{aligned}
$$
第4行使用了Jensen’s inequality. 上述推导过程,我们把 $q_{\phi}(z)$ 换成 amortized variational posterior,也就是 $q_{\phi}(z|x)$ 是不影响推导过程的,因此我们可以得到:

$$lnp(z)\ge E_{z\sim q_{\phi}(z|x)} [lnp(x|z)] - E_{z\sim q_{\phi}(z|x)}[lnq_{\phi}(z|x) - lnp(z)]$$

amortized variational posterior 非常有用,因为我们可以利用神经网络得到这样一个模型,给定输入 $x$,然后输出对应分布的参数。在 (Kim et al., 2018) 中,作者使用了一种 semi-amortized variational inference.

最后,我们得到一个auto encoder-like模型,其包括一个encoder, $q_{\phi}(z|x)$ 和 一个 decoder, $p(x|z)$. 我们用随机性来强调encoder和decoder其实就是概率分布,这与deterministic auto-encoder是有区别的(这里没太懂)。这个带有 amortized variational posterior的自编码模型就是变分自编码 Variational Auto-Encoder (Kingma & Welling, 2013; Rezende et al., 2014). 其似然函数的下界就是 Evidence LOwer Bound (ELBO).

ELBO的第一项是 reconstruction error,$E_{z\sim q_{\phi}(z|x)} [lnp(x|z)]$ ; 第二项是regularizer,$E_{z\sim q_{\phi}(z|x)}[lnq_{\phi}(z|x) - lnp(z)]$,恰好就是KL散度,所以也可以成为KL item. 但在更复杂的模型中,这一项并不一定是 KL term,因此称作regularizer更通用一点。

回顾下:熵 $\rightarrow$ 交叉熵 $\rightarrow$ 相对熵/KL散度

  1. 信息量与概率成反比,当概率越大时,信息量越小;并且信息量为非负数。因此,定义信息量为 $log\dfrac{1}{p}$. 熵是信息量的期望:$E_{x\sim p(x)}ln\dfrac{1}{p(x)}$

  2. 交叉熵指的是:根据真实分布p来衡量预测分布q的度量。同样的我们也用信息量的期望来衡量这个交叉熵度量,我们希望交叉熵越小时,预测的q越准确,也就是越接近于1。类似地,交叉熵可以写成 $H(p, q) = E_{x\sim p(x)}ln\dfrac{1}{q(x)}$. 实际上,我们是无法知道真实分布p的,只能依据现有的样本统计得到。

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    import torch
    import torch.nn as nn
    import torch.nn.functional as F

    size = 3
    input = torch.randn(2, size)
    target = torch.Tensor([0, 2]).long()

    # use loss function
    loss_fn = nn.CrossEntropyLoss(reduction="mean")
    loss = loss_fn(input, target)

    # computer nll loss step by step
    score = torch.log_softmax(input, dim=1)
    my_nll = torch.sum(-score * F.one_hot(target, size)) / target.size(0)
    # use nll loss

    nll_loss_fn = nn.NLLLoss()
    nll_loss = nll_loss_fn(score, target)

    print(nll_loss == loss == my_nll) # True
  3. 相对熵: 根据Gibbs’ inequality上述例子中的 $H(p,q) >= H(p)$ 恒成立。当且仅当q=p时,这个等号才成立。那么熵H(p,q)相比熵H(q)多出来的部分就是相对熵 ,也称为KL散度(Kullback–Leibler divergence,KLD).

$$
D(p||q)=H(p,q)-H(p)=\sum_{x\sim p(x)}[ln\dfrac{1}{q(x)} - ln\dfrac{1}{p(x)}]=E_{x\sim p(x)}[ln\dfrac{p(x)}{q(x)}]=\int_{x}p(x)[ln\dfrac{p(x)}{q(x)}]
$$

A different perspective on the ELBO

下面提供一种新的推导方法,个人觉得更能理解ELBO:
$$
\begin{align}
\ln p(\mathbf{x}) &= E_{\mathbf{z} \sim q_{\phi}(\mathbf{z}|\mathbf{x})} \left[ \ln p(\mathbf{x}) \right] \\
&= E_{\mathbf{z} \sim q_{\phi}(\mathbf{z}|\mathbf{x})} \left[ \ln \frac{p(\mathbf{z}|\mathbf{x}) p(\mathbf{x})}{p(\mathbf{z}|\mathbf{x})} \right] \\
&= E_{\mathbf{z} \sim q_{\phi}(\mathbf{z}|\mathbf{x})} \left[ \ln \frac{p(\mathbf{x}|\mathbf{z}) p(\mathbf{z})}{p(\mathbf{z}|\mathbf{x})} \right] \quad \text{根据贝叶斯公式p(z|x)p(x)=p(x|z)p(z)}\\
&= E_{\mathbf{z} \sim q_{\phi}(\mathbf{z}|\mathbf{x})} \left[ \ln \frac{p(\mathbf{x}|\mathbf{z}) p(\mathbf{z})}{p(\mathbf{z}|\mathbf{x})} \frac{q_{\phi}(\mathbf{z}|\mathbf{x})}{q_{\phi}(\mathbf{z}|\mathbf{x})}\right] \quad \text{这一步很关键, } 1 = \frac{q_{\phi(z|x)}}{q_{\phi(z|x)}}, q_{\phi(z|x)} 是用来近似真实后验概率的变分后验概率 \\
&= E_{\mathbf{z} \sim q_{\phi}(\mathbf{z}|\mathbf{x})} \left[ \ln p(\mathbf{x}|\mathbf{z}) \frac{p(\mathbf{z})}{q_{\phi}(\mathbf{z}|\mathbf{x})} \frac{q_{\phi}(\mathbf{z}|\mathbf{x})}{p(\mathbf{z}|\mathbf{x})} \right] \\
&= E_{\mathbf{z} \sim q_{\phi}(\mathbf{z}|\mathbf{x})} \left[ \ln p(\mathbf{x}|\mathbf{z}) - \ln \frac{q_{\phi}(\mathbf{z}|\mathbf{x})}{p(\mathbf{z})} + \ln \frac{q_{\phi}(\mathbf{z}|\mathbf{x})}{p(\mathbf{z}|\mathbf{x})} \right] \quad \text{到这里能看出KL散度了}\\
&= E_{\mathbf{z} \sim q_{\phi}(\mathbf{z}|\mathbf{x})} \left[ \ln p(\mathbf{x}|\mathbf{z}) \right] - KL\left[ q_{\phi}(\mathbf{z}|\mathbf{x}) | p(\mathbf{z}) \right] + KL \left[ q_{\phi}(\mathbf{z}|\mathbf{x}) |p(\mathbf{z}|\mathbf{x}) \right] .
\end{align}
$$
前两项就是 ELBO,最后一项中 $p(z|x)$ 这个表示真实的后验概率 real posterior, $q_{\phi}(z|x)$ 表示变分后验概率 variational posterior. 我们并不知道真实的后验概率,但是我们可以跳过这一项,因为KL散度一定大于等于0. 去掉最后一项,我们得到ELBO,同时我们知道 ELBO和真实的对数似然之间的间隔是 $KL \left[ q_{\phi}({z}|{x}) |p({z}|{x}) \right] $.
$$
\begin{align}
\ln p(\mathbf{x}) &= E_{\mathbf{z} \sim q_{\phi}(\mathbf{z}|\mathbf{x})} \left[ \ln p(\mathbf{x}|\mathbf{z}) \right] - KL\left[ q_{\phi}(\mathbf{z}|\mathbf{x}) | p(\mathbf{z}) \right] + KL \left[ q_{\phi}(\mathbf{z}|\mathbf{x}) |p(\mathbf{z}|\mathbf{x}) \right] \\
&\ge E_{\mathbf{z} \sim q_{\phi}(\mathbf{z}|\mathbf{x})} \left[ \ln p(\mathbf{x}|\mathbf{z}) \right] - KL\left[ q_{\phi}(\mathbf{z}|\mathbf{x}) | p(\mathbf{z}) \right]
\end{align}
$$
Beautiful! 但同样我们能看出一些问题,如果 $q_{\phi}({z}|{x})$ 和 $p(z|x)$ 距离很大,那么ELBO优化的再好,ELBO和真实的对数似然之间的差距依然很大。也就是说,如果我们采用很简单的后验概率,我们会得到一个不太好的VAE模型。

Figure 2. ELBO是对数似然的下界。ELBO最大时对应的 $\hat \theta$ 并不一定也能让对数似然最大。ELBO 越losser,这对模型参数的最大似然估计的偏差越大。

Components of VAEs

  • 我们使用amortized variational posteriors ${q_{\phi}(z|x)}_{\phi}$ 来近似真实的后验分布 $p(z|x)$. 这个概率分布可以看作是 encoder.
  • conditional likelihood p(x|z) 对应的概率分布可以看作是 decoder
  • $p(z)$ 是对于潜在变量的边缘分布,也可以看作是 prior 模型

至此,还有两个问题需要解决:

  1. 如何参数化这些distributions?
  2. 如何计算这些期望呢?也就是积分。

Parameterization of distributions

显然,我们用Neural Networks来表达上述两个distributions: encoder 和 decoder。在VAE框架下,大部分情况下,我们可以使用任何distributions. 但是,我们也必须满足对应的任务。

  1. 对于decoder分布 $p_{\theta}(x|z)$ 显然不可能是正态分布,因为image的pixel values是离散的。一个可能的distribution可以是 categorical distribution:

​ $$p_{\theta}(x|𝐳)=\text{Categorical}(𝐱|\theta(𝐳))$$

​ 其中 $\theta(z) = \text{softmax}(NN(z))$. 这里把重构当作分类任务来做,当然也可以是回归任务。

  1. 对于潜在变量的分布,为了方便,通常 $z$ 可视为连续随机变量的向量,$z\in \mathbb{R}^M$. 因此,我们可以使用 Gaussians 来表示 variational posterior 和 prior.
    $$
    \begin{align}
    q_{\phi}(\mathbf{z}|\mathbf{x}) &= \mathcal{N}\left(\mathbf{z} | \mu_{\phi}(\mathbf{x}), \mathrm{diag}\left[ \sigma_{\phi}^2(\mathbf{x}) \right] \right) \\
    p(\mathbf{z}) &= \mathcal{N}\left(\mathbf{z} | 0, \mathbf{I} \right)
    \end{align}
    $$
    其中 $\mu_{phi}(x), \sigma_{\phi}(x)$ 是神经网络的输出。在实际使用中,我们使用NN得到 $2M$ 的 values $\in R^{1\times 2M}$,其中 $R^{1\times M}$ 表示 means,$R^{1\times M}$ 表示 variances.

Reparameterization trick

到目前为止,我们学习了log-likelihood和ELBO。但是仍然有一个问题,我们用encoder $q_{\phi}(z|x)$ 得到关于潜在变量的分布,我们该如何计算$E_{z\sim q_{\phi}(z|x)}(x|z)$这个积分呢?显然,$z\sim q_{\phi}(z|x)$这个采样过程是不可导的,我们可以采用MC-approximation,但是这样仍然有个问题,从这个变分后验sample得到的z,在ELBO的训练过程中,在计算关于$\phi$的梯度时,梯度的方差特别大。

因此,另一个可能的方法是reparameterizing这个分布(Devroye, 1996). 具体地,我们可以将随机变量 $z$ 表示为具有简单分布的独立随机变量的原始变换(例如算术运算、对数等)的组合。换句话说,我们使用重参数技巧表达为确定性的变量:
$$
z = \mu + \sigma \cdot \epsilon
$$
其中 $\epsilon \sim {N}(\epsilon|0,1)$ .

使用Reparameterization方法能够显著减小梯度的方法。Why?因为随机性来自独立的分布$p(\epsilon)$,我们计算梯度是关于确定性函数(即神经网络),而不是随机的对象$z\sim q_{\phi}(z|x)$。 更棒的是,由于我们使用随机梯度下降来学习 VAE,因此在训练期间仅采样一次 $z$ 就足够了!

综上,VAE框架主要包括:

  1. variational posterior $q_{\phi}(z|x)$ using encoder
  2. sample z from $q_{\phi}(z|x)$ and feed it to decoder, using reparameterization trick
  3. conditional likelihood $p_{\theta}(x|z)$ using decoder
  4. reconstruction loss
  5. kl loss between variational posterior and prior $p_{\theta}$. 其中 $q_{\phi}(z|x)\sim {N}(z|\mu, \sigma^2I)$ 多维高斯分布,$p_{\theta}\sim {N}(0,1)$ 是正态分布。

KL 散度的推导如下:

image-20211209135812911

代码如下:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
class VariationalEncoder(nn.Module):
def __init__(self, latent_dims):
super(VariationalEncoder, self).__init__()
self.linear1 = nn.Linear(784, 512)
self.linear2 = nn.Linear(512, latent_dims)
self.linear3 = nn.Linear(512, latent_dims)

self.kl = 0

def forward(self, x):
x = torch.flatten(x, start_dim=1)
x = F.relu(self.linear1(x))
mu = self.linear2(x)
sigma = torch.exp(self.linear3(x))

# reparameterization trick
z = mu + sigma*torch.randn_like(sigma)

self.kl = 0.5*(sigma**2 + mu**2 - torch.log(sigma) - 1).sum()

return z

class Decoder(nn.Module):
def __init__(self, latent_dims):
super(Decoder, self).__init__()
self.linear1 = nn.Linear(latent_dims, 512)
self.linear2 = nn.Linear(512, 784)

def forward(self, z):
x_hat = F.relu(self.linear1(z))
x_hat = torch.sigmoid(self.linear2(x_hat))
return x_hat.reshape((-1, 1, 28, 28))

class VariationalAutoencoder(nn.Module):
def __init__(self, latent_dims):
super(VariationalAutoencoder, self).__init__()
self.encoder = VariationalEncoder(latent_dims)
self.decoder = Decoder(latent_dims)

def forward(self, x):
z = self.encoder(x)
return self.decoder(z)

def train(model, data, epochs=20):
opt = torch.optim.Adam(model.parameters())
for epoch in range(epochs):
print(f"epoch: {epoch+1}/{epochs}")
for x, y in data:
x = x.to(device) # GPU
opt.zero_grad()
x_hat = model(x)
loss = ((x - x_hat)**2).sum() + model.encoder.kl
loss.backward()
opt.step()
return model

latent_dims = 2
vae = VariationalAutoencoder(latent_dims).to(device) # GPU
vae = train(vae, data)

More about VAEs!

  • Estimation of the log-likelihood using importance weighting 我们前面提到过 ELBO 只是对数似然的下界,它不应该被用作对数似然的良好估计。(Burda et al., 2015; Rezende et al., 2014) 采用了一种 importance weighting procedure 方法。

  • Enhancing VAEs: Better encoders 意味着更好的后验概率, 一个很重要的方向是 conditional flow-based models (van den Berg et al., 2018; Hoogeboom et al., 2020; Kingma et al., 2016; Rezende & Mohamed, 2015; Tomczak & Welling, 2016; Tomczak & Welling, 2017).

  • Enhancing VAEs: Better decoders 可以使用不同的模型或者loss function来拟合原始的数据,比如pixel-CNN, transformer等。

  • Enhancing VAEs: Better priors 设定一个好的先验也是很重要的,能够减小与变分后验的gap。很多研究尝试解决这个问题,比如:using a multimodal prior mimicking the aggregated posterior (known as the VampPrior) (Tomczak & Welling, 2018), or a flow-based prior (e.g., (Gatopoulos & Tomczak, 2020)), an ARM-based prior (Chen et al., 2016) or using an idea of resampling (Bauer & Mnih, 2019).

  • VAEs for non-image data 不仅仅是图像数据,文本、序列数据等也可以。

  • Extending VAEs Here, we present the unsupervised version of VAEs. However, there is no restriction to that and we can introduce labels or other variables. In (Kingma et al., 2014) a semi-supervised VAE was proposed. This idea was further extended to the concept of fair representations (Louizos et al., 2015). In (Ilse et al., 2020), the authors proposed a specific latent representation that allows domain generalization in VAEs. In (Blundell et al., 2015) variational inference and the reparameterization trick were used for Bayesian Neural Nets. This paper is not necessarily introducing a VAE, but a VAE-like way of dealing with Bayesian neural nets.

  • Different latent spaces in (Davidson et al., 2018; Davidson et al., 2019) a hyperspherical latent-space was used, and in (Mathieu et al., 2019) the hyperbolic latent space was utilized.

  • **The posterior collapse ** There were many ideas proposed to deal with the posterior collapse. For instance, (He et al., 2019) propose to update variational posteriors more often than the decoder. In (Dieng et al., 2019) a new architecture of the decoder is proposed by introducing skip connection to avoiding the posterior collapse.

  • Various perspectives on the objective The core of the VAE is the ELBO. However, we can consider different objectives. For instance, (Dieng et al., 2017) propose an upper-bound to the log-likelihood that is based on the chi-square divergence (CUBO). In (Alemi et al., 2018) an information-theoretic perspective on the ELBO is presented. (Higgins et al., 2016) introduced the 𝛽β-VAE where the regularization term is weighted by a fudge factor 𝛽β. The objective does not correspond to the lowe-bound of the log-likelihood though.

  • Deterministic Regularized Auto-Encoders: We can take look at the VAE and the objective, as mentioned before, and think of it as a regularized version of an auto-encoder with a stochastic encoder and a stochastic decoder. (Ghosh et al., 2020) “peeled off” VAEs from all stochasticity and indicated similarities between deterministic regularized auto-encoders and VAEs, and highlited potential issues with VAEs. Moreover, they brilliantly pointed out that even with a deterministic encoders, due to stochasticity of the empirical distribution, we can fit a model to the aggregated posterior. As a result, the deterministic (regularized) auto-encoder could be turned into a generative model by sampling from our model, 𝑝𝜆(𝐳)pλ(z), and then, deterministically, mapping 𝐳z to the space of observable 𝐱x. In my opinion, this direction should be further explored and an important question is whether we indeed need any stochasticity at all.

  • Hierarchical VAEs Very recently, there are many VAEs with a deep, hierarchical structure of latent variables that achieved remarkable results! The most important ones are definitely BIVA (Maaløe et al., 2019), NVA (Vahdat & Kautz, 2020), and very deep VAEs (Child, 2020). Another interesting perspective on a deep, hierarchical VAE was presented in (Gatopoulos & Tomczak, 2020) where, additionally, a series of deterministic functions was used.

  • Adversarial Auto-Encoders Another interesting perspective on VAEs is presented in (Makhzani et al., 2015). Since learning the aggregated posterior as the prior is an important component mentioned in some papers (e.g., (Tomczak & Welling, 2018)), a different approach would be to train the prior with an adversarial loss. Further, (Makhzani et al., 2015) present various ideas how auto-encoders could benefit from adverarial learning.

Hierarchical Variational Auto-Encoders

交叉熵:$p_{data}(x)$ 是真实分布,$p_{\theta}(x)$是预测分布
$$
\begin{align}
\mathbb{CE}[p_{data}(x)|p_{\theta}(x)] &= E_{x\sim p_{data}(x)}[ln\dfrac{1}{p_{\theta}(x)}] \\
&= -\frac{1}{N}\sum_{1}^{N}lnp_{\theta}(x_n)
\end{align}
$$
相对熵:
$$
\begin{align}
\mathbb{KL}[p_{data}(x)|p_{\theta}(x)] &= -\mathbb{H}[p_{data}(x)] + \mathbb{CE}[p_{data}(x)|p_{\theta}(x)] \\
&= const + \mathbb{CE}[p_{data}(x)|p_{\theta}(x)]
\end{align}
$$

潜在变量模型存在一个问题: [后验崩塌**(posteroir collapse)**](python - What is “posterior collapse” phenomenon? - Data Science Stack Exchange). 当后验未出现崩塌时,$z_d$(d维潜在变量)通过encoder采样得到 $q_{\phi}(z|x)={N}(z|\mu, \sigma^2)$, 其中 $\mu, \sigma$是关于输入 $x$ 的稳定的函数。也就是encoder从输入中提取到了有用的信息赋予到 $\mu,\sigma$中。

当出现后验崩塌时,通过encoder从输入 $x$ 提取到的信息很弱或者噪声太大**(too weak or too noisy)**. 也就是说decoder在生成图像时,会忽略潜在变量 $z$ 的信息。too noisy意味着 $\mu,\sigma$不稳定, too weak意味着:
$$
q_{\phi}(z|x) \approx q_{\phi}(z) = \mathcal{N}(a,b)
$$
其中 $a,b$ 是常量。

从可视化的角度看这个问题:

image-20211209135812911

横轴是重构loss,纵轴是 潜在变量模型$p_{\theta}(z|x)$ 的参与度。当处于上图的左下角时,表明重构loss很小,但是后验概率并不重要,也就是出现了后验崩塌。所以问题来了,在优化过程中,一定存在这样的模型在完全抛弃潜在变量模型的情况下使得KL很小。但是在实际使用过程中,我们发现潜在变量是有用的,其原因是我们所选择的模型存在归纳偏置**(inductive bia)**,这使得潜在变量一定会学到有用的信息。

How to define a proper class of models? 该如何定义这一类模型,使得潜在变量一定能学到有用的信息呢?这定义了一个constrained optimization problem (Phuong et al., 2018; Rezende & Viola, 2018) 和正则化问题 an auxiliary regularizer (Sinha & Dieng, 2021; Tomczak, 2016) to (implicitly) define usefulness of the latents.

在下文中,我们使用层次化的结构 hierarchical architectures 来缓解这一问题。我们在前面提到潜在变量可以理解成构成目前的一些因素/概率,而这些因素可以是层次化结构的。如果潜在变量模型存在层次化结构,能够引入归纳偏置,进入限制我们的VAE模型,强制信息在潜在变量和观测目标中流动。

Hierarchical Variational Auto-Encoders

Two-level VAE:

这里有一个假设就是 $z_2\rightarrow z_1\rightarrow x$ 是一阶隐马尔可夫链,也就是 $x$ 不受 $x_2$的影响。那么有:
$$
p(x,z_1,z_2) = p(x|z_1,z_2)p(z_1,z_2) = p(x|z_1)p(z_1|z_2)p(z_2)
$$
那么对应的模型(概率分布)可以表示成:
$$
\begin{align}
p(z_1|z_2) &= N(z_1|\mu(z_2), \sigma^2(z_2)) \\
p(z_2) &= N(z_2|0,1) \\
q(z_1|z_2) &= N(z_1|\mu(x), \sigma^2(x)) \\
q(z_2|z_1) &= N(z_2|\mu(z_1), \sigma^2(z_1))
\end{align}
$$
其中 $p(z_1|z_2)$ 是用NN表示的模型,$p(z_2)$是先验模型,$q(z_1|x)$ 和 $q(z_2|z_1)$是后验概率模型。

我们使用层次模型的目的是保证潜在变量学到有用的信息,可是我们做到了吗?答案是NO,因为前面提到的问题依然还是存在。这里我们先推导下层次模型的ELBO,在进一步分析这个问题:

在推导之前有这么一个公式会用到,这一步的公式用到了一阶马尔可夫链的特性,也就是$z_2$与$x$无关:
$$
q(z_1,z_2|x) = q((z_2|z_1)|x)q(z_1|x) = q(z_2|z_1)q(z_1|x)
$$

$$
\begin{align}
p(x) &= \int p(x|z_1,z_2)p(z_1,z_2)d(z_1,z_2) \\
ELBO(x) &= E_{(z_1,z_2)\sim q_{\phi}(z_1,z_2|x)}[lnp(x|z_1,z_2)-KL[q_{\phi}(z_1,z_2|x)||p(z_1,z_2)]] \quad \text{利用一层VAE的推导可得到} \\
&= E_{(z_1,z_2)\sim q_{\phi}(z_1,z_2|x)}[lnp(x|z_1)-ln\dfrac{q(z_1,z_2|x)}{p(z_1,z_2)}] \quad x \text{不依赖于} z_2,以及KL散度公式 \\
&= E_{(z_1,z_2)\sim q_{\phi}(z_1,z_2|x)}[lnp(x|z_1) - ln\dfrac{q(z_1|x)q(z_2|z_1)}{p(z_2)p(z_1|z_2)}] \quad \text{贝叶斯公式,以及上一个公式} \\
&= E_{(z_1,z_2)\sim q_{\phi}(z_1,z_2|x)}[lnp(x|z_1) - ln\dfrac{q(z_1|x)}{p(z_1|z_2)} - ln\dfrac{q(z_2|z_1)}{p(z_2)} \\
&= E_{(z_1,z_2)\sim q_{\phi}(z_1,z_2|x)}[lnp(x|z_1) - KL[q(z_1|x)||p(z_1|z_2) - KL[q(z_2|z_1)||p(z_2)]]
\end{align}
$$

我们发现后验崩塌的问题依然存在,$z_2$ 大多数情况不会被使用(Burda et al., 2015; Maaloe et al., 2017).

Diffusion Model

学过 Hierarchical VAE之后再来看diffusion model就很容易理解这些公式背后的意义了。这里笔者参照H-VAE画了个图,与H-VAE的关系就一目了然了:

在VAE模型中,后验概率 $q_{\phi}(z|x)$ 是通过NN学习得到的,但是在diffusion model, (Sohl-Dickstein et al., 2015; Ho et al., 2020)将这个Gaussian diffusion process 定为确定性的,比如在 (Ho et al., 2020) 将其设置为 $\beta_i=10^{-4}$ to $\beta_T=0.02$.

Gaussian diffusion process 的样本可视化如下:

DDGM的目标函数和H-VAE的区别在于潜在变量的维度,但基本推导是完全一致的。
$$
ln p_{\theta}(x)=ln\int Q_{z_{\phi}(z_{1:T}|x)}\dfrac{p_{\theta}(x,z_{1:T})}{Q_{\phi(z_{1:T}|x)}}\mathrm{d}z_{1:T}
$$
根据H-VAE的推导过程,DDGM的ELBO可以直接写出如下:
$$
ELBO(x;\theta,\phi) = E_{Q_{\phi}(z_{1:T}|x)}[lnp_{\theta}(x|z_1)-KL[q(z_1|x)||p(z_1|z_2)] - \sum_{i=2}^{T-1}KL[q(z_i|z_{i-1})||p(z_i|z_{i+1})] - KL[q(z_T|z_{T-1})||p(z_T)]]
$$
其中:
$$
p_{\theta}(z_T) = {N}(z_T|0, I)
$$
到目前位置,我们学了H-VAE和DDGM,他们都是采用变分推断的方法来近似学习log-likelihood. 但是为什么这种层次化的方式能缓解我们在VAE中提到的后验崩塌问题: 如何保证潜在变量学到有用的信息?

Reference:

作者

Xie Pan

发布于

2021-12-07

更新于

2022-04-17

许可协议

评论