Basic Definition
Kullback-Leibler divergence
Monte Carlo method
SNR
Modalities
模态(Modality),就是数据的一种形式,例如图像、文本、语音等。 通常意义的多模态学习,就是利用模型去同时处理多个模态数据,例如同时处理图文,图生文本、文本生图等。
在概率论中,"变分"一词主要指的是优化或近似某个函数或过程的方式,尤其是通过变化某些参数来使某个量达到最优。它通常涉及到变分推断、变分原理等概念,尤其在机器学习和统计推断中具有重要应用。
Variational Inference
变分推断是一种用来近似复杂概率分布的方法。由于直接计算后验分布或某些复杂概率模型的推断可能非常困难(通常需要进行高维积分),变分推断通过将复杂的概率分布替换为某个简单的可处理的分布,来简化推断过程。变分推断的目标是最小化两个分布之间的差异,通常使用 KL散度(Kullback-Leibler Divergence) 来量化这种差异。
例如,在贝叶斯推断中,我们往往希望计算后验分布,但直接计算可能非常复杂。变分推断则通过引入一个简单的分布族(通常是参数化的)来近似真实的后验分布,并通过优化这个分布,使其尽可能接近真实后验。
变分推断常见步骤:
- 选择变分分布族:选择一个简单的分布族 q(θ)q(\theta),作为真实后验分布的近似。
- 定义目标函数:通常选择最小化KL散度,或者等价的最大化变分下界(ELBO,证据下界)。
- 优化变分参数:通过优化变分分布的参数,使得变分分布与真实后验分布尽可能接近。
Background: ELBO, VAE, and Hierarchical VAE
对于许多模态,我们可以将观察到的数据视为由相关的看不见的或无法直接获取的 潜在变量/隐变量(latent variables)$z$ 表示或生成。
在生成建模中,我们通常寻求学习低维的潜在表示,而不是高维的潜在表示。这是因为在没有强大先验的情况下,试图学习比观察更高维度的表示是徒劳的。另一方面,学习低维潜在也可能被视为一种压缩形式,并且有可能发现描述观察的语义上有意义的结构。
Evidence Lower Bound
假设潜在变量和观测到的数据是由联合分布 $p(x,z)$ 生成的。基于似然(likelihood-based) 的生成建模通过最大化观测数据 $x$ 的似然$p(x)$。考虑以下两种方式来表示似然$p(x)$:
- 显式地边缘化潜在变量 z:$$p(\boldsymbol{x})=\int p(\boldsymbol{x},\boldsymbol{z})d\boldsymbol{z}$$
- 或者采用条件概率形式:$$p(x)=\frac{p(x,z)}{p(z|x)}$$
直接计算和最大化似然 $p(x)$ 很困难,因为其要么涉及对方程 1 中的所有潜在变量 z 进行积分,这对于复杂模型来说很难处理,要么涉及访问方程 2 中的真实潜在编码器(ground truth latent encoder) $p(z|x)$ 。
但是我们有以下关系式成立:
$$
\log p(\boldsymbol{x})\geq\mathbb{E}_{q_\boldsymbol{\phi}(\boldsymbol{z}|\boldsymbol{x})}\left[\log\frac{p(\boldsymbol{x},\boldsymbol{z})}{q_\boldsymbol{\phi}(\boldsymbol{z}|\boldsymbol{x})}\right]
$$
称不等式右侧$\mathbb{E}_{q_\boldsymbol{\phi}(\boldsymbol{z}|\boldsymbol{x})}\left[\log\frac{p(\boldsymbol{x},\boldsymbol{z})}{q_\boldsymbol{\phi}(\boldsymbol{z}|\boldsymbol{x})}\right]$为证据下界 (ELBO) ,其中 $q_{\phi}(z|x)$ 为寻求优化的参数化模型,用于估计给定观测值$x$时潜变量$z$的真实分布,即我们的目标转化为通过调整参数$\phi$来最大化ELBO。
derive the ELBO
由条件期望定义以及Jesen不等式(log为凹函数)可以得到:
$$
\begin{aligned}\log p(\boldsymbol{x})&=\log\int p(\boldsymbol{x},\boldsymbol{z})d\boldsymbol{z}\\&=\log\int\frac{p(\boldsymbol{x},\boldsymbol{z})q_\boldsymbol{\phi}(\boldsymbol{z}|\boldsymbol{x})}{q_\boldsymbol{\phi}(\boldsymbol{z}|\boldsymbol{x})}d\boldsymbol{z}\\&=\log\mathbb{E}_{q_\boldsymbol{\phi}(\boldsymbol{z}|\boldsymbol{x})}\left[\frac{p(\boldsymbol{x},\boldsymbol{z})}{q_\boldsymbol{\phi}(\boldsymbol{z}|\boldsymbol{x})}\right]\\&\geq\mathbb{E}_{q_\boldsymbol{\phi}(\boldsymbol{z}|\boldsymbol{x})}\left[\log\frac{p(\boldsymbol{x},\boldsymbol{z})}{q_\boldsymbol{\phi}(\boldsymbol{z}|\boldsymbol{x})}\right]\end{aligned}
$$
或者从希望的目标出发,考虑似然$p(x)$的条件概率形式:
$$
\begin{aligned}\log p(\boldsymbol{x})&=\log p(\boldsymbol{x})\int q_\boldsymbol{\phi}(\boldsymbol{z}|\boldsymbol{x})dz\\&=\int q_\boldsymbol{\phi}(\boldsymbol{z}|\boldsymbol{x})(\log p(\boldsymbol{x}))dz\\&=\mathbb{E}_{q_\boldsymbol{\phi}(\boldsymbol{z}|\boldsymbol{x})}\left[\log p(\boldsymbol{x})\right]\\&=\mathbb{E}_{q_\boldsymbol{\phi}(\boldsymbol{z}|\boldsymbol{x})}\left[\log\frac{p(\boldsymbol{x},\boldsymbol{z})}{p(\boldsymbol{z}|\boldsymbol{x})}\right]\end{aligned}
$$
进一步变形并拆分期望引入KL Divergence:
$$
\begin{aligned}
&\log p(\boldsymbol{x})=\mathbb{E}_{q_{\boldsymbol{\phi}}(\boldsymbol{z}|\boldsymbol{x})}\left[\log\frac{p(\boldsymbol{x},\boldsymbol{z})q_{\boldsymbol{\phi}}(\boldsymbol{z}|\boldsymbol{x})}{p(\boldsymbol{z}|\boldsymbol{x})q_{\boldsymbol{\phi}}(\boldsymbol{z}|\boldsymbol{x})}\right] \\
& =\mathbb{E}_{q_{\boldsymbol{\phi}}(\boldsymbol{z}|\boldsymbol{x})}\left[\log\frac{p(\boldsymbol{x},\boldsymbol{z})}{q_{\boldsymbol{\phi}}(\boldsymbol{z}|\boldsymbol{x})}\right]+\mathbb{E}_{q_{\boldsymbol{\phi}}(\boldsymbol{z}|\boldsymbol{x})}\left[\log\frac{q_{\boldsymbol{\phi}}(\boldsymbol{z}|\boldsymbol{x})}{p(\boldsymbol{z}|\boldsymbol{x})}\right] \\
& =\mathbb{E}_{q_{\boldsymbol{\phi}}(\boldsymbol{z}|\boldsymbol{x})}\left[\log\frac{p(\boldsymbol{x},\boldsymbol{z})}{q_{\boldsymbol{\phi}}(\boldsymbol{z}|\boldsymbol{x})}\right]+D_{\mathrm{KL}}(q_{\boldsymbol{\phi}}(\boldsymbol{z}|\boldsymbol{x})\parallel p(\boldsymbol{z}|\boldsymbol{x})) \\
& \geq\mathbb{E}_{q_{\boldsymbol{\phi}}(\boldsymbol{z}|\boldsymbol{x})}\left[\log\frac{p(\boldsymbol{x},\boldsymbol{z})}{q_{\boldsymbol{\phi}}(\boldsymbol{z}|\boldsymbol{x})}\right]
\end{aligned}$$
从这个推导中,可以看到,对数似然 $log p(x)$ 等于 ELBO 加上近似后验 $q_{\phi}(z|x)$ 和真实后验(true posterior)$p(z|x)$ 之间的 KL 散度,而在对数似然固定的前提下(或者不依赖于参数$\phi$,通过调整$\phi$)最大化ELBO相当于最小化这个KL散度 —— 优化变分后验 $q_{\phi}(z|x)$的参数,以精确匹配真实后验$p(z|x)$。
因此,ELBO 可以最大化作为学习如何完美建模真实潜在后验分布的代理。
Variational Autoencoders
变分自动编码器 (VAE) 的默认公式直接最大化 ELBO。这种方法是变分的,因为我们在由 $\phi$ 参数化的潜在后验分布族(potential posterior distributions)中优化了最佳 $q_{\phi}(z|x)$ 。由于其中输入数据在经历中间瓶颈表示步骤后被训练为预测自身,所以它也是一种自动编码器(autoencoder,AE)。
为了深入理解上面这段话,可以进一步剖析ELBO:
$$
\begin{aligned} \mathbb{E}_{q_{\boldsymbol{\phi}}(\boldsymbol{z}|\boldsymbol{x})}\left[\log\frac{p(\boldsymbol{x},\boldsymbol{z})}{q_{\boldsymbol{\phi}}(\boldsymbol{z}|\boldsymbol{x})}\right] &= \mathbb{E}_{q_{\boldsymbol{\phi}}(\boldsymbol{z}|\boldsymbol{x})}\left[\log\frac{p_{\boldsymbol{\theta}}(\boldsymbol{x}|\boldsymbol{z})p(\boldsymbol{z})}{q_{\boldsymbol{\phi}}(\boldsymbol{z}|\boldsymbol{x})}\right] \\ &= \mathbb{E}_{q_{\boldsymbol{\phi}}(\boldsymbol{z}|\boldsymbol{x})}\left[\log p_{\boldsymbol{\theta}}(\boldsymbol{x}|\boldsymbol{z})\right] + \mathbb{E}_{q_{\boldsymbol{\phi}}(\boldsymbol{z}|\boldsymbol{x})}\left[\log\frac{p(\boldsymbol{z})}{q_{\boldsymbol{\phi}}(\boldsymbol{z}|\boldsymbol{x})}\right] \\ &= \underbrace{\mathbb{E}_{q_{\boldsymbol{\phi}}(\boldsymbol{z}|\boldsymbol{x})}\left[\log p_{\boldsymbol{\theta}}(\boldsymbol{x}|\boldsymbol{z})\right]}_{\text{reconstruction term}} - \underbrace{D_{\mathrm{KL}}(q_{\boldsymbol{\phi}}(\boldsymbol{z}|\boldsymbol{x})\parallel p(\boldsymbol{z}))}_{\text{prior matching term}} \end{aligned}
$$
- 学习了一个中间瓶颈分布 $q_{\phi}(z|x)$ ,它可以被视为 编码器(encoder) —— 它将 inputs 转换为可能的 latent 的分布。
- 同时学习了一个确定性函数 $p_{\theta}(x|z)$ ,它可悲视为 解码器(dencoder) —— 将给定的潜在向量 z 转换为观察值 x。
直观上可以这样理解:
- 第一项衡量解码器从我们的变分分布中重建的可能性;这可确保学习的分布正在对可以从中重新生成原始数据的有效潜在值进行建模。
- 第二项衡量学习到的变分分布与对潜在变量的先验的相似程度。最小化此项会鼓励编码器实际学习分布,而不是折叠成 Dirac delta 函数。因此,最大化 ELBO 等效于最大化其第一项和最小化其第二项。
VAE 的一个决定性特征是 ELBO 如何根据参数 φ 和 θ 进行联合优化。<u>通常选择 VAE 的编码器来模拟具有对角协方差的多元高斯,而先验通常被选为标准多元高斯</u>:
$$
q_{\boldsymbol{\phi}}(\boldsymbol{z}|\boldsymbol{x})=\mathcal{N}(\boldsymbol{z};\boldsymbol{\mu}_{\boldsymbol{\phi}}(\boldsymbol{x}),\boldsymbol{\sigma}_{\boldsymbol{\phi}}^2(\boldsymbol{x})\mathbf{I})
$$
$$
p(\boldsymbol{z})=\mathcal{N}(\boldsymbol{z};\boldsymbol{0},\mathbf{I})
$$
然后,可以通过解析方式计算 ELBO 的 KL 散度项,并且可以使用 Monte Carlo 估计 (对潜变量采样)来近似重建项。然后,我们的目标可以改写为:
$$
\arg\max_{\boldsymbol{\phi},\boldsymbol{\theta}}\mathbb{E}_{q_{\boldsymbol{\phi}}(\boldsymbol{z}|\boldsymbol{x})}\left[\log p_{\boldsymbol{\theta}}(\boldsymbol{x}|\boldsymbol{z})\right]-D_{\mathrm{KL}}(q_{\boldsymbol{\phi}}(\boldsymbol{z}|\boldsymbol{x})\parallel p(\boldsymbol{z}))\approx\arg\max_{\boldsymbol{\phi},\boldsymbol{\theta}}\sum_{l=1}^L\log p_{\boldsymbol{\theta}}(\boldsymbol{x}|\boldsymbol{z}^{(l)})-D_{\mathrm{KL}}(q_{\boldsymbol{\phi}}(\boldsymbol{z}|\boldsymbol{x})\parallel p(\boldsymbol{z}))
$$
对于数据集中的每个观测值 $x$,潜在变量 ${z^{(l)}}_{l=1}^L$ 是从 $q_\phi(z|x)$ 中采样的。
然而,这种默认设置存在一个问题:计算损失时,每个 $z^{(l)}$ 都是由随机采样过程生成的,这通常是不可微分的。幸运的是,当 $q_\phi(z|x)$ 被设计为对某些分布进行建模时,这个问题可以通过 重新参数化(reparameterization) 技巧来解决,特别是对于多元高斯分布。重新参数化技巧通过将随机变量重写为噪声变量的确定性函数来解决这个问题;这使得通过梯度下降优化非随机部分成为可能。
例如,来自正态分布 $x \sim \mathcal{N}(x; \mu, \sigma^2)$ 的样本,具有任意均值 $\mu$ 和方差 $\sigma^2$,可以重写为:
$$
x = \mu + \sigma \cdot \epsilon
$$
其中 $\epsilon \sim \mathcal{N}(0, I)$。
换句话说,通过重新参数化技巧,我们可以从任意高斯分布中采样,方法是首先从标准高斯分布中采样,按目标标准差缩放结果,然后按目标均值平移结果。
因此,在变分自编码器(VAE)中,每个 $z$ 都被计算为输入 $x$ 和辅助噪声变量的确定性函数:
$$
z=\mu_{\boldsymbol{\phi}}(x)+\sigma_{\boldsymbol{\phi}}(x)\odot\boldsymbol{\epsilon}\quad\mathrm{with~}\boldsymbol{\epsilon}\sim\mathcal{N}(\boldsymbol{\epsilon};\boldsymbol{0},\mathbf{I})
$$
其中 $\epsilon \sim \mathcal{N}(0, I)$,并且$\odot$表示元素乘积。在这个 $z$ 的重新参数化版本下,我们可以根据需要计算相对于 $\phi$ 的梯度,从而优化 $\mu_\phi$ 和 $\sigma_\phi$。因此,VAE 利用重新参数化技巧和蒙特卡洛估计值,在 $\phi$ 和 $\theta$ 上共同优化 ELBO(变分下界)。
训练完成后,VAE 可以通过直接从潜在空间 $p(z)$ 中采样,然后通过解码器生成新的数据。
Hierarchical Variational Autoencoders
层次变分自编码器(HVAE) 是变分自编码器(VAE)扩展到多个潜变量的层次结构的推广. 在这种形式下,潜变量本身被解释为从其他更高层次、更抽象的潜变量生成的。
- 具有 $T$ 层次潜在值的马尔可夫分层变分自动编码器。生成过程被建模为马尔可夫链,其中每个潜在 $z_t$仅从前一个潜在 $z_{t+1}$ 生成。
在一般的HVAE中,具有T个层次,每个潜变量可以条件化在所有之前的潜变量上。
接下来,我们专注于HVAE的特殊形式—— 马尔科夫变分自编码器(MHVAE)。在MHVAE中,<font color="#ffff00">生成过程是一个马尔科夫链</font>;也就是说,每一步从一个层次到另一个层次的过渡是马尔科夫过程,其中每个潜变量 $z_t$ 只条件化于之前的潜变量 $z_{t+1}$。
- DDPM中的定义
数学上,我们将马尔科夫变分自编码器(MHVAE)的联合分布和后验表示为(<font color="#ff0000">为什么后验也有这种表示,HMVAE是markov时间可逆的吗,不过扩散模型的前向和反向过程均建模为markov chain</font>):
$$
p(\boldsymbol{x},\boldsymbol{z}_{1:T})=p(\boldsymbol{z}_T)p_{\boldsymbol{\theta}}(\boldsymbol{x}|\boldsymbol{z}_1)\prod_{t=2}^Tp_{\boldsymbol{\theta}}(\boldsymbol{z}_{t-1}|\boldsymbol{z}_t)
$$
$$
q_{\boldsymbol{\phi}}(\boldsymbol{z}_{1:T}|\boldsymbol{x})=q_{\boldsymbol{\phi}}(\boldsymbol{z}_1|\boldsymbol{x})\prod_{t=2}^Tq_{\boldsymbol{\phi}}(\boldsymbol{z}_t|\boldsymbol{z}_{t-1})$$ ^bsqjxa
 ^pvf9zu
然后,我们可以轻松地将 ELBO 扩展为:
$$\begin{aligned}
\log p(\boldsymbol{x}) & =\log\int p(\boldsymbol{x},\boldsymbol{z}_{1:T})d\boldsymbol{z}_{1:T} \\
& =\log\int\frac{p(\boldsymbol{x},\boldsymbol{z}_{1:T})q_{\boldsymbol{\phi}}(\boldsymbol{z}_{1:T}|\boldsymbol{x})}{q_{\boldsymbol{\phi}}(\boldsymbol{z}_{1:T}|\boldsymbol{x})}d\boldsymbol{z}_{1:T} \\
& =\log\mathbb{E}_{q_{\boldsymbol{\phi}}(\boldsymbol{z}_{1:T}|\boldsymbol{x})}\left[\frac{p(\boldsymbol{x},\boldsymbol{z}_{1:T})}{q_{\boldsymbol{\phi}}(\boldsymbol{z}_{1:T}|\boldsymbol{x})}\right] \\
& \geq\mathbb{E}_{q_{\boldsymbol{\phi}}(\boldsymbol{z}_{1:T}|\boldsymbol{x})}\left[\log\frac{p(\boldsymbol{x},\boldsymbol{z}_{1:T})}{q_{\boldsymbol{\phi}}(\boldsymbol{z}_{1:T}|\boldsymbol{x})}\right]
\end{aligned}$$
然后,代入MHVAE的联合分布和后验表示,还可得到:
$$
\mathbb{E}_{q_{\boldsymbol{\phi}}(\boldsymbol{z}_{1:T}|\boldsymbol{x})}\left[\log\frac{p(\boldsymbol{x},\boldsymbol{z}_{1:T})}{q_{\boldsymbol{\phi}}(\boldsymbol{z}_{1:T}|\boldsymbol{x})}\right]=\mathbb{E}_{q_{\boldsymbol{\phi}}(\boldsymbol{z}_{1:T}|\boldsymbol{x})}\left[\log\frac{p(\boldsymbol{z}_T)p_{\boldsymbol{\theta}}(\boldsymbol{x}|\boldsymbol{z}_1)\prod_{t=2}^Tp_{\boldsymbol{\theta}}(\boldsymbol{z}_{t-1}|\boldsymbol{z}_t)}{q_{\boldsymbol{\phi}}(\boldsymbol{z}_1|\boldsymbol{x})\prod_{t=2}^Tq_{\boldsymbol{\phi}}(\boldsymbol{z}_t|\boldsymbol{z}_{t-1})}\right]
$$
正如我们将在下面展示的,当我们研究变分扩散模型时,这个目标可以进一步分解为可解释的组成部分。
Variational Diffusion Models
VDM相关假设
最简单的理解 变分扩散模型(VDM) 的方式是将其视为一个MHVAE(层次过渡是马尔科夫过程),并且有三个关键的限制条件:
- <u>潜变量维度与真实数据维度相等</u>(e.g. 真实数据为一张图像,隐变量为对应尺寸图像)
- <u>每个时间步的潜编码器结构$q(\boldsymbol{x}_t|\boldsymbol{x}_{t-1})$不是通过学习得到的, 而是预定义为线性高斯模型</u>。
- <u>潜编码器的高斯参数(与$\alpha_t$有关)随时间变化,使得在最终时间步 T 时,潜变量的分布为标准高斯分布</u>
根据第一个限制,略作符号上的调整,将真实数据样本和潜在变量都表示为 $x_t$,其中 $t = 0$ 表示真实数据样本,而 $t \in [1, T]$ 则表示相应的潜在变量,且该潜在变量的层次由 $t$ 索引。VDM 的后验与 MHVAE 的后验[[Understanding Diffusion Models#^bsqjxa|MHVAE后验]]相同,可重写为(在DDPM原文中扩散过程悲固定在Markov chain上):
$$
q(x_{1:T}|x_0)=\prod_{t=1}^Tq(x_t|x_{t-1})
$$
根据第二个假设,我们知道编码器中每个潜在变量的分布是一个高斯分布,中心位于其之前的层次潜在变量周围。与马尔科夫层次变分自编码器(MHVAE)不同,编码器在每个时间步 $t$ 的结构不是通过学习得到的;它被固定为一个线性高斯模型,其中均值和标准差可以作为<u>超参数预先设置(DDPM)</u>,或者<u>作为参数进行学习(Advances in neural information processing systems)</u>。
我们将高斯编码器的均值参数化为 $\mu_t(x_t) = \sqrt{\alpha_t} x_{t-1}$,方差参数化为 $\Sigma_t(x_t) = (1 - \alpha_t)I$,其中系数的形式被选择使得潜变量的方差保持在相似的尺度;换句话说,编码过程是保持方差不变的 (variance-preserving) (<font color="#ff0000">这里是如何保持方差不变的,可以见后续关于SDE的解释</font>)。请注意,也允许使用其他的高斯参数化形式,并且它们会导致类似的推导。主要的要点是,$\alpha_t$ 是一个(可能是可学习的)系数,它可以随层次深度 $t$ 变化,从而提供灵活性。编码器的 transitions 可以表示为:
$$
q(\boldsymbol{x}_t|\boldsymbol{x}_{t-1})=\mathcal{N}(\boldsymbol{x}_t;\sqrt{\alpha_t}\boldsymbol{x}_{t-1},(1-\alpha_t)\mathbf{I})$$ ^7j02d0
> PS:DDPM中表示方式(一样的):
> 扩散过程被固定为马尔可夫链,同时根据线性设置的方差表(variance schedule)$\beta_1 、... 、 \beta_T$ 往数据中加噪声:
> $$q(\mathbf{x}_t|\mathbf{x}_{t-1}):=\mathcal{N}(\mathbf{x}_t;\sqrt{1-\beta_t}\mathbf{x}_{t-1},\beta_t\mathbf{I})$$
$$
\alpha_{t}:=1-\beta_{t} \quad and \quad \bar{\alpha}_{t}:=\prod_{s=1}^{t}\alpha_{s}
$$
从第三个假设中,我们知道 ${\alpha}_t$ 根据固定或可学习的时间表结构随时间演变,使得<u>最终潜变量分布 $p(x_T)$ 的分布是标准的高斯分布</u>。然后,我们可以更新MHVAE 的联合分布,将 [[Understanding Diffusion Models#^bsqjxa|VDM 的联合分布]]写为:
$$
p(\boldsymbol{x}_{0:T})=p(\boldsymbol{x}_T)\prod_{t=1}^Tp_\boldsymbol{\theta}(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t)
$$
where $p(\boldsymbol{x}_T)=\mathcal{N}(\boldsymbol{x}_T;\mathbf{0},\mathbf{I})$
总的来说,<u>这组假设所描述的是图像输入随时间的稳定噪声化</u>;我们通过添加高斯噪声来逐渐破坏图像,直到最终它变得与纯高斯噪声完全相同。从视觉上看,这个过程如图所示:
PS:
- VDM中编码器分布 $q(x_t | x_{t-1})$ 不再由 $\phi$ 参数化,因为它们完全被建模为在每个时间步都有定义的均值和方差参数的高斯分布。因此,<u>在变分扩散模型(VDM)中,我们只关心学习条件概率 </u>$p_{\theta}(x_{t-1} | x_t)$,以便我们可以模拟新的数据。(设定扩散步,学习生成步)
- 优化 VDM 后,采样过程非常简单:从 $p(x_T)$ 中采样高斯噪声,然后迭代运行去噪过渡 $p_{\theta}(x_{t-1} | x_t)$ 进行 $T$ 步,生成一个新的 $x_0$。
ELBO推导
推导(一)
像任何层次变分自编码器(HVAE)一样,VDM 可以通过最大化 ELBO 来进行优化,ELBO 可以推导为:
ELBO 的推导形式可以根据其各个组成部分进行解释:
- $\mathbb{E}_q(x_1 | x_0) [\log p_{\theta}(x_0 | x_1)]$ 可以解释为 重建项 ,预测给定第一步潜在值的原始数据样本的对数概率。
- $\mathbb{E}_q(x_{T-1} | x_0) [D_{\mathrm{KL}}(q(x_T | x_{T-1}) \parallel p(x_T))]$ 是 先验匹配项 ;当最终潜在分布与高斯先验匹配时,它最小化。这一项不需要优化,因为它没有可训练的参数;此外,由于我们假设 $T$ 足够大,最终分布是高斯分布,因此这项实际上最终变为零。
- $\mathbb{E}_q(x_{t-1}, x_{t+1} | x_0) [D_{\mathrm{KL}}(q(x_t | x_{t-1}) \parallel p_{\theta}(x_t | x_{t+1}))]$ 是 一致性项 ;它的目标是使得在每个中间时间步长的分布一致,从前向过程和反向过程的角度来看。也就是说,对于每个中间时间步,来自噪声较多图像的降噪步骤应该与来自较清晰图像的相应噪声步骤匹配;这一点在数学上通过 KL 散度体现。
优化 VDM 的成本主要由第三项主导,因为我们必须在所有时间步长 t 上进行优化。从视觉上看,ELBO 的这种解释如图所示:
- VDM 可以通过确保对于每个中间$x_t$,其上方潜在变量的后验$p_{\theta}(x_t | x_{t+1})$与它之前的隐变量 $q(x_t|x_{t+1})$ 的高斯加噪相匹配来优化 VDM ——将粉红色和绿色箭头表示的分布之间的差异降至最低。
在此推导下,ELBO 的所有项均作为期望计算,因此可以使用蒙特卡洛估计进行近似计算。然而,实际上,使用上述推导出的ELBO格式可能不是最佳的;由于一致性项是作为每个时间步长对两个随机变量 $\{x_{t-1},x_{t+1}\}$的期望值计算的,因此其蒙特卡洛估计值的方差可能高于每个时间步长仅使用一个随机变量估计的项。由于 ELBO 是通过对 $T − 1$ 个一致性项求和来计算的,因此 ELBO 的最终估计值对于较大的 $T$ 值可能具有很大的方差。
推导(二)
为此推导ELBO便于操作的形式, 注意到在VDM中前向过程也设置为markov chain,由markov特性编码器 transitions 可改写为$q(x_t|x_{t−1}) = q(x_t|x_{t−1},x_0)$。然后,根据贝叶斯规则,我们可以将每个 transition 重写为:
$$
q(\boldsymbol{x}_t|\boldsymbol{x}_{t-1},\boldsymbol{x}_0)=\frac{q(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t,\boldsymbol{x}_0)q(\boldsymbol{x}_t|\boldsymbol{x}_0)}{q(\boldsymbol{x}_{t-1}|\boldsymbol{x}_0)}
$$
重新推导ELBO为:
因此,我们成功地推导出了 ELBO 的另一种形式,可以通过估算每个项的期望来降低方差,因为每个项最多只计算一个随机变量的期望。
这个形式也有一个优雅的解释,当我们检查每个单独的项时可以揭示出来:
- $\mathbb{E}_q(x_1 | x_0) [\log p_{\theta}(x_0 | x_1)]$ 可以解释为重建项;与标准 VAE 中 ELBO 的类似项一样,这个项可以通过蒙特卡洛估计进行近似和优化。
- $D_{\mathrm{KL}}(q(x_T | x_0) \parallel p(x_T))$ 表示最终噪声化输入的分布与标准高斯先验的接近程度。它没有可训练的参数,并且在我们的假设下也等于零。
- $\mathbb{E}_{q(x_t | x_0)} [D_{\mathrm{KL}}(q(x_{t-1} | x_t, x_0) \parallel p_{\theta}(x_{t-1} | x_t))]$ 是一个降噪匹配项。我们学习所需的降噪过渡步骤 $p_{\theta}(x_{t-1} | x_t)$,作为可处理的、真实的降噪过渡步骤 $q(x_{t-1} | x_t, x_0)$ 的近似。$q(x_{t-1} | x_t, x_0)$ 过渡步骤可以作为真实信号,因为它定义了如何使用最终完全降噪图像 $x_0$ 来对噪声图像 $x_t$ 进行降噪。因此,当这两个降噪步骤尽可能匹配时,该项会最小化,通过它们的 KL 散度来衡量匹配程度。
重参数化$q(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t,\boldsymbol{x}_0)$
在 ELBO 的这种推导中,大部分优化成本再次在于求和项,它主导了重建项。虽然每个 KL 散度项 $D_{\mathrm{KL}}(q(x_{t-1}|x_t,x_0)\parallel p_{\boldsymbol{\theta}}(x_{t-1}|x_t))$ 对于任意复杂的马尔可夫 HVAE 中的任意后验来说很难最小化,因为同时学习编码器的复杂性增加,而在 VDM 中,我们可以利用高斯过渡假设来使优化变得易于处理。
根据贝叶斯规则,我们有:
$$
q(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t,\boldsymbol{x}_0)=\frac{q(\boldsymbol{x}_t|\boldsymbol{x}_{t-1},\boldsymbol{x}_0)q(\boldsymbol{x}_{t-1}|\boldsymbol{x}_0)}{q(\boldsymbol{x}_t|\boldsymbol{x}_0)}
$$
正如我们已经知道的,从我们关于[[Understanding Diffusion Models#^7j02d0|编码器过渡的假设]](加噪过程为高斯分布)中,我们有:
$$
q(x_t | x_{t-1}, x_0) = q(x_t | x_{t-1}) = \mathcal{N}(x_t; \sqrt{\alpha_t} \boldsymbol{x}_{t-1}, (1-\alpha_t) \mathbf{I})
$$
接下来,我们需要推导 $q(\boldsymbol{x}_t | \boldsymbol{x}_0)$ 和 $q(\boldsymbol{x}_{t-1} | \boldsymbol{x}_0)$ 的形式。
在重新参数化技巧下,样本 $x_t \sim q(x_t | x_{t-1})$ 可以重写为:
$$
\boldsymbol{x}_t = \sqrt{\alpha_t} \boldsymbol{x}_{t-1} + \sqrt{1-\alpha_t} \boldsymbol{\epsilon} \quad \text{with}~ \boldsymbol{\epsilon} \sim \mathcal{N}(\boldsymbol{\epsilon}; \mathbf{0}, \mathbf{I})
$$
类似地,样本 $x_{t-1} \sim q(x_{t-1} | x_{t-2})$ 可以重写为:
$$
x_{t-1} = \sqrt{\alpha_{t-1}} x_{t-2} + \sqrt{1-\alpha_{t-1}} \boldsymbol{\epsilon} \quad \text{with}~ \boldsymbol{\epsilon} \sim \mathcal{N}(\boldsymbol{\epsilon}; \mathbf{0}, \mathbf{I})
$$
利用 VDM 中编码器过渡是线性的性质,假设我们有 $2T$ 个随机噪声变量$\{\epsilon_t^*,\epsilon_t\}_{t=0}^T\stackrel{\mathrm{iid}}{\sim}\mathcal{N}(\epsilon;\mathbf{0},\mathbf{I})$。那么,对于任意样本 $x_t \sim q(x_t | x_0)$,我们可以将其重写为:
$$
\begin{aligned}\boldsymbol{x}_{t}&=\sqrt{\alpha_t}\boldsymbol{x}_{t-1}+\sqrt{1-\alpha_t}\boldsymbol{\epsilon}_{t-2}^*\\&=\sqrt{\alpha_t}\left(\sqrt{\alpha_{t-1}}\boldsymbol{x}_{t-2}+\sqrt{1-\alpha_{t-1}}\boldsymbol{\epsilon}_{t-2}^*\right)+\sqrt{1-\alpha_t}\boldsymbol{\epsilon}_{t-1}^*\\&=\sqrt{\alpha_t\alpha_{t-1}}\boldsymbol{x}_{t-2}+\sqrt{\alpha_t-\alpha_t\alpha_{t-1}}\boldsymbol{\epsilon}_{t-2}^*+\sqrt{1-\alpha_t}\boldsymbol{\epsilon}_{t-1}^*\\&=\sqrt{\alpha_t\alpha_{t-1}}\boldsymbol{x}_{t-2}+\sqrt{\sqrt{\alpha_t-\alpha_t\alpha_{t-1}}^2+\sqrt{1-\alpha_t}^2}\boldsymbol{\epsilon}_{t-2}\\&=\sqrt{\alpha_t\alpha_{t-1}}\boldsymbol{x}_{t-2}+\sqrt{\alpha_t-\alpha_t\alpha_{t-1}+1-\alpha_t}\boldsymbol{\epsilon}_{t-2}\\&=\sqrt{\alpha_t\alpha_{t-1}}\boldsymbol{x}_{t-2}+\sqrt{1-\alpha_t\alpha_{t-1}}\boldsymbol{\epsilon}_{t-2}\\&=\ldots\\&=\sqrt{\prod_{i=1}^t\alpha_i}\boldsymbol{x}_0+\sqrt{1-\prod_{i=1}^t\alpha_i}\boldsymbol{\epsilon}_0\\&=\sqrt{\alpha_t}\boldsymbol{x}_0+\sqrt{1-\overline{\alpha}_t}\boldsymbol{\epsilon}_0\\&\sim\mathcal{N}(\boldsymbol{x}_t;\sqrt{\alpha_t}\boldsymbol{x}_0,(1-\bar{\alpha}_t)\mathbf{I})\end{aligned}$$ ^zst3xt
> PS:利用了两个独立高斯随机变量之和仍然是高斯分布的事实:其均值是两个均值之和,方差是两个方差之和。将 $\sqrt{1 - \alpha_t} \boldsymbol{\epsilon}_{t-1}^*$ 解释为来自高斯分布 $\mathcal{N}(\mathbf{0}, (1-\alpha_t) \mathbf{I})$ 的样本,将 $\sqrt{\alpha_t - \alpha_t \alpha_{t-1}} \boldsymbol{\epsilon}_{t-2}^*$ 解释为来自高斯分布 $\mathcal{N}(\mathbf{0}, (\alpha_t - \alpha_t \alpha_{t-1}) \mathbf{I})$ 的样本,然后可以将它们的和视为从高斯分布 $\mathcal{N}(\mathbf{0}, (1 - \alpha_t + \alpha_t - \alpha_t \alpha_{t-1}) \mathbf{I}) = \mathcal{N}(\mathbf{0}, (1 - \alpha_t \alpha_{t-1}) \mathbf{I})$ 中采样的随机变量。然后,可以使用重新参数化技巧将这个分布的样本表示为 $\sqrt{1 - \alpha_t \alpha_{t-1}} \boldsymbol{\epsilon}_{t-2}$。
因此,我们推导出了 $q(x_t|x_0)$ 的高斯形式。可以修改此推导,以产生描述 $q(x_{t-1}|x_0)$ 的高斯参数化。现在,知道了 $q(x_t|x_0)$ 和 $q(x_{t-1}|x_0)$ 的形式,我们可以通过代入贝叶斯规则展开来继续计算 $q(x_{t-1}|x_t,x_0)$ 的形式:
 ^d9tj94
- 在推导中 $C(x_t, x_0)$ 是一个关于 $x_{t-1}$ 的常数项,它是仅由 $x_t$, $x_0$ 和 $\alpha$ 值组合计算得到的;该项在方程 最后一行中隐式返回,用于完成平方。
因此,我们已经证明了,在每个步骤中,$x_{t-1} \sim q(x_{t-1}|x_t,x_0)$ 是正态分布的,均值为 $\boldsymbol{\mu}_q(x_t, \boldsymbol{x}_0)$,它是 $x_t$ 和 $x_0$ 的函数,方差为 $\boldsymbol{\Sigma}_q(t)$,可以将其重写为 $\Sigma_q(t) = \sigma_q^2(t)\mathbf{I}$,其中:
$$\sigma_q^2(t) = \frac{(1 - \alpha_t)(1 - \bar{\alpha}_{t-1})}{1 - \bar{\alpha}_t}$$
它是 $\alpha$ 系数的函数。这些 $\alpha$ 系数在每个时间步长上都是已知的且固定的;它们可以在建模时作为超参数永久设置,或者作为网络当前推断的输出进行处理。 ^znpntd
建模优化过程
为了使近似的去噪过渡步骤 $p_{\boldsymbol{\theta}}(x_{t-1}|x_t)$ 尽可能接近真实的去噪过渡步骤 $q(x_{t-1}|x_t, \boldsymbol{x}_0)$,我们也可以将其建模为高斯分布。此外,由于所有的 $\alpha$ 项在每个时间步长上都已知且被冻结,我们可以立即构造近似去噪过渡步骤的方差,使其也是 $\boldsymbol{\Sigma}_q(t) = \sigma_q^2(t)\mathbf{I}$。然而,我们必须将其均值 $\boldsymbol{\mu}_{\boldsymbol{\theta}}(\boldsymbol{x}_t,t)$ 参数化为 $x_t$ 的函数,因为 $p_{\boldsymbol{\theta}}(x_{t-1}|x_t)$ 并不以 $x_0$ 为条件。
回顾一下,两个高斯分布之间的KL散度为:
$$
D_{\mathrm{KL}}(\mathcal{N}(\boldsymbol{x};\boldsymbol{\mu}_x,\boldsymbol{\Sigma}_x) \parallel \mathcal{N}(\boldsymbol{y};\boldsymbol{\mu}_y,\boldsymbol{\Sigma}_y)) = \frac{1}{2}\left[\log \frac{|\boldsymbol{\Sigma}_y|}{|\boldsymbol{\Sigma}_x|} - d + \operatorname{tr}(\boldsymbol{\Sigma}_y^{-1}\boldsymbol{\Sigma}_x) + (\boldsymbol{\mu}_y - \boldsymbol{\mu}_x)^T \boldsymbol{\Sigma}_y^{-1} (\boldsymbol{\mu}_y - \boldsymbol{\mu}_x)\right]
$$
在我们的例子中,我们可以将两个高斯的方差设置为完全匹配,优化 KL 散度项可以减少到最小化两个分布的均值之间的差异:
$$
\begin{aligned}
& \arg\min_{\boldsymbol{\theta}}D_{\mathrm{KL}}(q(x_{t-1}|x_{t},x_{0})\parallel p_{\boldsymbol{\theta}}(x_{t-1}|x_{t})) \\
& =\arg\min_{\boldsymbol{\theta}}D_{\mathrm{KL}}(\mathcal{N}(x_{t-1};\boldsymbol{\mu}_{q},\Sigma_{q}(t))\parallel\mathcal{N}(x_{t-1};\boldsymbol{\mu}_{\boldsymbol{\theta}},\Sigma_{q}(t))) \\
& =\arg\min_{\boldsymbol{\theta}}\frac{1}{2}\left[\log\frac{|\Sigma_q(t)|}{|\Sigma_q(t)|}-d+\operatorname{tr}(\Sigma_q(t))+(\mu_{\boldsymbol{\theta}}-\mu_q)^T\Sigma_q(t)^{-1}(\mu_{\boldsymbol{\theta}}-\mu_q)\right] \\
& =\arg\min_{\boldsymbol{\theta}}\frac{1}{2}\left[\log1-d+d+(\mu_{\boldsymbol{\theta}}-\mu_q)^T\Sigma_q(t)^{-1}(\mu_{\boldsymbol{\theta}}-\mu_q)\right] \\
& =\arg\min_{\boldsymbol{\theta}}\frac{1}{2}\left[(\mu_{\boldsymbol{\theta}}-\mu_q)^T\Sigma_q(t)^{-1}(\mu_{\boldsymbol{\theta}}-\mu_q)\right] \\
& =\arg\min_{\boldsymbol{\theta}}\frac{1}{2}\left[(\mu_{\boldsymbol{\theta}}-\mu_q)^T\left(\sigma_q^2(t)\mathbf{I}\right)^{-1}\left(\mu_{\boldsymbol{\theta}}-\mu_q\right)\right] \\
& =\arg\min_{\boldsymbol{\theta}}\frac{1}{2\sigma_q^2(t)}\left[\|\mu_{\boldsymbol{\theta}}-\mu_q\|_2^2\right]
\end{aligned}$$
- 简写$\mu_q(x_t,x_0)$为$\mu_q$
也就是说,我们希望优化 $\mu_{\boldsymbol{\theta}}(x_t,t)$,使其与 $\mu_q(x_t,x_0)$ 匹配,根据[[Understanding Diffusion Models#^d9tj94|该分布形式]], $\mu_q(x_t,x_0)$ 具有以下形式:
$$
\boldsymbol{\mu}_q(\boldsymbol{x}_t, \boldsymbol{x}_0) = \frac{\sqrt{\alpha_t}(1 - \bar{\alpha}_{t-1}) \boldsymbol{x}_t + \sqrt{\bar{\alpha}_{t-1}}(1 - \alpha_t) \boldsymbol{x}_0}{1 - \bar{\alpha}_t}
$$
这表明我们希望通过优化模型的参数 $\boldsymbol{\theta}$,使得模型输出的均值尽可能接近真实的去噪步骤的均值 $\boldsymbol{\mu}_q$。
由于$\boldsymbol{\mu_{\boldsymbol{\theta}}}(\boldsymbol{x}_t,t)$也依赖于$x_t$,我们可以通过将其设置为以下形式来使$\boldsymbol{\mu}_q(\boldsymbol{x}_t,\boldsymbol{x}_0)$尽可能匹配:
$$
\mu_{\boldsymbol{\theta}}(x_t,t) = \frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1})\boldsymbol{x}_t + \sqrt{\bar{\alpha}_{t-1}}(1-\alpha_t)\hat{\boldsymbol{x}}_{\boldsymbol{\theta}}(\boldsymbol{x}_t,t)}{1-\bar{\alpha}_t}
$$
其中$\hat{\boldsymbol{x}}_{\boldsymbol{\theta}}(\boldsymbol{x}_t,t)$是由神经网络参数化的,旨在从噪声图像$\boldsymbol{x}_t$和时间索引$t$预测$\boldsymbol{x}_0$。接下来,优化问题简化为:
因此,优化一个VDM可以归结为学习一个神经网络,从任意噪声版本的原始图像$\boldsymbol{x}_0$中预测其原始图像$\boldsymbol{x}_0$。此外,通过最小化我们推导的[[Understanding Diffusion Models#^z1b6ji|ELBO]]中所有噪声级别的总和,可以通过最小化所有时间步的期望来近似:
$$
\arg\min_{\boldsymbol{\theta}} \mathbb{E}_{t \sim U\{2,T\}} \left[ \mathbb{E}_{q(\boldsymbol{x}_t|\boldsymbol{x}_0)} \left[ D_{\mathrm{KL}}(q(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t,\boldsymbol{x}_0) \parallel p_{\boldsymbol{\theta}}(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t)) \right] \right]
$$
然后可以使用对时间步的随机样本来优化此目标。
Learning Diffusion Noise Parameters
上文推导出了,扩散过程可以利用神经网络学习参数化的 $\hat{\boldsymbol{x}}_{\boldsymbol{\theta}}(\boldsymbol{x}_t,t)$ 进行优化,在DDPM原论文中将 时间步长的参数因子 $\beta_t$ 线性设置。
但我们也可以采用学习VDM 的噪声参数的方式。一种可能的方法是使用参数为 $η$ 的神经网络 $\hat{\alpha}_{\boldsymbol{\eta}}(t)$对 $\alpha_t$ 进行建模。但是,这是低效的,因为必须在每个时间步 t 执行多次推理才能计算 $\bar{\alpha_t}$。虽然缓存可以减轻这种计算成本,但我们也可以推导出另一种方法来学习扩散噪声参数。
通过将[[Understanding Diffusion Models#^znpntd|方差]]代入上述优化问题,我们得出的每时间步目标:
我们知道 $q(x_t|x_0)$ 是一个高斯分布,形式为 $\mathcal{N}(x_t;\sqrt{\bar{\alpha}_t}x_0,(1-\bar{\alpha}_t)\mathbf{I})$。然后,依据 信噪比(SNR)的定义 为 $\mathrm{SNR} = \frac{\mu^2}{\sigma^2}$,我们可以将每个时间步 $t$ 的 SNR 写成:
$$
\mathrm{SNR}(t) = \frac{\bar{\alpha}_t}{1-\bar{\alpha}_t}
$$
则优化目标可以简化为:
$$
\frac{1}{2\sigma_q^2(t)} \frac{\bar{\alpha}_{t-1}(1-\alpha_t)^2}{(1-\bar{\alpha}_t)^2} \left[ \left\| \hat{x}_{\boldsymbol{\theta}}(x_t,t) - x_0 \right\|_2^2 \right] = \frac{1}{2} \left( \mathrm{SNR}(t-1) - \mathrm{SNR}(t) \right) \left[ \left\| \hat{x}_{\boldsymbol{\theta}}(x_t,t) - x_0 \right\|_2^2 \right]
$$
这个公式表明,信噪比的变化可以与误差项的变化进行关联,从而在优化过程中帮助我们调整网络的预测。
顾名思义,<u>SNR 表示原始信号与存在的噪声量之间的比率;较高的 SNR 表示更多的信号,较低的 SNR 表示更多的噪声</u>。在扩散模型中,我们要求 SNR 随着时间步长 t 的增加而单调减小;这正式化了扰动输入 XT 随着时间的推移变得越来越嘈杂的概念,直到它在 $t = T$ 时变得与标准高斯相同。
进一步地,为了优化目标,我们可以使用神经网络直接参数化每个时间步的 SNR,并与扩散模型一起学习它。由于 SNR 必须随时间单调减小,因此我们可以将其表示为:
$$
\mathrm{SNR}(t)=\exp(-\omega_{\boldsymbol{\eta}}(t))
$$
其中,$\omega_\boldsymbol{\eta}(t)$ 被建模为一个单调递增的神经网络,参数为 $\boldsymbol{\eta}$。取负 $\omega_\boldsymbol{\eta}(t)$ 会得到一个单调递减的函数,而指数函数则强制结果项为正数。我们还可以显式地推导出 $\bar{\alpha}_t$ 以及 1-$\bar{\alpha}_t$ 的优雅形式:
$$
\begin{aligned} &\frac{\bar{\alpha}_t}{1-\bar{\alpha}_t}=\exp(-\omega_{\boldsymbol{\eta}}(t)) \\ &\therefore \bar{\alpha}_t = \mathrm{sigmoid}(-\omega_{\boldsymbol{\eta}}(t)) \\ &\therefore 1-\bar{\alpha}_t = \mathrm{sigmoid}(\omega_{\boldsymbol{\eta}}(t)) \end{aligned}
$$
Three Equivalent Interpretations
正如我们之前证明的那样,可以通过简单地学习神经网络来训练变分扩散模型,以从任意噪声版本 $x_t$ 及其时间索引 $t$ 中预测原始自然图像 $x_0$。但是 $x_0$ 还有另外两个等效的参数化,这导致了对 VDM 两种进一步解释。
Second
首先,我们可以利用 reparameterization 技巧。在我们推导 $q(x_t|x_0)$ 的形式时,我们可以重新排列[[Understanding Diffusion Models#^zst3xt|方程]]得到:
$$
x_0=\frac{x_t-\sqrt{1-\bar{\alpha}_t}\epsilon_0}{\sqrt{\bar{\alpha}_t}}
$$
将其代入我们之前推导的真降噪过渡平均值 $μ_q(x_t,x_0)$ 中,我们可以重新推导为:
因此,我们可以将近似降噪过渡平均值 $μ_{\theta}(x_t,t)$ 设置为:
$$
\mu_{\boldsymbol{\theta}}(x_t,t)=\frac1{\sqrt{\alpha_t}}x_t-\frac{1-\alpha_t}{\sqrt{1-\bar{\alpha}_t}\sqrt{\alpha_t}}\boldsymbol{\hat{\epsilon}}_{\boldsymbol{\theta}}(\boldsymbol{x}_t,t)
$$
相应的优化问题变为:
其中,$\boldsymbol{\hat{\epsilon}}_{\boldsymbol{\theta}}(\boldsymbol{x}_{t},t)$是一个神经网络,它学习预测从 $x_0$ 确定 $x_t$的源噪声 $\epsilon_{0}\sim\mathcal{N}(\boldsymbol{\epsilon};\boldsymbol{0},\mathbf{I})$。<u>因此,我们已经证明,通过预测原始图像 $x_0$ 来学习 VDM 等同于学习预测噪声</u>。
然而,从实证上讲,一些工作发现预测噪声会带来更好的性能。
Third
为了得出变分扩散模型的第三种常见解释,我们求助于 Tweedie 公式。在英语中,Tweedie 公式指出,给定从指数家庭分布中提取的样本,指数家庭分布的真实平均值可以通过样本的最大似然估计(又名经验平均值)加上一些涉及估计分数的修正项来估计。在只有一个观测样本的情况下,经验均值只是样本本身。它通常用于减轻样本偏倚;如果观察到的样本都位于基础分布的一端,则负分值会变大,并将样本的朴素最大似然估计值校正为真实均值。
从数学上讲,对于一个高斯变量 $z \sim \mathcal{N}(z; \boldsymbol{\mu}_z, \boldsymbol{\Sigma}_z)$,Tweedie公式表示为(这里别看成求和号):
$$
\mathbb{E}\left[\mu_z|z\right] = z + \Sigma_z \nabla_{\boldsymbol{z}} \log p(z)
$$
在这种情况下,我们将其应用于根据样本预测 $x_t$ 的真实后验均值:
$$
q(\boldsymbol{x}_t|\boldsymbol{x}_0) = \mathcal{N}(\boldsymbol{x}_t; \sqrt{\bar{\alpha}_t}\boldsymbol{x}_0, (1-\bar{\alpha}_t)\mathbf{I})
$$
通过Tweedie公式,可以得到:
$$
\mathbb{E}\left[\boldsymbol{\mu}_{x_t}|\boldsymbol{x}_t\right] = \boldsymbol{x}_t + (1-\bar{\alpha}_t) \nabla_{\boldsymbol{x}_t} \log p(\boldsymbol{x}_t)
$$
为了简化符号,我们<u>将 $\nabla_{\boldsymbol{x}_t} \log p(\boldsymbol{x}_t)$ 写作 $\nabla \log p(\boldsymbol{x}_t)$</u>。根据Tweedie公式,生成 $x_t$ 的真实均值的最佳估计,$\mu_{x_t} = \sqrt{\bar{\alpha}_t} x_0$,定义为:
$$
\sqrt{\bar{\alpha}_t} x_0 = \boldsymbol{x}_t + (1-\bar{\alpha}_t) \nabla \log p(\boldsymbol{x}_t)
$$
因此,
$$
x_0 = \frac{\boldsymbol{x}_t + (1-\bar{\alpha}_t) \nabla \log p(\boldsymbol{x}_t)}{\sqrt{\bar{\alpha}_t}}
$$
然后,我们可以将方程 其代入我们的真实降噪过渡平均值 $μ_q(x_t,x_0)$ 并推导出一种新形式:
$$
\begin{aligned}\mu_q(x_t,x_0)&=\frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1})x_t+\sqrt{\alpha_t-1}(1-\alpha_t)x_0}{1-\bar{\alpha}_t}\\&=\frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1})x_t+\sqrt{\bar{\alpha}_t-1}(1-\alpha_t)\frac{\boldsymbol{x}_t+(1-\bar{\alpha}_t)\nabla\log p(\boldsymbol{x}_t)}{\sqrt{\bar{\alpha}_t}}}{1-\bar{\alpha}_t}\\&=\frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1})x_t+(1-\alpha_t)\frac{\boldsymbol{x}_t+(1-\bar{\alpha}_t)\nabla\log p(\boldsymbol{x}_t)}{1-\bar{\alpha}_t}}{1-\bar{\alpha}_t}\\&=\frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1})x_t+(1-\alpha_t)\frac{\boldsymbol{x}_t+(1-\bar{\alpha}_t)\nabla\log p(\boldsymbol{x}_t)}{1-\bar{\alpha}_t}}{1-\bar{\alpha}_t}\\&=\left(\frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_t}+\frac{1-\alpha_t}{(1-\bar{\alpha}_t)\sqrt{\alpha_t}}\right)x_t+\frac{1-\alpha_t}{\sqrt{\alpha_t}}\nabla\log p(\boldsymbol{x}_t)\\&=\frac{\alpha_t-\bar{\alpha}_t+1-\alpha_t}{(1-\bar{\alpha}_t)\sqrt{\alpha_t}}x_t+\frac{1-\alpha_t}{(1-\bar{\alpha}_t)\sqrt{\alpha_t}}\\&=\frac{\alpha_t-\bar{\alpha}_t+1-\alpha_t}{(1-\bar{\alpha}_t)\sqrt{\alpha_t}}x_t+\frac{1-\alpha_t}{\sqrt{\alpha_t}}\nabla\log p(\boldsymbol{x}_t)\\&=\frac{1-\bar{\alpha}_t}{(1-\bar{\alpha}_t)\sqrt{\alpha_t}}x_t+\frac{1-\alpha_t}{\sqrt{\alpha_t}}\nabla\log p(\boldsymbol{x}_t)\\&=\frac1{\sqrt{\alpha_t}}x_t+\frac{1-\alpha_t}{\sqrt{\alpha_t}}\nabla\log p(\boldsymbol{x}_t)\\&=\frac1{\sqrt{\alpha_t}}x_t+\frac{1-\alpha_t}{\sqrt{\alpha_t}}\nabla\log p(\boldsymbol{x}_t)\end{aligned}
$$
因此,我们还可以将近似降噪过渡平均值 $μ_{\theta}(x_t,t)$ 设置为:
$$
\mu_{\boldsymbol{\theta}}(x_t,t)=\frac{1}{\sqrt{\alpha_t}}x_t+\frac{1-\alpha_t}{\sqrt{\alpha_t}}s_{\boldsymbol{\theta}}(x_t,t)
$$
相应的优化问题变为:
在这里,$s_{\boldsymbol{\theta}}(\boldsymbol{x}_t,t)$ 是一个神经网络,学习预测得分函数 $\nabla_{\boldsymbol{x}_t} \log p(\boldsymbol{x}_t)$,即数据空间中 $\boldsymbol{x}_t$ 的梯度,对于任意的噪声水平 $t$。
- PS:<u>得分函数 $\nabla \log p(x_t)$ 的形式与源噪声 $\epsilon_0$ 仅相差一个与时间成比例的常数因子</u>:$$x_0 = \frac{x_t + (1 - \bar{\alpha}_t) \nabla \log p(x_t)}{\sqrt{\bar{\alpha}_t}} = \frac{x_t - \sqrt{1 - \bar{\alpha}_t} \epsilon_0}{\sqrt{\bar{\alpha}_t}}$$因此,$$\nabla \log p(x_t) = - \frac{1}{\sqrt{1 - \bar{\alpha}_t}} \epsilon_0$$
<u>得分函数衡量的是如何在数据空间中移动,以最大化对数概率;直观上,由于源噪声是加到自然图像上来破坏它的,沿着它的反方向移动会“去噪”图像,并且是增加后续对数概率的最佳更新</u>。我们的数学证明证实了这一直觉;我们已经明确展示了,学习建模得分函数等同于建模源噪声的负值(加上一个缩放因子)。
因此,我们推导出了优化VDM的三种等价目标:
- 学习一个神经网络来预测原始图像 $x_0$
- 学习一个神经网络来预测源噪声 $\epsilon_0$
- 学习一个神经网络来预测图像在任意噪声水平下的score $\nabla \log p(\boldsymbol{x}_t)$
VDM 可以通过随机采样时间步 $t$ 并最小化预测与真实目标之间的范数来进行可扩展的训练。
Score-based Generative Models
可以通过优化神经网络 $s_{\boldsymbol{\theta}}(\boldsymbol{x}_t,t)$ 预测分数函数 $\nabla_{\boldsymbol{x}_t} \log p(\boldsymbol{x}_t)$ 来学习变分扩散模型。然而,在上述的推导中,分数函数的引入来自 Tweedie 公式的应用,接下来更加深入探讨为什么可以通过预测score function进行建模。
幸运的是,我们可以寻找另一类生成模型,即 基于分数的生成模型(Score-based Generative Modeling formulation)(<font color="#ff0000">可参考Yang Song教授的经典论文</font>),正是这种直觉。事实证明,我们可以证明我们之前推导出的 VDM 公式具有等效的基于分数的生成建模公式,允许我们在这两种解释之间随意灵活切换。
为了开始理解为什么优化评分函数是有意义的,我们绕道而行,重新审视 基于能量的模型(energy-based models)。任意灵活的概率分布可以写成以下形式:
$$
p_{\boldsymbol{\theta}}(\boldsymbol{x})=\frac{1}{Z_{\boldsymbol{\theta}}}e^{-f_{\boldsymbol{\theta}}(\boldsymbol{x})}
$$
在这里,$f_\boldsymbol{\theta}(x)$ 是一个任意灵活、可参数化的函数,称为能量函数(energy function),通常由神经网络建模,$Z_{\boldsymbol{\theta}}$ 是一个归一化常数,用于确保 $\int p_{\boldsymbol{\theta}}(x)dx = 1$。学习这样的分布的一种方法是最大似然估计(已知分布形式):然而,这需要可行地计算归一化常数 $Z_{\boldsymbol{\theta}} = \int e^{-f_{\boldsymbol{\theta}}(\boldsymbol{x})} , d\boldsymbol{x}$,对于复杂的 $f_\boldsymbol{\theta}(\boldsymbol{x})$ 函数,这可能是不可行的。
避免计算或建模归一化常数的一种方法是使用一个神经网络 $s_{\boldsymbol{\theta}}(x)$ 来学习分布 $p(\boldsymbol{x})$ 的得分函数 $\nabla \log p(\boldsymbol{x})$。这一方法的动机来自于以下观察:对能量函数两边取对数并求导可以得到:
$$
\begin{aligned}
\nabla_{\boldsymbol{x}}\log p_{\boldsymbol{\theta}}(\boldsymbol{x}) & =\nabla_{\boldsymbol{x}}\log(\frac{1}{Z_{\boldsymbol{\theta}}}e^{-f_{\boldsymbol{\theta}}(\boldsymbol{x})}) \\
& =\nabla_{\boldsymbol{x}}\log\frac{1}{Z_{\boldsymbol{\theta}}}+\nabla_{\boldsymbol{x}}\log e^{-f_{\boldsymbol{\theta}}(\boldsymbol{x})} \\
& =-\nabla_{\boldsymbol{x}}f_{\boldsymbol{\theta}}(\boldsymbol{x}) \\
& \approx s_{\boldsymbol{\theta}}(\boldsymbol{x})
\end{aligned}$$
而这可以自由表示为神经网络(拟合$\theta$),而不涉及任何归一化常数。
可以通过使用真值评分函数(ground truth score function)最小化 Fisher Divergence 来优化评分模型:
$$
\mathbb{E}_{p(\boldsymbol{x})}\left[\left\|s_\boldsymbol{\theta}(x)-\nabla\log p(\boldsymbol{x})\right\|_2^2\right]
$$
^k6f2kc
score function代表什么?对于每个 x,取其对数似然相对于 x 的梯度基本上描述了数据空间中的哪个方向移动,以进一步增加其似然。
直观地说,得分函数定义了一个向量场,覆盖数据 $x$ 所处的整个空间,指向数据分布的高峰(模态)。从视觉上看,如图 所示:
- 使用 Langevin 动力学生成的三个随机采样轨迹的可视化,所有轨迹都从同一个初始化点开始,用于高斯混合。左图在三维等值线上绘制了这些采样轨迹,而右图根据 groundtruth score 函数绘制了采样轨迹。从同一个初始化点,由于 Langevin 动力学采样过程中的随机噪声项,我们能够从不同模式生成样本;如果没有它,从固定点采样将始终确定性地在每次试验中将分数跟随到相同的模式。
然后,通过学习真实数据分布的 score 函数,我们可以从同一空间的任意点开始,并迭代地跟踪分数,直到达到某种模式,从而生成样本。此采样过程称为 Langevin dynamics,在数学上描述为:
$$
x_{i+1}\leftarrow x_i+c\nabla\log p(x_i)+\sqrt{2c}\epsilon,\quad i=0,1,...,K
$$
其中 $x_0$ 是从先验分布(例如 uniform)中随机采样的,$\epsilon\sim\mathcal{N}(\boldsymbol{\epsilon};\boldsymbol{0},\mathbf{I})$ 是一个额外的噪声项,以确保生成的样本不会总是折叠到某个模式上,而是悬停在它周围以实现分集。此外,由于学习到的分数函数是确定性的,因此涉及噪声项的采样增加了生成过程的随机性,使我们能够避免确定性轨迹。当从位于多个模式之间的位置初始化采样时,这特别有用。
请注意,[[Understanding Diffusion Models#^k6f2kc|方程]]中的目标依赖于能够访问真值分数函数,这对于复杂分布(例如建模自然图像的分布)不可用。幸运的是,已经衍生出称为 分数匹配(score matching) 的替代技术,可以在不知道 ground truth score 的情况下最小化这种 Fisher 散度,并且可以通过随机梯度下降进行优化。 总的来说,学习将分布表示为分数函数并使用它通过马尔可夫链蒙特卡洛技术(如朗之万动力学)生成样本,被称为 基于分数的生成建模(Score-based Generative Modeling) 。
原版分数匹配有三个主要问题。
- 首先,当 x 位于高维空间中的低维流形上时,分数函数定义不明确。这可以从数学上看出;所有不在低维流形上的点的概率为零,其对数是未定义的(定义域有限)。当尝试在自然图像上学习生成模型时,这尤其不方便,因为众所周知,自然图像位于整个环境空间的低维流形上。
- 其次,通过原版分数匹配训练的估计分数函数在低密度区域不准确。这从我们在方程 157 中最小化的目标中可以明显看出。因为它是对 p(x) 的期望,并且显式地从中训练了样本,所以模型不会收到很少见或未见过的样本的准确学习信号。这是有问题的,因为我们的采样策略涉及从高维空间中的随机位置开始,这很可能是随机噪声,并根据学习到的分数函数移动。由于我们遵循的是噪声或不准确的分数估计,因此最终生成的样本也可能是次优的,或者需要更多的迭代才能收敛到准确的输出。
- 最后,Langevin 动力学采样可能不会混合,即使它是使用真值分数执行的。假设真实数据分布是两个不相交分布的混合:
$$
p(\boldsymbol{x})=c_1p_1(\boldsymbol{x})+c_2p_2(\boldsymbol{x})
$$
然后,在计算分数时,这些混合系数会丢失,因为对数运算会从分布中拆分系数,而梯度运算会将其归零。
为了直观地说明这一点,请注意,上图所示的真值得分函数与三个分布之间的不同权重无关;来自所描述的初始化点的 Langevin 动力学采样获得每个模式的机会大致相等,尽管右下角的模式在实际的高斯混合中具有更高的权重。
事实证明,这三个缺点可以通过向数据中添加多级高斯噪声来同时解决。首先,由于高斯噪声分布的支持是整个空间,因此扰动数据样本将不再局限于低维流形。其次,添加较大的高斯噪声将增加每种模式在数据分布中覆盖的面积,从而在低密度区域增加更多的训练信号。最后,添加方差增加的多个级别的高斯噪声将导致遵循真实混合系数的中间分布。
正式的,我们可以选择一个噪声级 $\{\sigma_t\}_{t=1}^T$ 的正序列,并定义一个逐渐扰动的数据分布序列:
$$
p_{\sigma_t}(\boldsymbol{x}_t)=\int p(\boldsymbol{x})\mathcal{N}(\boldsymbol{x}_t;\boldsymbol{x},\sigma_t^2\mathbf{I})d\boldsymbol{x}
$$
然后,使用分数匹配学习神经网络$s_{\boldsymbol{\theta}}(\boldsymbol{x}_t,t)$,以同时学习所有噪声级别的分数函数:
$$
\arg\min_{\boldsymbol{\theta}}\sum_{t=1}^T\lambda(t)\mathbb{E}_{p_{\sigma_t}(\boldsymbol{x}_t)}\left[\left\|s_{\boldsymbol{\theta}}(\boldsymbol{x},t)-\nabla\log p_{\sigma_t}(\boldsymbol{x}_t)\right\|_2^2\right]
$$
其中 $\lambda(t)$ 是一个正加权函数,它以噪声水平 t 为条件。请注意,该目标几乎与[[Understanding Diffusion Models#^kwguln|方程]]中得出的用于训练变分扩散模型的目标完全匹配。
此外,作者提出 退火朗之万动力学采样(Langevin dynamics sampling) 作为一种生成程序,其中样本是通过按顺序对每个 $t = T、T −1、...、2、1$ 运行 Langevin 动力学来生成的。初始化是从一些固定的先验(例如 uniform)中选择的,并且每个后续采样步骤都从前一个 simulation 的最终样本开始。因为噪声水平随着时间步 t 而稳步降低,并且我们随着时间的推移减小步长,所以样本最终收敛到 true 模式。这直接类似于变分扩散模型的马尔可夫 HVAE 解释中执行的采样过程,其中随机初始化的数据向量在降低的噪声水平上迭代优化。
因此,我们在变分扩散模型和基于分数的生成模型之间建立了明确的联系,无论是在训练目标还是采样程序中。
一个问题是如何将扩散模型自然地推广到无限数量的时间步长。在马尔可夫 HVAE 视图下,这可以解释为将层次结构的数量扩展到无穷大 $T → ∞$。从等效的基于分数的生成模型的角度来表示这一点更清晰;在无限数量的噪声尺度下,图像在连续时间内的扰动可以表示为随机过程,因此可以用随机微分方程 (SDE) 来描述。然后通过反转 SDE 进行采样,这自然需要估计每个连续值噪声水平的得分函数。SDE 的不同参数化基本上描述了随时间变化的不同扰动方案,从而能够灵活地建模噪声过程。
Guidance
到目前为止,我们只专注于对数据分布 $p(x)$ 进行建模。然而,我们通常也对学习条件分布 $p(x|y)$ 感兴趣,这将使我们能够明确地控制通过条件信息 $y$ 生成的数据。这构成了图像超分辨率模型(如级联扩散模型 [18])以及最先进的图像文本模型(如 DALL-E 2 和 Imagen)的支柱。 添加条件信息的一种自然方法是在每次迭代时与时间步信息一起添加。回想生成过程的联合分布:
$$
p(\boldsymbol{x}_{0:T})=p(\boldsymbol{x}_T)\prod_{t=1}^Tp_{\boldsymbol{\theta}}(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t)
$$
然后,为了将其转换为条件扩散模型,我们可以简单地在每个转换步骤添加任意条件信息 y ,如下所示:
$$
p(\boldsymbol{x}_{0:T}|y)=p(\boldsymbol{x}_T)\prod_{t=1}^Tp_\boldsymbol{\theta}(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t,y)
$$
例如,$y$ 可以是图像文本生成中的文本编码,或者是用于超分辨率处理的低分辨率图像。因此,我们仍然能够像以前一样学习 VDM 的核心神经网络,通过预测$$$\hat{\boldsymbol{x}}_{\boldsymbol{\theta}}(x_t,t,y)\approx x_0 \quad \hat{\boldsymbol{\epsilon}}_{\boldsymbol{\theta}}(x_t,t,y)\approx\boldsymbol{\epsilon}_0,\quad \mathrm{or} \quad s_{\boldsymbol{\theta}}(x_t,t,y)\approx\nabla\log p(\boldsymbol{x}_t|y)$$
但是以这种方式训练的条件扩散模型可能会学会忽略或淡化任何给定的条件信息。Guidance建议作为一种更明确地控制模型赋予调节信息的权重的方法,但以样本多样性为代价。
两种最流行的guidance 形式
- Classifier Guidance
- Classifier-Free Guidance
Classifier Guidance
让我们从扩散模型的基于评分的形式开始,其中我们的目标是 在任意噪声水平 $t$ 下 $\log p(x_t|y)$,即条件模型的评分。根据贝叶斯法则,我们可以推导出以下等效形式($\nabla$ :$\nabla_{{\boldsymbol{x}_{t}}}$):
$$
\begin{aligned} \nabla\log p(x_t|y) &= \nabla\log\left(\frac{p(x_t)p(y|x_t)}{p(y)}\right) \\ &= \nabla\log p(x_t) + \nabla\log p(y|x_t) - \nabla\log p(y) \\ &= \underbrace{\nabla\log p(x_t)}_{\text{unconditional score}} + \underbrace{\nabla\log p(y|x_t)}_{\text{adversarial gradient}} \end{aligned}$$ ^dy67zt
- $\log p(y)$ 关于 $x_t$ 的梯度为零
我们最终的导出结果可以解释为学习一个无条件得分函数与分类器 $p(y|x_t)$ 的对抗梯度相结合。
因此,在Classifier Guidance中,无条件扩散模型的分数是按照先前导出的方式学习的,同时分类器接受任意噪声 $x_t$ 并尝试预测条件信息 $y$。
然后,在采样过程中,用于退火朗之万动力学的总体条件得分函数被计算为无条件得分函数与噪声分类器的对抗梯度之和。 为了引入细粒度控制来鼓励或阻止模型考虑调节信息,分类器指导通过 γ 超参数项来缩放噪声分类器的对抗梯度。在分类器指导下学习的得分函数可以总结为:
$$\nabla\log p(\boldsymbol{x}_t|y)=\nabla\log p(\boldsymbol{x}_t)+\gamma\nabla\log p(y|\boldsymbol{x}_t)$$
直观上,当 γ = 0 时,条件扩散模型学会完全忽略条件信息,而当 γ 较大时,条件扩散模型学会生成严重遵循条件信息的样本。这将以样本多样性为代价,因为它只会产生易于重新生成所提供的条件信息的数据,即使在噪声水平下也是如此。 Classifier Guidance的一个值得注意的缺点是它依赖于单独学习的分类器。由于分类器必须处理任意噪声输入,而大多数现有的预训练分类模型并未对此进行优化,因此必须与扩散模型一起进行临时学习。 ^kjpkik
Classifier-Free Guidance
在Classifier-Free Guidance(Classifier-free diffusion guidance)中,作者放弃了训练一个单独的分类器模型,而采用了无条件扩散模型和条件扩散模型。为了推导无分类器引导下的评分函数,我们可以首先重新排列[[Understanding Diffusion Models#^dy67zt|方程]]:
$$
\nabla\log p(y|x_t) = \nabla\log p(x_t|y) - \nabla\log p(x_t)
$$
然后,将其代入[[Understanding Diffusion Models#^kjpkik|方程]],我们得到:
$$
\begin{aligned} \nabla\log p(x_t|y) &= \nabla\log p(x_t) + \gamma\left(\nabla\log p(x_t|y) - \nabla\log p(x_t)\right) \\ &= \nabla\log p(x_t) + \gamma\nabla\log p(x_t|y) - \gamma\nabla\log p(x_t) \\ &= \underbrace{\gamma\nabla\log p(x_t|y)}_{\text{conditional score}} + \underbrace{(1-\gamma)\nabla\log p(x_t)}_{\text{unconditional score}} \end{aligned}
$$
γ 是一个控制我们学习的条件模型对条件信息的关心程度的术语。
- 当γ = 0时,学习到的条件模型完全忽略条件并学习无条件扩散模型。
- 当 γ = 1 时,模型在没有指导的情况下显式学习普通条件分布。
- 当γ > 1时,扩散模型不仅优先考虑条件得分函数,而且会向远离无条件得分函数的方向移动。换句话说,它降低了生成不使用条件信息的样本的概率,有利于明确使用条件信息的样本。这也会降低样本多样性,但代价是生成准确匹配调节信息的样本。
由于学习两个单独的扩散模型的成本很高,因此我们可以将条件扩散模型和无条件扩散模型一起学习为单个条件模型;可以通过用固定常数值(例如零)替换条件信息来查询无条件扩散模型。这本质上是对条件信息执行随机丢失。
同步过来有的markdown语法不支持,有时间改下
天呐
地啊