0%

论文阅读-image-generation-Denoising Diffusion Probabilistic Models

动机

  • 扩散模型定义简单,训练效率高,但还没有方法显示其能够生成高质量图像。

贡献

  • 提出了一种使用扩散概率模型(一种受非平衡热力学启发的隐变量模型)的图像合成方法;
  • 发现了扩散模型的某种参数化与训练期间在多个噪声水平上的去噪分数匹配和采样期间的annealed Langevin dynamics等价;
  • 发现了模型的大部分无损编码被用于描述难以察觉的图像细节,证明了扩散模型的采样过程是一种渐进解码类型,可以解释为自回归解码的泛化。

扩散模型

扩散模型是一种形式为 p_{\theta}\left(\mathbf{x}_{0}\right):=\int p_{\theta}\left(\mathbf{x}_{0: T}\right) d \mathbf{x}_{1: T}​ 的隐变量模型,其中 \mathbf{x}_1,…,\mathbf{x}_T​ 是与数据 \mathbf{x}_0\sim q(\mathbf{x}_0)​ 具有相同维度的隐变量,联合概率 p_{\theta}\left(\mathbf{x}_{0: T}\right)​ 被称为反向过程,定义为从 p\left(\mathbf{x}_{T}\right)=\mathcal{N}\left(\mathbf{x}_{T} ; \mathbf{0}, \mathbf{I}\right)​ 开始学习的高斯变换的马尔科夫链:

p_{\theta}\left(\mathbf{x}_{0: T}\right):=p\left(\mathbf{x}_{T}\right) \prod_{t=1}^{T} p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right), \quad p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right):=\mathcal{N}\left(\mathbf{x}_{t-1} ; \boldsymbol{\mu}_{\theta}\left(\mathbf{x}_{t}, t\right), \mathbf{\Sigma}_{\theta}\left(\mathbf{x}_{t}, t\right)\right)

扩散模型与其他类型的隐变量模型的区别在于,近似后验 q(\mathbf{x}_{1: T}|\mathbf{x}_0) ,称为前向过程或扩散过程,被固定在马尔可夫链上,该马尔可夫链逐步将高斯噪声的方差以一种策略 \beta_1,…,\beta_T 添加到数据中:

q\left(\mathbf{x}_{1: T} \mid \mathbf{x}_{0}\right):=\prod_{t=1}^{T} q\left(\mathbf{x}_{t} \mid \mathbf{x}_{t-1}\right), \quad q\left(\mathbf{x}_{t} \mid \mathbf{x}_{t-1}\right):=\mathcal{N}\left(\mathbf{x}_{t} ; \sqrt{1-\beta_{t}} \mathbf{x}_{t-1}, \beta_{t} \mathbf{I}\right)

训练则通过优化负对数似然的变分上界,即:

\mathbb{E}\left[-\log p_{\theta}\left(\mathbf{x}_{0}\right)\right] \leq \mathbb{E}_{q}\left[-\log \frac{p_{\theta}\left(\mathbf{x}_{0: T}\right)}{q\left(\mathbf{x}_{1: T} \mid \mathbf{x}_{0}\right)}\right]=\mathbb{E}_{q}\left[-\log p\left(\mathbf{x}_{T}\right)-\sum_{t \geq \mathbf{I}} \log \frac{p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right)}{q\left(\mathbf{x}_{t}\left[\mathbf{x}_{t-1}\right)\right.}\right]=: L

前向过程方差 \beta_t 可通过重参数化学习或者保持为常数作为超参数,反向过程的表达由 p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right) 中高斯条件的选择来保证。前向过程的一个显著特性是,它允许在任意时间步 t 上以闭合形式对 \mathbf{x}_t 进行采样:定义 \alpha_t:=1-\beta_t,\bar{\alpha}_t:=\prod_{s=1}^i\alpha_s ,则:

q\left(\mathbf{x}_{t} \mid \mathbf{x}_{0}\right)=\mathcal{N}\left(\mathbf{x}_{t} ; \sqrt{\bar{\alpha}_{t}} \mathbf{x}_{0},\left(1-\bar{\alpha}_{t}\right) \mathbf{I}\right)

因此,可以通过随机梯度下降法优化 L​ 的随机项,进一步的改进是通过重写 L​ 来减少方差:

\mathbb{E}_{q}[\underbrace{D_{\mathrm{KL}}\left(q\left(\mathbf{x}_{T} \mid \mathbf{x}_{0}\right) \| p\left(\mathbf{x}_{T}\right)\right)}_{L_{T}}+\sum_{t>1} \underbrace{D_{\mathrm{KL}}\left(q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}, \mathbf{x}_{0}\right) \| p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right)\right)}_{L_{t-1}} \underbrace{-\log p_{\theta}\left(\mathbf{x}_{0} \mid \mathbf{x}_{1}\right)}_{L_{0}}]

上式利用 KL 距离去比较 p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right) 和正向过程的后验,其在 \mathbf{x}_0 条件下易跟踪:

\begin{aligned} q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}, \mathbf{x}_{0}\right) &=\mathcal{N}\left(\mathbf{x}_{t-1} ; \tilde{\boldsymbol{\mu}}_{t}\left(\mathbf{x}_{t}, \mathbf{x}_{0}\right), \tilde{\beta}_{t} \mathbf{I}\right) \\ \text { where } \quad \tilde{\boldsymbol{\mu}}_{t}\left(\mathbf{x}_{t}, \mathbf{x}_{0}\right) &:=\frac{\sqrt{\bar{\alpha}_{t-1}} \beta_{t}}{1-\bar{\alpha}_{t}} \mathbf{x}_{0}+\frac{\sqrt{\alpha_{t}}\left(1-\bar{\alpha}_{t-1}\right)}{1-\bar{\alpha}_{t}} \mathbf{x}_{t} \quad \text { and } \quad \tilde{\beta}_{t}:=\frac{1-\bar{\alpha}_{t-1}}{1-\bar{\alpha}_{t}} \beta_{t} \end{aligned}

所以, KL 距离均是高斯分布之间的比较,故可用闭式表达式形式的Rao-Blackwellized方法代替高方差的Monte Carlo估计。

本文的方法

扩散模型可以看作是一种隐变量模型,在实现过程中包含了大量的自由度,如前向过程的方差 \beta_t ,反向过程的模型结构和反向过程的高斯分布参数化。本文在扩散模型和去噪分数匹配之间建立了一种显式的联系,从而为扩散模型建立了一个简化的加权变分边界目标。

前向过程和 L_T

本文将前向过程方差 \beta_t 固定为常数,故近似后验 q 没有可学习参数,故 L_T 是常数(因为 p\left(\mathbf{x}_{T}\right)=\mathcal{N}\left(\mathbf{x}_{T} ; \mathbf{0}, \mathbf{I}\right) , q\left(\mathbf{x}_{t} \mid \mathbf{x}_{0}\right)=\mathcal{N}\left(\mathbf{x}_{t} ; \sqrt{\bar{\alpha}_{t}} \mathbf{x}_{0},\left(1-\bar{\alpha}_{t}\right) \mathbf{I}\right) ),可忽略。

反向过程和 L_{1:T-1}

对于 p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right)=\mathcal{N}\left(\mathbf{x}_{t-1} ; \boldsymbol{\mu}_{\theta}\left(\mathbf{x}_{t}, t\right), \mathbf{\Sigma}_{\theta}\left(\mathbf{x}_{t}, t\right)\right),1<t\le T ,首先设置 \mathbf{\Sigma}_{\theta}\left(\mathbf{x}_{t}, t\right) 为未经训练的时间依赖常数,如 \sigma_{t}^{2}=\beta_t \sigma_{t}^{2}=\tilde{\beta}_{t}=\frac{1-\bar{\alpha}_{t-1}}{1-\bar{\alpha}_{t}} \beta_{t} ,第一种选择对于 \mathbf{x}_0\sim \mathcal{N}(\mathbf{0}, \mathbf{I}) 是最优的,第二种选择对于 \mathbf{x}_0 被设置为一个点来说是最优的,这是两种极端的选择,对应于具有单位方差的数据的反向过程的熵的上界和下界。

对于 \boldsymbol{\mu}_{\theta}\left(\mathbf{x}_{t}, t\right) ,基于如下的分析提出了一种参数化方法:对于 p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right)=\mathcal{N}\left(\mathbf{x}_{t-1} ; \boldsymbol{\mu}_{\theta}\left(\mathbf{x}_{t}, t\right), \sigma_{t}^{2} \mathbf{I}\right) ,有:

L_{t-1}=\mathbb{E}_{q}\left[\frac{1}{2 \sigma_{t}^{2}}\left\|\tilde{\boldsymbol{\mu}}_{t}\left(\mathbf{x}_{t}, \mathbf{x}_{0}\right)-\boldsymbol{\mu}_{\theta}\left(\mathbf{x}_{t}, t\right)\right\|^{2}\right]+C

其中 C 是不依赖 \theta 的常数,最简单的参数化是 \boldsymbol{\mu}_{\theta} 是一个模型,去预测前向过程的后验均值 \tilde{\boldsymbol{\mu}}_{t} ,通过上文的 q\left(\mathbf{x}_{t} \mid \mathbf{x}_{0}\right)=\mathcal{N}\left(\mathbf{x}_{t} ; \sqrt{\bar{\alpha}_{t}} \mathbf{x}_{0},\left(1-\bar{\alpha}_{t}\right) \mathbf{I}\right) 去扩展,即 \mathbf{x}_{t}\left(\mathbf{x}_{0}, \boldsymbol{\epsilon}\right)=\sqrt{\bar{\alpha}_{t}} \mathbf{x}_{0}+\sqrt{1-\bar{\alpha}_{t}} \boldsymbol{\epsilon},\boldsymbol{\epsilon}\sim\mathcal{N}(\mathbf{0}, \mathbf{I}) ,则得到:

\begin{aligned} L_{t-1}-C &=\mathbb{E}_{\mathbf{x}_{0}, \epsilon}\left[\frac{1}{2 \sigma_{t}^{2}}\left\|\tilde{\boldsymbol{\mu}}_{t}\left(\mathbf{x}_{t}\left(\mathbf{x}_{0}, \boldsymbol{\epsilon}\right), \frac{1}{\sqrt{\bar{\alpha}_{t}}}\left(\mathbf{x}_{t}\left(\mathbf{x}_{0}, \boldsymbol{\epsilon}\right)-\sqrt{1-\bar{\alpha}_{t}} \boldsymbol{\epsilon}\right)\right)-\boldsymbol{\mu}_{\theta}\left(\mathbf{x}_{t}\left(\mathbf{x}_{0}, \boldsymbol{\epsilon}\right), t\right)\right\|^{2}\right] \\ &=\mathbb{E}_{\mathbf{x}_{0}, \boldsymbol{\epsilon}}\left[\frac{1}{2 \sigma_{t}^{2}}\left\|\frac{1}{\sqrt{\alpha_{t}}}\left(\mathbf{x}_{t}\left(\mathbf{x}_{0}, \boldsymbol{\epsilon}\right)-\frac{\beta_{t}}{\sqrt{1-\bar{\alpha}_{t}}} \boldsymbol{\epsilon}\right)-\boldsymbol{\mu}_{\theta}\left(\mathbf{x}_{t}\left(\mathbf{x}_{0}, \boldsymbol{\epsilon}\right), t\right)\right\|^{2}\right] \end{aligned}

上式则显示了 \boldsymbol{\mu}_{\theta} 对于给定的 \mathbf{x}_t 必须预测 \frac{1}{\sqrt{\alpha_{t}}}\left(\mathbf{x}_{t}-\frac{\beta_{t}}{\sqrt{1-\bar{\alpha}_{t}}} \boldsymbol{\epsilon}\right) ,由于 \mathbf{x}_t 可以作为模型的输入,所以可以将 \boldsymbol{\mu}_{\theta}(\mathbf{x}_t,t)​ 参数化为:

\boldsymbol{\mu}_{\theta}\left(\mathbf{x}_{t}, t\right)=\overline{\boldsymbol{\mu}}_{t}\left(\mathbf{x}_{t}, \frac{1}{\sqrt{\bar{\alpha}_{t}}}\left(\mathbf{x}_{t}-\sqrt{1-\bar{\alpha}_{t}} \boldsymbol{\epsilon}_{\theta}\left(\mathbf{x}_{t}\right)\right)\right)=\frac{1}{\sqrt{\alpha_{t}}}\left(\mathbf{x}_{t}-\frac{\beta_{t}}{\sqrt{1-\bar{\alpha}_{t}}} \epsilon_{\theta}\left(\mathbf{x}_{t}, t\right)\right)

其中 \boldsymbol{\epsilon}_{\theta} 是一个函数,旨在从 \mathbf{x}_t 预测 \boldsymbol{\epsilon} 。采样 \mathbf{x}_{t-1}\sim p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_t) 是通过计算 \mathbf{x}_{t-1}=\frac{1}{\sqrt{\alpha_{t}}}\left(\mathbf{x}_{t}-\frac{\beta_{t}}{\sqrt{1-\bar{\alpha}_{t}}} \boldsymbol{\epsilon}\right)+\sigma_t\mathbf{z},\mathbf{z}\sim \mathcal{N}(\mathbf{0}, \mathbf{I}) (由 p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right)=\mathcal{N}\left(\mathbf{x}_{t-1} ; \boldsymbol{\mu}_{\theta}\left(\mathbf{x}_{t}, t\right), \sigma_{t}^{2} \mathbf{I}\right) 推导),完整的采样过程类似带数据密度的可学习的梯度 \boldsymbol{\epsilon}_{\theta} Langevin dynamics,由上式的参数化,故 L_{t-1}-C 简化为:

\mathbb{E}_{\mathbf{x}_{0}, \boldsymbol{\epsilon}}\left[\frac{\beta_{t}^{2}}{2 \sigma_{t}^{2} \alpha_{t}\left(1-\bar{\alpha}_{t}\right)}\left\|\boldsymbol{\epsilon}-\boldsymbol{\epsilon}_{\theta}\left(\sqrt{\bar{\alpha}_{t}} \mathbf{x}_{0}+\sqrt{1-\bar{\alpha}_{t}} \boldsymbol{\epsilon}, t\right)\right\|^{2}\right]

这类似于在多个噪声尺度(索引为 t )上的去噪分数匹配,由于上式等价于类似Langevin dynamics反向过程的变分边界,优化类似去噪分数匹配的目标等价于使用变分推断去拟合类似Langevin dynamics的采样链的有限时间边际。

总而言之,可以训练反向过程均值函数近似 \boldsymbol{\mu}_{\theta} 去预测 \tilde{\boldsymbol{\mu}}_t ,也可以去预测 \boldsymbol{\epsilon} (也可以预测 \mathbf{x}_0 ,但实验早期会导致较差的样本质量)。作者表明,预测 \boldsymbol{\epsilon} 类似于Langevin dynamics,且能够简化扩散模型的变分边界到一个类似去噪分数匹配的目标函数,注意它只是另一种 p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right) 的参数化方法。

数据缩放,反向过程解码器和 L_0

假设图像数据由 \{0,1,…,255\} 线性缩放到 [-1,1] ,这保证神经网络反向过程从标准正态先验开始,在缩放一致的输入上操作。为了得到离散对数似然,将反向过程的最后一项设置为一个由高斯函数 \mathcal{N}\left(\mathbf{x}_{0} ; \boldsymbol{\mu}_{\theta}\left(\mathbf{x}_{1}, 1\right), \sigma_{1}^{2} \mathbf{I}\right) 派生的独立的离散解码器:

\begin{aligned} p_{\theta}\left(\mathbf{x}_{0} \mid \mathbf{x}_{1}\right) &=\prod_{i=1}^{D} \int_{\delta_{-}\left(x_{0}^{i}\right)}^{\delta_{+}\left(x_{0}^{i}\right)} \mathcal{N}\left(x ; \mu_{\theta}^{i}\left(\mathbf{x}_{1}, 1\right), \sigma_{1}^{2}\right) d x \\ \delta_{+}(x) &=\left\{\begin{array}{ll} \infty & \text { if } x=1 \\ x+\frac{1}{255} & \text { if } x<1 \end{array} \quad \delta_{-}(x)=\left\{\begin{array}{ll} -\infty & \text { if } x=-1 \\ x-\frac{1}{255} & \text { if } x>-1 \end{array}\right.\right. \end{aligned}

其中 D 是数据维数,上标 i 表示提取一个坐标(可以合并一个更强大的解码器),类似于VAE的解码器和自回归模型的离散连续分布,本文的选择确保了变分边界是离散数据的无损编码长度,不需要在数据中添加噪声,也不需要在对数似然中纳入比例运算的Jacobian矩阵。在采样的最后,会得到无损的 \boldsymbol{\mu}_{\theta}\left(\mathbf{x}_{1}, 1\right)

简化训练目标

由上述的推导,实验中使用以下变分边界的变体(有利于图像质量且更容易实现):

L_{\text {simple }}(\theta):=\mathbb{E}_{t, \mathbf{x}_{0}, \epsilon}\left[\left\|\boldsymbol{\epsilon}-\boldsymbol{\epsilon}_{\theta}\left(\sqrt{\bar{\alpha}_{t}} \mathbf{x}_{0}+\sqrt{1-\bar{\alpha}_{t}} \boldsymbol{\epsilon}, t\right)\right\|^{2}\right]

其中, t 属于区间 [1,T] t=1 的情况对应于 L_0 ,前文离散解码器的定义中的积分近似于高斯概率密度函数乘以bin宽度(忽略 \sigma_{1}^2 和边缘效应), t>1 的情况对应于未加权的 \mathbb{E}_{\mathbf{x}_{0}, \boldsymbol{\epsilon}}\left[\frac{\beta_{t}^{2}}{2 \sigma_{t}^{2} \alpha_{t}\left(1-\bar{\alpha}_{t}\right)}\left|\boldsymbol{\epsilon}-\boldsymbol{\epsilon}_{\theta}\left(\sqrt{\bar{\alpha}_{t}} \mathbf{x}_{0}+\sqrt{1-\bar{\alpha}_{t}} \boldsymbol{\epsilon}, t\right)\right|^{2}\right] ,类似于NCSN去噪分数匹配模型中使用的加权损失( L_T 没有出现,因为前向过程的方差 \beta_t 是固定的)。

简化的目标函数不考虑前文推导的权重,与标准变分边界相比,加权变分边界强调重构的不同方面。本文实验设置使得简化目标降低对应小 t 的损失项的权重,这部分训练网络去降噪,所以在更大的 t 时,降低它们的权重有利于专注于困难的去噪任务。

训练过程和采样过程如下:

部分实验结果

采样质量:

比较了ISFID和负对数似然(无损编码长度),训练集的FID 3.17 ,测试集的FID 5.24

反向过程参数化和训练目标消融实验:

实验发现在基线模型中, \tilde{\boldsymbol{\mu}} 适用于训练真正的变分约束;学习反向过程方差(将参数化的 \mathbf{\Sigma}_{\theta}(\mathbf{x}_t) 加入到变分边界中)会导致不稳定的训练;固定方差时预测 \boldsymbol{\epsilon} \tilde{\boldsymbol{\mu}} 好;预测 \boldsymbol{\epsilon} 时,简化的变分边界有更好的样本质量。

渐进编码:

  • 渐进有损压缩

    如下图引入一个渐进有损编码,去进一步探究模型的rate-distortion behavior。它假设访问一个过程,对于任意分布 p q ,其中只有 p 预先对接收方可用,平均使用大约 D_{KL}(q(\mathbf{x})\left |\right|(\mathbf{x})) bit去传输一个样本 \mathbf{x}\sim q(\mathbf{x})

    失真在率失真图的低率区域急剧下降,这表明大多数比特分配给了难以察觉的失真。

  • 渐进生成

    通过使用Algorithm 2采样反向过程,去预测 \hat{\mathbf{x}}_0

    下图显示按照反向过程去采样的结果,大尺度图像特征最先出现,细节最后出现。

    下图显示了 \mathbf{x}_0\sim p_{\theta}(\mathbf{x}_0|\mathbf{x}_t) 在不同 t 时的随机预测,当 t 很小时,几乎所有细节都被保留,当 t 很大时,只有大尺度的特征被保留。

  • 与自回归解码的联系

    扩散模型的变分边界可以重写为:

    L=D_{\mathrm{KL}}\left(q\left(\mathrm{x}_{T}\right) \| p\left(\mathrm{x}_{T}\right)\right)+\mathbb{E}_{q}\left[\sum_{t \geq 1} D_{\mathrm{KL}}\left(q\left(\mathrm{x}_{t-1} \mid \mathrm{x}_{t}\right) \| p_{\theta}\left(\mathrm{x}_{t-1} \mid \mathrm{x}_{t}\right)\right)\right]+H\left(\mathrm{x}_{0}\right)

    作者证明在某些情况下,其等价于训练自回归模型。

插值:

q 作为随机编码器,对于源图像 \mathbf{x}_0 ,由 \mathbf{x}_0^{\prime}\sim q(\mathbf{x}_0) \mathbf{x}_t^{\prime}\sim q(\mathbf{x}_t|\mathbf{x}_0) ,然后对线性插值 \overline{\mathbf{x}}_{t}=(1-\lambda) \mathbf{x}_{0}+\lambda \mathbf{x}_{0}^{\prime} 通过反向过程解码到图像空间 \overline{\mathbf{x}}_0\sim p(\mathbf{x}_0|\overline{\mathbf{x}}_t)