【翻译】最近兴起的扩散模型

3,437 阅读19分钟

持续创作,加速成长!这是我参与「掘金日新计划 · 10 月更文挑战」的第22天,点击查看活动详情


本文翻译自:The recent rise of diffusion-based models | Maciej Domagała (maciejdomagala.github.io)

作者博客:Maciej Domagała | ML, DL, Data Science. (maciejdomagala.github.io)


简介

在过去的一年半里,每个搞生成模型的人都绝对像活在梦里一样,在过去的几个月里边,文本到图像(text-to-image)生成的论文层出不穷,新工作一个接一个,一篇比一篇更🐂。我们可以看到社交媒体上迅速出现大量的使用人工智能模型生成的图片。比如这只金毛猎犬一脸丧了吧唧地在竞选途中回答棘手的问题,或者这颗大脑乘坐火箭前往外太空。

Sources: openai.com/dall-e-2/ and imagen.research.google/

前一个是用openAI 的 DALL·E-2生成的,后一个是用的谷歌的Imagen。

在这篇文章中,我们将总结一下解决文本到图像生成问题的这一段简短的历史,然后带大家了解一下关于扩散模型的最新进展。为啥要讲扩散模型,因为刚才看到的那些图都是使用扩散模型生成的,扩散模型在是现在最新的SOTA中很重要的一个内容。

Short timeline of image generation and text-to-image solutions. Source: Maciej Domagała


这一切都是从DALL·E开始的

在2020年,OpenAl团队发布了GPT-3模型——一个在无所不能的巨型语言模型,它能够进行机器翻译、文本生成、语义分析等,该模型当时被认为是语言模型的SOTA,而DALL·E可以认为是Transformer在计算机视觉领域的应用。

先看一下DALL·E是什么:

自回归方法

作者提出了一种优雅的两阶段方法:

  • 训练离散的VAE模型将图像压缩为图像token,

  • 将编码的文本片段与图像token连接起来,训练自回归的Transformer学习文本和图像的联合分布。

最终训练用的数据集是从互联网上获得的2.5亿个文本-图像对。

CLIP

在推理过程中,模型可以直接输出一个batch的生成图像,那我们如何评估生成的哪个图像最好呢?在DALL·E发布的同时,OpenAI团队提出了一个打通图像和文本的模型CLIP,CLIP提供了一种将文本片段与其图像表示配对的方法。先抛开所有的技术实现方面不谈,仅看训练过程的话,训练这类模型的想法是很简单的:对一些文本片段进行编码,拿一些图片进行编码,就这样对四亿个图片-文本对进行编码。训练的时候使用对比训练的方法进行训练。

Visualisation of CLIP contrastive pre-training, source: arxiv.org/abs/2103.00…

这种映射可以让我们评估:在考虑到输入文本的情况下哪个生成的图像是最匹配的。DALL·E引起了AI领域内外人士的广泛关注。从此爆火,获得了众多关注,获得广泛讨论。但是即便如此,它也只是在这里被提及一下,因为此后涌现出了更多新奇的东西。


All you need is diffusion

Sohl-Dickstein等人提出了图像生成的新思路——扩散模型。

Generative models, source: lilianweng.github.io/posts/2021-…

扩散模型是受到非平衡热力学的启发产生的,它背后充满了有趣的数学概念。我们可以看到这里有我们熟悉的编码器-解码器(encoder-decoder)结构的概念,但其基本思想与我们在传统的变分自编码器中(variational autoencoders)观察到的略有不同。为了理解这个模型的基础,我们需要讲一下正向和反向扩散过程。

前向图像扩散过程

这一过程可以理解为我们逐渐给图片加高斯噪声,直到它变得完全无法辨认。这一过程在随机意义上是固定的:添加噪声这一步骤可以被认为是连续扩散步骤的马尔可夫链。为了稍微理清这些难懂的措辞,我们可以用一些公式来简洁地描述这一过程。

假设我们图像数据集的起始分布为q(x0)q(x_0),我们可以从这个分布中采样一张图x0x_0,我们要执行一个链式的扩散步骤x0x1x2...xTx_0 → x_1→ x_2→ ...→ x_T,每一步都将我们采样到的图片逐渐加噪破坏掉。

怎么给它加噪声?我们需要预定义一个噪声表{βt}t=1T\{\beta_t\}^T_{t=1},其中βt(0,1)\beta_t \in (0,1)。这样我们就可以将前向过程定义为:

q(xtxt1)=N(1βxxt1,βtI)q(x_t | x_{t-1}) = \mathcal N (\sqrt{1-\beta_x} x_{t-1},\beta_t \mathbf I)
  • 时间步增加(TT → \infty)最终的分布q(xT)q(x_T)非常接近球形高斯分布,之后你从噪声中采样就会非常方便。

  • 用高斯核噪声还有另外一个好处:

    你不需要一步一步计算潜变量,因为会用到参数重整化技巧,所以可以一步到位计算任意时间步tt的扩散结果:

    q(xtx0)=N(αˉtx0,(1αˉt)I)=αˉtx0+1αˉtϵq\left({x}_t \mid {x}_0\right)=\mathcal{N}\left(\sqrt{\bar{\alpha}_t} {x}_0,\left(1-\bar{\alpha}_t\right) \mathbf{I}\right)=\sqrt{\bar{\alpha}_t} {x}_0+\sqrt{1-\bar{\alpha}_t} \cdot \epsilon

    其中αt:=1βt,αˉt:=k=0tαk\alpha_t:=1-\beta_t, \bar{\alpha}_t:=\prod_{k=0}^t \alpha_kϵN(0,I)\epsilon \sim \mathcal{N}(0, \mathbf{I}). ϵ\epsilon是一个高斯噪声,这个公式对扩散模型的训练很重要。

反向重建过程

我们已经定义好了前向过程,那会有人问,然后嘞?为什么我们不能直接定义一个反向过程q(xt1xt)q(x_{t-1}|x_t)直接从噪声回溯到原图? 首先,这个从概念上就是错的,因为我们想要一个神经网络学会怎么解决问题,所以我们不能给它提供解决方案。 其次,想要从噪声样本回到初始的分布q(x0)q(x_0),我们必须计算所有可能从噪声中产生x0x_0的方法的边缘概率,包括所有的潜变量。这意味着要计算q(x0:T)dx1:T\int q(x_{0:T})dx_{1:T},但是这样玩意儿是算不出来的。所以我们要去近似它。

所以反向过程的核心思想是找一个可靠的解决方案,搞一个可学习的神经网络,近似这个反向扩散过程。实现这一过程就是要去估计每一步去噪的均值和方差。 pθ(xt1xt)=N(μθ(xt,t),Σθ(xt,t))p_\theta\left({x}_{t-1} \mid {x}_t\right)=\mathcal{N}\left(\mu_\theta\left({x}_t, t\right), \Sigma_\theta\left({x}_t, t\right)\right) μθ(xt,t)\mu_\theta\left({x}_t, t\right)可以用神经网络学出来,Σθ(xt,t)\Sigma_\theta\left({x}_t, t\right)可以固定为与噪声调度相关的某个常数,比如βtIβ_t \mathbf I

Forward and reverse diffusion processes, source: [14]

用神经网络估计μθ(xt,t)\mu_\theta\left({x}_t, t\right)是可以的,但是谷歌大脑那个人想出了另一个不同的训练方法 :训练一个神经网络ϵθ(xt,t)\epsilon_\theta\left({x}_t, t\right) 直接通过前边的公式 q(xtx0)q\left({x}_t \mid {x}_0\right)去学习噪声ϵ\epsilon

image.png

训练过程就变成了:

  1. 从分布q(x0)q(x_0)中采样一张图片x0x_0

  2. 从扩散步骤U({1,2,...,T})U(\{1,2,...,T\})中采样时间步tt

  3. 为其添加从标准正态分布中采样的噪声ϵ\epsilon

  4. 估计噪声ϵθ(xt,t)=ϵθ(αˉtx0+1αˉtϵ,t)\epsilon_\theta\left(\mathbf{x}_t, t\right)=\epsilon_\theta\left(\sqrt{\bar{\alpha}_t} \mathbf{x}_0+\sqrt{1-\bar{\alpha}_t} \cdot \epsilon, t\right)

  5. 梯度下降计算lossθϵϵθ(xt,t)2\nabla_\theta\left\|\epsilon-\epsilon_\theta\left(\mathbf{x}_t, t\right)\right\|^2

损失可以表示为: Ldiffusion =Et,x0,ϵ[ϵϵθ(xt,t)2]L_{\text {diffusion }}=\mathbb{E}_{t, \mathbf{x}_0, \epsilon}\left[\left\|\epsilon-\epsilon_\theta\left(\mathbf{x}_t, t\right)\right\|^2\right]

所有的公式、参数重整化、推导都涉及大量的数学知识,如果想继续深入了解,现在有一些很好的资源,比如

可以去这些大佬的博客看一下。

引导扩散 Guidance

上边我们讲了为什么扩散模型可以拿来生成图像。一旦模型ϵθ(xt,t)\epsilon_\theta\left({x}_t,t\right)被训练出来,我们可以使用它将噪声xtx_t重建回x0x_0,因为从球状高斯分布采样噪声很简单的,这样我们就可以获得无限的图像变化。

我们还可以在训练过程中通过向网络输入额外的信息来指导图像的生成。假设图像是有标签的,那么类别标签yy的信息可以被送入一个类条件扩散模型ϵθ(xt,ty)\epsilon_\theta\left({x}_t, t|y\right)中。 怎么给训练过程引入指导?

第一种方法是训练一个单独的模型作为噪声图像分类器,在去噪的每一步中都去检查图像是否在向正确的方向上去噪,在梯度下降的过程中将自己的损失函数组合到扩散模型的整体损失中。

Ho等人还提出另一个想法,就是不用训练额外的分类器情况下也能把类别信息传递给模型。就是在训练类条件扩散模型ϵθ(xt,ty)\epsilon_\theta\left({x}_t, t|y\right)的过程中,有时候不给它实际的类别信息yy而是传入一个空值ϕ\phi代替类别信息。这样模型就能学会在有和没有指导的情况下进行扩散。在推理过程中模型进行两次预测,一次是给定类的标签信息,一次是不给定标签信息。通过控制引导参数s1s \geqslant 1,模型最终的结果会靠近ϵθ(xt,ty)\epsilon_\theta\left({x}_t, t \mid y\right),远离ϵθ(xt,tϕ)\epsilon_\theta\left({x}_t, t \mid \phi\right)

ϵ^θ(xt,ty)=ϵθ(xt,tϕ)+s(ϵθ(xt,ty)ϵθ(xt,tϕ))\hat{\epsilon}_\theta\left({x}_t, t \mid y\right)=\epsilon_\theta\left({x}_t, t \mid \phi \right)+s \cdot\left(\epsilon_\theta\left({x}_t, t \mid y\right)-\epsilon_\theta\left({x}_t, t \mid \phi\right)\right)

这种无分类器指导,只使用扩散模型自己的理解的情况会产生更好的结果。

更详细可以看这里:Guidance,让扩散模型的指标更能打 - 掘金 (juejin.cn)

使用GLIDE用文本引导扩散模型

尽管在本文讨论的所有论文中,描述GLIDE体系结构的论文关注是最少的,但可以说这个工作提出了有新意度最有趣的想法。就是这个工作结合了扩散模型。

我们已经知道扩散模型是如何工作的,我们可以用它们来生成图像。我们现在想回答的两个问题是:

  1. 如何利用文本信息来指导扩散模型?

  2. 我们怎样才能保证模型的质量足够好呢?

框架选择

模型可以归纳为由三个部分组成:

  1. 一个基于U-Net的模型负责扩散学习的图像部分

  2. 基于transformer的模型负责从文本片段获得embedding

  3. 一个上采样扩散模型来提高输出图像的分辨率

前两者一起工作,以创建文本引导的图像输出,而最后一个用于放大图像,同时保持质量。

核心框架就是一个经典的U-Net,由上采样和下采样卷积组成。其中还包含一个注意力层,便于和文本处理衔接。这部分拥有约23亿个参数,数据集和DALL·E一样。

用于做指导的文本被编码到token中,然受输入到基于Transformer的模型中,这部分是12亿个参数,由24个宽度为2048的残差块构建而成。Transformer的输出有两个方法:

  1. 最后的token embedding 作为分类embedding yy传给条件扩散模型 ϵθ(xt,ty)\epsilon_\theta\left({x}_t, t \mid y\right)

  2. 最后一层输出的token embedding加到上边U-Net的每一层attention中

很明显,为了确保模型收到足够的文本相关信息以生成准确的图像,我们投入了大量的精力。该模型以文本片段embedding为条件,编码后的文本与注意力上下文相连接,在训练期间,使用类条件扩散模型ϵθ(xt,ty)\epsilon_\theta\left({x}_t, t \mid y\right)。 最后一部分,作者使用扩散模型,利用ImageNet的上采样器提高图像的分辨率。

GLIDE interpretation of ‘a corgi in a field’, source: arxiv.org/abs/2112.10…

GLIDE集合了近年来发展的一些新成就,并为文本引导图像生成的概念带来了新思路。考虑到和DALL·E模型基于不同的结构,可以说GLIDE的发表标志着基于扩散模型的文本到图像生成时代的到来

DALL·E 2

OpenAI团队似乎没有休息,因为在4月份,他们用DALL·E 2在互联网上又掀起了一场风暴。它借鉴了两个前辈的元素:

  • 严重依赖CLIP

  • 解决方案的很大一部分参考了GLIDE架构。

DALL-E 2有两个主要的底层组件,称为prior和图片解码器,二者组合在一起时能够产生图像输出。整个机制被命名为unCLIP,听这个名字想必了解CLIP的人就要已经能能到这部分做了什么工作了。

Visualization of DALL·E 2 two-stage mechanism. Source: Zero-Shot Text-to-Image Generation

prior模型

第一部分的目的是将caption(一个文本片段)转换为文本嵌入(如“柯基演奏喷火的小号”)。这部分用的是一个冻结的CLIP模型。

在文本嵌入之后就是一个有趣的部分了,我们现在需要获得一个图像嵌入,类似于通过CLIP模型获得的图像嵌入。我们希望它还能封装来自文本嵌入的所有重要信息,因为它将通过扩散生成图像。那你想问这不就和CLIP的作用一样吗? 如果我们想给输入的短语找到一个对应的图像嵌入,我们可以直接去查看在CLIP编码空间中与我们的文本嵌入接近的内容就行了。为什么还要这么麻烦搞一个prior?

DALL·E 2的一位作者发表了一篇很好的解释性的文章,说了一下为什么不能用这个决方案,以及为什么需要搞一个prior出来:

“An infinite number of images could be consistent with a given caption, so the outputs of the two encoders will not perfectly coincide. Hence, a separate prior model is needed to “translate” the text embedding into an image embedding that could plausibly match it”

作者大致意思是:对于同一个caption,可能会有无数的图片与之对应,因此CLIP两个编码器输出的内容并不是完全一致的,因此在这里需要一个单独的prior模型将文本嵌入“翻译”为可能与之相匹配的图像嵌入。

在此基础上,作者通过实验证明了prior在整个网络中的重要性。证明了使用prior出来的图像嵌入+文本信息的效果要优于仅使用文本信息或者使用CLIP的文本嵌入

Samples generated conditioned on: caption, text embedding and image embedding. Source: Zero-Shot Text-to-Image Generation

作者为了验证prior的效果分别测试了两类模型:自回归模型和扩散模型。这篇文章我们只讲扩散模型部分的prior,因为它的表现比自回归模型好很多(毕竟这篇文章就是说扩散模型的,将自回归干什么),特别是从数学计算的角度来看。对于先验的训练,选择了仅解码器的Transformer模型。它通过使用几个输入序列进行训练:

  • 编码的文本,

  • CLIP文本嵌入,

  • 扩散时间步嵌入,

  • 噪声图像嵌入,

以输出一个嵌入ziz_i的无噪声图像为目标,与Ho等人在前几节中提出的训练方式不同,直接预测无噪声图像嵌入比预测噪声更适合。记住前面关于导向模型中扩散损失的公式。

Ldiffusion =Et,x0,ϵ[ϵϵθ(xt,ty)2]L_{\text {diffusion }}=\mathbb{E}_{t, \mathbf{x}_{0, \epsilon}}\left[\left\|\epsilon-\epsilon_\theta\left(\mathbf{x}_t, t \mid y\right)\right\|^2\right]

将前半部分扩散的loss改为:

Lprior:diffusion =Et[zifθ(zit,ty)2]L_{\text {prior:diffusion }}=\mathbb{E}_t\left[\left\|z_i-f_\theta\left(z_i^t, t \mid y\right)\right\|^2\right]

其中fθf_\theta是prior模型,zitz_i^t是带噪图像嵌入,tt是时间步,yy是用于模型引导的文本信息(caption)。

解码器

我们已经讲了unCLIP的前半部分(prior模型),prior就是一个能够将文本中的所有重要信息封装到类似CLIP的图像嵌入中的模型。现在我们想使用前边获得的这个图像嵌入生成一个实际的图像作为输出。这就是unCLIP名字的来由:我们要用图像embedding获得对应的图像,这可以看做是CLIP模型图像编码器的反向过程。

俗话说得好啊,一个扩散模型后边该跟着另一个扩散模型了!

第二个扩散模型就是我们刚刚说的GLIDE模型,也不全是,是对GLIDE稍微改了一下子。主要的一点更改是将CLIP生成的图像嵌入加入到原始的GLIDE的文本编码器中。毕竟prior设计出来的目的就是为了给图像解码器提供信息的。至于Guidance技巧,DALL·E 2这里用的和GLIDE一样classifier-free Guidance方法,。将 10% 的CLIP embeddings 设置为 ϕ\phi,将50%的文本caption信息设置为ϕ\phi

另一个和GLIDE一样的是先生成一个小图然后使用上采样提高像素。把输出丢进扩散模型里边,生成一个64×64的小图,和GLIDE不同的是这里用了两个上采样。第一次是将像素提高到256×256,第二次是将像素提高到1024×1024。

Imagen

谷歌大脑也不甘落后,在DALL·E2之后不到两个月后,他们也展示了自己的劳动成果——Imagen。

Overview of Imagen architecture. Source: Zero-Shot Text-to-Image Generation

Imagen的从结构上看起来出奇的简单。使用一个预训练好的文本模型生成嵌入,用于扩散过程生成图片。然后通过一个提高分辨率的超分辨率扩散模型来提高生成图片的分辨率,这一点和DALL·E 2中是一样的。创新点分布在模型的不同地方,有的是在模型上,有的是在训练过程上。与前边的几个工作相比,它做出了一些创新。鉴于前边我们都说了好几个模型了,我们现在可以通过比较一下Imagen和前面工作的不同来解释这个模型:

使用预训练模型,而不是自己去训练Transformer

和OpenAI的工作相比,这一点可以认为是核心创新点。在GLIDE中,所有与文本嵌入有关的东西,用的都是一个新的,专门训练好的Transformer模型;但是在Imagen中作者用的是一个预训练的好的T5-XXL模型。这样做是考虑到之前的模型是仅对于图片的caption进行训练的,但是使用T5-XXL,它是已经预训练好的不需要再微调的模型,并且含有更为丰富的上下文信息,所以可以产生更有价值的嵌入。

让底层神经网络更有效

用于提高分辨率的超分辨率扩散模型, 把原来普通的U-net换成了Efficient U-net神经网络作为其模型的骨干。Efficient U-net比以前的版本更节省内存,更简单,收敛速度也更快。这些改进主要是由残差块和网络内部值的额外缩放带来的。如果你想深入了解这一部分可以去看一下Zero-Shot Text-to-Image Generation

使用条件增强来增强图像的保真度

因为这个模型是由一系列扩散模型组成的,因此在模型相关联的区域进行增强是有必要的。Ho等人提出了一种称为条件反射增强的解决方案。简单地说,这相当于在将低分辨率图像输入超分辨率模型之前,对其应用各种数据增强技术,如高斯模糊。


对低FID得分和高图像保真度(如动态阈值)改进还有一些很重要的方法,这些在Zero-Shot Text-to-Image Generation也中有详细解释。

Some of Imagen generations with captions. Source: Zero-Shot Text-to-Image Generation

效果怎么样?

截止本文写作时间为止,就文本到图像的生成而言,谷歌的Imagen被认为是最先进的。但这到底是为什么呢?我们如何评估模型并相互比较?

Imagen的作者选择了两种评估方法:

  • 一种是现在公认的文本到图像生成的指标,即在COCO验证数据集上计算FID score。作者说Imagen显示出最先进的性能,因为其zero-shot生成图像的FID优于其他模型,甚至超过那些专门在COCO数据集上训练的模型。

Comparison of several models. Source: [7]

  • 另一个更有趣的的评估手段是Imagen作者提出的新的评价指标DrawBench用来评估比较不同的文本到图像生成的模型。就是一组全面并且很有挑战性的文本提示(prompt),包含11个类别200个文本提示。评估过程是用275个没有偏见的评分人员完成的,每一类别25个人。把这些文本提示丢给不同的模型让它们生成图像,不是选出最优生成结果,是生成啥就拿啥。然后把生成的图像给评分人员,问他们两个问题。

    1. 你觉得哪组图片生成质量更高?

    2. 哪组图片和文本说明更贴合?

    这两个问题能清晰表达出一个好的文本到图像生成模型的两个最重要的特征:产生的图像的质量(保真度)和它到底能不能很好得反映输入的文本提示(对齐)。

    评分者只能说A模型更好、B模型更好或者两个一样。但是最后只能有一个模型是胜者。还有一个有趣的事情,在这个评价指标上GLIDE的效果比DALL·E 2要好。

    Imagen vs other models. Source: Zero-Shot Text-to-Image Generation

该论文的很大一部分都是在比较Imagen和GLIDE、DALL·E 2生成的图像,更多的可以去看附录。

未完待续

像以前的发展状况一样,随着新体系结构获得认可,大量有趣的工作涌现出来。由于发展速度快,我们不可能follow每一个工作。除了原始的生成能力之外,模型还有许多有趣的特性需要发现,例如图像修复、风格转换和图像编辑等等。

除了对生成性模型的新时代感到兴奋之外,我们还要知道扩散模型结构中还存在一些缺陷,例如与以前的模型相比,采样速度较慢等问题。我们也要在这些方向上不断进行改进。

Models comparison. Source:Improving Diffusion Models as an Alternative To GANs, Part 1