Unsupervised Representation Learning from Pre-trained Diffusion Probabilistic Models

论文链接：

https://arxiv.org/abs/2212.12990

代码链接：

https://github.com/ckczzj/PDAE

Abstract

扩散概率模型(DPM)已经显示出生成高质量图像样本的强大能力。最近，扩散自编码器(Diff-AE)被提出来探索通过自编码进行表示学习的DPM。他们的关键思想是联合训练一个编码器，用于从图像中发现有意义的表示，和一个条件DPM作为解码器，用于重建图像。

考虑到从无开始训练DPM需要很长时间，并且已有大量的预训练DPM，我们提出了一种将现有的预训练DPM适应解码器进行图像重建的通用方法——预训练DPM自动编码(pre-trained DPM AutoEncoding, PDAE)，该方法具有比Diff-AE更好的训练效率和性能。

具体来说，我们发现预训练的DPM无法从其潜在变量重建图像的原因是由于前向过程的信息丢失，导致其预测的后验均值与真实均值之间存在差距。从这个角度来看，分类器引导的抽样方法可以解释为计算一个额外的均值偏移来填补空白，重建样本中丢失的类信息。这意味着缺口对应于图像中丢失的信息，我们可以通过填充缺口来重建图像。从中获得灵感，我们使用一个可训练的模型来根据编码表示来预测均值偏移，并训练它来尽可能多地填充空白，这样，编码器就被迫从图像中学习尽可能多的信息来帮助填充。通过重用部分预训练的DPM网络，重新设计扩散损失的加权方案，PDAE可以有效地从图像中学习有意义的表示。

大量的实验证明了PDAE的有效性、高效性和灵活性。

1 Introduction

深度生成模型，如变分自编码器(VAE)[25,39]，生成对抗网络(GAN)[13]，自回归模型[50,48]，归一化流(NF)[38,23]和基于能量的模型(EM)[9,45]，在合成醒目的图像样本方面表现出卓越的能力。最近，另一种生成模型扩散概率模型(Diffusion probistic models, DPM)[43,14]得到了进一步发展，并因其稳定的训练过程和最先进的样本质量[8]而受到欢迎。虽然在实现中有大量的自由度，但本文讨论的DPM仅指DDPM[14]中提出的去噪方法训练的DPM。

基于生成建模的无监督表示学习是计算机视觉领域的热门课题。潜在变量生成模型，如GAN和VAE，是自然的候选模型，因为它们固有地涉及它们生成的数据的潜在表示。同样，DPM本身也是通过正向过程产生潜在变量。然而，这些潜在变量缺乏高级语义信息，因为它们只是一组空间损坏的图像。有鉴于此，扩散自编码器(Diff-AE)[36]探索了通过自编码进行表示学习的DPM。具体来说，他们使用编码器从图像中发现有意义的表示，并使用条件DPM作为解码器，通过将编码的表示作为输入条件进行图像重建。Diff-AE在图像重建方面与最先进的模型竞争，并能够完成各种下游任务。

遵循自编码器的范式，PDAE旨在使现有的预训练的DPM适应解码器进行图像重建并从中受益。一般情况下，由于正向过程的信息丢失，预训练的DPM在逆向过程中无法准确地从xt预测xt−1的后验均值，导致其预测的后验均值与真实的后验均值存在差距。这就是为什么他们无法从潜在变量( x t x_t xt)重建图像( x 0 x_0 x0)的原因。从这个角度来看，分类器引导的抽样方法[8]可以解释为通过用分类器梯度计算的额外项来移动预测的后测均值来填补空白，从而重建样本中丢失的类信息。从这种使用先验知识(类标签)来填充空白的方法中获得灵感，我们的目标是从空白中反向提取知识，即学习可以帮助填充空白的表示。鉴于此，我们采用了一种新的梯度估计器，根据编码表示来预测均值偏移，并训练它来尽可能多地填充空白，这样编码器就被迫从图像中学习尽可能多的信息来帮助填充。PDAE遵循这一原则，基于预训练的DPM构建自编码器。

此外，我们发现不同时间阶段的后验平均间隙包含不同层次的信息，因此我们重新设计了扩散损失的加权方案，以鼓励模型有效地学习丰富的表示。我们还重用了预先训练的DPM网络的一部分，以加速模型的收敛。基于预训练的DPM, PDAE完成表征学习所需的训练时间不到Diff-AE的一半，但仍优于Diff-AE。此外，PDAE还支持其他一些有趣的特性。

2 Background

2.1 Denoising Diffusion Probabilistic Models

DDPM[14]采用前向过程，从数据分布 q ( x 0 ) q(x_0) q(x0)开始，用固定方差表 { β t } t = 1 T {\lbraceβ_t\rbrace}^T_{t=1} {βt}t=1T定义的马尔可夫扩散核 q ( x t ∣ x t − 1 ) q(x_t|x_{t−1}) q(xt∣xt−1)依次将其破坏到 N ( 0 , I ) N (0, I) N(0,I)。该过程可以表示为:
q ( x t ∣ x t − 1 ) = N ( x t ; 1 − β t x t − 1 , β t I ) q ( x 1 : T ∣ x 0 ) = ∏ t = 1 T q ( x t ∣ x t − 1 ) (1) q\left(\boldsymbol{x}_t \mid \boldsymbol{x}_{t-1}\right)=\mathcal{N}\left(\boldsymbol{x}_t ; \sqrt{1-\beta_t} \boldsymbol{x}_{t-1}, \beta_t \mathbf{I}\right) \quad q\left(\boldsymbol{x}_{1: T} \mid \boldsymbol{x}_0\right)=\prod_{t=1}^T q\left(\boldsymbol{x}_t \mid \boldsymbol{x}_{t-1}\right) \tag{1} q(xt∣xt−1)=N(xt;1−βt xt−1,βtI)q(x1:T∣x0)=t=1∏Tq(xt∣xt−1)(1)
其中 { x t } t = 1 T {\lbrace{x_t}\rbrace}^T_{t=1} {xt}t=1T为dDDPM的潜在变量。根据正态分布随机变量和规则，对于任意t, q ( x t ∣ x 0 ) = N ( x t ; α ˉ t x 0 , ( 1 − α ˉ t ) I ) q(x_t|x_0) =N\left(\mathbf{x}_{t} ; \sqrt{\bar{\alpha}_{t}} \mathbf{x}_{0},\left(1-\bar{\alpha}_{t}\right) \mathbf{I}\right) q(xt∣x0)=N(xt;αˉt x0,(1−αˉt)I)，我们可以直接从 x 0 x_0 x0中对 x t x_t xt进行抽样，其中 α t = 1 − β t α_t=1−β_t αt=1−βt和 α t = ∏ s = 0 t α s α_t=\prod\limits_{s=0}\limits^{t}{α_s} αt=s=0∏tαs

反向(生成)过程被定义为另一个由θ参数化的马尔可夫链来描述相同但相反的过程，将任意高斯噪声去噪到干净的数据样本：
p θ ( x t − 1 ∣ x t ) = N ( x t − 1 ; μ θ ( x t , t ) , Σ θ ( x t , t ) ) p θ ( x 0 : T ) = p ( x T ) ∏ t = 1 T p θ ( x t − 1 ∣ x t ) (2) p_\theta\left(\boldsymbol{x}_{t-1} \mid \boldsymbol{x}_t\right)=\mathcal{N}\left(\boldsymbol{x}_{t-1} ; \boldsymbol{\mu}_\theta\left(\boldsymbol{x}_t, t\right), \boldsymbol{\Sigma}_\theta\left(\boldsymbol{x}_t, t\right)\right) \quad p_\theta\left(\boldsymbol{x}_{0: T}\right)=p\left(\boldsymbol{x}_T\right) \prod_{t=1}^T p_\theta\left(\boldsymbol{x}_{t-1} \mid \boldsymbol{x}_t\right) \tag{2} pθ(xt−1∣xt)=N(xt−1;μθ(xt,t),Σθ(xt,t))pθ(x0:T)=p(xT)t=1∏Tpθ(xt−1∣xt)(2)
其中 p ( x T ) = N ( x T ; 0 , I ) p(x_T)=\mathcal{N}(x_T;0,I) p(xT)=N(xT;0,I)；它采用高斯形式的 p θ ( x t − 1 ∣ x t ) p_θ(x_{t−1}|x_t) pθ(xt−1∣xt)，因为当 β t β_t βt较小时，扩散过程的反转与正扩散过程具有相同的函数形式[11,43]。生成分布可以表示为 p θ ( x 0 ) = ∫ p θ ( x 0 : T ) d x 1 : T p_\theta\left(\boldsymbol{x}_0\right)=\int p_\theta\left(\boldsymbol{x}_{0: T}\right) d \boldsymbol{x}_{1: T} pθ(x0)=∫pθ(x0:T)dx1:T。

训练是通过最小化负的变分上界来最大化模型对数似然 ∫ q ( x 0 ) log ⁡ p θ ( x 0 ) d x 0 \int q\left(\boldsymbol{x}_0\right) \log p_\theta\left(\boldsymbol{x}_0\right) d \boldsymbol{x}_0 ∫q(x0)logpθ(x0)dx0$。最终目标是通过一些参数化和简化[14]得到的：
L simple ( θ ) = E x 0 , t , ϵ [ ∥ ϵ − ϵ θ ( α ˉ t x 0 + 1 − α ˉ t ϵ , t ) ∥ 2 ] (3) \mathcal{L}_{\text {simple }}(\theta)=\mathbb{E}_{x_0, t, \epsilon}\left[\left\|\epsilon-\epsilon_\theta\left(\sqrt{\bar{\alpha}_t} x_0+\sqrt{1-\bar{\alpha}_t} \epsilon, t\right)\right\|^2\right] \tag{3} Lsimple (θ)=Ex0,t,ϵ[ ϵ−ϵθ(αˉt x0+1−αˉt ϵ,t) 2](3)
ϵ θ \epsilon_{\theta} ϵθ是一个函数逼近器从 x t x_t xt来预测 ϵ \epsilon ϵ。

2.2 Denoising Diffusion Implicit Models

DDIM[44]定义了一个非马尔可夫正向过程，该过程导致与DDPM相同的训练目标，但相应的反向过程可以更灵活和更快地进行采样。具体来说，可以使用一些预训练DDPM的 ϵ θ \epsilon_{\theta} ϵθ从 x t x_t xt中采样 x t − 1 x_{t−1} xt−1，方法如下：
x t − 1 = α ˉ t − 1 ( x t − 1 − α ˉ t ⋅ ϵ θ ( x t , t ) α ˉ t ) + 1 − α ˉ t − 1 − σ t 2 ⋅ ϵ θ ( x t , t ) + σ t ϵ t (4) \boldsymbol{x}_{t-1}=\sqrt{\bar{\alpha}_{t-1}}\left(\frac{\boldsymbol{x}_t-\sqrt{1-\bar{\alpha}_t} \cdot \boldsymbol{\epsilon}_\theta\left(\boldsymbol{x}_t, t\right)}{\sqrt{\bar{\alpha}_t}}\right)+\sqrt{1-\bar{\alpha}_{t-1}-\sigma_t^2} \cdot \boldsymbol{\epsilon}_\theta\left(\boldsymbol{x}_t, t\right)+\sigma_t \epsilon_t \tag{4} xt−1=αˉt−1 (αˉt xt−1−αˉt ⋅ϵθ(xt,t))+1−αˉt−1−σt2 ⋅ϵθ(xt,t)+σtϵt(4)
其中 ϵ t ∼ N ( 0 , I ) \epsilon_{t}\sim \mathcal{N} (0, I) ϵt∼N(0,I)， σt控制正向过程的随机性。跨步大于1允许加速采样。当 σ t = 0 σ_t = 0 σt=0时，生成过程具有确定性，称为DDIM过程。

2.3 Classifier-guided Sampling Method

分类器引导的采样方法[44,46,8]表明，可以在噪声数据上训练分类器 p φ ( y ∣ x t ) p_φ(y|x_t) pφ(y∣xt)，并使用其梯度 ∇ x t l o g p φ ( y ∣ x t ) ∇x_t log p_φ(y|x_t) ∇xtlogpφ(y∣xt)引导一些预先训练的无条件DDPM朝着指定的类别y采样。条件反向过程可以近似为类似于Eq.(2)中无条件的高斯过程，但具有移位的均值：
p θ , ϕ ( x t − 1 ∣ x t , y ) ≈ N ( x t − 1 ; μ θ ( x t , t ) + Σ θ ( x t , t ) ⋅ ∇ x t log ⁡ p ϕ ( y ∣ x t ) , Σ θ ( x t , t ) ) (5) p_{\theta, \phi}\left(\boldsymbol{x}_{t-1} \mid \boldsymbol{x}_t, \boldsymbol{y}\right) \approx \mathcal{N}\left(\boldsymbol{x}_{t-1} ; \boldsymbol{\mu}_\theta\left(\boldsymbol{x}_t, t\right)+\boldsymbol{\Sigma}_\theta\left(\boldsymbol{x}_t, t\right) \cdot \nabla_{\boldsymbol{x}_t} \log p_\phi\left(\boldsymbol{y} \mid \boldsymbol{x}_t\right), \boldsymbol{\Sigma}_\theta\left(\boldsymbol{x}_t, t\right)\right) \tag{5} pθ,ϕ(xt−1∣xt,y)≈N(xt−1;μθ(xt,t)+Σθ(xt,t)⋅∇xtlogpϕ(y∣xt),Σθ(xt,t))(5)
对于像DDIM这样的确定性抽样方法，可以使用基于分数的条件调节技巧[46,45]来为条件抽样定义一个新的函数逼近器：
ϵ ^ θ ( x t , t ) = ϵ θ ( x t , t ) − 1 − α ˉ t ⋅ ∇ x t log ⁡ p ϕ ( y ∣ x t ) (6) \hat{\boldsymbol{\epsilon}}_\theta\left(\boldsymbol{x}_t, t\right)=\boldsymbol{\epsilon}_\theta\left(\boldsymbol{x}_t, t\right)-\sqrt{1-\bar{\alpha}_t} \cdot \nabla_{\boldsymbol{x}_t} \log p_\phi\left(\boldsymbol{y} \mid \boldsymbol{x}_t\right) \tag{6} ϵ^θ(xt,t)=ϵθ(xt,t)−1−αˉt ⋅∇xtlogpϕ(y∣xt)(6)
更一般地，任何有噪声数据与条件之间的相似估计器都可以应用于引导采样，例如noise - clip引导[33,31]。

3 Method

3.1 Forward Process Posterior Mean Gap

一般来说，我们可以通过分别学习 p θ ( x t − 1 ∣ x t ) = N ( x t − 1 ; µ θ ( x t , t ) , Σ θ ( x t , t ) ) p_θ(x_{t−1}|x_t) = \mathcal{N} (x_{t−1};µ_θ(x_t, t),Σ_θ(x_t, t)) pθ(xt−1∣xt)=N(xt−1;µθ(xt,t),Σθ(xt,t))和 p θ ( x t − 1 ∣ x t , y ) = N ( x t − 1 ; μ θ ( x t , y , t ) , Σ θ ( x t , y , t ) ) p_θ(x_{t−1} | x_t, y) = \mathcal{N} (x_{t−1};μ_θ(x_t, y, t),Σ_θ(x_t, y, t)) pθ(xt−1∣xt,y)=N(xt−1;μθ(xt,y,t),Σθ(xt,y,t))近似相同的正向过程后验 q ( x t − 1 ∣ x t , x 0 ) = N ( x t − 1 ; μ ~ t ( x t , x 0 ) , 1 − α ˉ t − 1 1 − σ ˉ t β t I ) q\left(\boldsymbol{x}_{t-1} \mid \boldsymbol{x}_t, \boldsymbol{x}_0\right)=\mathcal{N}\left(\boldsymbol{x}_{t-1} ; \tilde{\boldsymbol{\mu}}_t\left(\boldsymbol{x}_t, \boldsymbol{x}_0\right), \frac{1-\bar{\alpha}_{t-1}}{1-\bar{\sigma}_t} \beta_t \mathbf{I}\right) q(xt−1∣xt,x0)=N(xt−1;μ~t(xt,x0),1−σˉt1−αˉt−1βtI)

这里y是一个条件，它包含了对应 x 0 x_0 x0的一些先验知识，比如类标签。假设两个 Σ θ Σ_θ Σθ都设置为未训练的时间相关常数，在相同的实验设置下，条件DPM将达到较低的优化扩散损失。图1中的实验可以证明这一事实，这意味着 μ θ ( x t , y , t ) \mu_θ(x_t, y, t) μθ(xt,y,t)比 μ θ ( x t , t ) \mu_θ(x_t, t) μθ(xt,t)更接近 μ t ~ ( x t , x 0 ) \tilde{\mu_t}(x_t, x_0) μt~(xt,x0)。

这意味着，在无条件DPM( μ θ ( x t , t ) \mu_θ(x_t, t) μθ(xt,t))而真正的 μ t ~ ( x t , x 0 ) \tilde{\mu_t}(x_t, x_0) μt~(xt,x0)后验平均缺口本质上是由于正向过程的信息丢失造成的，反向过程仅根据 x t x_t xt无法在 x t − 1 x_{t−1} xt−1中恢复。如果我们为DPM引入一些关于 x 0 x_0 x0的知识，比如这里的y，差距就会变小。y包含的 x 0 x_0 x0的信息越多，差距就越小。

此外，由式(5)可知，分类器引导条件逆向过程的高斯均值相对于无条件逆向过程的高斯均值多了一个移位项。从后验平均缺口来看，均值偏移项可以部分填补缺口，帮助逆向过程重建样本中丢失的类别信息。理论上，如果Eq.(5)中的y包含x0的全部信息，则均值漂移将完全填补空白，并引导逆向过程重构 x 0 x_0 x0。另一方面，如果我们使用一个模型来根据我们的编码表示z来预测均值偏移，并训练它来填充尽可能多的空白，编码器将被迫从 x 0 x_0 x0中学习尽可能多的信息来帮助填充。缺口填得越多，均值漂移越准确，重构越完美，z包含的 x 0 x_0 x0信息越多。PDAE遵循这一原则，基于预训练的DPM构建自编码器。

3.2 Unsupervised Representation Learning by Filling the Gap

遵循自编码器的范式，我们使用编码器 z = E ϕ ( x 0 ) z = E_ϕ(x_0) z=Eϕ(x0)从输入图像中学习紧凑且有意义的表示，并采用预先训练的无条件DPM p θ ( x t − 1 ∣ x t ) = N ( x t − 1 ; µ θ ( x t , t ) , Σ θ ( x t , t ) ) p_θ(x_{t−1}|x_t) = \mathcal{N} (x_{t−1};µ_θ(x_t, t),Σ_θ(x_t, t)) pθ(xt−1∣xt)=N(xt−1;µθ(xt,t),Σθ(xt,t))到解码器进行图像重建。

具体来说，我们使用梯度估计量 G ψ ( x t , z , t ) G_ψ(x_t, z, t) Gψ(xt,z,t)来模拟 ∇ x t l o g p ( z ∣ x t ) ∇x_t log p(z|x_t) ∇xtlogp(z∣xt)，其中 p ( z ∣ x t ) p(z|x_t) p(z∣xt)是一些我们不会显式使用的隐式分类器，并使用它来组装一个条件DPM p θ , ψ ( x t − 1 ∣ x t , z ) = N ( x t − 1 ; μ θ ( x t , t ) + Σ θ ( x t , t ) ⋅ G ψ ( x t , z , t ) , Σ θ ( x t , t ) ) p_{\theta, \psi}\left(\boldsymbol{x}_{t-1} \mid \boldsymbol{x}_t, \boldsymbol{z}\right)=\mathcal{N}\left(\boldsymbol{x}_{t-1} ; \boldsymbol{\mu}_\theta\left(\boldsymbol{x}_t, t\right)+\boldsymbol{\Sigma}_\theta\left(\boldsymbol{x}_t, t\right) \cdot \boldsymbol{G}_\psi\left(\boldsymbol{x}_t, \boldsymbol{z}, t\right), \boldsymbol{\Sigma}_\theta\left(\boldsymbol{x}_t, t\right)\right) pθ,ψ(xt−1∣xt,z)=N(xt−1;μθ(xt,t)+Σθ(xt,t)⋅Gψ(xt,z,t),Σθ(xt,t))译码器。然后，我们通过优化以下派生目标(假设采用 ϵ \epsilon ϵ-预测参数化)，像常规条件DPM一样训练它：
L ( ψ , φ ) = E x 0 , t , ϵ [ λ t ∥ ϵ − ϵ θ ( x t , t ) + α t 1 − α ˉ t β t ⋅ Σ θ ( x t , t ) ⋅ G ψ ( x t , E φ ( x 0 ) , t ) ∥ 2 ] (7) \mathcal{L}(\psi, \varphi)=\mathbb{E}_{\boldsymbol{x}_0, t, \epsilon}\left[\lambda_t\left\|\epsilon-\boldsymbol{\epsilon}_\theta\left(\boldsymbol{x}_t, t\right)+\frac{\sqrt{\alpha_t} \sqrt{1-\bar{\alpha}_t}}{\beta_t} \cdot \boldsymbol{\Sigma}_\theta\left(\boldsymbol{x}_t, t\right) \cdot \boldsymbol{G}_\psi\left(\boldsymbol{x}_t, \boldsymbol{E}_{\varphi}\left(\boldsymbol{x}_0\right), t\right)\right\|^2\right] \tag{7} L(ψ,φ)=Ex0,t,ϵ[λt ϵ−ϵθ(xt,t)+βtαt 1−αˉt ⋅Σθ(xt,t)⋅Gψ(xt,Eφ(x0),t) 2](7)
其中KaTeX parse error: \tag works only in display equations和 λ t λ_t λt是一种新的加权方案，我们将在3.4节中讨论。

注意，我们使用预训练的DPM，以便θ在优化过程中被冻结。通常我们将 Σ θ = 1 − α ˉ t − 1 1 − σ ˉ t β t I Σ_θ = \frac{1-\bar{\alpha}_{t-1}}{1-\bar{\sigma}_t} \beta_t \mathbf{I} Σθ=1−σˉt1−αˉt−1βtI设置为未训练的时间依赖常数。优化相当于最小化 ∥ Σ θ ( x t , t ) ⋅ G ψ ( x t , E φ ( x 0 ) , t ) − ( μ ~ t ( x t , x 0 ) − μ θ ( x t , t ) ) ∥ 2 \left\|\boldsymbol{\Sigma}_\theta\left(\boldsymbol{x}_t, t\right) \cdot \boldsymbol{G}_\psi\left(\boldsymbol{x}_t, \boldsymbol{E}_{\varphi}\left(\boldsymbol{x}_0\right), t\right)-\left(\tilde{\boldsymbol{\mu}}_t\left(\boldsymbol{x}_t, \boldsymbol{x}_0\right)-\boldsymbol{\mu}_\theta\left(\boldsymbol{x}_t, t\right)\right)\right\|^2 ∥Σθ(xt,t)⋅Gψ(xt,Eφ(x0),t)−(μ~t(xt,x0)−μθ(xt,t))∥2，迫使预测的均值偏移 Σ θ ( x t , t ) ⋅ G ψ ( x t , E φ ( x 0 ) , t ) \boldsymbol{\Sigma}_\theta\left(\boldsymbol{x}_t, t\right) \cdot \boldsymbol{G}_\psi\left(\boldsymbol{x}_t, \boldsymbol{E}_{\varphi}\left(\boldsymbol{x}_0\right), t\right) Σθ(xt,t)⋅Gψ(xt,Eφ(x0),t)填充后验平均间隙 ( μ ~ t ( x t , x 0 ) − μ θ ( x t , t ) ) \left(\tilde{\boldsymbol{\mu}}_t\left(\boldsymbol{x}_t, \boldsymbol{x}_0\right)-\boldsymbol{\mu}_\theta\left(\boldsymbol{x}_t, t\right)\right) (μ~t(xt,x0)−μθ(xt,t))。

对于训练好的 G ψ ( x t , z , t ) G_ψ(x_t, z, t) Gψ(xt,z,t)，我们可以将其视为最优分类器 p ( z ∣ x t ) p(z|x_t) p(z∣xt)的得分，并使用Eq.(5)中的分类器引导采样方法进行DDPM采样或使用改进的函数逼近器 ϵ ^ θ \hat{\epsilon}_θ ϵ^θ在Eq.(6)中进行DDIM采样，基于预训练 ϵ θ ( x t , t ) \epsilon_θ(x_t, t) ϵθ(xt,t)。我们在附录A中列出了详细的算法过程。

除了语义潜码z，我们还可以通过反向运行DDIM的确定性生成过程来推断出随机潜码 x T x_T xT [36]：
x t + 1 = α ˉ t + 1 ( x t − 1 − α ˉ t ⋅ ϵ ^ θ ( x t , t ) α ˉ t ) + 1 − α ˉ t + 1 ⋅ ϵ ^ θ ( x t , t ) (8) \boldsymbol{x}_{t+1}=\sqrt{\bar{\alpha}_{t+1}}\left(\frac{\boldsymbol{x}_t-\sqrt{1-\bar{\alpha}_t} \cdot \hat{\boldsymbol{\epsilon}}_\theta\left(\boldsymbol{x}_t, t\right)}{\sqrt{\bar{\alpha}_t}}\right)+\sqrt{1-\bar{\alpha}_{t+1}} \cdot \hat{\boldsymbol{\epsilon}}_\theta\left(\boldsymbol{x}_t, t\right) \tag{8} xt+1=αˉt+1 (αˉt xt−1−αˉt ⋅ϵ^θ(xt,t))+1−αˉt+1 ⋅ϵ^θ(xt,t)(8)
这个过程是可选的，但是在使用DDIM采样时，对于重建输入图像的微小细节，有助于接近精确的重建和实图像操作。

我们还训练了一个潜在的DPM p ω ( z t − 1 ∣ z t ) p_ω(z_{t−1}|z_t) pω(zt−1∣zt)来模拟学习到的语义潜在空间，与Diff-AE[36]的方法相同。通过训练好的潜在DPM，我们可以从中对z进行采样，以帮助预先训练好的DPM在 G ψ ( x t , z , t ) G_ψ(x_t, z, t) Gψ(xt,z,t)的指导下实现更快更好的无条件采样。

3.3 Network Design

图2显示了PDAE的网络和数据流程。对于编码器 E ϕ E_ϕ Eϕ，与使用U-Net[40]编码器部分的Diff-AE不同，我们发现简单的堆叠卷积层和线性层足以从 x 0 x_0 x0学习有意义的z。对于梯度估计器 G ψ G_ψ Gψ，我们使用类似于预训练DPM的函数逼近器 ϵ θ \epsilon_θ ϵθ的U-Net。考虑到 ϵ θ \epsilon_θ ϵθ也以 x t x_t xt和t作为输入，我们可以进一步利用预训练DPM的知识，重用其训练好的编码器部分和时间嵌入层，这样我们只需要对 G ψ G_ψ Gψ采用新的U-Net中间块、解码器部分和输出块。为了将z纳入其中，我们遵循[8]扩展组归一化[53]，在归一化特征映射上应用缩放和移动两次：
AdaGN ⁡ ( h , t , z ) = z s ( t s GroupNorm ⁡ ( h ) + t b ) + z b (9) \operatorname{AdaGN}(\boldsymbol{h}, t, \boldsymbol{z})=\boldsymbol{z}_s\left(t_s \operatorname{GroupNorm}(\boldsymbol{h})+t_b\right)+\boldsymbol{z}_b \tag{9} AdaGN(h,t,z)=zs(tsGroupNorm(h)+tb)+zb(9)

其中 [ t s , t b ] [t_s, t_b] [ts,tb]和 [ z s , z b ] [z_s, z_b] [zs,zb]分别由t和z的线性投影得到。请注意，我们仍然使用从重用的编码器到新的解码器的跳过连接。这样， G ψ G_ψ Gψ完全由预训练的DPM决定，可以普遍应用于不同的U-Net架构。

3.4 Weighting Scheme Redesign

我们最初在DDPM[14]中使用了类似的简化训练目标，即在Eq.(7)中设置 λ t = 1 λ_t = 1 λt=1，但发现训练非常不稳定，导致慢/不收敛，性能较差。受p2 -加权[7]的启发，扩散损失的加权方案可以极大地影响DPM的性能，我们将这一现象归因于加权方案，并在图3中进行了研究。

具体来说，我们在MNIST[28]上训练一个无条件DPM和一个有噪声分类器，并将扩散正向过程分为三个阶段:0 ~ t 1 t_1 t1之间的早期阶段、 t 1 t_1 t1 ~ t 2 t_2 t2之间的关键阶段和 t 2 t_2 t2 ~ T之间的晚期阶段，如上图所示。然后，我们设计了一个混合抽样程序，采用无条件抽样，但只在指定阶段切换到分类器引导抽样。下面三行显示了由三种不同的混合抽样程序生成的样本，其中每行只在右侧的指定阶段使用分类器引导的抽样。我们可以看到，只有在关键阶段被分类器引导的样本才能匹配输入的类标签。我们可以得出结论，与其他两个阶段相比，关键阶段的均值漂移包含了更多重建样本中输入类别标签的关键信息。从扩散轨迹的角度来看，采样轨迹在关键阶段是相互分离的，它们需要均值漂移来引导它们向指定的方向移动，否则将由朗之万动力学的随机性决定。因此，我们选择在早期和后期降低t的目标函数的权重，以鼓励模型从关键阶段学习丰富的表示形式。

受p2 -加权[7]的启发，我们根据信噪比[24] ( SNR ⁡ ( t ) = α ˉ t 1 − α ˉ t ) \left(\operatorname{SNR}(t)=\frac{\bar{\alpha}_t}{1-\bar{\alpha}_t}\right) (SNR(t)=1−αˉtαˉt)重新设计了扩散损失(λt, Eq.(7))的加权方案：
λ t = ( 1 1 + SNR ⁡ ( t ) ) 1 − γ ⋅ ( SNR ⁡ ( t ) 1 + SNR ⁡ ( t ) ) γ (10) \qquad \lambda_t=\left(\frac{1}{1+\operatorname{SNR}(t)}\right)^{1-\gamma} \cdot\left(\frac{\operatorname{SNR}(t)}{1+\operatorname{SNR}(t)}\right)^\gamma \tag{10} λt=(1+SNR(t)1)1−γ⋅(1+SNR(t)SNR(t))γ(10)
第一项是早期的，第二项是晚期的。γ是一个超参数，用于平衡两个项目之间的权重降低强度。根据经验，我们设定γ = 0.1。图4显示了不同DPM相对于真实变分下限损失的扩散损失的归一化加权方案。与其他DPM相比，我们的加权方案在低信噪比和高信噪比时都降低了扩散损失。

4 Experiments

为了比较PDAE和Diff-AE[36]，我们在相同的设置下进行了实验。此外，我们还展示了PDAE支持一些附加特性。为了公平比较，我们使用官方Diff-AE实现提供的基线dpm作为我们预训练的模型(也作为我们的基线)具有与Diff-AE模型相同的网络架构(超参数)。为了简单起见，我们使用“FFHQ128-130M-z512-64M”这样的符号来命名我们的模型，这意味着我们使用130M图像预训练的基线DPM，并利用它对64M图像进行PDAE训练，在128 × 128 FFHQ数据集[21]上，语义潜在代码z为512-d。我们将所有的实现细节放在附录B中，后续实验的附加样本放在附录C中。

4.1 Training Efficiency

我们从训练时间和次数两个方面证明了PDAE比Diff-AE的训练效率更高。对于训练时间，我们使用4个Nvidia A100- sxm4 GPU在128×128图像数据集上训练具有相同网络架构(超参数)的两个模型进行分布式训练，并将批大小设置为128(每个GPU 32个)，以计算它们的训练吞吐量(imgs/秒/A100)。

PDAE的吞吐量为81.57,Diff-AE的吞吐量为75.41。由于重用了预训练DPM的U-Net编码器部分，PDAE的可训练参数更少，训练吞吐量比Diff-AE更高。对于训练次数，我们发现PDAE需要的训练批次(图像)数量大约是Diff-AE需要的损失收敛的13 ~ 12倍。我们认为这是因为基于预训练的DPM建模后验平均差距比从头开始建模条件DPM更容易。网络复用和加权方案的重新设计也有一定的帮助。因此，基于预训练的DPM, PDAE完成表示学习所需的训练时间不到Diff-AE的一半。

4.2 Learned Mean Shift Fills Posterior Mean Gap

4.3 Autoencoding Reconstruction

我们用“FFHQ128-130M-z512-64M”分别用DDIM和DDPM的PDAE生成过程运行了一些自编码重构实例。如图6所示，两种方法都生成了与输入内容相似的样本。当引入随机性时，一些随机变化[36]出现在头发、眼睛和皮肤的小细节上。由于性能相似在DDPM和DDIM之间随机xT，我们将在以后的实验中始终使用DDIM抽样方法。利用由上述ODE推导出的随机潜码，可以得到近似精确的重构，进一步证明了随机潜码控制局部细节。

为了进一步评价PDAE的自编码重建质量，我们用Diff-AE进行了同样的定量实验。具体来说，我们使用“FFHQ128-130M-z512-64M”对CelebA-HQ[20]的所有30k图像进行编码和构建，并以其平均SSIM [52]， LPIPS[56]和MSE来评估重建质量。我们使用[36]中描述的相同基线，结果如表1所示。我们可以看到，PDAE与最先进的NV AE具有竞争力，即使潜在维度要低得多，并且在除随机xT的LPIPS之外的所有指标上都优于Diff-AE。此外，PDAE只需要大约Diff-AE进行表征学习所需训练时间的一半，这表明基于预训练的DPM, PDAE可以更有效地从图像中学习更丰富的表征。

4.4 Interpolation of Semantic Latent Codes and Trajectories

4.5 Attribute Manipulation

我们可以在监督的方式下进一步探索学习的语义潜空间。为了说明这一点，我们训练了一个“CelebA-HQ128-52M-z512-25M”模型，并利用CelebA-HQ数据集的属性注释进行属性操作实验。具体来说，我们首先将图像编码为其语义潜码，然后沿学习方向移动图像，最后解码为处理后的图像。与Diff-AE类似，我们训练一个线性分类器来分离具有不同属性标签的图像的语义潜码，并使用分离超平面的法向量(即线性分类器的权值)作为方向向量。我们在图8中展示了一些属性操作示例。我们可以看到，PDAE通过在不同的尺度上沿着期望属性的方向移动语义潜码来成功地操纵图像。与Diff-AE一样，如果使用输入图像的推断xT, PDAE可以改变属性相关的特征，同时保持其他不相关的细节几乎不变。

4.6 Truncation-like Effect

根据[8,15]，我们可以通过缩放分类器引导的强度在dpm中获得类似截断的效果。我们假设通过填充后验平均间隙训练的 G ψ ( x t , z , t ) G_ψ(x_t, z, t) Gψ(xt,z,t)模拟了某些隐式分类器的梯度，并且它实际上可以按预期工作。理论上是可以的也适用于截断效应。为了说明这一点，我们直接将类标签合并到 G ψ ( x t , y , t ) G_ψ(x_t, y, t) Gψ(xt,y,t)中，并训练它来填充空白。具体来说，我们训练了一个“ImageNet64-77M-y-38M”模型，并使用100步的DDIM采样方法生成50k个样本，以不同尺度的预测均值偏移为指导，达到类似截断的效果。图9显示了扫过尺度的样本质量效应。我们可以看到，它达到了类似于分类器引导抽样方法的截断效果，这有助于我们在填充后验均值间隙和分类器引导抽样方法之间建立联系。通过填充后验平均间隙训练的梯度估计器可以替代噪声分类器。

4.7 Few-shot Conditional Generation

4.8 Improved Unconditional Sampling

如4.2节所示，在z的帮助下，PDAE只需要一步就可以生成似是而非的图像。如果我们可以提前得到z, PDAE可以在相同的采样步数下获得比预训练的dpm更好的样本质量。与Diff-AE相似，我们在语义潜空间上训练一个潜在的DPM，并从中提取样本z，以改进预训练的DPM的无条件采样。与Diff-AE必须以z为输入进行采样不同，PDAE使用独立的梯度估计器作为预训练DPM的校正器进行采样。我们发现只有在最后几个采样步骤中使用预训练的dpm才能获得更好的样本质量，这可能是因为梯度估计器对最后几个采样步骤中的z和样本的随机性比较敏感Z会导致域外样本。asyp[27]也发现了类似的现象。根据经验，我们在最后30%的抽样步骤中执行这一策略。对“FFHQ128-130M-z512-64M”、“Horse128-130M-z512-64M”、“Bedroom128-120M-z512-70M”和“CelebA64-72M-z512-38M”采用不同步骤的DDIM抽样方法进行无条件抽样结果评价。对于每个数据集，我们计算50k生成的样本和50k从数据集随机选择的真实图像之间的FID分数。表2显示，PDAE显著提高了预训练dpm的样本质量，并优于Diff-AE。请注意，PDAE可以应用于任何预训练的dpm作为辅助助推器，以提高其样本质量。

5 Related Work

我们的工作基于一种新兴的潜变量生成模型，称为扩散概率模型(Diffusion probability Models, DPM)[43,14]，该模型目前因其稳定的训练过程和具有竞争力的样本质量而流行。大量的研究[34,24,8,15,44,19,46,30]和应用[5,26,18,32,57,6,29,41,3,16,17]进一步显著地改进和扩展了DPM。

基于生成建模的无监督表示学习是计算机视觉领域的热门课题。潜在变量生成模型，如GAN [13]， VAE[25,39]和DPM，是自然的候选模型，因为它们固有地涉及它们生成的数据的潜在表示。对于GANs，由于其缺乏推理功能，人们必须通过一种称为GAN反转[54]的额外技术来提取任何给定真实样本的表示，该技术将样本反转回训练过的GANs的潜在空间。现有反演方法[58,35,4,1,2,51]要么重构质量有限，要么计算成本明显较高。vae显式地学习样本的表示，但仍然面临表示生成权衡的挑战[49,42]。VQ-VAE[49, 37]和D2C[42]通过对潜在变量进行不同的事后建模来克服这些问题。DPM还通过正向过程产生潜在变量。然而，这些潜在变量缺乏高级语义信息，因为它们只是一组空间损坏的图像。有鉴于此，扩散自编码器(Diff-AE)[36]探索了通过自编码进行表示学习的dpm。具体来说，他们联合训练了一个编码器，用于从图像中发现有意义的表示，并训练了一个有条件的DPM作为图像重建的解码器，将表示作为输入条件。Diff-AE在图像重建方面与最先进的模型竞争，并能够完成各种下游任务。与Diff-AE相比，PDAE同样通过自动编码利用现有的预训练DPM进行表示学习，但具有更好的训练效率和性能。

具有类似思想的并行工作是预训练的文本到图像DPM[12]的文本反转。具体来说，对于用户提供的概念(如对象或样式)，只给他们3-5张图像，他们就学会在冻结的文本到图像DPM的嵌入空间中通过新的“单词”来表示它。这些习得的“词”可以进一步组合成自然语言句子，直观地指导个性化创作。从后验平均间隙的角度，对于给定的新概念，语篇反演优化其对应的新“词”嵌入向量，以找到最佳的语篇条件( c \mathcal{c} c)，这样就可以输入预先训练好的文本-图像DPM ( ϵ θ ( x t , c , t ) \epsilon_θ(x_t,c,t) ϵθ(xt,c,t))为了填补空白( ϵ − ϵ θ ( x t , ∅ , t ) \epsilon-\epsilon_θ(x_t,∅,t) ϵ−ϵθ(xt,∅,t))越好。

6 Conclusion

总之，我们提出了一种称为PDAE的通用方法，该方法利用预训练的DPM通过自动编码进行表示学习，并获得了比Diff-AE更好的训练效率和性能。我们的主要思想是基于后验平均差的概念及其与分类器引导抽样方法的联系。同时进行的工作，预训练的文本到图像DPM的文本反转，也可以从这个角度解释。我们认为这个想法可以进一步探索，从预训练的DPM中提取知识，例如可解释的方向发现[51]，我们把它留给未来的工作。

A Algorithm

B Implementation Details

B.1 Network Architecture

B.2 Experimental Details

C Additional Samples

C.1 Learned Mean Shift Fills Posterior Mean Gap

C.2 Autoencoding Reconstruction

C.3 Interpolation of Semantic Latent Codes and Trajectories

C.4 Attribute Manipulation

C.5 Few-shot Conditional Generation

C.6 Visualization of Mean Shift

D Limitations and Potential Negative Societal Impacts

尽管PDAE训练效率更高，但由于额外的梯度估计器，PDAE的推理速度比Diff-AE慢，而且需要更多的内存和存储空间。生成速度慢是基于pm的作品普遍存在的问题。尽管许多研究已经能够通过很少的反向步骤实现良好的性能，但它们仍然落后于只需要一次网络通过的VAEs和GANs。此外，几乎完美的PDAE重建需要数百个额外的前向步骤来推断随机潜码。此外，我们发现扩散损失的加权方案是PDAE不可缺少的，但我们还没有探索其机制，这可能有助于进一步提高PDAE的效率和性能。我们把这方面的实证和理论研究留作以后的工作。我们工作的潜在负面影响主要涉及深度造假，它利用机器学习和人工智能的强大生成技术来创建合成媒体，这些合成媒体可能被用于欺诈、欺诈、欺凌或报复。尽管一些合成样本很难区分，但研究人员已经开发出了类似于用于构建deepfake的算法，以高精度检测它们。区块链和数字签名等其他技术可以帮助平台验证媒体的来源。