扩散模型

社区文章 发布于2024年5月19日

** github 代码 **

image/png

首先,我们需要定义什么是生成模型?

生成模型学习数据集的概率分布,以便我们可以从该分布中采样以创建新的数据实例。例如,如果我们有很多猫的图片,并且我们训练了一个生成模型,那么我们就可以从这个分布中采样来创建新的猫的图像。

现在,什么是扩散模型?扩散模型是生成模型,在过去几年中获得了显著的流行,这有充分的理由。仅在2020年代发布的一些开创性论文就向世界展示了扩散模型的能力,例如在图像合成方面击败[GANs]。最近,从业者会看到DALL-E 2中使用了扩散模型,这是OpenAI上个月发布的图像生成模型。

扩散模型——介绍

扩散模型是**生成**模型,这意味着它们用于生成与训练数据相似的数据。从根本上说,扩散模型通过连续添加高斯噪声来**破坏训练数据**,然后通过**逆转**这种去噪过程来**学习恢复**数据。训练后,我们可以通过简单地**将随机采样的噪声通过学习到的去噪过程**来生成数据。

扩散模型受**非平衡热力学**的启发。它们定义了扩散步骤的**马尔可夫链**,以缓慢地向数据添加随机噪声,然后学习逆转扩散过程,从噪声中构建所需的数据样本。与**VAE或流模型**不同,扩散模型采用固定过程学习,并且潜在变量具有高维度(与原始数据相同)。


现在深入了解扩散模型:

扩散模型由以下两个过程组成,如下图所示

  • 前向过程(红线)。
  • 逆向过程(蓝线)。

image/png

如前所述,扩散模型由**前向过程**(或**扩散过程**)组成,其中数据(通常是图像)逐渐被噪声化,以及**逆向过程**(或**逆向扩散过程**),其中噪声被转换回目标分布中的样本。

对于图像来说,更详细的设置包括两个过程

  • 一个我们选择的固定(或预定义)前向扩散过程 `q`,它逐渐向图像添加高斯噪声,直到最终变成纯噪声。
  • 一个学习到的逆向去噪扩散过程 `p_θ`,其中训练一个神经网络,从纯噪声开始逐渐去噪图像,直到最终得到一个真实的图像。

1. 前向过程(固定):

当前向过程中的噪声水平足够低时,采样链的转换可以设置为条件高斯分布。将这一事实与马尔可夫假设相结合,可以得到前向过程的简单参数化:

q(xtxt1)=N(xt;1βtxt1,βtI)q(x1:Tx0)=t=1Tq(xtxt1) 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})

其中 `beta_1, ..., beta_T` 是方差调度(可以是学习到的,也可以是固定的),如果表现良好,它**确保当 T 足够大时,x_T 几乎是一个各向同性高斯分布**。数据样本 `f{x}_0` 随着步长 `t` 的增大逐渐失去其可辨别的特征。最终,当 `T` 趋于无穷时,它等同于一个各向同性高斯分布。

image/png

2. 逆向过程(学习)

如前所述,扩散模型的“魔力”在于**逆向过程**。在训练过程中,模型学习逆转这个扩散过程,以生成新的数据。从纯高斯噪声开始,模型学习联合分布,如下所示:

最终,图像渐近地转换为纯高斯噪声。训练扩散模型的**目标**是学习**逆向**过程——即训练。通过沿着这条链反向遍历,我们可以生成新的数据。

pθ(x0:T)=p(xT)t=1Tpθ(xt1xt)pθ(xt1xt)=N(xt1;μθ(xt,t),Σθ(xt,t)) 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))

image/png

其中高斯转换的时间相关参数是学习得到的。特别值得注意的是,马尔可夫公式断言,给定的逆向扩散转换分布仅取决于前一个时间步(或后一个时间步,取决于如何看待)。

前向和逆向过程都以 `t` 为索引,在有限时间步 `T` 内发生(DDPM 作者使用 `T=1000`)。你从 `t=0` 开始,从数据分布中采样一个真实图像 `x_0`(比如 ImageNet 中的猫图片),然后前向过程在每个时间步 `t` 从高斯分布中采样一些噪声,并将其添加到前一个时间步的图像中。给定足够大的 `T` 和一个良好行为的噪声添加调度,你最终会通过一个渐进过程得到一个名为各向同性高斯分布的 `t=T`。


去噪扩散概率模型

论文背景:

image/png

红色的数学方程是前向过程,黄色的方程是逆向过程,中间用数字 (3) 表示的方程是逆向过程的学习过程。

本文介绍了扩散概率模型 [53] 的进展。扩散概率模型(为简洁起见,我们称之为“扩散模型”)是一种参数化的马尔可夫链,通过变分推理进行训练,以在有限时间后生成与数据匹配的样本。该链的转换被学习以逆转扩散过程,扩散过程是一个马尔可夫链,它逐渐向数据中添加噪声,与采样方向相反,直到信号被破坏。当扩散由少量高斯噪声组成时,将采样链转换设置为条件高斯分布也足够,从而允许特别简单的神经网络参数化。扩散模型易于定义且训练高效,但据我们所知,尚未证明它们能够生成高质量样本。我们证明扩散模型实际上能够生成高质量样本,有时甚至优于其他类型生成模型已发表的结果(第 4 节)。此外,我们展示了扩散模型的特定参数化揭示了训练期间去噪分数匹配与多个噪声水平的等效性,以及采样期间与退火朗之万动力学的等效性(第 3.2 节)[55, 61]。我们使用这种参数化获得了最佳的样本质量结果(第 4.2 节),因此我们认为这种等效性是我们主要贡献之一。尽管样本质量高,但我们的模型的对数似然与其他基于似然的模型相比不具有竞争力(然而,我们的模型的对数似然优于退火重要性采样对基于能量的模型和分数匹配报告的较大估计 [11, 55])。我们发现我们的模型的大部分无损码长都用于描述难以察觉的图像细节(第 4.3 节)。我们用有损压缩的语言对这一现象进行了更精细的分析,并且我们表明扩散模型的采样过程是一种渐进式解码,类似于沿位序的自回归解码,这极大地概括了自回归模型通常可能实现的功能。

请注意,前向过程是固定的,我们只是使用公式向图像添加噪声,但逆向过程是扩散模型的主要公式,扩散模型实际上是在这里学习的。但是我们如何才能通过仅仅使用逆向过程来使模型学习呢?扩散模型通过**寻找使训练数据似然最大化的逆马尔可夫转换**进行训练。在实践中,训练等效地包括最小化负对数似然的变分上界。

论文实验:

image/png


数学解释:

image/png q(xtxt1)=N(xt;1βtxt1,βtI)q(x1:Tx0)=t=1Tq(xtxt1) 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})

上述过程的一个很好的特性是,我们可以使用重新参数化技巧在任意时间步 `t` 以封闭形式进行采样

xt=αtxt1+1αtϵt1 ;where ϵt1,ϵt2,N(0,I)=αtαt1xt2+1αtαt1ϵˉt2 ;where ϵˉt2 merges two Gaussians (*).==αˉtx0+1αˉtϵq(xtx0)=N(xt;αˉtx0,(1αˉt)I) \begin{aligned} \mathbf{x}_t &= \sqrt{\alpha_t}\mathbf{x}_{t-1} + \sqrt{1 - \alpha_t}\boldsymbol{\epsilon}_{t-1} & \text{ ;where } \boldsymbol{\epsilon}_{t-1}, \boldsymbol{\epsilon}_{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{\boldsymbol{\epsilon}}_{t-2} & \text{ ;where } \bar{\boldsymbol{\epsilon}}_{t-2} \text{ merges two Gaussians (*).} \\ &= \dots \\ &= \sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\boldsymbol{\epsilon} \\ 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}

如前所述,逆向过程是模型必须学习执行的任务,那么如何做到这一点呢?通过**找到使训练数据似然最大化的逆向马尔可夫转换**。在实践中,训练等效地包括最小化负对数似然的变分上界。

[!提示] 首先,您应该了解什么是证据下界**ELBO**,**变分下界**(VLB),什么是**变分自编码器**(VAE)模型,它如何工作,以及什么是**Kull-back散度**(D_kl),因为变分自编码器是扩散模型的核心。

需要了解的基本数学概念:

image/png image/png image/png image/png image/png


现在我们理解了ELBO和D_{KL}的推导,接下来将推导反向过程,即找到能够最大化训练数据似然的反向马尔可夫转移。在实践中,训练等效于最小化负对数似然的变分上界。

请注意,Lvlb 在技术上是一个我们试图最小化的**上**界(ELBO的负值),但为了与文献保持一致,我们仍称其为Lvlb。

因此,我们可以使用变分下界来优化负对数似然。

logpθ(x0)logpθ(x0)+DKL(q(x1:Tx0)pθ(x1:Tx0))=logpθ(x0)+Ex1:Tq(x1:Tx0)[logq(x1:Tx0)pθ(x0:T)/pθ(x0)]=logpθ(x0)+Eq[logq(x1:Tx0)pθ(x0:T)+logpθ(x0)]=Eq[logq(x1:Tx0)pθ(x0:T)]Let LVLB=Eq(x0:T)[logq(x1:Tx0)pθ(x0:T)]Eq(x0)logpθ(x0) \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, the following is **optional: LCE=Eq(x0)logpθ(x0)=Eq(x0)log(pθ(x0:T)dx1:T)=Eq(x0)log(q(x1:Tx0)pθ(x0:T)q(x1:Tx0)dx1:T)=Eq(x0)log(Eq(x1:Tx0)pθ(x0:T)q(x1:Tx0))Eq(x0:T)logpθ(x0:T)q(x1:Tx0)=Eq(x0:T)[logq(x1:Tx0)pθ(x0:T)]=LVLB \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}

now here is the Variational lower bound derivation for the revers process: LVLB=Eq(x0:T)[logq(x1:Tx0)pθ(x0:T)]=Eq[logt=1Tq(xtxt1)pθ(xT)t=1Tpθ(xt1xt)]=Eq[logpθ(xT)+t=1Tlogq(xtxt1)pθ(xt1xt)]=Eq[logpθ(xT)+t=2Tlogq(xtxt1)pθ(xt1xt)+logq(x1x0)pθ(x0x1)]=Eq[logpθ(xT)+t=2Tlog(q(xt1xt,x0)pθ(xt1xt)q(xtx0)q(xt1x0))+logq(x1x0)pθ(x0x1)]=Eq[logpθ(xT)+t=2Tlogq(xt1xt,x0)pθ(xt1xt)+t=2Tlogq(xtx0)q(xt1x0)+logq(x1x0)pθ(x0x1)]=Eq[logpθ(xT)+t=2Tlogq(xt1xt,x0)pθ(xt1xt)+logq(xTx0)q(x1x0)+logq(x1x0)pθ(x0x1)]=Eq[logq(xTx0)pθ(xT)+t=2Tlogq(xt1xt,x0)pθ(xt1xt)logpθ(x0x1)]=Eq[DKL(q(xTx0)pθ(xT))LT+t=2TDKL(q(xt1xt,x0)pθ(xt1xt))Lt1logpθ(x0x1)L0] \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}

让我们分别标注变分下界损失中的每个组成部分

LVLB=LT+LT1++L0where LT=DKL(q(xTx0)pθ(xT))Lt=DKL(q(xtxt+1,x0)pθ(xtxt+1)) for 1tT1L0=logpθ(x0x1) \begin{aligned} L_\text{VLB} &= L_T + L_{T-1} + \dots + L_0 \\ \text{where } 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}

如需更好地理解推导过程,请参阅**手写推导**。


以下是训练和采样算法:

image/png

训练算法:

换句话说,

  • 我们从真实未知且可能复杂的数据分布 q(x₀) 中抽取一个随机样本 x₀。
  • 我们从 1 到 T 之间均匀采样一个噪声水平 t(即一个随机时间步)。
  • 我们从高斯分布中抽取一些噪声,并使用此噪声在 t 水平(使用上面定义的良好属性)上破坏输入。
  • 神经网络被训练用于根据受损图像 x_t 预测此噪声(即基于已知调度 β_t 对 x_0 施加的噪声)。

[!提示] 实际上,所有这些都在数据批次上完成,因为我们使用随机梯度下降法优化神经网络。要了解为什么我们使用 SGD 而不是 BGD,请阅读本节随机梯度下降与批量梯度下降。

采样算法:

image/png


模型架构:

image/png

通过在较低维度的*潜在*空间而不是实际像素空间应用扩散过程,潜在扩散可以降低内存和计算复杂性。这是标准扩散模型与潜在扩散模型之间的关键区别:**在潜在扩散模型中,模型被训练以生成图像的潜在(压缩)表示。**

image/png

模型管道:**它们在Stable Diffusion管道中的作用是什么**

潜在扩散模型主要包含三个部分。

  1. 一个自动编码器 (VAE)。自动编码器包含两个部分 -
    1. `编码器` 接收图像作为输入,并将其转换为低维潜在表示。
    2. `解码器` 接收潜在表示,并将其转换回图像。

    image/jpeg

    image/png

如上所示,编码器就像一个压缩器,将图像压缩到更低的维度,而解码器则从压缩版本中重新创建原始图像。

p = FastDownload().download('https://lafeber.com/pet-birds/wp-content/uploads/2018/06/Scarlet-Macaw-2.jpg') 
img = load_image(p) 
print(f"Dimension of this image: {np.array(img).shape}") img	
Dimension of this image: (512, 512, 3)

image/png 现在我们用 VAE 编码器压缩这张图像,我们将使用 `pil_to_latents` 辅助函数。

latent_img = pil_to_latents(img) 
print(f"Dimension of this latent representation: {latent_img.shape}")
Dimension of this latent representation: torch.Size([1, 4, 64, 64])

我们可以看到 VAE 如何将 3 x 512 x 512 维的图像压缩成 4 x 64 x 64 的图像。这达到了 48 倍的压缩率!让我们可视化这四个潜在表示通道。

fig, axs = plt.subplots(1, 4, figsize=(16, 4)) 
for c in range(4):
    axs[c].imshow(latent_img[0][c].detach().cpu(), cmap='Greys')

image/png

这个潜在表示理论上应该捕获了关于原始图像的大量信息。让我们在这个表示上使用解码器,看看能得到什么。为此,我们将使用 `latents_to_pil` 辅助函数。

decoded_img = latents_to_pil(latent_img) 
decoded_img[0]

image/png 从上图中我们可以看出,VAE解码器能够从48倍压缩的潜在表示中恢复原始图像。这令人印象深刻!

如果您仔细观察解码后的图像,会发现它与原始图像并不完全相同,注意眼睛周围的差异。这就是为什么VAE编码器/解码器不是无损压缩。

  1. 一个U-Net

image/png

U-Net 模型接收两个输入:
1. `带噪声的潜在` 或 `噪声` - 带噪声的潜在是 VAE 编码器生成的潜在(如果提供了初始图像),并添加了噪声;或者在只想根据文本描述创建随机新图像的情况下,它可以接收纯噪声输入。
2. `文本嵌入` - 由输入的文本提示生成的基于 CLIP 的嵌入。

U-Net 模型的输出是预测的噪声残差,即输入噪声潜在中包含的噪声。换句话说,它预测了从噪声潜在中减去的噪声,以返回原始去噪潜在。

image/png

from diffusers import UNet2DConditionModel, LMSDiscreteScheduler [][]
## Initializing a scheduler 
scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000) 
## Setting number of sampling steps 
scheduler.set_timesteps(51)

unet = UNet2DConditionModel.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="unet", torch_dtype=torch.float16).to("cuda")

正如您可能从上面的代码中注意到的,我们不仅导入了 `unet`,还导入了 `scheduler`。`scheduler` 的目的是确定在扩散过程的给定步骤中要向潜在中添加多少噪声。让我们可视化调度器函数 -

image/png

扩散过程遵循这个采样调度,我们从高噪声开始,逐渐去噪图像。让我们可视化这个过程 -

image/png 让我们看看 U-Net 如何从图像中去除噪声。我们首先向图像添加一些噪声。

image/png

prompt = [""]

text_input = tokenizer(prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")

with torch.no_grad(): 
text_embeddings = text_encoder(text_input.input_ids.to("cuda"))[0]

latent_model_input =torch.cat([encoded_and_noised.to("cuda").float()]).half() with torch.no_grad(): 
noise_pred = unet(latent_model_input, 40, encoder_hidden_states=text_embeddings)["sample"]

latents_to_pil(encoded_and_noised- noise_pred)[0]

image/png

如上所示,U-Net 的输出比传入的原始带噪输入更清晰。

  1. 一个文本编码器,*例如* CLIP的文本编码器

任何机器学习模型都无法理解文本数据。为了让模型理解文本数据,我们需要将文本转换为数字,这些数字承载着文本的含义,称为`嵌入`。将文本转换为数字的过程可以分为两个部分:
1. **分词器** - 将每个单词分解为子词,然后使用查找表将其转换为数字。
2. **词元到嵌入编码器** - 将这些数字子词转换为包含该文本表示的表示。

image/png

image/png

Stable Diffusion 只使用 CLIP 训练的编码器进行文本到嵌入的转换。这成为 U-net 的输入之一。从高层次上看,CLIP 使用图像编码器和文本编码器来创建在潜在空间中相似的嵌入。这种相似性更精确地定义为一种对比目标

tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14", torch_dtype=torch.float16) 

text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14", torch_dtype=torch.float16).to("cuda")

prompt = ["a dog wearing hat"] []tok =tokenizer(prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt") []print(tok.input_ids.shape) []tok
torch.Size([1, 77])
{'input_ids': tensor([[49406,   320,  1929,  3309,  3801, 49407, 49407, 49407, 49407, 49407,
         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
         49407, 49407, 49407, 49407, 49407, 49407, 49407]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
         0, 0, 0, 0, 0]])}

它们在 Stable Diffusion 管道中的作用

潜在扩散模型使用U-Net在潜在空间中逐步去除噪声,经过多个步骤最终达到所需的输出。在每个步骤中,添加到潜在的噪声量会逐渐减少,直到我们达到最终的去噪输出。U-Net最初由这篇论文提出,用于生物医学图像分割。U-Net包含一个编码器和一个解码器,由ResNet块组成。Stable Diffusion的U-Net还具有交叉注意力层,使其能够根据提供的文本描述来调节输出。交叉注意力层被添加到U-Net的编码器和解码器部分,通常在ResNet块之间。您可以在此处了解更多关于此U-Net架构的信息。

image/png

稳定扩散模型接收文本输入和种子。文本输入随后通过 CLIP 模型生成大小为 77x768 的文本嵌入,种子用于生成大小为 4x64x64 的高斯噪声,作为第一个潜在图像表示。

接下来,U-Net 在文本嵌入的条件下,迭代地对随机潜在图像表示进行去噪。U-Net 的输出是预测的噪声残差,然后通过调度器算法用于计算条件潜在。去噪和文本条件化的过程重复 N 次(我们将使用 50 次),以获得更好的潜在图像表示。一旦此过程完成,潜在图像表示 (4x64x64) 将由 VAE 解码器解码,以检索最终输出图像 (3x512x512)。


基于潜在扩散模型的高分辨率图像合成

1. 论文导言:

扩散模型 [82] 是一种概率模型,旨在通过逐渐去噪正态分布变量来学习数据分布 p(x),这对应于学习固定长度为 T 的马尔可夫链的逆过程。对于图像合成,最成功的模型 [15,30,72] 依赖于 p(x) 变分下限的重加权变体,这与去噪分数匹配 [85] 相呼应。这些模型可以解释为去噪自动编码器 θ(xt, t) 的等权重序列;t = 1 . . . T,它们被训练来预测其输入 xt 的去噪变体,其中 xt 是输入 x 的噪声版本。相应的目标可以简化为 (Sec. B)

image/png

其中 t 从 {1, . . . , T} 中均匀采样。潜在表示的生成建模 借助我们训练好的感知压缩模型(由 E 和 D 组成),我们现在可以访问一个高效、低维的潜在空间,其中高频、难以察觉的细节被抽象掉。与高维像素空间相比,这个空间更适合基于似然的生成模型,因为它们现在可以 (i) 专注于数据中重要的语义位,以及 (ii) 在低维、计算效率更高的空间中进行训练。与之前依赖于高度压缩、离散潜在空间中的自回归、基于注意力的变压器模型 [23,66,103] 的工作不同,我们可以利用我们模型提供的图像特定归纳偏差。这导致我们能够使用相对温和的压缩率,并实现非常好的重建效果,因为我们能够利用学习到的潜在空间 z = E(x) 的二维结构。这与之前的工作 [23, 66] 形成对比,后者依赖于学习空间 z 的任意一维排序来自动回归地建模其分布,从而忽略了 z 的许多固有结构。因此,我们的压缩模型更好地保留了 x 的细节(参见表 8)。完整的训练目标和细节可以在补充材料中找到。 3.2. 潜在扩散模型 扩散模型 [82] 是一种概率模型,旨在通过逐渐去噪正态分布变量来学习数据分布 p(x),这对应于学习固定长度为 T 的马尔可夫链的逆过程。对于图像合成,最成功的模型 [15,30,72] 依赖于 p(x) 变分下限的重加权变体,这与去噪分数匹配 [85] 相呼应。这些模型可以解释为去噪自动编码器 θ(xt, t) 的等权重序列;t = 1 . . . T,它们被训练来预测其输入 x_t 的去噪变体,其中 x_t 是输入 x 的噪声版本。相应的目标可以简化为 (Sec. B) LDM = Ex,∼N(0,1),th k − θ(xt, t)k 2^2 i , (1) 其中 t 从 {1, . . . , T} 中均匀采样。潜在表示的生成建模 借助我们训练好的感知压缩模型(由 E 和 D 组成),我们现在可以访问一个高效、低维的潜在空间,其中高频、难以察觉的细节被抽象掉。与高维像素空间相比,这个空间更适合基于似然的生成模型,因为它们现在可以 (i) 专注于数据中重要的语义位,以及 (ii) 在低维、计算效率更高的空间中进行训练。与之前依赖于高度压缩、离散潜在空间中的自回归、基于注意力的变压器模型 [23,66,103] 的工作不同,我们可以利用我们模型提供的图像特定归纳偏差。这包括构建底层 U-Net 主要由 2D 卷积层组成的能力,并使用重加权界限将目标进一步集中在感知上最相关的位上,该界限现在为 LLDM := EE(x),∼N(0,1),th k − θ(zt, t)k 2^2 i . (2) 我们模型的神经骨干 θ(◦, t) 实现为时间条件 U-Net [71]。由于前向过程是固定的,因此在训练期间可以从 E 高效地获得 z_t,并且可以通过 D 的一次传递将 p(z) 的样本解码到图像空间。

image/png

2. 论文实验:

image/png


随机梯度下降与批量梯度下降:

image/png

  1. 批量梯度下降 (BGD):

    • 在 BGD 中,模型参数是使用在整个数据集上计算的梯度来更新的。
    • 这意味着对于每次迭代,梯度都是通过考虑整个数据集来计算的,这可能导致更新速度较慢。
    • BGD 确保了更精确的梯度估计,因为它考虑了完整的数据集。
    • 然而,对于大型数据集而言,BGD 可能在计算上和内存上都非常昂贵,因为它需要一次性存储和处理整个数据集。
  2. 随机梯度下降 (SGD):

    • 在SGD中,模型参数是使用从单个随机选择的数据点或一小部分数据点(小批量)计算的梯度进行更新的。
    • 这导致更新速度更快,因为每次迭代只考虑一小部分数据集。
    • SGD 由于依赖单个或小批量数据点,在参数更新中引入了更多噪声。
    • 尽管存在噪声,SGD 仍然可以更容易地逃离局部最小值并更快地收敛,尤其是在嘈杂或高维数据集中。
    • 此外,SGD 的内存占用更小,因为它每次只需处理单个数据点或一小部分数据。

在扩散模型中,SGD 通常优于 BGD,原因如下:

  • **效率:**扩散模型通常处理大型数据集或高维数据。SGD 在内存使用和计算速度方面的效率使其比 BGD 更适用于这些场景。

  • **对噪声的鲁棒性:**扩散模型通常涉及噪声数据或复杂的相互作用。SGD 的随机性有助于它在噪声中导航并收敛到合理的解决方案。

  • **可扩展性:**随着数据集的增大,BGD 的计算和内存需求变得 prohibitive。SGD 处理大型数据集的可扩展性使其成为首选。


变分自编码器 (VAE):

image/png

变分自编码器(VAE)是一种用于机器学习和人工智能的生成模型。它是一个概率模型,旨在以无监督的方式学习输入数据的低维表示。VAE在生成与训练数据相似的新数据样本等任务中特别受欢迎。

VAE 通常的工作原理如下:

  1. **编码器网络(识别模型):**编码器网络接收输入数据(例如图像、文本等),并将其映射到潜在空间中的概率分布(通常是高斯分布)。该网络学习将输入数据编码为潜在表示。编码器网络可以是用于图像数据的卷积神经网络(CNN),也可以是用于文本等序列数据的循环神经网络(RNN)。

  2. 潜在空间:潜在空间是一个低维空间,其中每个点都代表输入数据的潜在编码。为了便于采样,该空间通常被限制为具有某些属性,例如高斯分布。

  3. 采样:一旦编码器网络生成了潜在空间中概率分布的参数(均值和方差),就会从该分布中采样一个点。该采样点表示输入数据的潜在表示。

  4. **解码器网络(生成模型):**解码器网络接收采样到的潜在点,并从中重建输入数据。它学习将潜在表示解码回原始数据空间。与编码器一样,解码器可以是 CNN、RNN 或其他适合数据类型的架构。

  5. **重构损失:**重构损失衡量输入数据与解码器重构数据之间的差异。对于二元数据(例如图像),此损失通常是交叉熵损失;对于连续数据,则是均方误差。

  6. **正则化损失(KL散度):**除了重建损失,VAE还包含一个称为KL散度的正则化项。该项鼓励潜在空间遵循特定分布,通常是单位高斯分布。它有助于确保潜在空间是连续且结构良好的,从而便于有意义的插值和新数据样本的生成。

  7. 训练:在训练期间,VAE 优化重建损失和 KL 散度正则化项的组合。目标是最小化重建损失,同时使潜在空间接近所需分布。

  8. **生成:**一旦训练完成,解码器网络就可以通过从潜在空间采样点并对其进行解码来生成新的数据样本。通过探索潜在空间的不同区域,VAE 可以生成多样且逼真的数据样本。

image/png

就像你用 Python 生成一个 1 到 100 之间的随机数一样,你是在从 1 到 100 之间的均匀(伪)随机分布中采样。同样地,我们可以从潜在空间中采样,以生成一个随机向量,然后将其输入解码器,从而生成新的数据。

image/png

理解潜在空间:

通过柏拉图的洞穴寓言

image/png

在这个寓言中,我们看到囚犯们看到雕像的影子,他们相信他们所看到的就是实际的物体(可观察到的数据)。然而,与此同时,实际的物体在他们身后(潜在的、隐藏的数据)。

我不会深入探讨 VAE,但我将解释我们需要了解的关于扩散模型的最重要的内容:ELBO、库尔巴克散度、VLB 以及我们需要了解的推导。

ELBO:

在变分自编码器(VAE)的背景下,证据下界(ELBO)是用于这些模型训练和评估的关键概念。VAE是生成模型,它们在潜在空间中学习输入数据(通常是图像或序列)的概率表示。它们由两个主要组件组成:一个编码器网络和一个解码器网络。

ELBO作为数据对数似然的下界,在训练过程中的优化中发挥着关键作用。它源于变分推断框架,其目标是近似给定观测数据的潜在变量的真实后验分布。

以下是ELBO在VAE中如何公式化的:

令 p(x,z) 表示数据 x 和潜在变量 $z$ 的联合分布,令 $q(z∣x)$ 为变分分布,它近似于真实的后验 p(z∣x)。ELBO 定义为

ELBO=Eq(zx)[logp(xz)]KL[q(zx)p(z)] ELBO=Eq(z∣x)​[logp(x∣z)]−KL[q(z∣x)∣∣p(z)]

其中

  • Eq(z|x)​[logp(x|z)] 是重构项,表示解码器分布下数据的预期对数似然。
  • KL[q(z|x) || p(z)] 是变分分布与潜在空间上的先验分布之间的 Kullback-Leibler (KL) 散度。此项鼓励变分分布接近先验分布,起到正则化项的作用。

ELBO 也可以解释为证据间隙的负值,即数据边际似然与变分分布和先验分布之间的 KL 散度之差。

在训练过程中,VAE旨在最大化ELBO相对于模型参数(编码器和解码器网络)的值。最大化ELBO鼓励模型学习数据在潜在空间中的良好表示,同时使变分分布接近先验。

证据下界 (ELBO) 的拼图类比:

好的,想象一下你有一个很大的拼图要解决,但你没有完整的图片。你有一些碎片,但不是全部。现在,你想弄清楚完整的图片可能是什么样子。这有点像变分自编码器(VAEs)所做的事情,但它们处理的是数据而不是拼图。

我们来分解一下

  1. **拼图块(数据):**在我们的例子中,这些是图片,比如动物或房子的照片,任何你想学习的东西。

  2. **完整图片(潜在空间):**这就像拼图的大图。我们不能直接看到它,但我们想弄清楚它可能是什么样子。在 VAE 中,我们称之为“潜在空间”。它是一个空间,其中每个点代表一个不同的可能的完整图片。

  3. **解码器(拼凑):**就像你可能把拼图碎片拼凑起来看到完整的画面一样,VAE 中的解码器试图从潜在空间中取一个点并将其转化为一幅画。

  4. **编码器(寻找碎片):**现在,如果我们有一张图片,我们可能想弄清楚我们用了哪些拼图碎片来制作它。VAE 中的编码器试图做到这一点。它接收一张图片,并试图找到潜在空间中可能用于制作该图片的点。

  5. **证据下界(ELBO):**现在,ELBO 就像一个乐于助人的朋友,它告诉我们我们的猜测有多好。它帮助我们确保我们的猜测接近现实。在我们的拼图类比中,ELBO 就像一个度量,它检查我们猜测的完整图片与真实完整图片的匹配程度,以及我们认为用于制作图片的碎片与实际碎片的匹配程度。

因此,当我们训练 VAE 时,我们希望调整我们的猜测(解码器和编码器),以使 ELBO 尽可能高。这意味着我们猜测的完整图片看起来更像真实的完整图片,并且我们认为用于制作图片的碎片与实际碎片更好地匹配。

简单来说,ELBO 帮助我们确保 VAE 能够学习生成优质图片,并找出用于生成这些图片的碎片,即使我们无法直接看到完整的图片。

Kull-back 散度:

库尔巴克-莱布勒(KL)散度是一种衡量两个概率分布之间差异的度量。具体来说,在 VAE 中,它用于量化两个重要分布之间的差异:

  1. **变分分布 (q(z|x)):**这是给定输入数据 (x) 的潜在变量 (z) 的分布。简单来说,它告诉我们给定特定输入数据点时,潜在变量不同值的可能性有多大。

  2. **先验分布 (p(z)):**这是我们在看到任何特定数据之前假定的潜在变量的分布。它代表我们对潜在空间的初始信念或期望。

现在,VAE 中这两个分布之间的 KL 散度(KL[q(z|x) || p(z)])具有重要作用:

  • 它衡量了当我们使用变分分布 (q(z|x)) 来近似真实分布 (p(z)) 时丢失了多少信息。
  • 如果KL散度很低,则表示变分分布与先验分布非常相似,这很好,因为它意味着我们的模型正在学习数据中有意义的表示。
  • 如果 KL 散度很高,则表示变分分布与先验分布差异很大,这表明我们的模型可能未能很好地捕捉数据的潜在结构。

本质上,最小化 KL 散度鼓励变分分布接近先验分布,这有助于在潜在空间中学习数据的良好表示。

柏拉图对 Kull-back 散度的类比:

image/png

柏拉图的洞穴寓言可以为理解变分自编码器(VAE)中变分分布和先验分布的作用提供有益的类比。

  1. **变分分布 (q(z|x)):**在洞穴寓言中,想象囚犯们被锁在洞穴里,只能看到墙壁上的影子。这些影子在我们的类比中代表观察到的数据 (x)。现在,变分分布 (q(z|x)) 对应于囚犯们对真实世界中可能投射这些影子的物体的解释或猜测。这是他们试图仅根据看到的影子来理解物体真实本质的尝试。

  2. **先验分布 (p(z)):**在柏拉图的寓言中,囚犯们有一种由他们在洞穴内的经历塑造的特定世界观或信念系统。这种世界观代表了我们对潜在空间的先验信念。这就像囚犯们对洞穴外真实世界可能是什么样子的假设,即使他们没有直接体验过。

现在,让我们将此与 VAE 联系起来:

  • 洞穴墙壁上的影子代表观测数据,类似于 VAE 中的输入数据 (x)。
  • 囚犯们对这些影子的解释代表了变分分布 (q(z|x)),这是模型试图根据观测输入来理解数据的潜在结构的尝试。
  • 囚犯们的世界观或信念系统对应于先验分布 (p(z)),它代表了我们在观察任何数据之前对潜在空间的初始假设。

在寓言的背景下,VAE 的目标类似于囚犯们渴望了解投射影子的物体的真实性质。最小化变分分布和先验分布之间的 KL 散度有助于模型收敛到与我们先验信念相符且能准确解释观测数据的潜在空间表示,就像囚犯们对影子的解释应该与其对真实世界的看法相符一样。

[!提示] 现在我们已经理解了这些术语,我们将进行一些推导**数学解释**。


手写推导:

  • 推导变分自编码器(负 ELBO)

image/jpeg

  • 推导 L_T + L_t-1 + L_0

image/jpeg

  • 正向和反向过程

image/jpeg


让我们深入了解 Stable Diffusion 的超参数。

1. `prompt` - 这是我们用于生成图像的文本提示。类似于我们在第一部分中看到的 `pipe(prompt)` 函数。
2. `g` 或 `guidance scale` - 这是一个值,用于确定图像与文本提示的接近程度。这与一种称为无分类器引导的技术有关,该技术可以提高生成图像的质量。引导尺度值越高,图像与文本提示的接近程度越高。
3. `seed` - 这设置了生成初始高斯噪声潜在的种子。
4. `steps` - 生成最终潜在图像所需的去噪步数。
5. `dim` - 图像的尺寸,为简单起见,我们目前生成的是正方形图像,因此只需一个值。
6. `save_int` - 这是可选的布尔标志,如果我们想保存中间的潜在图像,这有助于可视化。


资源:

研究论文:

去噪扩散概率模型
潜在扩散模型的高分辨率图像合成

网站:

lilianweng_diffusion-models
assemblyai
参数 韩吉
HuggingFace-DMs
HuggingFace_2-DMs
HuggingFace-github_blog

YouTube:

异常值
乌玛·贾米尔
乔治·霍茨_part1
乔治·霍茨_part3
计算机爱好者
伯克利机器学习


以下是此研究的 BibTeX 引文,可用于学术作品中的正确引用:

引文:

@article{Gumaan2024diffusion,
  title   = "Diffusion Models?",
  author  = "Gumaan, Esmail",
  blog = "https://huggingface.co/blog/Esmail-AGumaan/",
  year    = "2024",
  month   = "April",
  url     = "https://huggingface.co/blog/Esmail-AGumaan/diffusion-models#diffusion-models"
}

社区

注册登录 以评论