生成模型VAE

likelihood-based

likelihood-based generative models是生成模型的一类范式，通过最大化所有观测数据的似然函数来学习模型参数。

该怎么去理解likelihood-based，基于似然的生成模型，优化的目标函数到底是什么？

如图所示，我们希望待学习分布$p_{\theta }(x)$和真实数据分布$p_{data}(x)$尽可能近。那该如何度量两个分布之间的距离？自然想到KL散度。
$$\begin{array}{rl}{D}_{KL}(P_{data}||P_{\theta })& =E_{x\sim P_{data}}[log\frac{P_{data}(x)}{P_{\theta }(x)}]\\ & =E_{x\sim P_{data}}[log{P}_{data}(x)]-E_{x\sim P_{data}}[log{P}_{\theta }(x)]\end{array}$$
RHS第一项是常量，不依赖待学习参数$\theta$。因此最小化KL散度等价于最大化log-likelihood
$$\begin{array}{rl}\underset{P_{\theta }}{\mathrm{argmin}}{D}_{KL}(P_{data}||P_{\theta })& =\underset{P_{\theta }}{\mathrm{argmin}}-E_{x\sim P_{data}}[log{P}_{\theta }(x)]\\ & =\underset{P_{\theta }}{\mathrm{argmax}}{E}_{x\sim P_{data}}[log{P}_{\theta }(x)]\end{array}$$

ELBO

假设数据集中有N个数据点 { x ( 1 ) , . . . , x ( N ) } \{x^{(1)}, ..., x^{(N)}\} {x(1),...,x(N)}，定义在其上的概率似然可以写作：
$$E_{x\sim P_{data}}[log{P}_{\theta }(x)]=\sum _{i=1}^N[log{p}_{\theta }(x^{(i)})]$$

对于$log{p}_{\theta }(x^{(i)})$我们有：
$$log{p}_{\theta }(x^{(i)})=D_{KL}(q_{\varphi }(z|x^{(i)}||p_{\theta }(z|x^{(i)}))+L(\theta ,\varphi ;x^{(i)})$$

RHS的第一项是近似分布和真实后验分布之间的KL散度，因为KL散度的非负特性，RHS的第二项被称为数据点$x^{(i)}$边际似然的(variational) lower bound，简称为ELBO。因此：
$$log{p}_{\theta }(x^{(i)})\ge L(\theta ,\varphi ;x^{(i)})=E_{z\sim q_{\varphi }(z|x^i)}[log{p}_{\theta }(x^{(i)}|z)]-D_{KL}(q_{\varphi }(z|x^{(i)})||p(z))$$

对于上述结论推导如下。
为了书写简化，下文以$p(x)$代称$p_{\theta }(x^{(i)})$。正常上式中的$p(x)$非常复杂，无法求解。因此采用变分推断的方式，通过最大化ELBO，达到间接优化似然函数的目的，通常也称为proxy或者surrogate。

对于上面的“$p(x)$非常复杂，无法求解。”做进一步的解释：
我们有两种方式去利用联合概率分布$p(x,z)$恢复出观测数据的概率似然$p(x)$，一种是边际概率似然$p(x)=\int p(x,z)dz$，还有一个是贝叶斯法则$p(x)=\frac{p(x,z)}{p(z|x)}$。但这两个方法在计算过程中，方法一是对所有隐变量$z$做积分，方法二需要知道隐变量编码器$p(z|x)$，对于复杂模型，均无法直接求解。因此我们可以利用上述两个公式，推导出ELBO，做为对数似然的代理优化目标函数。

对于上面的“通过最大化ELBO，达到间接优化似然函数的目的”做进一步解释：
引入隐变量$z$，我们的目标是学习可以描述观测数据的隐变量结构。换句话说，我们希望优化变分后验$q_{\varphi }(z|x)$的参数，使其与真实后验分布$p(z|x)$完全一致。这可以通过最小化两者之间的KL散度实现。但不幸的是，我们不知道真实的$p(z|x)$分布，没法直接求解两者之间的KL散度项。但我们注意到数据的对数似然$logp(x)$始终是常值，ELBO + KL等于常值，最大化ELBO项等价于最小化KL散度项，ELBO最大时，模型参数达到最优。也因此ELBO可以做为建模并学习后验分布的proxy。

$$\begin{array}{rl}logp(x)& =logp(x)\int q_{\varphi }(z|x)dz\\ & =\int q_{\varphi }(z|x)logp(x)dz\\ & =E_{z\sim q_{\varphi }(z|x)}[logp(x)]\\ & =E_{z\sim q_{\varphi }(z|x)}[log\frac{p(x,z)}{p(z|x)}]\\ & =E_{z\sim q_{\varphi }(z|x)}[log\frac{p(x,z)q_{\varphi }(z|x)}{p(z|x)q_{\varphi }(z|x)}]\\ & =E_{z\sim q_{\varphi }(z|x)}[log\frac{p(x,z)}{q_{\varphi }(z|x)}]+E_{z\sim q(z|x)}[log\frac{q_{\varphi }(z|x)}{p(z|x)}]\\ & =E_{z\sim q_{\varphi }(z|x)}[log\frac{p(x,z)}{q_{\varphi }(z|x)}]+KL(q_{\varphi }(z|x)||p(z|x))\\ & \ge E_{z\sim q_{\varphi }(z|x)}[log\frac{p(x,z)}{q_{\varphi }(z|x)}]=ELBO\end{array}$$
推导过程用到了如下特性：
$$\int q_{\varphi }(z|x)dz=1$$
$$KL(q_{\varphi }(z|x)||p(z|x))\ge 0$$
evidence和ELBO之间的差值是非负的KL项，因此ELBO的值不会超过evidence，且当参数化模型近似的后验分布$q(z|x)$和真实后验分布$p(z|x)$一致时，等号成立，因此ELBO可以认为是evidence的严格下界。上述ELBO的推导过程也可以用Jensen不等式推导一遍，但从ELBO是evidence严格下界这个角度去理解，没有上述过程那么直观。
$$\begin{array}{rl}logp(x)& =log\int p(x,z)dz\\ & =log\int p(x,z)\frac{q_{\varphi }(z|x)}{q_{\varphi }(z|x)}dz\\ & =log{E}_{z\sim q_{\varphi }(z|x)}[\frac{p(x,z)}{q_{\varphi }(z|x)}]\\ & \ge E_{z\sim q_{\varphi }(z|x)}[\frac{p(x,z)}{q_{\varphi }(z|x)}]=ELBO\end{array}$$

VAE

将ELBO变换形式可得：
$$ELBO=E_{z\sim q(z|x)}[logp(x|z)]-KL(q_{\varphi }(z|x)||p(z))$$
我们用$\varphi$参数化$p(z|x)$部分，可称之为decoder；用$\theta$参数化$p(x|z)$部分，可称之为encoder，ELBO可以进一步写为：
$$ELBO=E_{z\sim q_{\varphi }(z|x)}[log{p}_{\theta }(x|z)]-KL(q_{\varphi }(z|x)||p(z))$$

ELBO的第一部分是重建loss，第二部分是对隐变量分布做约束，让其尽可能贴近先验分布$p(z)$，最小化该项可以鼓励encoder学习一个正常的分布，而不致于坍塌为Dirac delta函数。

SGVB估计器和AEVB算法

SGVB: Stochastic Gradient Variational Bayes。
AEVB: Auto-Encoding Variational Bayes。

参考下文的重参数化技巧，我们可以将随机变量$z\sim q_{\varphi }(z|x)$使用一个辅助噪声$ϵ$的可微变换$g_{\varphi }(ϵ,x)$重参数化。
$$z=g_{\varphi }(ϵ,x),ϵ\in p(ϵ)$$
对于一个函数$f(z)$相对与$q_{\varphi }(z|x^{(i)})$分布的期望，我们可以使用Monte Carlo估计：
$$E_{z\sim q_{\varphi }(z|x^{(i)})}[f(z)]=E_{p(ϵ)}[f(g_{\varphi }(ϵ,x^{(i)}))]\approx \frac{1}{L}\sum _{l=1}^{L}f(g_{\varphi }({ϵ}^{(l)},x^{(i)})$$
where ${ϵ}^{(l)}\sim p(ϵ)$
我们可以该技术用在ELBO中，改写下式：
$$ELBO=E_{z\sim q_{\varphi }(z|x)}[\frac{p(x,z)}{q_{\varphi }(z|x)}]$$
可以得到通用的SGVB估计器：
$${\tilde{L}}^A(\theta ,\varphi ;x^{(i)})=\frac{1}{L}\sum _{l=1}^{L}log{p}_{\theta }(x^{(i)},z^{(i,l)})-log{q}_{\varphi }(z^{(i,l)}|x^{(i)})$$
where $z^{(i,l)}=g_{\varphi }({ϵ}^{(i,l)},x^{(i)})={\mu }^{(i)}+{\sigma }^{(i)}\odot {ϵ}^{(l)}$, and ${ϵ}^{(l)}\sim p(ϵ)$。

不过网络训练的时候也不会用上式。而且基于下式做进一步改写：
$$ELBO=E_{z\sim q_{\varphi }(z|x)}[log{p}_{\theta }(x|z)]-KL(q_{\varphi }(z|x)||p(z))$$
得到第二版SGVB估计器，因为第二项的KL散度有解析形式，因此只有第一项的重建误差需要通过采样进行估计。KL散度项可理解为对参数$\varphi$的正则化，鼓励近似后验分布接近先验分布$p(z)$。
$${\tilde{L}}^B(\theta ,\varphi ;x^{(i)})=\frac{1}{L}\sum _{l=1}^{L}log{p}_{\theta }(x^{(i)}|z^{(i,l)})-D_{KL}(q_{\varphi }(z|x^{(i)})||p(z)))$$

这里需要大白话翻译一下：输入第$i$个数据点$x^{(i)}$，经过$\varphi$参数化模型，得到后验分布$p(z|x)$的均值${\mu }^{(i)}$和方差${\sigma }^{(i)}$。利用该均值方差，我们用蒙特卡洛估计去估算$E_{z\sim q_{\varphi }(z|x)}[log{p}_{\theta }(x|z)]$，也就是$log{p}_{\theta }(x|z)$在隐变量$z$上的期望值。具体做法就是随机采样$L$个辅助噪声，利用重参数化技巧构建出隐变量$z^{(i,l)}$，输入到$\theta$参数化的生成模型中，计算$L$个输出的loss均值，作为最终估计。

假定数据集$X$有N个数据点，从中随机取M个数据点 X M = { x ( i ) } i = 1 M X^M=\{x^{(i)}\}_{i=1}^M XM={x(i)}i=1M​作为minibatch，去估计上述ELBO，则有：
$$L(\theta ,\varphi ;X)\approx {\tilde{L}}^M(\theta ,\varphi ;X^M)=\frac{N}{M}\tilde{L}(\theta ,\varphi ;x^{(i)})$$
实验过程中，我们发现只要minibatch够大，比如M=100，做蒙特卡洛估计时L可以取1，而不用大量采样，拉长学习时长。

重建损失

第一部分的$p_{\theta }(x|z)$属于重建损失，但其分布形式不确定，根据数据集特性，我们可以假设其为多元高斯分布/伯努利分布/拉普拉斯分布。

高斯分布

如果假设其为高斯分布，$p_{\varphi }(x|z)=N(x;{\mu }_{\varphi }(z),{\sigma }_{\varphi }^2(z))$，则：
$$\begin{array}{rl}log{p}_{\varphi }(x|z)& =-\frac{1}{2}log(2\pi {\sigma }_{\varphi }^2(z))-\frac{1}{2{\sigma }_{\varphi }^2(z)}(x-{\mu }_{\varphi }(z){)}^2\\ & =-\frac{1}{2}log(2\pi )-log{\sigma }_{\varphi }(z)-\frac{1}{2}(\frac{x-{\mu }_{\varphi }(z)}{{\sigma }_{\varphi }(z)}{)}^2\end{array}$$
如果假设方差为定值$I$，网络只学习${\mu }_{\varphi }(z)$，这样就简化为最一般的MSE Loss（这里还要展开讲一讲，实际训练时，只用MSE Loss监督的话，重建效果会blur，进一步引出了VQGAN，用GAN Loss起到锐化作用）。

伯努利分布

如果假设分布是伯努利分布，那么最后将得到一个交叉熵损失。 比如用MNIST数据集做实验，我们可以假设decoder输出服从伯努利分布。
$$log{p}_{\varphi }(x|z)=\sum _{i=1}^D{x}_{i}log{y}_i+(1-x_i)log(1-y_i)$$
这里$y=f_{\sigma }(W_{2}tanh(W_{1}z+b_1)+b_2)$，$f_{\sigma }$是逐元素的sigmoid激活函数。原论文做实验用的是最简单的MLPs结构。

拉普拉斯分布

如果假设分布是拉普拉斯分布，那么得到的是L1损失。
$$log{p}_{\varphi }(x|z)=-log(2)-log({\sigma }_{\varphi }(z))-\frac{|x-{\mu }_{\varphi }(z)|}{{\sigma }_{\varphi }(z)}$$

目前基于LDM框架的diffusion模型中用的VAE都是基于此假设分布训练得到。

rec_loss = torch.abs(inputs.contiguous() - reconstructions.contiguous())
if self.perceptual_weight > 0:
	p_loss = self.perceptual_loss(inputs.contiguous(), reconstructions.contiguous())
	rec_loss = rec_loss + self.perceptual_weight * p_loss
	
nll_loss = rec_loss / torch.exp(self.logvar) + self.logvar
weighted_nll_loss = nll_loss
if weights is not None:
	weighted_nll_loss = weights*nll_loss
	weighted_nll_loss = torch.sum(weighted_nll_loss) / weighted_nll_loss.shape[0]

可以再延伸下，如果假设输出服从高斯分布，代码是不是可以改为：

rec_loss = (inputs.contiguous() - reconstructions.contiguous()) / torch.exp(self.logvar)
nll_loss = 0.5 * torch.pow(rec_loss, 2) + self.logvar

隐变量分布约束

对于第二部分的KL约束，可理解为是对隐变量分布的约束or正则化，希望其能学成多元高斯分布，方便采样生成。进一步假设多元变量之间相互独立，每个位置等价为一元高斯分布，每个位置的KL散度，参考数理基础部分有详细推导：
$$KL(q_{\theta }(z|x)||p(z))=-log{\sigma }_{\theta }(x)+\frac{1}{2}({\sigma }_{\theta }^2(x)+{\mu }_{\theta }^2(x))-\frac{1}{2}$$

重参数化

$z\sim p_{\theta }(z|x)$可以想象为将样本$x$，输入到参数化为$\theta$的encoder中，得到隐变量$z$概率分布的参数估计，假设隐变量$z$服从高斯分布$z\sim N(z;{\mu }_{\theta }(x),{\sigma }_{\theta }^2(x))$。需要从该分布中采样出样本$z$，输入到decoder部分。这里的采样需要用到重参数化。如果直接利用torch.normal，输入encoder估计出的均值和方差，生成样本，训练过程中梯度无法反传，因此需要用到重参数化，绕开torch.normal采样，采用$mean+std\ast ϵ$, $ϵ\sim N(0,I)$的方式采样，得到样本$z$，这样训练过程中，梯度就可以正常反传，做参数优化。

import torch
mean = torch.rand(3, requires_grad=True)
std = torch.rand(3, requires_grad=True)
val = torch.normal(mean, std)
loss_1 = sum(val)
loss_1.backward()  # backpropgation计算梯度
print(mean.grad) # tensor([0., 0., 0.])
print(std.grad)  # tensor([0., 0., 0.])

# 重参数化采样
noise = torch.randn(3)
val = mean + std * noise
loss_2 = sum(val)
loss_2.backward()
print(mean.grad) # tensor([1., 1., 1.])
print(std.grad)  # noise value