生成模型(四):扩散模型(Diffusion Models)

本文大纲如下:
生成模型(四):扩散模型(Diffusion Models)_第1张图片

生成模型种类

到目前为止,我已经写了三种类型的生成模型,[[生成模型-GAN]]、[[生成模型-VAE]]和[[生成模型-Flow based model]]。 它们在生成高质量样本方面显示出巨大的成功,但每一种都有其自身的一些局限性。下面是对GAN、VAE和基于流的生成模型之间区别的一个快速总结。

生成对抗网络。GAN提供了一个聪明的解决方案,将数据生成这一无监督的学习问题建模为有监督的问题。鉴别器模型学习从生成器模型产生的虚假样本中区分出真实数据。两个模型同时训练,进行最大最小博弈。但是训练不稳定,多样性较少。
变量自动编码器。VAE通过最大化证据下限(ELBO),不明确地优化了数据的对数可能性。但是依赖于surrogate loss。
基于流的生成模型。基于流的生成模型是由一连串的可逆变换构建的。与其他两种不同,该模型明确地学习了数据分布,因此损失函数只是负对数似然。需要使用专门的架构来构建可逆变换。
扩散模型是受非平衡热力学的启发。它们定义一个扩散步骤的马尔可夫链,逐渐向数据添加随机噪声,然后学习逆扩散过程,从噪声中构建所需的数据样本。与VAE或流动模型不同,扩散模型是用固定的程序学习的,而且隐变量具有高维度(与原始数据相同)。

现有的生成模型可以根据如何表示概率分布分为以下两类:

基于似然的模型,通过(近似)最大似然直接学习分布的概率密度(或质量)函数。典型的基于似然的模型包括自回归模型,归一化流模型, 基于能量的模型 (EBM)和变分自动编码器 (VAE)。
隐式生成模型,其概率分布由其采样过程的模型隐式表示。最突出的例子是生成对抗网络(GAN),其中来自数据分布的新样本是通过用神经网络转换随机高斯向量来合成的。

然而,基于似然的模型和隐性生成模型都有很大的局限性。基于似然的模型要么需要对模型结构进行严格的限制,以确保似然计算有一个可操作的归一化常数,要么必须依靠代理目标(surrogate objectives)来接近最大似然训练。另一方面,隐式生成模型往往需要对抗性训练,这是众所周知的不稳定,可能导致模式崩溃.

我们将介绍另一种表示概率分布的方法,它可能会规避一些限制。其关键思想是对对数概率密度函数的梯度进行建模,这个量通常被称为(Stein)score function . 这种基于分数的模型不需要有一个可操作的归一化常数,并且可以通过分数匹配直接学习。

假设我们得到一个数据集 x 1 , x 2 , ⋯   , x N {\mathbf{x}_1, \mathbf{x}_2, \cdots, \mathbf{x}_N} x1,x2,,xN,其中每个点都是从一个基础数据分布p(\mathbf{x})独立抽取的。考虑到这个数据集,生成式建模的目标是为数据分布拟合一个模型,这样我们就可以通过从分布中取样,随意合成新的数据点。

为了建立这样的生成模型,我们首先需要一种方法来表示概率分布。其中一种方法,如基于似然的模型,是直接对概率密度函数(p.d.f.)或概率质量函数(p.m.f.)建模。让 f θ ( x ) ∈ R f_\theta(\mathbf{x}) \in \mathbb{R} fθ(x)R 是一个实值函数,由一个可学习的参数\theta作为参数。我们可以通过以下方式定义一个p.d.f.

p_\theta(\mathbf{x}) = \frac{e^{-f_\theta(\mathbf{x})}}{Z_\theta} \

其中Z_\theta \gt 0是依赖于\theta的归一化常数,从而\int p_\theta(\mathbf{x}) \textrm{d} \mathbf{x} = 1。这里的函数f_\theta(\mathbf{x})通常被称为非标准化的概率模型,或基于能量的模型。 我们可以通过最大化数据的对数似然来训练p_\theta(\mathbf{x}):

\begin{align} \max_\theta \sum_{i=1}^N \log p_\theta(\mathbf{x}_i). \end{align} \

然而,方程式要求p_\theta(\mathbf{x})是一个归一化的概率密度函数。这是不可取的,因为为了计算p_\theta(\mathbf{x}),我们必须评估归一化常数Z_\theta。对于任何一般的f_\theta(\mathbf{x})来说都是难以解决的数量。因此,为了使最大似然训练可行,基于似然的模型必须限制其模型结构(例如自回归模型中的因果卷积,归一化流模型中的可逆网络)以使Z_\theta具有可操作性,或者近似归一化常数(例如VAE中的变分推理,或contrastive divergence中使用的MCMC采样),这可能在计算上变得很昂贵。

通过对分数函数而不是密度函数进行建模,我们可以避开难以解决的归一化常数的困难。分布p(\mathbf{x})的得分函数被定义为

\begin{equation} \nabla_\mathbf{x} \log p(\mathbf{x}), \notag \end{equation} \

关于分数函数的模型被称为基于分数的模型,我们把它表示为\mathbf{s}\theta(\mathbf{x})。基于分数的模型是这样学习的:\mathbf{s}\theta(\mathbf{x})\approx \nabla_\mathbf{x} log p(\mathbf{x}),并且可以在不担心归一化常数的情况下进行参数化。例如,我们可以很容易地将基于分数的模型与方程中定义的基于能量的模型进行参数化,方法是

\begin{equation} \mathbf{s}\theta (\mathbf{x}) = \nabla{\mathbf{x}} \log p_\theta (\mathbf{x} ) = -\nabla_{\mathbf{x}} f_\theta (\mathbf{x}) - \underbrace{\nabla_\mathbf{x} \log Z_\theta}{=0} = -\nabla\mathbf{x} f_\theta(\mathbf{x}). \end{equation} \

请注意,基于分数的模型\mathbf{s}\theta(\mathbf{x})与归一化常数Z\theta无关! 这极大地扩展了我们可以实际使用的模型家族,因为我们不需要任何特殊的架构来使归一化常数变得可操作。

什么是扩散模型
一些基于扩散的生成模型被提出,其下有类似的想法,包括:

  • 扩散概率模型(diffusion probabilistic modelsSohl-Dickstein et al., 2015),
  • 噪声条件得分网络(NCSNYang & Ermon, 2019),
  • 去噪扩散概率模型(DDPMHo et al. 2020)。

前向扩散过程
给出从真实数据分布中取样的数据点\mathbf{x}_0\sim q(\mathbf{x}),定义一个前向扩散过程,在这个过程中我们以T的步骤向样本添加少量的高斯噪声,产生一连串的噪声样本\mathbf{x}_1, \dots, \mathbf{x}T。步长由方差{\beta_t\in(0,1)}{t=1}^t控制。

q(\mathbf{x}t \vert \mathbf{x}{t-1}) = \mathcal{N}(\mathbf{x}t; \sqrt{1 - \beta_t} \mathbf{x}{t-1}, \beta_t\mathbf{I}) \quad q(\mathbf{x}_{1:T} \vert \mathbf{x}0) = \prod^T{t=1} q(\mathbf{x}t \vert \mathbf{x}{t-1}) \

数据样本\mathbf{x}_0随着步长t的增大而逐渐失去其可辨别的特征。最终当T\to \infty时,\mathbf{x}_T相当于一个各向同性的高斯分布。

上述过程的一个很好的特性是,我们可以在任何任意时间步长t处用重参数化技巧对\mathbf{x}_t进行封闭式采样。让\alpha_t = 1 - \beta_t和\bar{\alpha}t = \prod{i=1}^T \alpha_i。

x t = α t x t − 1 + 1 − α t z t − 1  ;where  z t − 1 , z t − 2 , ⋯ ∼ N ( 0 , I ) = α t α t − 1 x t − 2 + 1 − α t α t − 1 z ˉ t − 2  ;where  z ˉ t − 2  merges two Gaussians (*). = … = α ˉ t x 0 + 1 − α ˉ t z q ( x t ∣ x 0 ) = N ( x t ; α ˉ t x 0 , ( 1 − α ˉ t ) I ) \begin{aligned} \mathbf{x}_t &= \sqrt{\alpha_t}\mathbf{x}_{t-1} + \sqrt{1 - \alpha_t}\mathbf{z}_{t-1} & \text{ ;where } \mathbf{z}_{t-1}, \mathbf{z}_{t-2}, \dots \sim \mathcal{N}(\mathbf{0}, \mathbf{I}) \\ &= \sqrt{\alpha_t \alpha_{t-1}} \mathbf{x}_{t-2} + \sqrt{1 - \alpha_t \alpha_{t-1}} \bar{\mathbf{z}}_{t-2} & \text{ ;where } \bar{\mathbf{z}}_{t-2} \text{ merges two Gaussians (*).} \\ &= \dots \\ &= \sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\mathbf{z} \\ q(\mathbf{x}_t \vert \mathbf{x}_0) &= \mathcal{N}(\mathbf{x}_t; \sqrt{\bar{\alpha}_t} \mathbf{x}_0, (1 - \bar{\alpha}_t)\mathbf{I}) \end{aligned} xtq(xtx0)=αt xt1+1αt zt1=αtαt1 xt2+1αtαt1 zˉt2==αˉt x0+1αˉt z=N(xt;αˉt x0,(1αˉt)I) ;where zt1,zt2,N(0,I) ;where zˉt2 merges two Gaussians (*).

回顾一下,当我们合并两个具有不同方差的高斯,\mathcal{N}(\mathbf{0}, \sigma_1^2\mathbf{I})和\mathcal{N}(\mathbf{0}, \sigma_2^2\mathbf{I}),新的分布是\mathcal{N}(\mathbf{0}, (\sigma_1^2 + \sigma_2^2)\mathbf{I})。这里的合并标准差是\sqrt{(1-\alpha_t)+ \alpha_t (1-\alpha_{t-1})} = \sqrt{1-\alpha_t\alpha_{t-1}}。
通常情况下,当样本变得更加随机时,可以尝试更大的更新步骤,因此\beta_1 \lt; \beta_2 \lt; \dots \lt; \beta_T , \bar{\alpha}_1 \gt; \dots \gt; \bar{\alpha}_T.

与随机梯度Langevin动力学的联系
Langevin动力学是一个来自物理学的概念,为统计学上的分子系统建模而开发。与随机梯度下降相结合,随机梯度Langevin动力学(Welling & Teh 2011)可以从概率密度p(\mathbf{x})中产生样本,只使用梯度\nabla_\mathbf{x} \log p(\mathbf{x}) (前文提到的score function)的马尔科夫更新链。

\mathbf{x}t = \mathbf{x}{t-1} + \frac{\epsilon}{2} \nabla_\mathbf{x} p(\mathbf{x}_{t-1}) + \sqrt{\epsilon} \mathbf{z}_t ,\quad\text{where } \mathbf{z}_t \sim \mathcal{N}(\mathbf{0}, \mathbf{I}) \

其中\epsilon是步长。当T\to \infty, \epsilon\to 0时,\mathbf{x}_T等于真实概率密度p(\mathbf{x})。

与标准SGD相比,随机梯度Langevin动力学将高斯噪声注入参数更新中,以避免陷入局部最小值。

逆向扩散过程
把上述过程倒过来,从q(\mathbf{x}_{t-1} \vert \mathbf{x}_t)中取样,能从高斯噪声输入中重现真实样本,\mathbf{x}T \sim \mathcal{N}(\mathbf{0}, \mathbf{I}) 。注意,如果\beta_t足够小,q(\mathbf{x}{t-1} \vert \mathbf{x}t)也将是高斯的。不幸的是,我们很难估计q(\mathbf{x}{t-1} \vert \mathbf{x}t),因为它需要整个数据集,因此需要学习一个模型p\theta来近似这些条件概率,以便运行反向扩散过程。

p_\theta(\mathbf{x}{0:T}) = p(\mathbf{x}T) \prod^T{t=1} p\theta(\mathbf{x}{t-1} \vert \mathbf{x}t) \quad p\theta(\mathbf{x}{t-1} \vert \mathbf{x}t) = \mathcal{N}(\mathbf{x}{t-1}; \boldsymbol{\mu}_\theta(\mathbf{x}t, t), \boldsymbol{\Sigma}\theta(\mathbf{x}_t, t)) \

值得注意的是,当以\mathbf{x}_0为条件时,反向条件概率是可控的。

q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}0) = \mathcal{N}(\mathbf{x}{t-1}; \color{blue}{\tilde{\boldsymbol{\mu}}}(\mathbf{x}_t, \mathbf{x}_0), \color{red}{\tilde{\beta}_t} \mathbf{I}) \

使用贝叶斯法则:

\begin{aligned} q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0) &= q(\mathbf{x}t \vert \mathbf{x}{t-1}, \mathbf{x}0) \frac{ q(\mathbf{x}{t-1} \vert \mathbf{x}_0) }{ q(\mathbf{x}t \vert \mathbf{x}0) } \ &\propto \exp \Big(-\frac{1}{2} \big(\frac{(\mathbf{x}t - \sqrt{\alpha_t} \mathbf{x}{t-1})^2}{\beta_t} + \frac{(\mathbf{x}{t-1} - \sqrt{\bar{\alpha}{t-1}} \mathbf{x}0)^2}{1-\bar{\alpha}{t-1}} - \frac{(\mathbf{x}_t - \sqrt{\bar{\alpha}_t} \mathbf{x}0)^2}{1-\bar{\alpha}t} \big) \Big) \ &= \exp\Big( -\frac{1}{2} \big( \color{red}{(\frac{\alpha_t}{\beta_t} + \frac{1}{1 - \bar{\alpha}{t-1}})} \mathbf{x}{t-1}^2 - \color{blue}{(\frac{2\sqrt{\alpha_t}}{\beta_t} \mathbf{x}_t + \frac{2\sqrt{\bar{\alpha}_t}}{1 - \bar{\alpha}_t} \mathbf{x}0)} \mathbf{x}{t-1} + C(\mathbf{x}_t, \mathbf{x}_0) \big) \Big) \end{aligned} \

其中C(\mathbf{x}_t, \mathbf{x}0)是某个不涉及\mathbf{x}{t-1}的函数。按照标准的高斯密度函数,均值和方差可以被参数化如下。

\begin{aligned} \tilde{\beta}t &= 1/(\frac{\alpha_t}{\beta_t} + \frac{1}{1 - \bar{\alpha}{t-1}}) = \frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \cdot \beta_t \ \tilde{\boldsymbol{\mu}}_t (\mathbf{x}_t, \mathbf{x}_0) &= (\frac{\sqrt{\alpha_t}}{\beta_t} \mathbf{x}_t + \frac{\sqrt{\bar{\alpha}_t}}{1 - \bar{\alpha}t} \mathbf{x}0)/(\frac{\alpha_t}{\beta_t} + \frac{1}{1 - \bar{\alpha}{t-1}}) = \frac{\sqrt{\alpha_t}(1 - \bar{\alpha}{t-1})}{1 - \bar{\alpha}_t} \mathbf{x}t + \frac{\sqrt{\bar{\alpha}{t-1}}\beta_t}{1 - \bar{\alpha}_t} \mathbf{x}_0\ \end{aligned} \

可以表示\mathbf{x}_0=\frac{1}{\sqrt{\bar{\alpha}_t}}(\mathbf{x}_t - \sqrt{1 - \bar{\alpha}_t}\mathbf{z}_t)并将其插入上述方程式,得:

\begin{aligned} \tilde{\boldsymbol{\mu}}t &= \frac{\sqrt{\alpha_t}(1 - \bar{\alpha}{t-1})}{1 - \bar{\alpha}_t} \mathbf{x}t + \frac{\sqrt{\bar{\alpha}{t-1}}\beta_t}{1 - \bar{\alpha}_t} \frac{1}{\sqrt{\bar{\alpha}_t}}(\mathbf{x}_t - \sqrt{1 - \bar{\alpha}_t}\mathbf{z}_t) \ &= \color{cyan}{\frac{1}{\sqrt{\alpha_t}} \Big( \mathbf{x}_t - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}_t}} \mathbf{z}_t \Big)} \end{aligned} \

这样的设置与VAE非常相似,因此我们可以使用变分下界来优化负对数似然:

\begin{aligned} - \log p_\theta(\mathbf{x}0) &\leq - \log p\theta(\mathbf{x}0) + D\text{KL}(q(\mathbf{x}{1:T}\vert\mathbf{x}0) | p\theta(\mathbf{x}{1:T}\vert\mathbf{x}0) ) \ &= -\log p\theta(\mathbf{x}0) + \mathbb{E}{\mathbf{x}{1:T}\sim q(\mathbf{x}{1:T} \vert \mathbf{x}0)} \Big[ \log\frac{q(\mathbf{x}{1:T}\vert\mathbf{x}0)}{p\theta(\mathbf{x}{0:T}) / p\theta(\mathbf{x}0)} \Big] \ &= -\log p\theta(\mathbf{x}0) + \mathbb{E}q \Big[ \log\frac{q(\mathbf{x}{1:T}\vert\mathbf{x}0)}{p\theta(\mathbf{x}{0:T})} + \log p_\theta(\mathbf{x}0) \Big] \ &= \mathbb{E}q \Big[ \log \frac{q(\mathbf{x}{1:T}\vert\mathbf{x}0)}{p\theta(\mathbf{x}{0:T})} \Big] \ \text{Let }L_\text{VLB} &= \mathbb{E}{q(\mathbf{x}{0:T})} \Big[ \log \frac{q(\mathbf{x}{1:T}\vert\mathbf{x}0)}{p\theta(\mathbf{x}{0:T})} \Big] \geq - \mathbb{E}_{q(\mathbf{x}0)} \log p\theta(\mathbf{x}_0) \end{aligned} \

It is also straightforward to get the same result using Jensen’s inequality. Say we want to minimize the cross entropy as the learning objective,

\begin{aligned} L_\text{CE} &= - \mathbb{E}{q(\mathbf{x}0)} \log p\theta(\mathbf{x}0) \ &= - \mathbb{E}{q(\mathbf{x}0)} \log \Big( \int p\theta(\mathbf{x}{0:T}) d\mathbf{x}{1:T} \Big) \ &= - \mathbb{E}{q(\mathbf{x}0)} \log \Big( \int q(\mathbf{x}{1:T} \vert \mathbf{x}0) \frac{p\theta(\mathbf{x}{0:T})}{q(\mathbf{x}{1:T} \vert \mathbf{x}{0})} d\mathbf{x}{1:T} \Big) \ &= - \mathbb{E}{q(\mathbf{x}0)} \log \Big( \mathbb{E}{q(\mathbf{x}{1:T} \vert \mathbf{x}0)} \frac{p\theta(\mathbf{x}{0:T})}{q(\mathbf{x}{1:T} \vert \mathbf{x}{0})} \Big) \ &\leq - \mathbb{E}{q(\mathbf{x}{0:T})} \log \frac{p\theta(\mathbf{x}{0:T})}{q(\mathbf{x}{1:T} \vert \mathbf{x}{0})} \ &= \mathbb{E}{q(\mathbf{x}{0:T})}\Big[\log \frac{q(\mathbf{x}{1:T} \vert \mathbf{x}{0})}{p\theta(\mathbf{x}{0:T})} \Big] = L\text{VLB} \end{aligned} \

为了将方程中的每一个项转换为可分析的计算,目标可以进一步改写为几个KL散度和熵项的组合:

\begin{aligned} L_\text{VLB} &= \mathbb{E}{q(\mathbf{x}{0:T})} \Big[ \log\frac{q(\mathbf{x}{1:T}\vert\mathbf{x}0)}{p\theta(\mathbf{x}{0:T})} \Big] \ &= \mathbb{E}q \Big[ \log\frac{\prod{t=1}^T q(\mathbf{x}t\vert\mathbf{x}{t-1})}{ p_\theta(\mathbf{x}T) \prod{t=1}^T p_\theta(\mathbf{x}{t-1} \vert\mathbf{x}t) } \Big] \ &= \mathbb{E}q \Big[ -\log p\theta(\mathbf{x}T) + \sum{t=1}^T \log \frac{q(\mathbf{x}t\vert\mathbf{x}{t-1})}{p\theta(\mathbf{x}{t-1} \vert\mathbf{x}t)} \Big] \ &= \mathbb{E}q \Big[ -\log p\theta(\mathbf{x}T) + \sum{t=2}^T \log \frac{q(\mathbf{x}t\vert\mathbf{x}{t-1})}{p\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)} + \log\frac{q(\mathbf{x}_1 \vert \mathbf{x}0)}{p\theta(\mathbf{x}_0 \vert \mathbf{x}1)} \Big] \ &= \mathbb{E}q \Big[ -\log p\theta(\mathbf{x}T) + \sum{t=2}^T \log \Big( \frac{q(\mathbf{x}{t-1} \vert \mathbf{x}t, \mathbf{x}0)}{p\theta(\mathbf{x}{t-1} \vert\mathbf{x}_t)}\cdot \frac{q(\mathbf{x}_t \vert \mathbf{x}0)}{q(\mathbf{x}{t-1}\vert\mathbf{x}_0)} \Big) + \log \frac{q(\mathbf{x}_1 \vert \mathbf{x}0)}{p\theta(\mathbf{x}_0 \vert \mathbf{x}1)} \Big] \ &= \mathbb{E}q \Big[ -\log p\theta(\mathbf{x}T) + \sum{t=2}^T \log \frac{q(\mathbf{x}{t-1} \vert \mathbf{x}t, \mathbf{x}0)}{p\theta(\mathbf{x}{t-1} \vert\mathbf{x}t)} + \sum{t=2}^T \log \frac{q(\mathbf{x}_t \vert \mathbf{x}0)}{q(\mathbf{x}{t-1} \vert \mathbf{x}_0)} + \log\frac{q(\mathbf{x}_1 \vert \mathbf{x}0)}{p\theta(\mathbf{x}_0 \vert \mathbf{x}1)} \Big] \ &= \mathbb{E}q \Big[ -\log p\theta(\mathbf{x}T) + \sum{t=2}^T \log \frac{q(\mathbf{x}{t-1} \vert \mathbf{x}t, \mathbf{x}0)}{p\theta(\mathbf{x}{t-1} \vert\mathbf{x}_t)} + \log\frac{q(\mathbf{x}_T \vert \mathbf{x}_0)}{q(\mathbf{x}_1 \vert \mathbf{x}_0)} + \log \frac{q(\mathbf{x}_1 \vert \mathbf{x}0)}{p\theta(\mathbf{x}_0 \vert \mathbf{x}1)} \Big]\ &= \mathbb{E}q \Big[ \log\frac{q(\mathbf{x}T \vert \mathbf{x}0)}{p\theta(\mathbf{x}T)} + \sum{t=2}^T \log \frac{q(\mathbf{x}{t-1} \vert \mathbf{x}t, \mathbf{x}0)}{p\theta(\mathbf{x}{t-1} \vert\mathbf{x}t)} - \log p\theta(\mathbf{x}0 \vert \mathbf{x}1) \Big] \ &= \mathbb{E}q [\underbrace{D\text{KL}(q(\mathbf{x}T \vert \mathbf{x}0) \parallel p\theta(\mathbf{x}T))}{L_T} + \sum{t=2}^T \underbrace{D\text{KL}(q(\mathbf{x}{t-1} \vert \mathbf{x}t, \mathbf{x}0) \parallel p\theta(\mathbf{x}{t-1} \vert\mathbf{x}t))}{L{t-1}} \underbrace{- \log p\theta(\mathbf{x}_0 \vert \mathbf{x}1)}{L_0} ] \end{aligned} \

整理得:

\begin{aligned} L_\text{VLB} &= L_T + L_{T-1} + \dots + L_0 \ \text{其中 } L_T &= D_\text{KL}(q(\mathbf{x}T \vert \mathbf{x}0) \parallel p\theta(\mathbf{x}T)) \ L_t &= D\text{KL}(q(\mathbf{x}t \vert \mathbf{x}{t+1}, \mathbf{x}0) \parallel p\theta(\mathbf{x}t \vert\mathbf{x}{t+1})) \text{ for }1 \leq t \leq T-1 \ L_0 &= - \log p\theta(\mathbf{x}_0 \vert \mathbf{x}_1) \end{aligned} \

L_\text{VLB}中的每个KL项(除了L_0)都是比较两个高斯分布,因此它们可以以封闭形式计算出来。L_T是常数,在训练中可以忽略,因为q没有可学习的参数,\mathbf{x}_T是高斯噪声。Ho等人在2020年使用一个单独的离散解码器对L_0进行建模,该解码器由\mathcal{N}(\mathbf{x}0; \boldsymbol{\mu}\theta(\mathbf{x}1, 1), \boldsymbol{\Sigma}\theta(\mathbf{x}_1, 1))导出。

训练损失的L_t的参数化
回顾一下,我们需要学习一个神经网络来近似逆向扩散过程中的条件概率分布,p_\theta(\mathbf{x}{t-1} \vert \mathbf{x}t) = \mathcal{N}(\mathbf{x}{t-1}; \boldsymbol{\mu}\theta(\mathbf{x}t, t), \boldsymbol{\Sigma}\theta(\mathbf{x}t, t))。我们希望训练\boldsymbol{\mu}\theta来预测\tilde{\boldsymbol{\mu}}_t = \frac{1}{\sqrt{\alpha_t}}\Big(\mathbf{x}_t - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}_t}} \mathbf{z}_t \Big)。由于\mathbf{x}_t在训练时可作为输入,我们可以重新参数化高斯噪声项,而使其在时间步长t时从输入\mathbf{x}_t预测 \mathbf{z}_t。

\begin{aligned} \boldsymbol{\mu}_\theta(\mathbf{x}_t, t) &= \color{cyan}{\frac{1}{\sqrt{\alpha_t}} \Big( \mathbf{x}t - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}t}} \mathbf{z}\theta(\mathbf{x}t, t) \Big)} \ \text{因此 }\mathbf{x}{t-1} &= \mathcal{N}(\mathbf{x}{t-1}; \frac{1}{\sqrt{\alpha_t}} \Big( \mathbf{x}_t - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}t}} \mathbf{z}\theta(\mathbf{x}t, t) \Big), \boldsymbol{\Sigma}\theta(\mathbf{x}_t, t)) \end{aligned} \

损失项L_t的参数化是为了最小化与\tilde{\boldsymbol{\mu}}的差异。

\begin{aligned} L_t &= \mathbb{E}_{\mathbf{x}0, \mathbf{z}} \Big[\frac{1}{2 | \boldsymbol{\Sigma}\theta(\mathbf{x}_t, t) |^2_2} | \color{blue}{\tilde{\boldsymbol{\mu}}_t(\mathbf{x}_t, \mathbf{x}0)} - \color{green}{\boldsymbol{\mu}\theta(\mathbf{x}t, t)} |^2 \Big] \ &= \mathbb{E}{\mathbf{x}0, \mathbf{z}} \Big[\frac{1}{2 |\boldsymbol{\Sigma}\theta |^2_2} | \color{blue}{\frac{1}{\sqrt{\alpha_t}} \Big( \mathbf{x}_t - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}_t}} \mathbf{z}_t \Big)} - \color{green}{\frac{1}{\sqrt{\alpha_t}} \Big( \mathbf{x}_t - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}t}} \boldsymbol{\mathbf{z}}\theta(\mathbf{x}t, t) \Big)} |^2 \Big] \ &= \mathbb{E}{\mathbf{x}_0, \mathbf{z}} \Big[\frac{ \beta_t^2 }{2 \alpha_t (1 - \bar{\alpha}t) | \boldsymbol{\Sigma}\theta |^2_2} |\mathbf{z}t - \mathbf{z}\theta(\mathbf{x}t, t)|^2 \Big] \ &= \mathbb{E}{\mathbf{x}_0, \mathbf{z}} \Big[\frac{ \beta_t^2 }{2 \alpha_t (1 - \bar{\alpha}t) | \boldsymbol{\Sigma}\theta |^2_2} |\mathbf{z}t - \mathbf{z}\theta(\sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\mathbf{z}_t, t)|^2 \Big] \end{aligned} \

Simplification
根据经验,Ho等人(2020)发现,用忽略加权项的简化目标训练扩散模型效果更好。

L_t^\text{simple} = \mathbb{E}_{\mathbf{x}_0, \mathbf{z}_t} \Big[|\mathbf{z}t - \mathbf{z}\theta(\sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\mathbf{z}_t, t)|^2 \Big] \

最终目标函数为:

L_\text{simple} = L_t^\text{simple} + C \

其中C是一个不取决于\theta的常数。

Connection with noise-conditioned score networks (NCSN)
Song ; Ermon(2019)提出了一种基于分数的生成式建模方法,其中样本是通过connecting-with-stochastic-gradient-langevin-dynamicsLangevin动力学产生的,使用分数匹配估算的数据分布的梯度。每个样本\mathbf{x}的密度概率的得分被定义为其梯度\nabla_{\mathbf{x}}\log p(\mathbf{x})。一个分数网络s_\theta: \to \mathbb{R}^D 被训练来估计它。为了使其在深度学习环境下可扩展到高维数据,他们提议使用去噪分数匹配(向数据添加预先指定的小噪声;Vincent,2011)或滑动分数匹配(使用随机投影;Yang等人,2019)。

回顾一下,Langevin动力学可以从概率密度分布中抽出数据点,在一个迭代过程中只使用分数\nabla_{\mathbf{x}} log p(\mathbf{x})。

然而,根据流形假说,大部分数据预计会集中在一个低维流形中,尽管观察到的数据可能看起来只是任意的高维。由于数据点不能覆盖整个空间,这对分数估计带来了负面影响。在数据密度低的区域,分数估计的可靠性较低。在加入一个小的高斯噪声,使扰动的数据分布覆盖整个空间\mathbb{R}^D后,分数估计器网络的训练变得更加稳定。Song, Ermon(2019)通过对数据进行不同程度的噪声扰动,并训练一个噪声条件下的分数网络来联合估计所有扰动数据在不同噪声水平下的分数,对其进行改进。

噪声水平增加的机制类似于前向扩散过程。

The schedule of increasing noise levels resembles the forward diffusion process.

\beta_t 的参数化
forward variances被Ho等人(2020)设定为一系列线性增加的常数,从\beta_1=10^{-4}到\beta_T=0.02。与[-1, 1]之间的归一化图像像素值相比,它们相对较小。扩散模型在其实验中显示了高质量的样本,但仍然不能像其他生成模型那样取得有竞争力的对数似然模型。

Nichol & Dhariwal (2021)提出了几种改进技术来帮助扩散模型获得更低的NLL。其中一项改进是使用cosine-based variance schedule。schedule函数的选择可以是任意的,只要它在训练过程中间可以以近乎线性的方式下降,并在t=0和t=T附近有细微的变化。

\beta_t = \text{clip}(1-\frac{\bar{\alpha}t}{\bar{\alpha}{t-1}}, 0.999) \quad\bar{\alpha}_t = \frac{f(t)}{f(0)}\quad\text{其中 }f(t)=\cos\Big(\frac{t/T+s}{1+s}\cdot\frac{\pi}{2}\Big) \

其中小的偏移量s是为了防止\beta_t在接近t=0时太小。

逆向过程方差 \boldsymbol{\Sigma}\theta的参数化
Ho等人 (2020)选择将\beta_t固定为常数,而不是可学习参数,并设置\boldsymbol{\Sigma}
\theta(\mathbf{x}_t t) = \sigma^2_t \mathbf{I} ,其中 \sigma_t 不是学习的参数,而是设置为 \beta_t 或 \tilde{\beta}t = \frac{1 - \bar{\alpha}{t-1}}{1 - \bar{\alpha}t} \cdot \beta_t。因为他们发现,学习对角线方差\boldsymbol{\Sigma}\theta会导致训练不稳定,样本质量较差。

Nichol; Dhariwal (2021)提出通过模型预测混合向量\mathbf{x}t,学习\boldsymbol{\Sigma}\theta(\mathbf{x}_t, t)作为\beta_t和\tilde{\beta}_t之间的内插值。

\boldsymbol{\Sigma}_\theta(\mathbf{x}_t, t) = \exp(\mathbf{v} \log \beta_t + (1-\mathbf{v}) \log \tilde{\beta}_t) \

然而,简单目标函数L_\text{simple}并不依赖于\boldsymbol{\Sigma}\theta。为了增加依赖性,他们构建了一个混合目标L\text{hybrid}=L_\text{simple} + \lambda L_\text{VLB} 其中\lambda=0.001是在L_\text{VLB}项中对\boldsymbol{\mu}\theta的停止梯度,这样L\text{VLB}只指导\boldsymbol{\Sigma}\theta的学习。根据经验,他们观察到L\text{VLB}可能由于噪声梯度的原因,优化起来相当具有挑战性,因此他们建议使用时间平均平滑版本的带有重要性抽样的L_\text{VLB}。

加速扩散模型采样
以马尔可夫链为主的反向扩散过程从DDPM中生成样本是非常缓慢的,因为T可以达到一或几千步。Song等人2020年的一个数据:从DDPM中对大小为32×32的50k图像进行采样需要约20个小时,但在Nvidia 2080 Ti GPU上从GAN中进行采样则不到一分钟。

一个简单的方法是运行一个strided采样时间表(Nichol ; Dhariwal, 2021),通过每\lceil T/S \rceil步进行采样更新,将该过程从T减少到S步。生成的新抽样计划是{\tau_1, \dots, \tau_S},其中\tau_1 \lt; \tau_2 \lt; \dots \lt;\tau_S\in [1, T] 和 S \lt; T。

对于另一种方法,让我们重写q_\sigma(\mathbf{x}_{t-1}\vert \mathbf{x}_t, \mathbf{x}_0),根据nice属性,用期望的标准偏差\sigma_t作为参数。

\begin{aligned} \mathbf{x}{t-1} &= \sqrt{\bar{\alpha}{t-1}}\mathbf{x}0 + \sqrt{1 - \bar{\alpha}{t-1}}\mathbf{z}{t-1} \ &= \sqrt{\bar{\alpha}{t-1}}\mathbf{x}0 + \sqrt{1 - \bar{\alpha}{t-1} - \sigma_t^2} \mathbf{z}t + \sigma_t\mathbf{z} \ &= \sqrt{\bar{\alpha}{t-1}}\mathbf{x}0 + \sqrt{1 - \bar{\alpha}{t-1} - \sigma_t^2} \frac{\mathbf{x}_t - \sqrt{\bar{\alpha}_t}\mathbf{x}0}{\sqrt{1 - \bar{\alpha}t}} + \sigma_t\mathbf{z} \ q\sigma(\mathbf{x}{t-1} \vert \mathbf{x}t, \mathbf{x}0) &= \mathcal{N}(\mathbf{x}{t-1}; \sqrt{\bar{\alpha}{t-1}}\mathbf{x}0 + \sqrt{1 - \bar{\alpha}{t-1} - \sigma_t^2} \frac{\mathbf{x}_t - \sqrt{\bar{\alpha}_t}\mathbf{x}_0}{\sqrt{1 - \bar{\alpha}_t}}, \sigma_t^2 \mathbf{I}) \end{aligned} \

回顾 q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}0) = \mathcal{N}(\mathbf{x}{t-1}; \tilde{\boldsymbol{\mu}}(\mathbf{x}_t, \mathbf{x}_0), \tilde{\beta}_t \mathbf{I}), 有:

\tilde{\beta}t = \sigma_t^2 = \frac{1 - \bar{\alpha}{t-1}}{1 - \bar{\alpha}_t} \cdot \beta_t \

让\sigma_t^2 = \eta \cdot \tilde{\beta}_t,我们可以调整\eta \in \mathbb{R}^+作为超参数来控制采样的随机性。在\eta = 0的特殊情况下,采样过程是确定的。这样的模型被命名为去噪扩散隐含模型(DDIM;Song等人,2020)。DDIM具有相同的边际噪声分布,但确定性地将噪声映射回原始数据样本。

在生成过程中,我们只对S扩散步骤的子集{\tau_1, \dots, \tau_S}进行采样,推理过程变成:

q_{\sigma, \tau}(\mathbf{x}{\tau{i-1}} \vert \mathbf{x}{\tau_t}, \mathbf{x}0) = \mathcal{N}(\mathbf{x}{\tau{i-1}}; \sqrt{\bar{\alpha}{t-1}}\mathbf{x}0 + \sqrt{1 - \bar{\alpha}{t-1} - \sigma_t^2} \frac{\mathbf{x}{\tau_i} - \sqrt{\bar{\alpha}_t}\mathbf{x}_0}{\sqrt{1 - \bar{\alpha}_t}}, \sigma_t^2 \mathbf{I}) \

虽然在实验中所有的模型都是用T=1000的扩散步骤训练的,但他们观察到,当S较小时,DDIM(\eta=0)可以产生质量最好的样本,而DDPM(\eta=1)在小S上表现得更差。当我们有能力运行完整的反向马尔科夫扩散步骤(S=T=1000)时,DDPM确实表现更好。使用DDIM,可以训练扩散模型到任意数量的正向步骤,但只从生成过程中的一个步骤子集中取样。

与DDPM相比,DDIM能够:

用更少的步骤产生更高质量的样本。
具有一致性属性,因为生成过程是确定性的,这意味着以同一潜在变量为条件的多个样本应该具有类似的高层特征。
由于具有一致性,DDIM可以在隐变量中做有语义的插值。
条件性生成
在ImageNet数据上训练生成式模型时,通常是以类别标签为条件生成样本。为了明确地将类别信息纳入扩散过程,Dhariwal & Nichol (2021)在嘈杂的图像\mathbf{x}t上训练了一个分类器f\phi(y \vert \mathbf{x}t, t)并使用梯度\nabla\mathbf{x} log f_\phi(y \vert \mathbf{x}_t, t)来引导扩散采样过程向目标类别标签y移动。他们的消融扩散模型(ADM)和带有额外分类器指导的模型(ADM-G)能够取得比SOTA生成模型(如BigGAN)更好的结果。

此外,Dhariwal & Nichol(2021)在UNet架构上做了一些修改,显示出比GAN的扩散模型更好的性能。架构的修改包括更大的模型深度/宽度、更多的注意力头、多分辨率注意力、用于上/下采样的BigGAN残差块、以1/\sqrt{2}重新划分残差连接和自适应组归一化(AdaGN)。

小结
优点:可操作性和灵活性是生成模型中两个相互冲突的目标。可操作的模型可以被分析评估并简单的拟合数据数据(例如通过高斯或拉普拉斯),但它们不容易描述丰富数据集中的结构。灵活的模型可以适合数据中的任意结构,但评估、训练或从这些模型中采样通常是昂贵的。扩散模型同时具有可操作性和灵活性。

缺点:扩散模型依赖于长的马尔可夫链的扩散步骤来产生样本,所以在时间和计算上可能相当昂贵。新的方法已经被提出来,使这个过程快得多,但采样仍然比GAN慢。

参考文献
[1] Jascha Sohl-Dickstein et al. “Deep Unsupervised Learning using Nonequilibrium Thermodynamics.” ICML 2015.

[2] Max Welling & Yee Whye Teh. “Bayesian learning via stochastic gradient langevin dynamics.” ICML 2011.

[3] Yang Song & Stefano Ermon. “Generative modeling by estimating gradients of the data distribution.” NeurIPS 2019.

[4] Yang Song & Stefano Ermon. “Improved techniques for training score-based generative models.” NeuriPS 2020.

[5] Jonathan Ho et al. “Denoising diffusion probabilistic models.” arxiv Preprint arxiv:2006.11239 (2020). [code]

[6] Jiaming Song et al. “Denoising diffusion implicit models.” arxiv Preprint arxiv:2010.02502 (2020). [code]

[7] Alex Nichol & Prafulla Dhariwal. “ Improved denoising diffusion probabilistic models” arxiv Preprint arxiv:2102.09672 (2021). [code]

[8] Prafula Dhariwal & Alex Nichol. “Diffusion Models Beat GANs on Image Synthesis." arxiv Preprint arxiv:2105.05233 (2021). [code]




参考资料:
生成模型(四):扩散模型

你可能感兴趣的:(#,CV/生成模型,人工智能,算法)