Understanding Diffusion Model[上]

Understanding Diffusion Model

1. Introduction: Generative Models

给定真实样本 x x xgenerative model的目标是建模出其真实数据分布 p ( x ) p(x) p(x)

生成对抗网络 (GAN) 对复杂分布的采样过程进行建模,该过程以对抗方式学习。 另一类生成模型,称为“likelihood-based”,旨在学习一个模型,该模型将高可能性分配给观察到的数据样本。 这包括autoregressive
models, normalizing flows, and Variational Autoencoders (VAEs)

2. Background: ELBO, VAE, and Hierarchical VAE

对于大多数的modalities,我们可以将我们观察到的数据认为是由相关的隐变量 z z z表示或者生成的。

在生成建模中,我们通常寻求学习低维潜在表示而不是高维表示。一方面,如果没有强大的先验信息,学习比观察更高维的表示是徒劳的;另一方面,学习低维的潜在表示也可以被视为一种压缩形式,可以潜在的揭示描述观察的语义上更有意义的结构。

Evidence Lower Bound

我们可以将隐变量 z z z和我们观察到的数据 x x x认为是由联合分布 p ( x , z ) p(x,z) p(x,z)建模的。

回想一下一种"likelihood-based"的生成建模方法,是学习一个模型以最大化所有观察到的 x x x的似然概率 p ( x ) p(x) p(x)。有两种方法可以从联合分布 p ( x , z ) p(x,z) p(x,z)中获取我们的似然概率 p ( x ) p(x) p(x)。方法一是对隐变量 z z z进行积分,得到 x x x的边缘概率分布:
p ( x ) = ∫ z p ( x , z )    d z ( 1 ) p(x) = \int_z p(x,z) \; dz \qquad (1) p(x)=zp(x,z)dz(1)
或者使用链式法则:
p ( x ) = p ( x , z ) p ( z ∣ x ) ( 2 ) p(x) = \frac{p(x,z)}{p(z \mid x)} \qquad (2) p(x)=p(zx)p(x,z)(2)
直接计算和最大化似然概率 p ( x ) p(x) p(x)是非常困难的,因为它要么涉及公式1中的所有隐变量 z z z,这对复杂模型来说是难以处理的;要么涉及访问ground truth latent encoder p ( z ∣ x ) p(z \mid x) p(zx)。( p ( z ∣ x ) p(z \mid x) p(zx)是我们要建模的编码器)。但是我们可以利用上面两个公式推导出生成模型的置信下界,即Evidence Lower Bound (ELBO)

在这种情况下,evidence被量化为观测数据的对数似然;然后最大化ELBO成为优化潜变量模型的代理目标;在最好的情况下,当ELBO被很好的参数化和完美优化的时候,它变得和evidence等价。从数学形式上看,ELBO的公式如下:
E q ϕ ( z ∣ x ) [ l o g    p ( x , z ) q ϕ ( z ∣ x ) ] ( 3 ) \mathbf E_{q_\phi(z \mid x)}\left[log \; \frac{p(x,z)}{q_\phi(z \mid x)}\right] \qquad (3) Eqϕ(zx)[logqϕ(zx)p(x,z)](3)
ELBO与似然概率 p ( x ) p(x) p(x)之间的关系可以写成下面的公式:
l o g    p ( x ) ≥ E q ϕ ( z ∣ x ) [ l o g    p ( x , z ) q ϕ ( z ∣ x ) ] ( 4 ) \begin{aligned} log \; p(x) \ge \mathbf E_{q_\phi(z \mid x)}\left[log \; \frac{p(x,z)}{q_\phi(z \mid x)}\right] && (4) \end{aligned} logp(x)Eqϕ(zx)[logqϕ(zx)p(x,z)](4)
其中, q ϕ ( z ∣ x ) q_\phi(z \mid x) qϕ(zx)是一种近似变分分布,其中参数 ϕ \phi ϕ是我们待优化的参数。直观来说,可以认为它是一种参数化的模型,用来学习给定观测 x x x来估计潜变量 z z z的真实分布,也就是说用来寻求近似真实的后验概率 p ( z ∣ x ) p(z \mid x) p(zx)

正如我们在探寻VAE中所看到的,当我们通过调整参数以最大化ELBO来增加下限的时候,我们获得可用于建模真实数据分布的components并可以从中采样,从而学习了生成模型。下面将推导为什么ELBO是我们希望最大化的目标。
log ⁡ p ( x ) = log ⁡ ∫ p ( x , z ) d z  (Apply Equation 1)  (5) = log ⁡ ∫ p ( x , z ) q ϕ ( z ∣ x ) q ϕ ( z ∣ x ) d z  (Multiply by  1 = q ϕ ( z ∣ x ) q ϕ ( z ∣ x ) ) (6) = log ⁡ E q ϕ ( z ∣ x ) [ p ( x , z ) q ϕ ( z ∣ x ) ]  (Definition of Expectation)  (7) ≥ E q ϕ ( z ∣ x ) [ log ⁡ p ( x , z ) q ϕ ( z ∣ x ) ]  (Apply Jensen’s Inequality)  ( 8 ) \begin{aligned} \log p(\boldsymbol{x}) &=\log \int p(\boldsymbol{x}, \boldsymbol{z}) d \boldsymbol{z} & & \text { (Apply Equation 1) } && \text{(5)}\\ &=\log \int \frac{p(\boldsymbol{x}, \boldsymbol{z}) q_{\boldsymbol{\phi}}(\boldsymbol{z} \mid \boldsymbol{x})}{q_{\boldsymbol{\phi}}(\boldsymbol{z} \mid \boldsymbol{x})} d \boldsymbol{z} & &\text { (Multiply by } \left.1=\frac{q_{\boldsymbol{\phi}}(\boldsymbol{z} \mid \boldsymbol{x})}{q_{\boldsymbol{\phi}}(\boldsymbol{z} \mid \boldsymbol{x})}\right) && \text{(6)}\\ &=\log \mathbb{E}_{q_\phi(\boldsymbol{z} \mid \boldsymbol{x})}\left[\frac{p(\boldsymbol{x}, \boldsymbol{z})}{q_{\boldsymbol{\phi}}(\boldsymbol{z} \mid \boldsymbol{x})}\right] & & \text { (Definition of Expectation) } && \text{(7)}\\ & \geq \mathbb{E}_{q_\phi(\boldsymbol{z} \mid \boldsymbol{x})}\left[\log \frac{p(\boldsymbol{x}, \boldsymbol{z})}{q_{\boldsymbol{\phi}}(\boldsymbol{z} \mid \boldsymbol{x})}\right] & & \text { (Apply Jensen's Inequality) } && \text(8) \end{aligned} logp(x)=logp(x,z)dz=logqϕ(zx)p(x,z)qϕ(zx)dz=logEqϕ(zx)[qϕ(zx)p(x,z)]Eqϕ(zx)[logqϕ(zx)p(x,z)] (Apply Equation 1)  (Multiply by 1=qϕ(zx)qϕ(zx)) (Definition of Expectation)  (Apply Jensen’s Inequality) (5)(6)(7)(8)
在上面的公式推导中,我们使用Jensen不等式直接得到了置信下界,但是这个证明没有给出明确的说明,为什么ELBO实际上是evidence的下界(Jensen不等式将一部分丢弃了),并且仅仅知道ELBO是evidence的下界并不能真正告诉我们为什么要将ELBO作为最大化的目标。于是我们需要再进行一次推导,如下所示:
l o g    p ( x ) = l o g    p ( x ) ∫ q ϕ ( z ∣ x ) d z ( 1 = ∫ q ϕ ( z ∣ x ) d z ) ( 9 ) = ∫ q ϕ ( z ∣ x ) ( l o g    p ( x ) ) d z 将p(x)放到积分号里面 ( 10 ) = E q ϕ ( z ∣ x ) [ l o g    p ( x ) ] 数学期望的定义 ( 11 ) = E q ϕ ( z ∣ x ) [ l o g    p ( x , z ) p ( z ∣ x ) ] 公式2链式法则 ( 12 ) = E q ϕ ( z ∣ x ) [ l o g    p ( x , z ) q ϕ ( z ∣ x ) p ( z ∣ x ) q ϕ ( z ∣ x ) ] 上下同乘以 q ϕ ( z ∣ x ) ( 13 ) = E q ϕ ( z ∣ x ) [ l o g    p ( x , z ) q ϕ ( z ∣ x ) ] + E q ϕ ( z ∣ x ) [ l o g    q ϕ ( z ∣ x ) p ( z ∣ x ) ] log乘法拆成加法 ( 14 ) = E q ϕ ( z ∣ x ) [ l o g    p ( x , z ) q ϕ ( z ∣ x ) ] + D K L ( q ϕ ( z ∣ x )    ∣ ∣    p ( z ∣ x ) ) KL散度定义 ( 15 ) ≥ E q ϕ ( z ∣ x ) [ l o g    p ( x , z ) q ϕ ( z ∣ x ) ] KL散度 ≥ 0 ( 16 ) \begin{aligned} log \; p(x) &= log \; p(x) \int q_\phi(z \mid x)dz &&(1 = \int q_\phi(z \mid x)dz) && (9) \\ & = \int q_\phi(z \mid x)(log \; p(x))dz && \text{将p(x)放到积分号里面} && (10) \\ & = \mathbf E_{q_\phi(z \mid x)}[log \; p(x)] && \text{数学期望的定义} && (11) \\ & = \mathbf E_{q_\phi(z \mid x)} \left[log \; \frac{p(x,z)}{p(z \mid x)}\right] && \text{公式2链式法则} && (12) \\ & = \mathbf E_{q_\phi(z \mid x)}\left[log \; \frac{p(x,z)q_\phi(z \mid x)}{p(z \mid x)q_\phi(z \mid x)}\right] && \text{上下同乘以} q_\phi(z \mid x) && (13) \\ & = \mathbf E_{q_\phi(z \mid x)} \left[log \; \frac{p(x,z)}{q_\phi(z \mid x)}\right] + \mathbf E_{q_\phi(z \mid x)} \left[log \; \frac{q_\phi(z \mid x)}{p(z \mid x)}\right] && \text{log乘法拆成加法} && (14) \\ & = \mathbf E_{q_\phi(z \mid x)} \left[log \; \frac{p(x,z)}{q_\phi(z \mid x)}\right] + D_{KL}(q_\phi(z \mid x) \; || \; p(z \mid x)) && \text{KL散度定义} &&(15) \\ & \ge \mathbf E_{q_\phi(z \mid x)} \left[log \; \frac{p(x,z)}{q_\phi(z \mid x)}\right] && \text{KL散度} \ge 0 && (16) \end{aligned} logp(x)=logp(x)qϕ(zx)dz=qϕ(zx)(logp(x))dz=Eqϕ(zx)[logp(x)]=Eqϕ(zx)[logp(zx)p(x,z)]=Eqϕ(zx)[logp(zx)qϕ(zx)p(x,z)qϕ(zx)]=Eqϕ(zx)[logqϕ(zx)p(x,z)]+Eqϕ(zx)[logp(zx)qϕ(zx)]=Eqϕ(zx)[logqϕ(zx)p(x,z)]+DKL(qϕ(zx)p(zx))Eqϕ(zx)[logqϕ(zx)p(x,z)](1=qϕ(zx)dz)p(x)放到积分号里面数学期望的定义公式2链式法则上下同乘以qϕ(zx)log乘法拆成加法KL散度定义KL散度0(9)(10)(11)(12)(13)(14)(15)(16)
从公式(15)我们可以看出,evidence项等于ELBO加上一个KL散度项,这个KL散度项在第一次推导的时候被Jensen不等式去掉了。

首先由于evidence项(即log p(x))跟ELBO之间差了一个严格非负的KL散度项,因此ELBO确实是evidence的一个下限。

其次,我们要探寻为什么要去优化ELBO项。在引入了我们想要建模的隐变量 z z z之后,我们的目标是学习描述我们观察到的数据 x x x的潜在结构。换句话说,我们希望去优化variational posterior(变分后验) q ϕ ( z ∣ x ) q_\phi(z \mid x) qϕ(zx)的参数去精确的匹配**true posterior(真实后验)**分布 p ( z ∣ x ) p(z \mid x) p(zx),这是通过最小化这两者之间的KL散度来实现的(详见附录[最小化KL散度])。然而,直接优化KL散度项是很困难的,因为我们没有办法得到ground truth p ( z ∣ x ) p(z \mid x) p(zx)的分布。不过,在公式15的左侧,观测数据(and therefore our evidence term log p(x)这句话不知道怎么翻译)的likelihood始终是相对于 ϕ \phi ϕ的常数,因为 p ( x ) p(x) p(x)是通过从联合分布 p ( x , z ) p(x,z) p(x,z)边缘积分 z z z来得到的,并不依赖任何的参数 ϕ \phi ϕ(这句话可以简单理解为,自然世界的真实图像分布是不受任何后验参数影响的,它本身就在那儿)。由于ELBO和KL项的总和为常数,因此最大化ELBO项就相当于最小化KL项。因此,我们越优化ELBO,我们的变分后验就越接近真实后验。并且,一旦进行训练,ELBO也可以用于估计观测或者生成data的likelihood,因为ELBO被训练建模成为 l o g    p ( x ) log \; p(x) logp(x)的近似

Variational Autoencoders

在这里插入图片描述

图1:图形表示VAE。其中,encoder q ( z ∣ x ) q(z \mid x) q(zx)定义了已知观察值 x x x潜变量 z z z的分布,decoder从已知的潜变量 z z z中解码出 x x x

在变分自动编码器(VAE)的默认公式中,我们直接最大化ELBO。这种方法是variational(应该翻译成可变分的),因为我们优化了由 ϕ \phi ϕ参数化的一系列潜在后验分布得到最优的 q ϕ ( z ∣ x ) q_\phi(z \mid x) qϕ(zx)。这被称为autoencoder,其中input data在经历中间bottlenecking表示步骤之后用于预测自身。下面清晰的展示出ELBO项:
E q ϕ ( z ∣ x ) [ log ⁡ p ( x , z ) q ϕ ( z ∣ x ) ] = E q ϕ ( z ∣ x ) [ log ⁡ p θ ( x ∣ z ) p ( z ) q ϕ ( z ∣ x ) ]  (Chain Rule of Probability)  ( 17 ) = E q ϕ ( z ∣ x ) [ log ⁡ p θ ( x ∣ z ) ] + E q ϕ ( z ∣ x ) [ log ⁡ p ( z ) q ϕ ( z ∣ x ) ]  (Split the Expectation)  ( 18 ) = E q ϕ ( z ∣ x ) [ log ⁡ p θ ( x ∣ z ) ] ⏟ reconstruction term  − D K L ( q ϕ ( z ∣ x ) ∥ p ( z ) ) ⏟ prior matching term   (Definition of KL Divergence)  ( 19 ) \begin{array}{rll} \mathbb{E}_{q_\phi(\boldsymbol{z} \mid \boldsymbol{x})}\left[\log \frac{p(\boldsymbol{x}, \boldsymbol{z})}{q_\phi(\boldsymbol{z} \mid \boldsymbol{x})}\right] & =\mathbb{E}_{q_\phi(\boldsymbol{z} \mid \boldsymbol{x})}\left[\log \frac{p_{\boldsymbol{\theta}}(\boldsymbol{x} \mid \boldsymbol{z}) p(\boldsymbol{z})}{q_{\boldsymbol{\phi}}(\boldsymbol{z} \mid \boldsymbol{x})}\right] & \text { (Chain Rule of Probability) } & (17)\\ & =\mathbb{E}_{q_\phi(\boldsymbol{z} \mid \boldsymbol{x})}\left[\log p_{\boldsymbol{\theta}}(\boldsymbol{x} \mid \boldsymbol{z})\right]+\mathbb{E}_{q_\phi(\boldsymbol{z} \mid \boldsymbol{x})}\left[\log \frac{p(\boldsymbol{z})}{q_{\boldsymbol{\phi}}(\boldsymbol{z} \mid \boldsymbol{x})}\right] & \text { (Split the Expectation) } & (18) \\ & =\underbrace{\mathbb{E}_{q_\phi(\boldsymbol{z} \mid \boldsymbol{x})}\left[\log p_{\boldsymbol{\theta}}(\boldsymbol{x} \mid \boldsymbol{z})\right]}_{\text {reconstruction term }}-\underbrace{D_{\mathrm{KL}}\left(q_\phi(\boldsymbol{z} \mid x) \| p(\boldsymbol{z})\right)}_{\text {prior matching term }} & \text { (Definition of KL Divergence) } & (19) \end{array} Eqϕ(zx)[logqϕ(zx)p(x,z)]=Eqϕ(zx)[logqϕ(zx)pθ(xz)p(z)]=Eqϕ(zx)[logpθ(xz)]+Eqϕ(zx)[logqϕ(zx)p(z)]=reconstruction term  Eqϕ(zx)[logpθ(xz)]prior matching term  DKL(qϕ(zx)p(z)) (Chain Rule of Probability)  (Split the Expectation)  (Definition of KL Divergence) (17)(18)(19)
上面是将ELBO项做展开,不同于公式(9)是将evidence l o g    p ( x ) log \; p(x) logp(x)展开

在这种情况下,我们学习了一个中间bottlenecking分布 q ϕ ( z ∣ x ) q_\phi(z \mid x) qϕ(zx),可以将其视为一个encoder;它可以将输入转换成可能的潜在分布。同时我们还学习了一个决策函数(deterministic function) p θ ( x ∣ z ) p_\theta(x \mid z) pθ(xz)用于将给定的潜在变量 z z z转换成观测数据 x x x,可以将其视为一个decoder

公式(19)中的两个项都有一个直观的描述:第一项从我们的变分分布 q ϕ ( z ∣ x ) q_\phi(z \mid x) qϕ(zx)中测量了decoder的重建likelihood;这个保证了所学习的分布正在对可以从中重新生成原始数据的有效隐变量进行建模。第二项测量了学习的变分分布于隐变量的先验有多相似。最小化这一项会鼓励编码器实际学习一个分布,而不是塌陷为Dirac delta function.最大化ELBO等同于最大化第一项和最小化第二项。

VAE的一个定义上的特征就是如何在参数 ϕ , θ \phi,\theta ϕ,θ上联合优化ELBO。VAE的encoder通常被选择用来建模具有对角协方差的多元高斯,并且先验通常被选择为标准多元高斯

The encoderof the VAE is commonly chosen to model a multivariate Gaussian with diagonal covariance, and the prior is often selected to be a standard multivariate Gaussian:

q ϕ ( z ∣ x ) = N ( z ; μ ϕ ( x ) , σ ϕ 2 ( x ) I ) ( 20 ) p ( z ) = N ( z ; 0 , I ) ( 21 ) \begin{aligned} q_\phi(z \mid x) = \mathcal N(z;\mu_\phi(x),\sigma^2_\phi(x)I) && (20) \\ p(z) = \mathcal N(z;0,I) && (21) \end{aligned} qϕ(zx)=N(z;μϕ(x),σϕ2(x)I)p(z)=N(z;0,I)(20)(21)

然后,可以分析地计算ELBO的KL项,并且可以使用Monte Carlo estimate来估计重建项。我们的目标函数可以写为:
arg ⁡ max ⁡ ϕ , θ E q ϕ ( z ∣ x ) [ log ⁡ p θ ( x ∣ z ) ] − D K L ( q ϕ ( z ∣ x ) ∥ p ( z ) ) ≈ arg ⁡ max ⁡ ϕ , θ ∑ l = 1 L log ⁡ p θ ( x ∣ z ( l ) ) − D K L ( q ϕ ( z ∣ x ) ∥ p ( z ) ) ( 22 ) \underset{\boldsymbol{\phi}, \boldsymbol{\theta}}{\arg \max } \mathbb{E}_{q_{\boldsymbol{\phi}}(\boldsymbol{z} \mid \boldsymbol{x})}\left[\log p_{\boldsymbol{\theta}}(\boldsymbol{x} \mid \boldsymbol{z})\right]-D_{\mathrm{KL}}\left(q_{\boldsymbol{\phi}}(\boldsymbol{z} \mid \boldsymbol{x}) \| p(\boldsymbol{z})\right) \approx \underset{\boldsymbol{\phi}, \boldsymbol{\theta}}{\arg \max } \sum_{l=1}^L \log p_{\boldsymbol{\theta}}\left(\boldsymbol{x} \mid \boldsymbol{z}^{(l)}\right)-D_{\mathrm{KL}}\left(q_{\boldsymbol{\phi}}(\boldsymbol{z} \mid \boldsymbol{x}) \| p(\boldsymbol{z})\right) \qquad (22) ϕ,θargmaxEqϕ(zx)[logpθ(xz)]DKL(qϕ(zx)p(z))ϕ,θargmaxl=1Llogpθ(xz(l))DKL(qϕ(zx)p(z))(22)
上面的公式中,隐变量 { z ( l ) } l = 1 L \{z^{(l)}\}_{l=1}^L {z(l)}l=1L是对于每一个观测 x x x从分布 q ϕ ( z ∣ x ) q_\phi(z \mid x) qϕ(zx)中采样得到的。然而,这个默认设置会出现问题:我们用来计算损失的每一个 z ( l ) z^{(l)} z(l)都是由随机抽样生成的,这个过程是不可微分的

each z(l) that our loss is computed on is generated by a stochastic sampling procedure, which is generally non-differentiable.

幸运的是,当 q ϕ ( z ∣ x ) q_\phi(z \mid x) qϕ(zx)被设计为建模某些确定的分布(包括多元高斯分布)时,这可以使用reparameterization
trick
来解决。

重参数化技巧将随机变量重写成噪声的确定性函数(deterministic function);这允许通过梯度下降来优化非随机项。例如,从具有任意均值 μ \mu μ和方差 σ 2 \sigma^2 σ2的标准分布 x ∼ N ( x ; μ , σ 2 ) x \sim \mathcal N(x;\mu,\sigma^2) xN(x;μ,σ2)中采样可以改写成:
x = μ + σ ϵ w i t h    ϵ ∼ N ( ϵ ; 0 , I ) x = \mu + \sigma \epsilon \quad with \; \epsilon \sim \mathcal N(\epsilon;0,I) x=μ+σϵwithϵN(ϵ;0,I)

换句话说,任意的高斯分布都可以被解释为标准高斯分布( ϵ \epsilon ϵ为一个样本),高斯分布的均值通过加法从0偏移到 μ \mu μ,其方差拉伸到目标方差 σ 2 \sigma^2 σ2。(任意高斯分布都可以由标准高斯分布通过拉伸和平移得到)因此,通过重参数化技巧,可以通过从标准高斯中采样,按照目标均值和方差进行平移和缩放来达到从任意高斯分布中采样的结果。

在VAE中,每个隐变量 z z z因此被计算为输入 x x x和辅助噪声变量 ϵ \epsilon ϵ的确定性函数:
z = μ ϕ ( x ) + σ ϕ ( x ) ⊙ ϵ  with  ϵ ∼ N ( ϵ ; 0 , I ) \boldsymbol{z}=\boldsymbol{\mu}_\phi(\boldsymbol{x})+\boldsymbol{\sigma}_\phi(\boldsymbol{x}) \odot \boldsymbol{\epsilon} \quad \text { with } \boldsymbol{\epsilon} \sim \mathcal{N}(\boldsymbol{\epsilon} ; \mathbf{0}, \mathbf{I}) z=μϕ(x)+σϕ(x)ϵ with ϵN(ϵ;0,I)
其中, ⊙ \odot 表示按元素乘积。在重参数化下的隐变量 z z z,然后可以根据 ϕ \phi ϕ的需要计算梯度,去优化 μ ϕ , σ ϕ \mu_\phi,\sigma_\phi μϕ,σϕ。因此VAE利用重参数化技巧和蒙特卡洛估计来联合优化ELBO通过优化参数 θ , ϕ \theta,\phi θ,ϕ

**训练完VAE之后,可以通过直接从隐变量空间 p ( z ) p(z) p(z)中采样,然后利用解码器来生成新数据。**当隐变量 z z z的维度小于输入 x x x的维度时,VAE可能在学习一种紧凑、有用的表示。此外,当学习到语义上有意义的潜在空间时,可以在讲潜在向量传递给解码器之前对其进行修改,以便更加精准地控制生成的数据。

Hierarchical Variational Autoencoders

在这里插入图片描述

图2:具有T层潜在层的马尔科夫分层自编码器。生成过程被建模为马尔科夫链,每个隐变量 z t z_t zt仅从其前一个隐变量 z t + 1 z_{t+1} zt+1生成。

上图中, q q q是encoder, p p p是decoder

多层VAE是VAE的推广,扩展到了潜变量的多个层次。根据VAE的公式,隐变量本身被解释为是由其它更高层次、更加抽象的隐变量产生的。

在具有T个层级的HVAE中,每个latent都被允许以之前所有的latent为条件,在本文中我们关注一种特殊情况称之为马尔科夫HVAE。在MHVAE中,生成过程是马尔科夫链;也就是说,层次结构下的每个步骤都是马尔科夫的,即每个隐变量 z t z_t zt的decoding仅仅取决于先前的隐变量 z t + 1 z_{t+1} zt+1;在数学上,我们可以将数据 x x x和隐变量 z 1 : T z_{1:T} z1:T的联合分布和MHVAE的后验表示为如下形式:
p ( x , z 1 : T ) = p ( z T ) p θ ( x ∣ z 1 ) ∏ t = 2 T p θ ( z t − 1 ∣ z t ) ( 23 ) q ϕ ( z 1 : T ∣ x ) = q ϕ ( z 1 ∣ x ) ∏ t = 2 T q ϕ ( z t ∣ z t − 1 ) ( 24 ) \begin{aligned} p(x,z_{1:T}) = p(z_T)p_\theta(x \mid z_1) \prod_{t=2}^T p_\theta(z_{t-1} \mid z_t) && (23) \\ q_\phi(z_{1:T} \mid x) = q_\phi(z_1 \mid x) \prod_{t=2}^T q_\phi(z_t \mid z_{t-1}) && (24) \end{aligned} p(x,z1:T)=p(zT)pθ(xz1)t=2Tpθ(zt1zt)qϕ(z1:Tx)=qϕ(z1x)t=2Tqϕ(ztzt1)(23)(24)
*见附录详尽的推导

然后我们就可以将ELBO进行推广:
log ⁡ p ( x ) = log ⁡ ∫ p ( x , z 1 : T ) d z 1 : T  (Apply Equation 1)  ( 25 ) = log ⁡ ∫ p ( x , z 1 : T ) q ϕ ( z 1 : T ∣ x ) q ϕ ( z 1 : T ∣ x ) d z 1 : T  (Multiply by  1 = q ϕ ( z 1 : T ∣ x ) q ϕ ( z 1 : T ∣ x ) ( 26 ) = log ⁡ E q ϕ ( z 1 : T ∣ x ) [ p ( x , z 1 : T ) q ϕ ( z 1 : T ∣ x ) ]  (Definition of Expectation)  ( 27 ) ≥ E q ϕ ( z 1 : T ∣ x ) [ log ⁡ p ( x , z 1 : T ) q ϕ ( z 1 : T ∣ x ) ]  (Apply Jensen’s Inequality)  ( 28 ) \begin{aligned} \log p(\boldsymbol{x}) &=\log \int p\left(\boldsymbol{x}, \boldsymbol{z}_{1: T}\right) d \boldsymbol{z}_{1: T} & & \text { (Apply Equation 1) } && (25)\\ &=\log \int \frac{p\left(\boldsymbol{x}, \boldsymbol{z}_{1: T}\right) q_{\boldsymbol{\phi}}\left(\boldsymbol{z}_{1: T} \mid \boldsymbol{x}\right)}{q_{\boldsymbol{\phi}}\left(\boldsymbol{z}_{1: T} \mid \boldsymbol{x}\right)} d \boldsymbol{z}_{1: T} & & \text { (Multiply by } 1=\frac{q_{\boldsymbol{\phi}}\left(\boldsymbol{z}_{1: T} \mid \boldsymbol{x}\right)}{q_{\boldsymbol{\phi}}\left(\boldsymbol{z}_{1: T} \mid \boldsymbol{x}\right)} && (26)\\ &=\log \mathbb{E}_{q_\phi\left(\boldsymbol{z}_{1: T} \mid \boldsymbol{x}\right)}\left[\frac{p\left(\boldsymbol{x}, \boldsymbol{z}_{1: T}\right)}{q_{\boldsymbol{\phi}}\left(\boldsymbol{z}_{1: T} \mid \boldsymbol{x}\right)}\right] & & \text { (Definition of Expectation) } &&(27) \\ & \geq \mathbb{E}_{q_{\boldsymbol{\phi}}\left(\boldsymbol{z}_{1: T} \mid \boldsymbol{x}\right)}\left[\log \frac{p\left(\boldsymbol{x}, \boldsymbol{z}_{1: T}\right)}{q_{\boldsymbol{\phi}}\left(\boldsymbol{z}_{1: T} \mid \boldsymbol{x}\right)}\right] & & \text { (Apply Jensen's Inequality) } && (28) \end{aligned} logp(x)=logp(x,z1:T)dz1:T=logqϕ(z1:Tx)p(x,z1:T)qϕ(z1:Tx)dz1:T=logEqϕ(z1:Tx)[qϕ(z1:Tx)p(x,z1:T)]Eqϕ(z1:Tx)[logqϕ(z1:Tx)p(x,z1:T)] (Apply Equation 1)  (Multiply by 1=qϕ(z1:Tx)qϕ(z1:Tx) (Definition of Expectation)  (Apply Jensen’s Inequality) (25)(26)(27)(28)
然后,我们可以将联合分布(公式23)和后验概率分布(公式24)带入公式28,形成另一种形式:
E q ϕ ( z 1 : T ∣ x ) [ log ⁡ p ( x , z 1 : T ) q ϕ ( z 1 : T ∣ x ) ] = E q ϕ ( z 1 : T ∣ x ) [ log ⁡ p ( z T ) p θ ( x ∣ z 1 ) ∏ t = 2 T p θ ( z t − 1 ∣ z t ) q ϕ ( z 1 ∣ x ) ∏ t = 2 T q ϕ ( z t ∣ z t − 1 ) ] ( 29 ) \mathbb{E}_{q_{\boldsymbol{\phi}}\left(\boldsymbol{z}_{1: T} \mid \boldsymbol{x}\right)}\left[\log \frac{p\left(\boldsymbol{x}, \boldsymbol{z}_{1: T}\right)}{q_{\boldsymbol{\phi}}\left(\boldsymbol{z}_{1: T} \mid \boldsymbol{x}\right)}\right]=\mathbb{E}_{q_{\boldsymbol{\phi}}\left(\boldsymbol{z}_{1: T} \mid \boldsymbol{x}\right)}\left[\log \frac{p\left(\boldsymbol{z}_T\right) p_{\boldsymbol{\theta}}\left(\boldsymbol{x} \mid \boldsymbol{z}_1\right) \prod_{t=2}^T p_{\boldsymbol{\theta}}\left(\boldsymbol{z}_{t-1} \mid \boldsymbol{z}_t\right)}{q_{\boldsymbol{\phi}}\left(\boldsymbol{z}_1 \mid \boldsymbol{x}\right) \prod_{t=2}^T q_{\boldsymbol{\phi}}\left(\boldsymbol{z}_t \mid \boldsymbol{z}_{t-1}\right)}\right] \qquad (29) Eqϕ(z1:Tx)[logqϕ(z1:Tx)p(x,z1:T)]=Eqϕ(z1:Tx)[logqϕ(z1x)t=2Tqϕ(ztzt1)p(zT)pθ(xz1)t=2Tpθ(zt1zt)](29)
正如我们将在下面展示的,当我们研究变分扩散模型时,公式29可以进一步分解成为可解释的部分。

3. Variational Diffusion Models

最简单的考虑变分扩散模型(VDM)的方法是作为具有三个关键限制的MHVAE(马尔科夫分层变分自编码器):

  • 隐变量维度和输入数据维度完全相等
  • **每个时间步的隐变量的encoder结构没有被学习;它被预定义为线性高斯模型。**换句话说,它是以前一个时间步的输出为中心的高斯分布
  • 潜变量encoder的高斯分布参数随时间变化,使得在最终的时间步T的latent分布是标准的高斯分布

在这里插入图片描述

图3:变分扩散模型的视觉表示。其中 x 0 x_0 x0是真实的数据观测,例如自然图像; x T x_T xT表示纯的高斯噪声, x t x_t xt x 0 x_0 x0的中间噪声版本。每个 q ( x t ∣ x t − 1 ) q(x_t \mid x_{t-1}) q(xtxt1)都被建模为使用前一状态的输出作为其平均值的高斯分布。

此外,我们明确保持标准MHVAE的分层转换之间的马尔科夫性。

根据第一个限制(隐变量维度和真实数据维度完全相同),我们将真实的数据样本和隐变量都表示为 x t x_t xt,其中 t = 0 , 即 x 0 t=0,即x_0 t=0x0表示真实数据样本, t ∈ [ 1 , T ] t \in [1,T] t[1,T]代表对应的隐变量。VDM与MHVAE的后验概率相同(公式24),但是现在可以改写为:
q ( x 1 : T ∣ x 0 ) = ∏ t = 1 T    q ( x t ∣ x t − 1 ) ( 30 ) q(x_{1:T} \mid x_0) = \prod_{t=1}^T \; q(x_t \mid x_{t-1}) \qquad (30) q(x1:Tx0)=t=1Tq(xtxt1)(30)
从第二个限制中,我们可以知道encoder中的每个隐变量的分布都是以其先前层的隐变量为中心的高斯分布。与MHVAE不同,在每个时间步t的encoder是不被学习的;它被固定为线性高斯模型,其中的平均值和标准差可以被预设为超参数或者可学习的参数。我们将高斯encoder参数化为均值 μ t ( x t ) = α t x t − 1 \mu_t(x_t) = \sqrt{\alpha_t} x_{t-1} μt(xt)=αt xt1和方差 ∑ t ( x t ) = ( 1 − α t ) I \sum_t(x_t) = (1 - \alpha_t)I t(xt)=(1αt)I,其中系数的形式被选择为使得隐变量的方差保持在相似的尺度上;换句话说,encoding的过程是variance-preserving。主要的结论是,为了保持灵活性, α t \alpha_t αt是一个(潜在可学习)的系数,随着分层深度t来变化。数学上,encoder的转换可以表示为:
q ( x t ∣ x t − 1 ) = N ( x t ; α t x t − 1 , ( 1 − α t ) I ) ( 31 ) q(x_t \mid x_{t-1}) = \mathcal N(x_t;\sqrt{\alpha_t}x_{t-1},(1-\alpha_t)I) \qquad (31) q(xtxt1)=N(xt;αt xt1,(1αt)I)(31)
从第三个限制中,我们知道 α t \alpha_t αt根据固定或者可学习的策略来随着时间进行演变,该策略被构造为使得最终的隐变量分布 p ( x T ) p(x_T) p(xT)是标准高斯分布。然后,我们可以更新MHVAE的联合分布(公式23),可以将VDM的联合分布写成:
p ( x 0 , T ) = p ( x T ) ∏ t = 1 T    p θ ( x t − 1 ∣ x t ) ( 32 ) p ( x T ) = N ( x T ; 0 , I ) ( 33 ) p(x_0,T) = p(x_T) \prod_{t=1}^T \; p_\theta(x_{t-1} \mid x_t) \qquad (32) \\ p(x_T) = \mathcal N(x_T;0,I) \qquad (33) p(x0,T)=p(xT)t=1Tpθ(xt1xt)(32)p(xT)=N(xT;0,I)(33)

总之,VDM的三个限制描述的是图像输入随时间的稳定噪声化;我们通过添加高斯噪声来逐渐破坏图像,直到它最终与纯的高斯噪声完全相同。从视觉上看,这个过程与图3所示。

注意到,我们encoder的分布 q ( x t ∣ x t − 1 ) q(x_t \mid x_{t-1}) q(xtxt1)不在由 ϕ \phi ϕ参数化,因为它们在每个时间步都被完全建模为具有确定均值和方差的高斯分布。因此在VDM中,我们只对学习**条件概率 p θ ( x t − 1 ∣ x t ) p_\theta(x_{t-1} \mid x_t) pθ(xt1xt)**感兴趣,以便我们可以模拟新数据。在优化完VDM之后,采样的过程很简单,从高斯噪声 p ( x T ) p(x_T) p(xT)中采样,并迭代进行去噪转换 p θ ( x t − 1 ∣ x t ) p_\theta(x_{t-1} \mid x_t) pθ(xt1xt)去生成新的 x 0 x_0 x0

类似于任何的HVAE,VDM可以通过最大化ELBO来进行优化,推导如下:
在这里插入图片描述

在这里插入图片描述

整体的推导跟着步骤来一遍是不难的,需要注意的点主要是:

  1. 公式39到40是对连乘的下标做了一个统一
  2. 公式43到44有一个期望服从分布的变化(期望的下标),这是使用了马尔科夫假设
  3. 公式44到45期望的下标又一次发生了变化,这是由于KL散度本身就是计算一个期望,更加具体的推导可以把公式45按照定义完全展开就可以得到44的形式

上述ELBO的衍生形式可以根据各个组成部分进行解释:

  1. E q ( x 1 ∣ x 0 ) [ l o g    p θ ( x 0 ∣ x 1 ) ] \mathbb E_{q(x_1 \mid x_0)}[log \; p_\theta(x_0 \mid x_1)] Eq(x1x0)[logpθ(x0x1)]可以被解释为重建项(reconstruction term),该项预测了给定第一步潜在变量 x 1 x_1 x1的原始数据样本 x 0 x_0 x0的对数概率。这一项也出现在普通的VAE中,可以进行类似的训练
  2. E q ( x T − 1 ∣ x 0 ) [ D K L ( q ( x T ∣ x T − 1 ) ∣ ∣ p ( x T ) ) ] \mathbb E_{q(x_{T-1} \mid x_0)[D_{KL}(q(x_T \mid x_{T-1}) || p(x_T))]} Eq(xT1x0)[DKL(q(xTxT1)p(xT))]是先验匹配项(prior matching term);当最终的潜变量分布于高斯先验匹配时被最小化。这一项不需要进行优化,因为没有可训练的参数;并且由于我们假设了一个足够大的T,使得潜变量的最终分布为高斯分布,因此这个项实际上会变成0
  3. E q ( x t − 1 , x t + 1 ∣ x 0 ) [ D K L ( q ( x t ∣ x t − 1 ) ∣ ∣ p θ ( x t ∣ x t + 1 ) ) ] \mathbb E_{q(x_{t-1,x_{t+1}} \mid x_0)}[D_{KL}(q(x_t \mid x_{t-1}) || p_\theta(x_t \mid x_{t+1}))] Eq(xt1,xt+1x0)[DKL(q(xtxt1)pθ(xtxt+1))]是一致性项(consistency term);这项使得 x t x_t xt在前向和后向时分布保持一致。对于每个中间时间步,来自噪声较大的图像的去噪步骤应该与来自较干净图像的响应的加噪步骤相匹配,这在数学上由KL散度反映。当我们训练 p θ ( x t ∣ x t + 1 ) p_\theta(x_t \mid x_{t+1}) pθ(xtxt+1)去匹配高斯分布 q ( x t ∣ x t − 1 ) q(x_t \mid x_{t-1}) q(xtxt1)时,此项会最小化。

[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-PbjeUAt2-1669026624879)(D:\研究生资料\扩散模型\1.文章笔记\img\image-20221111145810710.png)]

图4:根据我们的推导,可以通过确保每个中间 x t x_t xt,使其后验概率 p θ ( x t ∣ x t + 1 ) p_\theta(x_t \mid x_{t+1}) pθ(xtxt+1)匹配其前面的 q ( x t ∣ x t − 1 ) q(x_t \mid x_{t-1}) q(xtxt1)的高斯衰减来优化VDM。在上面的图中,对于每个中间 x t x_t xt,我们将粉色和绿色箭头表示的分布之间的差异最小化。

优化VDM的成本主要由第三项决定,因为我们必须在所有时间步t上进行优化。

根据这种推导,ELBO的所有项都作为期望值计算,因此可以使用蒙特卡罗估计进行近似。然而实际使用我们推导出的项去优化ELBO可能是次优的;因为一致性项(公式45中的第三项)被计算为每个时间步的两个随机变量 { x t − 1 , x t + 1 } \{x_{t-1},x_{t+1}\} {xt1,xt+1}的期望,其蒙特卡罗估计的方差可能会高于每一个时间步仅使用一个随机变量估计的项。因为要对T-1个一致性项进行求和计算,对于较大的T值,ELBO的最终估计值可能具有高方差。

作为代替,让我们尝试推导出ELBO的一种形式,其中的每个项都能被计算为一次仅对一个随机变量的期望值。关键的步骤就是我们可以改写encoder转换的步骤: q ( x t ∣ x t − 1 ) = q ( x t ∣ x t − 1 , x 0 ) ) q(x_t \mid x_{t-1}) = q(x_t \mid x_{t-1},x_0)) q(xtxt1)=q(xtxt1,x0)),其中由于马尔科夫特性,额外的条件项是多余的(这句话的意思是多出来的 x 0 x_0 x0项)。然后根据贝叶斯规则,我们可以将每一步的encoder转换步骤写成:
q ( x t ∣ x t − 1 , x 0 ) = q ( x t − 1 ∣ x t , x 0 ) q ( x t ∣ x 0 ) q ( x t − 1 ∣ x 0 ) ( 46 ) q(x_t \mid x_{t-1},x_0) = \frac{q(x_{t-1} \mid x_t,x_0)q(x_t \mid x_0)}{q(x_{t-1} \mid x_0)} \qquad (46) q(xtxt1,x0)=q(xt1x0)q(xt1xt,x0)q(xtx0)(46)
有了公式46,我们可以尝试从公式37中重新进行ELBO项的推导:

在这里插入图片描述

在这里插入图片描述

上面的这个公式推导也是不难的,需要注意的点就是:

  1. 请务必耐着性子跟着公式一步一步推,不难的
  2. 公式48到公式49的连乘项的展开同公式38-39有所不同
  3. 公式52是使用了联合概率的贝叶斯公式
  4. 公式53的消去化简,可以将连乘符号展开,然后会有对项相消,对项相消之后就剩了公式54中的项
  • TODO:这部分的含义不是很懂

我们成功推导了ELBO的解释项,该解释可以用较低的方差进行估计,因为公式58中的每一项都是每次最多一个随机变量的期望值。公式58中的每一项也都有一个直观的解释:

  1. 重建项就不用解释了,可以看一看普通VAE的重建项;
  2. D K L ( q ( x T ∣ x 0 ) ∣ ∣ p ( x T ) ) D_{KL}(q(x_T \mid x_0) || p(x_T)) DKL(q(xTx0)p(xT))表示最终噪声化输入的分布于标准高斯先验的接近程度。它没有可训练的参数,并且在我们的假设情况下也等于零;
  3. E q ( x t ∣ x 0 ) [ D K L ( q ( x t − 1 ∣ x t , x 0 ) ∥ p θ ( x t − 1 ∣ x t ) ) ] \mathbb{E}_{q\left(\boldsymbol{x}_t \mid \boldsymbol{x}_0\right)}\left[D_{\mathrm{KL}}\left(q\left(\boldsymbol{x}_{t-1} \mid \boldsymbol{x}_t, \boldsymbol{x}_0\right) \| p_{\boldsymbol{\theta}}\left(\boldsymbol{x}_{t-1} \mid \boldsymbol{x}_t\right)\right)\right] Eq(xtx0)[DKL(q(xt1xt,x0)pθ(xt1xt))]是去噪匹配项(denoising matching term)。我们期望去噪转化步骤 p θ ( x t − 1 ∣ x t ) p_\theta(x_{t-1} \mid x_t) pθ(xt1xt)作为一种对于ground truth去噪步骤 q ( x t − 1 ∣ x t , x 0 ) q(x_{t-1} \mid x_t,x_0) q(xt1xt,x0)可控制的近似。(这里我们要注意,从图4中我们可以看出 q ( x t ∣ x t − 1 ) q(x_t \mid x_{t-1}) q(xtxt1)是我们加入高斯噪声的步骤,那么 q ( x t − 1 ∣ x t ) q(x_{t-1} \mid x_t) q(xt1xt)就是去噪声的步骤,并且是去噪步骤的ground truth。 p θ ( x t − 1 ∣ x t ) p_\theta(x_{t-1} \mid x_t) pθ(xt1xt)是我们要学习的去噪模型(decoder))。因此当这两个去噪步骤尽可能匹配的时候,KL散度项被最小化。

作为旁注,我们主要到在ELBO的推导过程中只使用了马尔科夫假设,因此这些公式将适用于任何任意的MHVAE中。此外,当我们设置T=1时,VDM的两个ELBO推导都能够清楚的重建出原始VAE的ELBO公式,即公式19.

在ELBO的第二种推导中,优化成本的主要部分还是集中于求和项,它主导着重建项。由于同时学习encoder增加的复杂性,第三项中的KL散度很难在任意复杂的马尔科夫HVAE中对任意后验进行最小化。在VDM中,我们可以利用高斯变换假设使得优化变得容易处理。
q ( x t − 1 ∣ x t , x 0 ) = q ( x t ∣ x t − 1 , x 0 ) q ( x t − 1 ∣ x 0 ) q ( x t ∣ x 0 ) q(x_{t-1} \mid x_t,x_0) = \frac{q(x_t \mid x_{t-1},x_0)q(x_{t-1}\mid x_0)}{q(x_t \mid x_0)} q(xt1xt,x0)=q(xtx0)q(xtxt1,x0)q(xt1x0)
并且根据encoder加入高斯噪声的步骤中(公式31)我们已经知道了 q ( x t ∣ x t − 1 , x 0 ) = q ( x t ∣ x t − 1 ) = N ( x t ; α t x t − 1 , ( 1 − α t ) I ) q(x_t \mid x_{t-1},x_0) = q(x_t \mid x_{t-1}) = \mathcal N(x_t;\sqrt{\alpha_t}x_{t-1},(1-\alpha_t)I) q(xtxt1,x0)=q(xtxt1)=N(xt;αt xt1,(1αt)I),剩下的就是 q ( x t ∣ x 0 ) , q ( x t − 1 ∣ x 0 ) q(x_t \mid x_0),q(x_{t-1} \mid x_0) q(xtx0),q(xt1x0)的公式推导。(这里注意, q ( x t ∣ x t − 1 ) q(x_t \mid x_{t-1}) q(xtxt1)是我们加入高斯噪声的步骤,这个步骤是符合马尔科夫假设的,但是反过来就不行了,这也就是我们需要利用贝叶斯公式去优化 q ( x t − 1 ∣ x t , x 0 ) q(x_{t-1} \mid x_t,x_0) q(xt1xt,x0)的原因)。幸运的是,通过利用VDM的encoder转换是线性高斯模型这一事实,上面的两个项变得易于处理起来。回忆起重参数化技巧,采样 x t ∼ q ( x t ∣ x t − 1 ) x_t \sim q(x_t \mid x_{t-1}) xtq(xtxt1)的步骤可以重写为(逐步加入高斯噪声的过程):
x t = α t x t − 1 + 1 − α t ϵ w i t h    ϵ ∼ N ( ϵ ; 0 , I ) ( 59 ) x_t = \sqrt{\alpha_t} x_{t-1} + \sqrt{1-\alpha_t} \epsilon \quad with \; \epsilon \sim \mathcal N(\epsilon;0,I) \qquad (59) \\ xt=αt xt1+1αt ϵwithϵN(ϵ;0,I)(59)
然后, q ( x t ∣ x 0 ) q(x_t \mid x_0) q(xtx0)的形式就可以通过重复应用重参数化的技巧递归推导出来。假设我们可以随机得到2T个随机噪声变量 { ϵ t ∗ , ϵ t } t = 0 T ∼ N ( ϵ ; 0 , I ) \{\epsilon^*_t,\epsilon_t\}_{t=0}^T \sim \mathcal N(\epsilon;0,I) {ϵt,ϵt}t=0TN(ϵ;0,I)。然后,对于任意采样 x t ∼ q ( x t ∣ x 0 ) x_t \sim q(x_t \mid x_0) xtq(xtx0),我们可以重写成下面的形式:

在这里插入图片描述

公式64中使用了高斯分布求和的一个定理

在这里插入图片描述

根据上面的推导,我们推导出了 q ( x t ∣ x 0 ) q(x_t \mid x_0) q(xtx0)的高斯形式。这个推导过程可以被修改以扩展到 q ( x t − 1 ∣ x 0 ) q(x_{t-1} \mid x_0) q(xt1x0)。现在我们知道了这两个项的形式,我们就可以通过带入贝叶斯规则展开来计算 q ( x t − 1 ∣ x t , x 0 ) q(x_{t-1} \mid x_t,x_0) q(xt1xt,x0)的形式。下面的公式我没有仔细推了:

在这里插入图片描述

在这里插入图片描述

在公式75中, C ( x t , x 0 ) C(x_t,x_0) C(xt,x0)相对于 x t − 1 x_{t-1} xt1l来说是一个常数;该项在公式84中隐式的返回以完成平方项。

因此我们展示出在每个步骤中, x t − 1 ∼ q ( x t − 1 ∣ x t , x 0 ) x_{t-1} \sim q(x_{t-1} \mid x_t,x_0) xt1q(xt1xt,x0)是正态分布的,其均值 μ q ( x t , x 0 ) \mu_q(x_t,x_0) μq(xt,x0) x t , x 0 x_t,x_0 xt,x0的函数,其方差 ∑ q ( t ) \sum_q(t) q(t)是系数 α \alpha α的函数。这些 α \alpha α系数是已知的,并且在每个时间步中都是固定的;它们要么是永久固定的(被建模为超参数),要么被视为对他们进行建模的网络的当前的推理输出。根据公式84,我们可以重写我们的方差公式为 ∑ q ( t ) = σ q 2 ( t ) I \sum_q(t) = \sigma_q^2(t)I q(t)=σq2(t)I,其中:
σ q 2 ( t ) = ( 1 − α t ) ( 1 − α ^ t − 1 ) 1 − α ^ t ( 85 ) \sigma_q^2(t) = \frac{(1-\alpha_t)(1-\hat{\alpha}_{t-1})}{1-\hat{\alpha}_t} \qquad (85) σq2(t)=1α^t(1αt)(1α^t1)(85)
为了让两个去噪步骤 p θ ( x t − 1 ∣ x t ) 同 q ( x t − 1 ∣ x t , x 0 ) p_\theta(x_{t-1} \mid x_t)同q(x_{t-1} \mid x_t,x_0) pθ(xt1xt)q(xt1xt,x0)进行匹配,我们也可以将其建模为高斯分布。此外,由于已知所有项在每个时间步中都被冻结,我们可以立即构造 p θ ( x t − 1 ∣ x t ) p_\theta(x_{t-1} \mid x_t) pθ(xt1xt)每步中的方差为 ∑ q ( t ) = σ q 2 ( t ) I \sum_q(t) = \sigma^2_q(t)I q(t)=σq2(t)I。我们必须将均值参数化 μ θ ( x t , t ) \mu_\theta(x_t,t) μθ(xt,t) x t x_t xt的参数,因为 p θ ( x t − 1 ∣ x t ) p_\theta(x_{t-1} \mid x_t) pθ(xt1xt)不以 x 0 x_0 x0为条件。

回忆两个高斯分布之间的KL散度为:

在这里插入图片描述

在我们的情况下,我们可以将两个高斯分布的方差是精确匹配上的,因此优化KL散度项就可以变成最小化两个分布之间的均值之间的差异:

在这里插入图片描述

为了简洁起见,我们将 μ q ( x t , x 0 ) \mu_q(x_t,x_0) μq(xt,x0)简写为 μ q \mu_q μq μ θ ( x t , t ) \mu_\theta(x_t,t) μθ(xt,t)简写为 μ θ \mu_\theta μθ。换句话说,我们希望去优化 μ θ ( x t , t ) \mu_\theta(x_t,t) μθ(xt,t)去匹配 μ q ( x t , x 0 ) \mu_q(x_t,x_0) μq(xt,x0),从公式84的推导中,我们可以得到如下的形式:

在这里插入图片描述

因为 μ θ ( x t , t ) \mu_\theta(x_t,t) μθ(xt,t)是同样以 x t x_t xt为条件,我们可以通过将 μ q ( x t , x 0 ) \mu_q(x_t,x_0) μq(xt,x0)设置为如下的形式去匹配该项:

在这里插入图片描述

其中 x ^ θ ( x t , t ) \hat x_\theta(x_t,t) x^θ(xt,t)由神经网络参数化用于从噪声图像 x t x_t xt中预测 x 0 x_0 x0。然后,这个优化问题被简化为:

在这里插入图片描述

因此,优化VDM就变成了训练一个神经网络,用于从任意噪声图像中去预测原始无修改的真实图像。此外,可以通过最小化所有时间步上的期望值来近似最小化我们推导出的ELBO(公式58)在所有噪声水平上的求和项:

在这里插入图片描述

然后可以使用随时间步变化的随机样本来对其进行优化。

第3节整理

整个第3节公式推导的逻辑如下所示:

1️⃣ 推导VDM的ELBO形式
l o g    p ( x ) ≥ E L B O = E q ( x 1 ∣ x 0 ) [ l o g    p θ ( x 0 ∣ x 1 ) ] ⏟ a . − E q ( x T − 1 ∣ x 0 ) [ D K L ] ( q ( x T ∣ x T − 1 ) ∣ ∣ p ( x T ) ) ⏟ b . − ∑ t = 1 T E q ( x t − 1 , x t + 1 ∣ x 0 ) [ D K L ( q ( x t ∣ x t − 1 ) ∣ ∣ p θ ( x t ∣ x t + 1 ) ) ] ⏟ c . \begin{aligned} log \; p(x) \ge ELBO &= \underbrace {\mathbb E_{q(x_1 \mid x_0)}[log \; p_\theta(x_0 \mid x_1)]} _{a.} - \underbrace{\mathbb E_{q(x_{T-1\mid x_0})}[D_{KL}](q(x_T\mid x_{T-1}) || p(x_T))}_{b.} \\ & - \underbrace{\sum_{t=1}^T \mathbb E_{q(x_{t-1},x_{t+1} \mid x_0)}[D_{KL}(q(x_t \mid x_{t-1}) || p_\theta(x_t \mid x_{t+1}))]}_{c.} \end{aligned} logp(x)ELBO=a. Eq(x1x0)[logpθ(x0x1)]b. Eq(xT1x0)[DKL](q(xTxT1)p(xT))c. t=1TEq(xt1,xt+1x0)[DKL(q(xtxt1)pθ(xtxt+1))]

2️⃣ 其中1️⃣中ELBO形式中的c.项,因为需要同时优化 { x t − 1 , x t + 1 } \{x_{t-1},x_{t+1}\} {xt1,xt+1}的期望,会导致蒙特卡罗估计的高方差,因此我们需要尝试推导ELBO项只对单一变量进行优化的形式

重写ELBO项,有:
E L B O = E q ( x 1 ∣ x 0 ) [ l o g    p θ ( x 0 ∣ x 1 ) ] ⏟ a . 重 建 项 − D K L [ q ( x T ∣ x 0 ) ∣ ∣ p ( x T ) ] ⏟ b . 先 验 匹 配 项 − ∑ t = 2 T E q ( x t ∣ x 0 ) [ D K L [ q ( x t − 1 ∣ x t , x 0 ) ∣ ∣ p θ ( x t − 1 ∣ x t ) ] ] ⏟ c . 去 噪 匹 配 项 ( 58 ) \begin{aligned} ELBO &= \underbrace{\mathbb E_{q(x_1 \mid x_0)}[log \; p_\theta(x_0 \mid x_1)]}_{a.重建项} - \underbrace{D_{KL}[q(x_T \mid x_0) || p(x_T)]}_{b.先验匹配项} \\ & - \underbrace{\sum_{t=2}^T \mathbb E_{q(x_t \mid x_0)}[D_{KL}[q(x_{t-1} \mid x_t,x_0) || p_\theta(x_{t-1} \mid x_t)]]}_{c.去噪匹配项} && (58) \end{aligned} ELBO=a. Eq(x1x0)[logpθ(x0x1)]b. DKL[q(xTx0)p(xT)]c. t=2TEq(xtx0)[DKL[q(xt1xt,x0)pθ(xt1xt)]](58)
针对公式(58),a.重建项好优化,b.匹配项根据我们的假设,最终会变成0;c.去噪匹配项代表的含义是我们建模的decoder p θ ( x t − 1 ∣ x t ) p_\theta(x_{t-1} \mid x_t) pθ(xt1xt)期望同真实的去噪步骤保持一致。

但是由于c项中的 q ( x t − 1 ∣ x t , x 0 ) q(x_{t-1} \mid x_t,x_0) q(xt1xt,x0)这个分布难以处理(并不符合马尔科夫假设),因此我们需要对其进行优化处理,使得变得容易处理。

3️⃣ 这一步的主要工作就是对第2️⃣步骤中的 q ( x t − 1 ∣ x t , x 0 ) q(x_{t-1} \mid x_t,x_0) q(xt1xt,x0)利用重参数化技巧进行处理,最终得到其符合正态分布的结论

首先使用贝叶斯公式有:
q ( x t − 1 ∣ x t , x 0 ) = q ( x t ∣ x t − 1 , x 0 ) q ( x t − 1 ∣ x 0 ) q ( x t ∣ x 0 ) q(x_{t-1} \mid x_t,x_0) = \frac{q(x_t \mid x_{t-1},x_0)q(x_{t-1}\mid x_0)}{q(x_t \mid x_0)} q(xt1xt,x0)=q(xtx0)q(xtxt1,x0)q(xt1x0)
其中, q ( x t ∣ x t − 1 , x 0 ) = q ( x t ∣ x t − 1 ) = N ( x t ; α t x t − 1 , ( 1 − α t ) I ) q(x_t \mid x_{t-1},x_0) = q(x_t \mid x_{t-1}) = \mathcal N(x_t;\sqrt{\alpha_t}x_{t-1},(1-\alpha_t)I) q(xtxt1,x0)=q(xtxt1)=N(xt;αt xt1,(1αt)I)

然后,我们假设我们可以随机得到2T个随机噪声变量 q ( x t ∣ x t − 1 , x 0 ) = q ( x t ∣ x t − 1 ) = N ( x t ; α t x t − 1 , ( 1 − α t ) I ) q(x_t \mid x_{t-1},x_0) = q(x_t \mid x_{t-1}) = \mathcal N(x_t;\sqrt{\alpha_t}x_{t-1},(1-\alpha_t)I) q(xtxt1,x0)=q(xtxt1)=N(xt;αt xt1,(1αt)I)。然后对于任意 x t ∼ q ( x t ∣ x 0 ) x_t \sim q(x_t \mid x_0) xtq(xtx0),我们可以使用重参数化技巧得到:
x t ∼ q ( x t ∣ x 0 ) = N ( x t ; α ˉ t x 0 , ( 1 − α ˉ t ) I ) ( 70 ) \begin{aligned} x_t \sim q(x_t \mid x_0) & = \mathcal N(x_t;\sqrt{\bar \alpha_t}x_0,(1-\bar \alpha_t)I )&& (70) \end{aligned} xtq(xtx0)=N(xt;αˉt x0,(1αˉt)I)(70)
其中, α ˉ t = ∏ t = 1 T α t \bar \alpha_t = \prod_{t=1}^T \alpha_t αˉt=t=1Tαt,同理 q ( x t − 1 ∣ x 0 ) q(x_{t-1} \mid x_0) q(xt1x0)也可以从公式(70)得到,这样我们就可以重新整理 q ( x t − 1 ∣ x t , x 0 ) q(x_{t-1} \mid x_t,x_0) q(xt1xt,x0)的形式:
x t − 1 ∼ q ( x t − 1 ∣ x t , x 0 ) = q ( x t ∣ x t − 1 , x 0 ) q ( x t − 1 ∣ x 0 ) q ( x t ∣ x 0 ) = N ( x t ; α t x t − 1 , ( 1 − α t ) I ) N ( x t − 1 ; α ˉ t − 1 x 0 , ( 1 − α ˉ t − 1 ) I ) N ( x t ; α ˉ t x 0 , ( 1 − α ˉ t ) I ) ∝ N ( x t − 1 ; α t ( 1 − α ˉ t − 1 ) x t + α ˉ t − 1 ( 1 − α t ) x 0 1 − α ˉ t ⏟ μ q ( x t , x 0 ) , ( 1 − α t ) ( 1 − α ˉ t − 1 ) 1 − α ˉ t I ⏟ Σ q ( t ) ( 84 ) \begin{aligned} x_{t-1} \sim q(x_{t-1} \mid x_t,x_0) &= \frac{q(x_t \mid x_{t-1},x_0)q(x_{t-1}\mid x_0)}{q(x_t \mid x_0)} \\ & = \frac{\mathcal N(x_t;\sqrt{\alpha_t}x_{t-1},(1-\alpha_t)I) \mathcal N(x_{t-1};\sqrt{\bar \alpha_{t-1}}x_0,(1-\bar \alpha_{t-1})I )}{\mathcal N(x_t;\sqrt{\bar \alpha_t}x_0,(1-\bar \alpha_t)I )} \\ & \propto \mathcal N(x_{t-1};\underbrace{\frac{\sqrt{\alpha_t} (1-\bar \alpha_{t-1})x_t + \sqrt{\bar \alpha_{t-1}} (1-\alpha_t)x_0}{1-\bar \alpha_t}}_{\mu_q(x_t,x_0)},\underbrace{\frac{(1-\alpha_t)(1-\bar \alpha_{t-1})}{1-\bar \alpha_t}I}_{\Sigma_q(t)} && (84) \end{aligned} xt1q(xt1xt,x0)=q(xtx0)q(xtxt1,x0)q(xt1x0)=N(xt;αˉt x0,(1αˉt)I)N(xt;αt xt1,(1αt)I)N(xt1;αˉt1 x0,(1αˉt1)I)N(xt1;μq(xt,x0) 1αˉtαt (1αˉt1)xt+αˉt1 (1αt)x0,Σq(t) 1αˉt(1αt)(1αˉt1)I(84)
从上面的推导中可以看出, q ( x t − 1 ∣ x t , x 0 ) q(x_{t-1} \mid x_t,x_0) q(xt1xt,x0)服从正态分布,其中均值 μ q ( x t , x 0 ) \mu_q(x_t,x_0) μq(xt,x0) x t , x 0 x_t,x_0 xt,x0的函数,方差 Σ q ( t ) \Sigma_q(t) Σq(t) α \alpha α的函数,是固定已知的。

4️⃣ 将公式(84)带入公式(58)中的c项,并使用高斯分布的KL散度公式进行化简,得到我们要优化的目标函数
KaTeX parse error: Got function '\min' with no arguments as argument to '\underset' at position 40: …nderset \theta \̲m̲i̲n̲ ̲D_{KL}(q(x_{t-1…
其中, μ θ = μ θ ( x t , t ) , μ q = μ q ( x t , x 0 ) \mu_\theta = \mu_\theta(x_t,t),\mu_q = \mu_q(x_t,x_0) μθ=μθ(xt,t),μq=μq(xt,x0),从公式(92)中我们看出,我们希望优化 μ θ \mu_\theta μθ去匹配 μ q \mu_q μq,因此我们可以将 μ θ \mu_\theta μθ设置成下面的形式来进行匹配:
μ θ ( x t , t ) = α t ( 1 − α ˉ t − 1 ) x t + α ˉ t − 1 ( 1 − α t ) x ^ θ ( x t , t ) 1 − α ˉ t ( 94 ) \mu_\theta(x_t,t) =\frac{\sqrt{\alpha_t} (1-\bar \alpha_{t-1})x_t + \sqrt{\bar \alpha_{t-1}} (1-\alpha_t)\hat x_\theta(x_t,t)}{1-\bar \alpha_t} \quad (94) μθ(xt,t)=1αˉtαt (1αˉt1)xt+αˉt1 (1αt)x^θ(xt,t)(94)
其中, x ^ θ ( x t , t ) \hat x_\theta(x_t,t) x^θ(xt,t)由神经网络参数化用于从噪声图像 x t x_t xt中预测 x 0 x_0 x0

5️⃣ 将公式(94)带入公式(92)并进行化简得到我们最终的优化目标
在这里插入图片描述

因此,优化VDM就变成了训练一个神经网络,用于从任意噪声图像中去预测原始无修改的真实图像。此外,可以通过最小化所有时间步上的期望值来近似最小化我们推导出的ELBO(公式58)在所有噪声水平上的求和项:
a r g m θ i n E t ∼ U { 2 , T } [ E q ( x t ∣ x 0 ) [ D K L [ ( q ( x t − 1 ∣ x t , x 0 ) ∣ ∣ p θ ( x t − 1 ∣ x t ) ) ] ] ( 100 ) arg \underset \theta min \mathbb E_{t \sim U\{2,T\}}[\mathbb E_{q(x_t \mid x_0)}[D_{KL}[(q(x_{t-1} \mid x_t,x_0)||p_\theta(x_{t-1} \mid x_t))]] \quad (100) argθminEtU{2,T}[Eq(xtx0)[DKL[(q(xt1xt,x0)pθ(xt1xt))]](100)


4. Learning Diffusion Noise Parameters

让我们研究如何联合学习VDM的噪声参数。一个可选的方法是一个神经网络模型 α ^ η ( t ) \hat \alpha_\eta(t) α^η(t)去建模 α t \alpha_t αt。然而,这是低效的,因为必须要在每个时间步t中执行多次推理才能计算 α ˉ t \bar \alpha_t αˉt。尽管使用缓存能够减轻这种计算成本,但是我们也可以推导出另一种方法来学习扩散噪声的参数。通过将公式85中的方差方程带入公式99中推导出的每个时间步的目标,我们可以简化公式如下:

在这里插入图片描述

回忆公式70, q ( x t ∣ x 0 ) q(x_t \mid x_0) q(xtx0)是高斯形式 N ( x t ; α ˉ t x 0 , ( 1 − α ˉ t ) I ) \mathcal N(x_t;\sqrt{\bar \alpha_t} x_0,(1-\bar \alpha_t)I) N(xt;αˉt x0,(1αˉt)I)。然后,根据SNR(信噪比)的定义 S N R = μ 2 σ 2 SNR = \frac{\mu^2}{\sigma^2} SNR=σ2μ2,我们可以写出每个时间步的SNR,如下:
S N R ( t ) = α ˉ t 1 − α ˉ t SNR(t) = \frac{\bar \alpha_t}{1 - \bar \alpha_t} SNR(t)=1αˉtαˉt

然后,我们的推导的公式108和公式99就可以简化为下面的形式:

在这里插入图片描述

顾名思义,SNR表示原始信号与存在的噪声量之间的比率;SNR越高表示信号越多,SNR越低表示噪声越多。在扩散模型中,我们要求SNR要随着时间步的t的增加而单调减少;这使得扰动输入 x t x_t xt变得越来越噪声化这个概念正式化起来,直到最终 x t x_t xt变得与标准高斯噪声相同。

在公式110中的目标简化之后,我们可以使用神经网络在每个时间步中直接参数化SNR,并同扩散模型一起优化它。由于SNR必须随着时间单调下降,
S N R ( t ) = e x p ( − ω η ( t ) ) ( 111 ) SNR(t) = exp(-\omega_\eta(t)) \qquad (111) SNR(t)=exp(ωη(t))(111)

其中 ω η ( t ) \omega_\eta(t) ωη(t)被建模为参数为 η \eta η的单调递增神经网络。不管 ω η \omega_\eta ωη是一个单调递减的函数,而SNR的指数项将强迫最终的结果项为正值。注意到,公式100中的目标现在也必须要通过参数 η \eta η来进行优化。通过将我们在公式111中的SNR参数化与我们在公式109中对SNR的定义相结合,我们可以显式地推导出如下的公式:
α ˉ t 1 − α ˉ t = e x p ( − ω η ( t ) ) ( 112 ) ∴ α ˉ t = s i g m o i d ( − ω η ( t ) ) ( 113 ) ∴ 1 − α ˉ t = s i g m o i d ( ω η ( t ) ) ( 114 ) \begin{aligned} \frac{\bar \alpha_t}{1 - \bar \alpha_t} = exp(-\omega_\eta(t)) && (112) \\ ∴ \bar \alpha_t = sigmoid(-\omega_\eta(t)) && (113) \\ ∴ 1- \bar \alpha_t = sigmoid(\omega_\eta(t)) && (114) \end{aligned} 1αˉtαˉt=exp(ωη(t))αˉt=sigmoid(ωη(t))1αˉt=sigmoid(ωη(t))(112)(113)(114)
这些技巧项对于各种计算是非常必要的;例如在优化期间,它们被用于利用重参数技巧来从输入 x 0 x_0 x0创建任意的噪声 x t x_t xt,如等式69所示。

第4节整理

  • TODO:个人理解待更正

首先我们回顾一下我们要优化的最终目标函数,即公式(99):
a r g m θ i n D K L ( q ( x t − 1 ∣ x t , x 0 ) ∣ ∣ p θ ( x t − 1 ∣ x t ) ) = a r g m θ i n D K L ( N ( x t − 1 ; μ q , Σ q ( t ) ) ∣ ∣ N ( x t − 1 ; μ θ , Σ q ( t ) ) ) = a r g m θ i n 1 2 σ q 2 ( t ) α ˉ t − 1 ( 1 − α t ) 2 ( 1 − α ˉ t ) 2 [ ∣ ∣ x ^ θ ( x t , t ) − x 0 ∣ ∣ 2 2 ] ( 99 ) σ q 2 ( t ) = ( 1 − α t ) ( 1 − α ˉ t − 1 ) 1 − α ˉ t ( 65 ) \begin{aligned} arg \underset \theta min D_{KL}(q(x_{t-1} \mid x_t,x_0) || p_\theta(x_{t-1} \mid x_t)) &= arg \underset \theta min D_{KL}(\mathcal N(x_{t-1};\mu_q,\Sigma_q(t))||\mathcal N(x_{t-1};\mu_\theta,\Sigma_q(t))) \\ & = arg \underset \theta min \frac{1}{2\sigma_q^2(t)} \frac{\bar \alpha_{t-1}(1-\alpha_t)^2}{(1-\bar\alpha_t)^2}[||\hat x_\theta(x_t,t)-x_0||^2_2] \quad (99) \\ \sigma_q^2(t) = \frac{(1-\alpha_t)(1-\bar \alpha_{t-1})}{1- \bar \alpha_t} \quad (65) \end{aligned} argθminDKL(q(xt1xt,x0)pθ(xt1xt))σq2(t)=1αˉt(1αt)(1αˉt1)(65)=argθminDKL(N(xt1;μq,Σq(t))N(xt1;μθ,Σq(t)))=argθmin2σq2(t)1(1αˉt)2αˉt1(1αt)2[x^θ(xt,t)x022](99)
从公式(99)中可以看出来,我们需要训练一个模型 x ^ θ ( x t , t ) \hat x_\theta(x_t,t) x^θ(xt,t)用于从噪声图像中预测出 x 0 x_0 x0。另外我们还需要有 α t \alpha_t αt来满足我们的假设:即不断向原始图像中加入高斯噪声,在T步之后变成标准高斯噪声。

因此这一节的核心idea是怎么去得到这个系数 α t \alpha_t αt

用一个神经网络去建模这个 α t \alpha_t αt存在一定的弊端,因此文章从信噪比SNR出发,最终推导出一个简洁的形式来进行参数化建模。(在扩散模型的训练过程中, x ^ θ , α t \hat x_\theta,\alpha_t x^θ,αt的是要被建模训练的。)

将公式(65)带入(99)中得到:
1 2 σ q 2 ( t ) α ˉ t − 1 ( 1 − α t ) 2 ( 1 − α ˉ t ) 2 [ ∣ ∣ x ^ θ ( x t , t ) − x 0 ∣ ∣ 2 2 ] = 1 2 ( α ˉ t − 1 1 − α ˉ t − 1 − α ˉ t 1 − α ˉ t ) [ ∣ ∣ x ^ θ ( x t , t ) − x 0 ∣ ∣ 2 2 ] ( 108 ) \frac{1}{2\sigma_q^2(t)} \frac{\bar \alpha_{t-1}(1-\alpha_t)^2}{(1-\bar\alpha_t)^2}[||\hat x_\theta(x_t,t)-x_0||^2_2] = \frac{1}{2}(\frac{\bar \alpha_{t-1}}{1-\bar \alpha_{t-1}}-\frac{\bar \alpha_t}{1-\bar \alpha_t}) [||\hat x_\theta(x_t,t)-x_0||^2_2] \quad (108) 2σq2(t)1(1αˉt)2αˉt1(1αt)2[x^θ(xt,t)x022]=21(1αˉt1αˉt11αˉtαˉt)[x^θ(xt,t)x022](108)
根据SNR的定义,有:
S N R ( t ) = α ˉ t 1 − α ˉ t SNR(t) = \frac{\bar \alpha_t}{1 - \bar \alpha_t} SNR(t)=1αˉtαˉt
最后公式(108)就可以化简为:

在这里插入图片描述

于是我们的目标就是,在训练扩散模型时使得信噪比SNR单调减小(噪声越来越多),因此我们可以直接参数化SNR,并在扩散模型训练的时候一起优化它。

后面的内容就直接看第4节的后半部分就好了。

  • 3
    点赞
  • 10
    收藏
    觉得还不错? 一键收藏
  • 0
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值