Sora相关技术解读

15 阅读36分钟

概述

2024年2月16日,OpenAI发布了视频生成模型Sora,其能基于已有的文本描述、图片和视频生成新的视频,且生成的视频可支持不同尺寸和时长(最长可达1分钟),并可模拟真实物理世界的相关特性(如镜头移动时物体可保持三维透视原理)。Sora生成的视频的视觉效果相比其他已有的视频生成模型有较大的提升。

OpenAI发布的关于Sora的技术报告——《Video Generation Models as World Simulators》对Sora的技术方案、生成效果作了简要的介绍,其核心技术基于Sora项目负责人之一——William Peebles在2023年发表的论文《Scalable Diffusion Models with Transformers》中提出的Diffusion Transformer(可简称为DiT),而Diffusion Transformer则是在扩散模型近几年不断发展的基础上,将扩散模型进一步与Transformer结合,从而在图片生成领域取得了更好的效果。

因此,Sora的成功并不是一蹴而就的,而是基于近几年扩散模型、Transformer等方面不断发展的技术成果。本文对Sora相关技术基于已有的论文和技术报告进行解读,如有不足之处,请指正。

Transformer

《Attention Is All You Need》是Google于2017年发表的一篇经典论文,其开创性的设计了Transformer模型结构,对于序列问题,通过编码器输出隐向量序列,再通过解码器输出目标序列,模型引入了注意力机制,充分挖掘序列各节点之间的深度信息,并通过矩阵计算的并行化加速模型训练和推理性能。Transformer模型结构的提出当初主要为了解决机器翻译问题,但由于其模型结构的通用性,目前已广泛应用于NLP、计算机视觉、计算广告等多个领域,更成为NLP大模型的底层基础结构。

模型结构

Transformer整体模型结构为编码器和解码器(Encoder-Decoder)架构,如图1所示。对于机器翻译问题,编码器将原句词元序列(x1,...,xn)(x_1,...,x_n)作为输入,输出向量序列z=(z1,...,zn)z=(z_1,...,z_n),如图1左侧所示;而解码器再将向量序列zz作为输入,输出翻译后的结果句词元序列(y1,...,ym)(y_1,...,y_m),如图1右侧所示。过程中,解码器每步输出一个词元,即第ii步输出yiy_i,实际输出时通过Softmax层计算每个词元在当前位置的概率,选取最大概率的词元作为结果。同时,解码器每步计算时,会将之前已输出的结果句词元序列也作为输入,即第ii步会使用(y1,...,yi1)(y_1,...,y_{i-1})作为输入,论文称之为自回归(Auto-Regressive)。编码器和解码器内部使用了注意力机制、全连接网络和残差连接充分挖掘词元序列的语义信息,论文下面对这些结构进行了详细的介绍。

图1 模型结构

编码器

编码器在Embedding和位置编码后。编码器由多层组成,每层的结构相同,均又包含两个子层:多头注意力层(Multi-Head Attention)和全连接网络层。编码器中多头注意力层输入的查询、键值对相同,均采用原句词元序列的Embedding和位置编码后的输出,因此实现自注意力(Self Attention)机制,挖掘原句词元序列两两词元之间的相关性。每个子层周围有一个残差连接,通过残差连接避免梯度消失问题,然后进行层归一化,通过层归一化在样本粒度进行归一化,而非批次粒度进行归一化,避免不同句子样本的长短不一问题。另外,为便于残差连接,Embedding和位置编码后的输出(即编码器层的输入)和编码器各子层的输出的维度dmodeld_{\text{model}}相等,论文中基线模型dmodeld_{\text{model}}值为512。每个子层的上述计算过程可以通过以下公式表示:

LayerNorm(x+Sublayer(x))\text{LayerNorm}(x +\text{Sublayer}(x))

其中,xx表示每个子层的输入,Sublayer(x)\text{Sublayer}(x)表示多头注意力层或全连接网络层的输出,x+Sublayer(x)x +\text{Sublayer}(x)表示残差连接,LayerNorm\text{LayerNorm}表示层归一化。

解码器

解码器也是由多层组成(论文中基线模型层数NN为6),每层的结构相同,均包含三个子层:带掩码的多头注意力层(Masked Multi-Head Attention)、多头注意力层(Multi-Head Attention)和全连接层,并在每个子层周围有一个残差连接,然后进行层归一化。可以看到解码器和编码器的结构基本相同,不同之处主要有以下几点:

  • 增加带掩码的多头注意力层,并将解码器之前已输出的结果句词元序列也作为输入,即第ii步会使用(y1,...,yi1)(y_1,...,y_{i-1})作为输入,其中掩码的作用是为了在训练阶段中,当第ii步预测yiy_i时,只使用训练样本中的(y1,...,yi1)(y_1,...,y_{i-1}),而抹去(yi,...,ym)(y_{i},...,y_m),保证训练和推理阶段数据的一致性;
  • 多头注意力层输入的查询采用带掩码的多头注意力层的输出,键值对采用编码器的输出,因此实现目标注意力(Target Attention)机制,挖掘原句词元序列和结果句词元序列两两词元之间的相关性;
  • 解码器的输出通过线性层和Softmax层计算每个词元在当前位置的概率,选取最大概率的词元作为结果。

注意力机制

注意力机制的输入包括查询和多个键值对,通过注意力评分函数计算查询和各个键的权重,基于权重对相应的各个值进行加权求和,从而突出和查询相关的值,实现注意力机制,类比人类观察事物,会着重关注自己感兴趣的事物部分而非整体。

缩放点积注意力

图2 缩放点击注意力

论文中的注意力机制底层实现是缩放点积注意力(Scaled Dot-Product Attention),其结构如图2所示。令多个查询构成矩阵QQ,其中每个查询的维度为dkd_k,多个键、值对分别构成矩阵KKVV,其中每个键的维度为dkd_k,每个值的维度为dvd_v,则图2的计算过程可以通过以下公式表示:

Attention(Q,K,V)=softmax(QKTdk)V\text{Attention}(Q,K,V)=\text{softmax}(\frac{QK^T}{\sqrt{d_k}})V

其中,通过QQKTK^T这两个矩阵相乘,实现查询和键两两之间的点积计算,作为点积注意力,并且各点积计算可以并行化加速计算过程。计算后的点积注意力会乘以1dk\frac{1}{\sqrt{d_k}}进行缩放,这是为了避免当dkd_k较大时,原始点积注意力值也会较大,Softmax函数输出值会落入梯度较小的区域,收敛慢。缩放后的点积注意力通过Softmax函数映射到取值范围为[0,1][0,1]、且和为1的概率空间。最后再通过一个矩阵相乘,得到各个查询下、各个值的加权和。

另外,图中还有一个可选项——掩码,其作用已在解码器部分介绍。

多头注意力

图3 多头注意力

如果仅使用上一节的缩放点积注意力,其查询、键值对在论文中均是词元向量,维度均为dmodeld_{\text{model}},论文中提到,若采用线性层,将查询、键值对映射到不同的维度空间,能够自动挖掘原始输入中的深度信息,带来效果的提升,因此引入了多头注意力层,其结构如图3所示。

多头注意力层分为多层(论文中基线模型层数hh为8),每个层之间可以并行计算。图3的计算过程可以通过以下公式表示:

MultiHead(Q,K,V)=Concat(head1,...,headh)WO\text{MultiHead}(Q,K,V)=\text{Concat}(\text{head}_1,...,\text{head}_{\text{h}})W^O
where headi=Attention(QWiQ,KWiK,VWiV)\text{where head}_i=\text{Attention}(QW_i^Q,KW_i^K,VW_i^V)

其中,对于当前层ii,先通过三个线性层矩阵WiQRdmodel×dkW_i^Q \in \mathbb{R}^{d_{\text{model}} \times d_k}WiKRdmodel×dkW_i^K \in \mathbb{R}^{d_{\text{model}}\times d_k}WiVRdmodel×dvW_i^V \in \mathbb{R}^{d_{\text{model}}\times d_v}分别将原始的词元序列转化为向量维度为dkd_k的查询矩阵、键矩阵和向量维度为dvd_v的值矩阵,再通过上一节的缩放点积注意力计算各个查询下、各个值的加权和headi\text{head}_i。得到各层向量维度为dvd_v的输出headi\text{head}_i后,再拼接在一起得到向量维度为h×dvh \times d_v的总输出,最后通过一个线性层矩阵WORhdv×dmodelW^O \in \mathbb{R}^{hd_v \times d_{\text{model}}}将总输出映射为向量维度为dmodeld_{\text{model}}的结果输出。论文中基线模型dk=dv=dmodel/h=64d_k=d_v=d_{\text{model}}/h=64

注意力计算示例

下面通过如图4所示的示例说明如何计算得到词元之间的注意力,并通过注意力对词元进行加权求和。

image.png

图4中,输入XX即词元向量序列,每一行为一个词元向量。输入XX使用多个“Self-Attention”进行处理(即多头注意力)。

令词元向量序列中只有两个词元向量aia_iaja_j,在第二个头中,通过权重矩阵W2QW_2^Q将其分别映射为qi,2q_{i,2}qj,2q_{j,2},通过权重矩阵W2KW_2^K将其分别映射为ki,2k_{i,2}kj,2k_{j,2},通过权重矩阵W2VW_2^V将其分别映射为vi,2v_{i,2}vj,2v_{j,2},然后在点积缩放注意力中,对qi,2q_{i,2}ki,2k_{i,2}计算点积得到aia_i和其自身的注意力ai,ia_{i,i},对qi,2q_{i,2}kj,2k_{j,2}计算点积得到aia_iaja_j的注意力ai,ja_{i,j}(此处暂不考虑缩放和归一化),再基于注意力对vi,2v_{i,2}vj,2v_{j,2}进行加权求和得到第二个头输出结果中的第ii行,即:bi,2=ai,ivi,2+ai,jvj,2b_{i,2}=a_{i,i}v_{i,2}+a_{i,j}v_{j,2}

image.png

图5是原Transformer论文列出的一个多头注意力示例,其中只列出了词元“making”和部分其他词元之间的注意力,不同头的注意力采用不同的颜色。

注意力机制在模型中的应用

这里论文进一步总结了注意力机制在模型中的应用:

  • 在编码器部分,多头注意力层输入的查询、键值对相同,均采用原句词元序列的Embedding和位置编码后的输出,因此实现自注意力(Self Attention)机制,挖掘原句词元序列两两词元之间的相关性;
  • 在解码器部分,带掩码多头注意力层输入的查询、键值对相同,均采用结果句词元序列的Embedding和位置编码后的输出,因此实现自注意力(Self Attention)机制,挖掘结果句词元序列两两词元之间的相关性;另外,增加掩码,在训练阶段中,当第ii步预测yiy_i时,只使用训练样本中的(y1,...,yi1)(y_1,...,y_{i-1}),而(yi,...,ym)(y_{i},...,y_m)在Softmax层的输入被重置为负无强大,使其在Softmax层的输出逼近0,抹去其影响,保证训练和推理阶段数据的一致性;
  • 在编码器和解码器之间,查询采用解码器的输入,键值对采用编码器的输出,因此实现目标注意力(Target Attention)机制,挖掘原句词元序列和结果句词元序列两两词元之间的相关性。

逐个位置的前馈神经网络

多头注意力层输出是句子序列各词元的向量序列,维度为dmodeld_{\text{model}}。经过残差连接和归一化后,输入前馈神经网络,由于多头注意力层已通过注意力机制考虑句子序列各词元之间的相关性,包含序列信息,因此,前馈神经网络对每个词元的向量作为输入,进行单独处理,论文称之为“Position-Wise”。

前馈神经网络的模型结构可以用以下公式表示:

FFN(x)=max(0,xW1+b1)W2+b2\text{FFN}(x)=\text{max}(0,xW_1+b_1)W_2+b_2

从中可以看出,其结构就是带有一层隐藏层的全连接网络,输入到隐层的权重矩阵为W1Rdmodel×dffW_1 \in \mathbb{R}^{d_{\text{model}} \times d_{ff}},偏移向量是b1b_1,而隐层中每个神经元的激活函数采用ReLU,隐层到输出的权重矩阵为W2Rdff×dmodelW_2 \in \mathbb{R}^{d_{ff} \times d_{\text{model}}},偏移向量是b2b_2,输入和输出的维度仍是dmodeld_{\text{model}}(512),隐层的维度是dffd_{ff}(2048)。

Embedding和Softmax

对于输入的词元,使用了统一的Embedding矩阵进行向量化,转化为维度为dmodeld_{\text{model}}的向量。Softmax层不再详述。

位置编码

对于一个句子,如果构造该句子的词元本身不变,只是词元顺序变化,那么句子所表达的含义也会变化,因此需要考虑词元在句子中的位置。在Transformer中,输入的词元序列在Embedding计算后,还会对每个词元在序列中的位置进行编码,充分考虑位置信息。

词元位置编码后输出的向量维度也是dmodeld_{\text{model}},和词元通过Embedding层后输出的向量维度相同,这样便于两者直接相加作为编码器层的输入。那么如何将长短不一的句子中的词元位置编码为固定维度的向量,论文中采用的方法是正余弦函数,公式如下:

PE(pos,2i)=sin(pos/100002i/dmodel)PE_{(pos,2i)}=sin(pos/10000^{2i/d_{\text{model}}})
PE(pos,2i+1)=cos(pos/100002i/dmodel)PE_{(pos,2i+1)}=cos(pos/10000^{2i/d_{\text{model}}})

其中,PEPE表示位置编码,pospos表示当前计算的是第pospos个位置的编码向量,2i2i2i+12i+1分别表示该位置编码向量中第2i2i和第2i+12i+1个值。由于论文中的向量维度dmodeld_{\text{model}}值为512,因此,ii取值范围为[0,255][0,255],即将每个位置向量的维度分为256组,第ii组的两个值分别使用波长为100002i/dmodel10000^{2i/d_{\text{model}}}的正余弦函数对位置pospos进行计算,映射到[1,1][-1,1]。也就是说,随着ii的取值,正余弦函数的波长范围在2π2 \pi20000π20000\pi,只要句子中词元序列的长度不超过20000π20000\pi,通过上述方法计算的位置编码能保证唯一。

论文中解释采用这种位置编码的原因,一是可以将将长短不一的句子中的词元位置编码为固定维度的向量,二是可以将离散的位置值转化为连续值,三是可以借助正余弦函数的特性,对于某个偏移量kk,可以基于PEposPE_{pos}通过线性变换快速计算出相对位置的PEpos+kPE_{pos+k},即:

[sin((pos+k)w_i)cos((pos+k)w_i)]=[cos(kw_i)sin(kw_i)sin(kw_i)cos(kw_i)][sin(posw_i)cos(posw_i)]\begin{bmatrix} sin((pos+k) * w\_i) \\ cos((pos+k) * w\_i) \end{bmatrix} = \begin{bmatrix} cos(k * w\_i) & sin(k * w\_i) \\ -sin(k * w\_i) & cos(k * w\_i) \end{bmatrix} \begin{bmatrix} sin(pos * w\_i) \\ cos(pos * w\_i) \end{bmatrix}

其中,w_i=1100002i/d_modelw\_i = \frac{1}{10000^{2i/d\_{\text{model}}}}

为什么采用自注意力机制

图6 自注意力机制和RNN、CNN在计算性能上的比较

论文在这部分,针对输入向量序列为(x1,...,xn)(x_1,...,x_n)、输出向量序列为(z1,...,zn)(z_1,...,z_n)、其中xix_iziRdz_i \in \mathbb{R}^d这一类序列问题场景,对自注意力机制,和RNN、CNN在计算性能上作比较。考察每一层的计算量、其中串行操作量(多个串行操作可以并行,以此来衡量并行度),以及数据在网络中进行前向推理和反向梯度更新遍历的最长路径这三种指标,分析结果如图6所示。其中:

  • 在每一层的计算量上,自注意力机制需要计算序列中任意两个节点之间的相关性,因此计算复杂度是O(n2d)O(n^2 \cdot d),而RNN需要依赖序列中前序节点的输出,因此计算复杂度是O(nd2)O(n \cdot d^2),CNN需要进行卷积层计算,因此若卷积核宽度为kk,则计算复杂度是O(knd2)O(k \cdot n \cdot d^2)。对于向量维度远大于序列长度的场景(比如论文中的向量维度是512,若句子的词元数量远小于512),则自注意力机制的计算复杂度更优;
  • 在串行操作量上,由于RNN需要依赖序列中前序节点的输出,因此计算复杂度是O(n)O(n),而自注意力机制和CNN均是O(1)O(1);
  • 在网络遍历最长路径上,同样由于RNN需要依赖序列中前序节点的输出,因此计算复杂度是O(n)O(n),而自注意力机制是O(1)O(1)

从中可以看出,自注意力机制可以并行地计算两两节点之间的相关性,从而提升模型训练和推理效率。

Vision Transformer(ViT)

起初Transformer被用于NLP领域,而Google于2021年发表的论文《An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale》提出了Vision Transformer(可简称为ViT),其参考NLP中将句子拆分为词元(Token)向量序列作为Transformer的输入,将图片拆分为块(Patch)向量序列作为Transformer的输入,并将Transformer用于图片分类识别任务,取得了很好的效果,并论证了Transformer在计算机视觉领域也满足Scaling Law,即模型规模越大,数据集越大,其效果越好

方法

图7 ViT的整体架构

ViT的整体架构如图7所示,其核心结构采用Transformer的编码器。

在输入的处理上,令图片为xRH×W×C\mathbf{x}\in\mathbb{R}^{H\times W\times C},其中HH为图片高,WW为图片宽,CC为图片像素的通道数。ViT将图片划分为多个块,每一个块被称为一个Patch,而图片的块,可以类比于句子中的词元。令每个块的长和宽均为PP,块的数量为NN,则N=HW/P2N=HW/P^2

类比NLP中将句子转化为词元向量序列,ViT将图片转化为块向量序列。令块向量序列为xpRN×(P2C)\mathbf{x}_p\in\mathbb{R}^{N\times(P^2\cdot C)},序列的长度即块的数量,为NN,序列中的每个块使用向量xpi\mathrm{x}_p^ii=1,,Ni=1,\dots,N)表示,向量中的每个元素表示块中某个像素某个通道的值,因此,向量的维度为P2CP^2\cdot C。块向量序列可表示为:

xp=[xp1;xp2;;xpN]\mathbf{x}_p=[\mathrm{x}_p^1;\mathrm{x}_p^2;\cdots;\mathrm{x}_p^N]

将块向量序列作为Transformer的输入,而Transformer中隐向量的维度为DD,因此通过矩阵ER(P2C×D)\mathbf{E}\in\mathbb{R}^{(P^2\cdot C\times D)}将每个块的原始向量xpi\mathrm{x}_p^i线性映射为DD维向量xpiE\mathrm{x}_p^i\mathbf{E}。块向量序列可进一步表示为:

[xp1E;xp2E;;xpNE][\mathrm{x}_p^1\mathbf{E};\mathrm{x}_p^2\mathbf{E};\cdots;\mathrm{x}_p^N\mathbf{E}]

类似于BERT在词元序列中增加特殊词元“[class]”,ViT在上述块向量序列的头部增加特殊的DD维向量xclass\mathrm{x}_{\text{class}},该向量最终经过Transformer编码器的输出作为图片的表征y\mathbf{y}。块向量序列可进一步表示为:

[xclass;xp1E;xp2E;;xpNE][\mathrm{x}_{\text{class}};\mathrm{x}_p^1\mathbf{E};\mathrm{x}_p^2\mathbf{E};\cdots;\mathrm{x}_p^N\mathbf{E}]

和原始Transformer类似,ViT也会对块序列中的各位置进行位置编码,原块序列长度为NN,增加xclass\mathrm{x}_{\text{class}}后,块序列长度为N+1N+1,各位置编码的向量维度也为DD,令所有N+1N+1个位置的位置编码为EposR(N+1)×D\mathbf{E}_{pos}\in\mathbb{R}^{(N+1)\times D}。因此,最终ViT的输入可以表示为:

z0=[xclass;xp1E;xp2E;;xpNE]+Epos, ER(P2C)×D,EposR(N+1)×D\mathbf{z}_0=[\mathrm{x}_{\text{class}};\mathrm{x}_p^1\mathbf{E};\mathrm{x}_p^2\mathbf{E};\cdots;\mathrm{x}_p^N\mathbf{E}]+\mathbf{E}_{pos},\space\mathbf{E}\in\mathbb{R}^{(P^2\cdot C)\times D},\mathbf{E}_{pos}\in\mathbb{R}^{(N+1)\times D}

获得输入后,将z0\mathbf{z}_0输入Transformer的编码器,如图7左侧所示。Transformer的编码器结构如图7右侧所示。编码器共有LL层,对于第\ell层,将上一层的输出z1\mathbf{z}_{\ell-1}作为该层的输入,进行层归一化(Layer Normalization)后再输入多头自注意力层(Multi-Head Self-Attention),多头自注意力层的输出再与输入z1\mathbf{z}_{\ell-1}进行残差连接得到z\mathbf{z}'_{\ell},上述操作可表示为:

z=MSA(LN(z1))+z1\mathbf{z}'_{\ell}=\text{MSA}(\text{LN}(\mathbf{z}_{\ell-1}))+\mathbf{z}_{\ell-1}

其中,LN\text{LN}表示层归一化,MSA\text{MSA}表示多头自注意力层。

z\mathbf{z}'_{\ell}进行层归一化后再输入多层感知机,多层感知机的输出再与输入z\mathbf{z}'_{\ell}进行残差连接得到第\ell层的输出z\mathbf{z}_{\ell},上述操作可表示为:

z=MLP(LN(z))+z\mathbf{z}_{\ell}=\text{MLP}(\text{LN}(\mathbf{z}'_{\ell}))+\mathbf{z}'_{\ell}

其中,MLP\text{MLP}表示多层感知机,每个多层感知机包含两层,使用GELU作为激活函数。

在输出的处理上,将编码器最后一层输出的第一个块(即对应编码器输入中的xclass\mathrm{x}_{\text{class}})的编码zL0\mathbf{z}_L^0,再进行一次层归一化,得到图片表征向量y\mathbf{y}

y=LN(zL0)\mathbf{y}=\text{LN}(\mathbf{z}_L^0)

综上,ViT的整体网络结构可表示为:

z0=[xclass;xp1E;xp2E;;xpNE]+Epos,ER(P2C)×D,EposR(N+1)×Dz=MSA(LN(z1))+z1,=1Lz=MLP(LN(z))+z,=1Ly=LN(zL0)\begin{align} \mathbf{z}_0&=[\mathrm{x}_{\text{class}};\mathrm{x}_p^1\mathbf{E};\mathrm{x}_p^2\mathbf{E};\cdots;\mathrm{x}_p^N\mathbf{E}]+\mathbf{E}_{pos},&\mathbf{E}\in\mathbb{R}^{(P^2\cdot C)\times D},\mathbf{E}_{pos}\in\mathbb{R}^{(N+1)\times D} \\ \mathbf{z}'_{\ell}&=\text{MSA}(\text{LN}(\mathbf{z}_{\ell-1}))+\mathbf{z}_{\ell-1},&\ell=1\dots L\\ \mathbf{z}_{\ell}&=\text{MLP}(\text{LN}(\mathbf{z}'_{\ell}))+\mathbf{z}'_{\ell},&\ell=1\dots L\\ \mathbf{y}&=\text{LN}(\mathbf{z}_L^0) \end{align}

Transformer编码器中多头注意力层的实现细节之前已介绍,此处不再详述。

在ViT的预训练和微调阶段,ViT会在上述网络结构的基础上,增加分类头,分类头的输入是zL0\mathbf{z}_L^0,输出是图片各分类的概率。预训练阶段的分类头是包含一个隐层的多层感知机,微调阶段的分类头是线性映射矩阵。

效果

ViT通过使用不同的超参配置,构建了三种规模的模型,分别是ViT-Base、ViT-Large、ViT-Huge,如图8所示。三个模型的规模逐渐增长,编码器的层数从12增大到32,编码器中向量的维度从768增大到1280,模型参数量从8千万增大到6亿。

图8 不同规模的ViT模型

图9 各模型在各分类任务上的准确率

各模型在各分类任务上的准确率如图9所示。从中可以看出,同一个模型(ViT-L/16)在较大数据集(JFT-300M)上预训练的效果优于在较小数据集(ImageNet-21k)上预训练的效果,同一个预训练数据集(JFT-300M)下较大模型(ViT-H/14)的效果优于较小模型的效果(ViT-L/16),同一个预训练数据集(JFT-300M)下ViT(ViT-H/14和ViT-L/16)的效果优于之前基于卷积神经网络的SOTA模型BiT-L。上述结果论证了Transformer在计算机视觉领域也满足Scaling Law,即模型规模越大,数据集越大,其效果越好。

DDPM

扩散模型(Diffusion Model)是继GAN、VAE后的一种生成式模型,而目前在文生图领域比较流行的工具,如OpenAI的DALL-E系列、Stability AI的Stable Diffusion等,均是以上述扩散模型为基础,不断进行算法优化、迭代,取得了令人惊艳的效果

扩散模型于2015年在论文《Deep Unsupervised Learning Using Nonequilibrium Thermodynamics》中被提出,并于2020年在论文《Denoising Diffusion Probabilistic Models》中被改进、用于图片生成。《Denoising Diffusion Probabilistic Models》中提出的扩散模型算法被称为DDPM

正向扩散过程

令原始图片样本为x0x_0,其满足分布x0q(x0)x_0 \sim q(x_0)。定义前向扩散过程,在TT步内,每步给样本增加一个小的满足高斯分布的噪声,从而产生TT个带噪声的样本x1,...,xTx_1,...,x_T,整个过程为一个一阶马尔可夫过程,xtx_t只与xt1x_{t-1}有关,可用以下公式表示:

q(xtxt1)=N(xt;1βtxt1,βtI)q(x_t|x_{t-1})=\mathcal{N}(x_t;\sqrt{1-\beta_t}x_{t-1},\beta_t\mathbf{I})

其中,q(xtxt1)q(x_t|x_{t-1})表示给定xt1x_{t-1}时,xtx_{t}的条件概率,即均值为1βtxt1\sqrt{1-\beta_t}x_{t-1}、方差为βtI\beta_t\mathbf{I}的高斯分布,集合{βt(0,1)}t=1T\{\beta_t \in (0,1)\}_{t=1}^{T}用于控制每步的噪声大小。进一步给定x0x_0时,整个马尔科夫过程的条件概率为各步条件概率的连乘,可用以下公式表示:

q(x1:Tx0)=t=1Tq(xtxt1)q(x_{1:T}|x_0)=\prod_{t=1}^{T}{q(x_t|x_{t-1})}

正向扩散过程可由图10从右到左的过程表示,其中x0x_0为原始图片,随着每步增加噪声,图片逐渐变得模糊。

图10 从右到左为正向扩散过程

对于上述正向扩散过程,可进一步令αt=1βt\alpha_t=1-\beta_t,且αˉt=i=1tαi\bar{\alpha}_t=\prod_{i=1}^{t}{\alpha_i},则xtx_t可用以下公式表示:

xt=αtxt1+1αtϵt1;where ϵt1,ϵt2,...N(0,I)=αtαt1xt2+1αtαt1ϵˉt2;where ϵˉt2 merge two Gaussians ().=...=αˉtx0+1αˉtϵ\begin{aligned} x_t&=\sqrt{\alpha_t}x_{t-1}+\sqrt{1-\alpha_t}\epsilon_{t-1} &;\text{where }\epsilon_{t-1},\epsilon_{t-2},...\sim\mathcal{N}(0,\mathbf{I})\\ &=\sqrt{\alpha_t \alpha_{t-1}}x_{t-2}+\sqrt{1-\alpha_t\alpha_{t-1}}\bar{\epsilon}_{t-2} &;\text{where }\bar{\epsilon}_{t-2}\text{ merge two Gaussians }(*). \\ &=... \\ &=\sqrt{\bar{\alpha}_t}x_0+\sqrt{1-\bar{\alpha}_t}\epsilon \end{aligned}

xtx_t是在xt1x_{t-1}的基础上,增加一个满足高斯分布的噪声ϵt1\epsilon_{t-1},循环递归,即xtx_t可进一步推导为在x0x_0的基础上,增加一个满足高斯分布的噪声ϵ\epsilon。这里使用了高斯分布的一个特性,即两个高斯分布合并后仍是一个高斯分布,例如分布N(0,σ12I)\mathcal{N}(0,\sigma_1^2\mathbf{I})N(0,σ22I)\mathcal{N}(0,\sigma_2^2\mathbf{I}),合并后的分布为N(0,(σ12+σ22)I)\mathcal{N}(0,(\sigma_1^2+\sigma_2^2)\mathbf{I})

反向扩散过程

以上介绍了正向扩散过程,即图10从右到左,对原始图片逐步增加噪声,如果将过程逆向,即图1从左到右,那么就能从满足高斯分布的噪音xTN(0,I)x_T \sim \mathcal{N}(0,\mathbf{I})逐步还原原始图片样本,这就是基于扩散模型生成图片的基本思想,即从xTx_Tx0x_0的每一步,在给定xtx_t时,根据条件概率q(xt1xt)q(x_{t-1}|x_t)采样求解xt1x_{t-1},直至最终得到x0x_0

而当正向扩散过程每步增加的噪声很小时,反向扩散过程的条件概率q(xt1xt)q(x_{t-1}|x_t)也可以认为满足高斯分布,但实际上,我们不能直接求解该条件概率,因为直接求解需要整体数据集合。除直接求解外,另一个方法是训练一个模型pθp_\theta近似预估上述条件概率,可用以下公式表示:

pθ(xt1xt)=N(xt1;μθ(xt,t),Σθ(xt,t))p_\theta(x_{t-1}|x_t)=\mathcal{N}(x_{t-1};\mu_\theta(x_t,t),\Sigma_\theta(x_t,t))
pθ(x0:T)=p(xT)t=1Tpθ(xt1xt)p_{\theta}(x_{0:T})=p(x_T)\prod_{t=1}^{T}{p_\theta(x_{t-1}|x_t)}

xTx_Tx0x_0的每一步,通过模型θ\theta,输入xtx_ttt,预测xt1x_{t-1}的高斯分布的均值μθ(xt,t)\mu_\theta(x_t,t)和方差Σθ(xt,t)\Sigma_\theta(x_t,t),基于预测值,可以从xt1x_{t-1}的高斯分布pθ(xt1xt)p_\theta(x_{t-1}|x_t)中进行采样,从而得到xt1x_{t-1}的一个可能取值,如此循环,直至最终得到x0x_0的一个可能取值。通过上述反向扩散过程,即可以实现从一个满足高斯分布N(0,I)\mathcal{N}(0,\mathcal{I})的随机噪声,生成一张图片。而由于每次预测均是从一个概率密度函数中进行采样,因此,可以保证生成图片的多样性。

更进一步,论文进一步将模型预测pθ(xt1xt)p_\theta(x_{t-1}|x_t)的均值μθ(xt,t)\mu_\theta(x_t,t)和方差Σθ(xt,t)\Sigma_\theta(x_t,t)转化为预测噪声ϵθ(xt,t)\epsilon_\theta(x_t,t),并推导出μθ(xt,t)\mu_\theta(x_t,t)ϵθ(xt,t)\epsilon_\theta(x_t,t)的关系:

μθ(xt,t)=1αt(xt1αt1αˉtϵθ(xt,t))\mu_\theta(x_t,t)=\frac{1}{\sqrt{\alpha_t}}\left(x_t-\frac{1-\alpha_t}{\sqrt{1-\bar{\alpha}_t}}\epsilon_\theta(x_t,t)\right)

因此,pθ(xt1xt)p_\theta(x_{t-1}|x_t)可表示为:

pθ(xt1xt)=N(xt1;1αt(xt1αt1αˉtϵθ(xt,t)),Σθ(xt,t))p_\theta(x_{t-1}|x_t)=\mathcal{N}(x_{t-1};\frac{1}{\sqrt{\alpha_t}}\left(x_t-\frac{1-\alpha_t}{\sqrt{1-\bar{\alpha}_t}}\epsilon_\theta(x_t,t)\right),\Sigma_\theta(x_t,t))

论文将Σθ(xt,t)\Sigma_\theta(x_t,t)固定为常量,通过模型预测ϵθ(xt,t)\epsilon_\theta(x_t,t),并使用上述公式的概率密度函数进行采样,从xtx_t降噪得到xt1x_{t-1},这也是论文标题中“Denosing”的由来。

模型结构

DDPM中预测ϵθ(xt,t)\epsilon_\theta(x_t,t)的模型基于OpenAI于2017年发布的一个U-Net形式的网络结构PixelCNN++。U-Net于2015年在论文《U-Net: Convolutional Networks for Biomedical Image Segmentation》中发布,起初主要用于医学图像的切割,目前作为常用的去噪结构,广泛应用于扩散模型中。

图11 U-Net网络结构

U-Net的网络结构如图11所示,因整体结构形似字母U而得名,U型左侧是4层编码器层,对图片进行降维,U型右侧是4层解码器层,对图片进行升维。编码器的每一层,先是连续的两个卷积层(卷积核维度为3×3)和ReLU层,再接一个池化层进行下采样,然后输入下一层,卷积层的通道数逐层加倍,例如,第一层的输入是572×572的单通道图片,经过两个卷积核维度为3×3、通道为64、无padding的卷积层,输出张量维度分别为570×570×64、568×568×64,再经过一个维度为2×2的最大池化层后,输出张量维度为284×284×64,如此循环,最后一层输出的张量维度为32×32×512。在编码器层和解码器层之间的中间层,经过两个卷积层(卷积核维度为3×3、通道为1024)和ReLU层,输出的张量维度为28×28×1024。解码器的每一层,和编码器类似,也先是连续的两个卷积层(卷积核维度为3×3)和ReLU层,和编码器不同的是,解码器从下层到上层,通过一个上卷积层(卷积核维度为3×3)进行上采样(长、宽维度加倍,但通道缩小),同时,解码器每层的输入除上一层上采样的输出外,还包括同层编码器输出的裁剪。例如,解码器第一层的输入,包括中间层上采样的输出,张量维度为56×56×512,和同层编码器输出的裁剪,张量维度为56×56×512,合并后的张量维度为56×56×1024,经过两个卷积核维度为3×3、通道为512的卷积层,输出张量维度分别为54×54×512、52×52×512。解码器最后一层的输出,再通过一个1×1的卷积层,将原先的64通道映射为指定的通道,因为原始U-Net用于图像切割,即对图像每个像素做分类,所以有多少个分类,即有多少个最终的通道。

图12 PixelCNN++网络结构

而PixelCNN++于2017年在论文《PixelCNN++: Improving the PixelCNN with Discretized Logistic Mixture Likelihood and Other Modifications》中发布,其网络结构如图12所示。图中,矩形区块对应于U-Net中的编码器或解码器层,共3个编码器层、3个解码器层。在每个编码器或解码器中,PixelCNN++在原U-Net两个卷积层的基础上,增加了一个残差连接。DDPM进一步进行网络结构的改进,包括:使用Group Normalization进行归一化;在残差卷积块后增加自注意力层;使用Transformer中Sinusoidal Position Embedding对步数tt编码成Embedding向量作为模型输入。

DDPM的代码开源,代码地址是:github.com/hojonathanh…,其深度学习框架采用Tensorflow,计算资源采用Google Cloud TPU v3-8。diffusion_tf/models/unet.py中定义了网络结构,核心代码如下所示(增加了部分注释):

  with tf.variable_scope(name, reuse=reuse):
    # Timestep embedding
    # 将步数t编码成Embedding向量  
    with tf.variable_scope('temb'):
      temb = nn.get_timestep_embedding(t, ch)
      temb = nn.dense(temb, name='dense0', num_units=ch * 4)
      temb = nn.dense(nonlinearity(temb), name='dense1', num_units=ch * 4)
      assert temb.shape == [B, ch * 4]

    # Downsampling
    # 多层编码器层
    hs = [nn.conv2d(x, name='conv_in', num_units=ch)] 
    for i_level in range(num_resolutions):
      with tf.variable_scope('down_{}'.format(i_level)):
        # Residual blocks for this resolution
        # 构造编码器层,残差卷积块+自注意力层,并进行下采样   
        for i_block in range(num_res_blocks):
          h = resnet_block(
            hs[-1], name='block_{}'.format(i_block), temb=temb, out_ch=ch * ch_mult[i_level], dropout=dropout)
          if h.shape[1] in attn_resolutions:
            h = attn_block(h, name='attn_{}'.format(i_block), temb=temb)
          hs.append(h)
        # Downsample
        if i_level != num_resolutions - 1:
          hs.append(downsample(hs[-1], name='downsample', with_conv=resamp_with_conv))

    # Middle
    # 中间层,残差卷积块+自注意力层+残差卷积块       
    with tf.variable_scope('mid'):
      h = hs[-1]
      h = resnet_block(h, temb=temb, name='block_1', dropout=dropout)
      h = attn_block(h, name='attn_1'.format(i_block), temb=temb)
      h = resnet_block(h, temb=temb, name='block_2', dropout=dropout)

    # Upsampling
    # 多层解码器层    
    for i_level in reversed(range(num_resolutions)):
      with tf.variable_scope('up_{}'.format(i_level)):  
        # Residual blocks for this resolution
        # 构造解码器层,残差卷积块+自注意力层,并进行上采样  
        for i_block in range(num_res_blocks + 1):
          h = resnet_block(tf.concat([h, hs.pop()], axis=-1), name='block_{}'.format(i_block),
                           temb=temb, out_ch=ch * ch_mult[i_level], dropout=dropout)
          if h.shape[1] in attn_resolutions:
            h = attn_block(h, name='attn_{}'.format(i_block), temb=temb)
        # Upsample
        if i_level != 0:
          h = upsample(h, name='upsample', with_conv=resamp_with_conv)
    assert not hs

    # End
    # 最后再经过一个卷积层输出  
    h = nonlinearity(normalize(h, temb=temb, name='norm_out'))
    h = nn.conv2d(h, name='conv_out', num_units=out_ch, init_scale=0.)
    assert h.shape == x.shape[:3] + [out_ch]
    return h

其中,残差卷积块的代码如下(增加了部分注释):

def resnet_block(x, *, temb, name, out_ch=None, conv_shortcut=False, dropout):
  B, H, W, C = x.shape
  if out_ch is None:
    out_ch = C

  with tf.variable_scope(name):
    h = x

    # 对图片进行归一化和非线性转化  
    h = nonlinearity(normalize(h, temb=temb, name='norm1'))
    # 对图片进行卷积  
    h = nn.conv2d(h, name='conv1', num_units=out_ch)

    # add in timestep embedding
    # 对步数t的embedding向量进行非线性转化,并合并至图片  
    h += nn.dense(nonlinearity(temb), name='temb_proj', num_units=out_ch)[:, None, None, :]

    # 对合并图片和步数后的输入再进行归一化和非线性转化,并再进行卷积  
    h = nonlinearity(normalize(h, temb=temb, name='norm2'))
    h = tf.nn.dropout(h, rate=dropout)
    h = nn.conv2d(h, name='conv2', num_units=out_ch, init_scale=0.)

    # 对两次卷积后的输出和原始输入进行残差连接  
    if C != out_ch:
      if conv_shortcut:
        x = nn.conv2d(x, name='conv_shortcut', num_units=out_ch)
      else:
        x = nn.nin(x, name='nin_shortcut', num_units=out_ch)

    assert x.shape == h.shape
    print('{}: x={} temb={}'.format(tf.get_default_graph().get_name_scope(), x.shape, temb.shape))
    return x + h

自注意力层的代码如下(即Transformer中的缩放点积注意力层):

def attn_block(x, *, name, temb):
  B, H, W, C = x.shape
  with tf.variable_scope(name):
    h = normalize(x, temb=temb, name='norm')
    q = nn.nin(h, name='q', num_units=C)
    k = nn.nin(h, name='k', num_units=C)
    v = nn.nin(h, name='v', num_units=C)

    w = tf.einsum('bhwc,bHWc->bhwHW', q, k) * (int(C) ** (-0.5))
    w = tf.reshape(w, [B, H, W, H * W])
    w = tf.nn.softmax(w, -1)
    w = tf.reshape(w, [B, H, W, H, W])

    h = tf.einsum('bhwHW,bHWc->bhwc', w, v)
    h = nn.nin(h, name='proj_out', num_units=C, init_scale=0.)

    assert h.shape == x.shape
    print(tf.get_default_graph().get_name_scope(), x.shape)
    return x + h

训练采样

训练

通过模型预测误差ϵθ(xt,t)\epsilon_\theta(x_t,t),损失函数采用均方误差(MSE,Mean-Squared Error):

Lsimple:=Et[1,T],x0q(x0),ϵN(0,I)[ϵϵθ(xt,t)2]L_{\text{simple}}:=E_{t\sim[1,T],x_0\sim q(x_0),\epsilon\sim\mathcal{N}(0,\mathbf{I})}[\parallel\epsilon-\epsilon_\theta(x_t,t)\parallel^2]

模型训练的目标即最小化上述损失函数,即使模型预测出的噪声ϵθ(xt,t)\epsilon_\theta(x_t,t)和真实噪声ϵ\epsilon尽可能接近。

图13 训练算法

训练算法如图13所示,采用梯度下降算法,循环下述过程直至模型收敛:

  • 对于样本x0x_0,从{1,...,T}\{1,...,T\}中随机采样步数tt
  • 从高斯分布N(0,I)\mathcal{N}(0,\mathbf{I})中采样真实噪声ϵ\epsilon
  • 根据样本x0x_0和真实噪声ϵ\epsilon,使用前面推导出的公式xt=αˉtx0+1αˉtϵx_t=\sqrt{\bar{\alpha}_t}x_0+\sqrt{1-\bar{\alpha}_t}\epsilon计算第tt步正向扩散后带噪声的图片xtx_t
  • 根据带噪声的图片xtx_t和步数tt,使用模型预测噪声ϵθ(xt,t)\epsilon_\theta(x_t,t),即ϵθ(αˉtx0+1αˉtϵ,t)\epsilon_\theta(\sqrt{\bar{\alpha}_t}x_0+\sqrt{1-\bar{\alpha}_t}\epsilon,t)
  • 根据真实噪声ϵ\epsilon和预测噪声ϵθ(αˉtx0+1αˉtϵ,t)\epsilon_\theta(\sqrt{\bar{\alpha}_t}x_0+\sqrt{1-\bar{\alpha}_t}\epsilon,t)计算损失函数的梯度,即θϵϵθ(αˉtx0+1αˉtϵ)2\nabla_\theta\parallel\epsilon-\epsilon_\theta(\sqrt{\bar{\alpha}_t}x_0+\sqrt{1-\bar{\alpha}_t}\epsilon)\parallel^2
  • 根据梯度和学习率超参更新模型参数。

采样

图14 采样算法

采样算法如图14所示,过程如下:

  • 从高斯分布N(0,I)\mathcal{N}(0,\mathbf{I})中采样完全噪声图片xTx_T
  • 循环TT步,步数从TT到1,直至计算得到x0x_0,生成最终的图片,对于其中的某一步tt
    • 根据带噪声的图片xtx_t和步数tt,使用模型预测噪声ϵθ(xt,t)\epsilon_\theta(x_t,t)
    • 前面已推导出概率密度函数pθ(xt1xt)p_\theta(x_{t-1}|x_t)满足高斯分布N(xt1;1αt(xt1αt1αˉtϵθ(xt,t)),Σθ(xt,t))\mathcal{N}(x_{t-1};\frac{1}{\sqrt{\alpha_t}}\left(x_t-\frac{1-\alpha_t}{\sqrt{1-\bar{\alpha}_t}}\epsilon_\theta(x_t,t)\right),\Sigma_\theta(x_t,t)),使用公式μθ(xt,t)=1αt(xt1αt1αˉtϵθ(xt,t))\mu_\theta(x_t,t)=\frac{1}{\sqrt{\alpha_t}}\left(x_t-\frac{1-\alpha_t}{\sqrt{1-\bar{\alpha}_t}}\epsilon_\theta(x_t,t)\right)由噪声ϵθ(xt,t)\epsilon_\theta(x_t,t)计算均值μθ(xt,t)\mu_\theta(x_t,t)
    • 对于上述概率密度函数,指定方差Σθ(xt,t)\Sigma_\theta(x_t,t)为常量,根据该分布进行采样,从xtx_t得到xt1x_{t-1},即xt1=1αt(xt1αt1αˉtϵθ(xt,t))+σtzx_{t-1}=\frac{1}{\sqrt{\alpha_t}}\left(x_t-\frac{1-\alpha_t}{\sqrt{1-\bar{\alpha}_t}}\epsilon_\theta(x_t,t)\right) +\sigma_tz

Improved DDPM

2021年OpenAI发表的论文《Improved Denoising Diffusion Probabilistic Models》对DDPM算法进行改进,包括噪声Schedule采用余弦函数、对方差进行学习(原DDPM算法将方差固定为常量)等

改进

噪声Schedule采用余弦函数

原始DDPM算法,使用公式xt=αˉtx0+1αˉtϵx_t=\sqrt{\bar{\alpha}_t}x_0+\sqrt{1-\bar{\alpha}_t}\epsilon计算第tt步正向扩散后带噪声的图片xtx_t,公式中的αˉt=i=1tαi\bar{\alpha}_t=\prod_{i=1}^{t}{\alpha_i}βt=1αt\beta_t=1-\alpha_t,即βt=1αˉtαˉt1\beta_t=1-\frac{\bar{\alpha}_t}{\bar{\alpha}_{t-1}}βt\beta_t表示每步噪声的大小,原始DDPM算法,令βt\beta_ttt线性增长,从β1=104\beta_1=10^{-4}增长到βT=0.02\beta_T=0.02,改进的DDPM算法,令αˉt\bar{\alpha}_ttt的变化采用余弦函数:

f(t)=cos(t/T+s1+sπ2)2αˉt=f(t)f(0)βt=clip(1αˉtαˉt1,0.999)\begin{aligned} f(t)&=\cos{\left(\frac{t/T+s}{1+s}\cdot\frac{\pi}{2}\right)^2} \\ \bar{\alpha}_t&=\frac{f(t)}{f(0)}\\ \beta_t&=\text{clip}(1-\frac{\bar{\alpha}_t}{\bar{\alpha}_{t-1}},0.999) \end{aligned}

两种噪声Schedule下,αˉt\bar{\alpha}_ttt的变化曲线如图15所示,相比线性函数,余弦函数的αˉt\bar{\alpha}_t下降相对较平缓,因而βt\beta_t相对较小,加噪相对较慢,不会过快地对原始图片加入过多的噪声。

图15 两种噪声Schedule下α随t的变化

Improved DDPM的代码开源,代码地址是:github.com/openai/impr…,其深度学习框架采用PyTorch。训练和采样分别执行以下脚本:

# 训练脚本
python scripts/image_train.py --data_dir path/to/images $MODEL_FLAGS $DIFFUSION_FLAGS $TRAIN_FLAGS

# 采样脚本
python scripts/image_sample.py --model_path /path/to/model.pt $MODEL_FLAGS $DIFFUSION_FLAGS

其中,MODEL_FLAGS、DIFFUSION_FLAGS、TRAIN_FLAGS分别表示模型结构(U-Net)、扩散过程和训练的配置,而基线模型(DDPM)的配置如下:

MODEL_FLAGS="--image_size 64 --num_channels 128 --num_res_blocks 3"
DIFFUSION_FLAGS="--diffusion_steps 4000 --noise_schedule linear"
TRAIN_FLAGS="--lr 1e-4 --batch_size 128"

如果采用余弦函数作为噪声Schedule,可以将DIFFUSION_FLAGS中的noise_schedule设置为cosine,而相应的计算噪声βt\beta_t的代码在improved_diffusion/gaussian_diffusion.py中,如下:

def get_named_beta_schedule(schedule_name, num_diffusion_timesteps):
    """
    Get a pre-defined beta schedule for the given name.

    The beta schedule library consists of beta schedules which remain similar
    in the limit of num_diffusion_timesteps.
    Beta schedules may be added, but should not be removed or changed once
    they are committed to maintain backwards compatibility.
    """
    if schedule_name == "linear":
        # Linear schedule from Ho et al, extended to work for any number of
        # diffusion steps.
        scale = 1000 / num_diffusion_timesteps
        beta_start = scale * 0.0001
        beta_end = scale * 0.02
        return np.linspace(
            beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64
        )
    elif schedule_name == "cosine":
        return betas_for_alpha_bar(
            num_diffusion_timesteps,
            lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2,
        )
    else:
        raise NotImplementedError(f"unknown beta schedule: {schedule_name}")

# 入参alpha_bar即公式推导中的f(t) 
def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):
    """
    Create a beta schedule that discretizes the given alpha_t_bar function,
    which defines the cumulative product of (1-beta) over time from t = [0,1].

    :param num_diffusion_timesteps: the number of betas to produce.
    :param alpha_bar: a lambda that takes an argument t from 0 to 1 and
                      produces the cumulative product of (1-beta) up to that
                      part of the diffusion process.
    :param max_beta: the maximum beta to use; use values lower than 1 to
                     prevent singularities.
    """
    betas = []
    for i in range(num_diffusion_timesteps):
        t1 = i / num_diffusion_timesteps
        t2 = (i + 1) / num_diffusion_timesteps
        betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))
    return np.array(betas)        

对方差进行学习

原始DDPM算法,满足高斯分布的概率密度函数pθ(xt1xt)p_\theta(x_{t-1}|x_t)中的方差Σθ(xt,t)\Sigma_\theta(x_t,t)被固定为常量σt2I\sigma_t^2\mathbf{I},其中σt2\sigma_t^2直接取值βt\beta_t。改进的DDPM算法通过模型对Σθ(xt,t)\Sigma_\theta(x_t,t)也进行了学习,这样可以使用更少的步数、且生成更高质量的图片。

图16 两个β的比值和t之间的关系

具体如何学习Σθ(xt,t)\Sigma_\theta(x_t,t)呢?DDPM的论文已推导Σθ(xt,t)\Sigma_\theta(x_t,t)的取值在βt\beta_tβ~t\tilde{\beta}_t之间,而β~t:=1αˉt11αˉtβt\tilde{\beta}_t:=\frac{1-\bar{\alpha}_{t-1}}{1-\bar{\alpha}_t}\beta_t,图16表示了β~t/βt\tilde{\beta}_t/\beta_ttt之间的关系,从中可见,除了t=0t=0外,其他tt取值下,βt\beta_tβ~t\tilde{\beta}_t近似相等,所以原始DDPM算法将方差直接取值βt\beta_t,而改进的DDPM算法设计了中间向量vv,由模型预测vv,并将Σθ(xt,t)\Sigma_\theta(x_t,t)表示为:

Σθ(xt,t)=exp(vlogβt+(1v)logβ~t)\Sigma_\theta(x_t,t)=\exp(v\log{\beta_t}+(1-v)\log{\tilde{\beta}_t})

原始DDPM算法的损失函数为:

Lsimple:=Et[1,T],x0q(x0),ϵN(0,I)[ϵϵθ(xt,t)2]L_{\text{simple}}:=E_{t\sim[1,T],x_0\sim q(x_0),\epsilon\sim\mathcal{N}(0,\mathbf{I})}[\parallel\epsilon-\epsilon_\theta(x_t,t)\parallel^2]

其中并不包含Σθ(xt,t)\Sigma_\theta(x_t,t),因此改进的DDPM算法设计了新的损失函数为:

Lhybrid=Lsimple+λLvlbL_\text{hybrid}=L_\text{simple}+\lambda L_\text{vlb}

LvlbL_\text{vlb}的定义如下:

Lvlb:=L0+L1++LT1+LTwhere L0:=logpθ(x0x1)Lt1:=DKL(q(xt1xt,x0)pθ(xt1xt)) for 2tTLT:=DKL(q(xTx0)p(xT))\begin{aligned} L_{\text{vlb}}&:=L_0+L_1+\cdots+L_{T-1}+L_T \\ \text{where }L_0&:=-\log{p_\theta(x_0|x_1)} \\ L_{t-1}&:=D_{\text{KL}}(q(x_{t-1}|x_{t},x_0)\parallel p_\theta(x_{t-1}|x_{t}))\text{ for } 2\le t\le T \\ L_T&:=D_{\text{KL}}(q(x_T|x_0)\parallel p(x_T)) \end{aligned}

损失函数中λ\lambda被设置为0.001,以减少LvlbL_\text{vlb}LsimpleL_\text{simple}的影响。另外梯度更新时,LvlbL_\text{vlb}部分不更新涉及μθ(xt,t)\mu_\theta(x_t,t)的参数,只更新涉及Σθ(xt,t)\Sigma_\theta(x_t,t)的参数。 如果需要对方差进行学习,可以将训练脚本参数MODEL_FLAGS中的learn_sigma设置为True,这样,模型结构(U-Net)的输出维度增加,增加的部分作为vv,相关代码在improved_diffusion/script_util.py的create_model方法中,如下:

    return UNetModel(
        in_channels=3,
        model_channels=num_channels,
        # 模型结构(U-Net)的输出维度增加,增加的部分作为v
        out_channels=(3 if not learn_sigma else 6),
        num_res_blocks=num_res_blocks,
        attention_resolutions=tuple(attention_ds),
        dropout=dropout,
        channel_mult=channel_mult,
        num_classes=(NUM_CLASSES if class_cond else None),
        use_checkpoint=use_checkpoint,
        num_heads=num_heads,
        num_heads_upsample=num_heads_upsample,
        use_scale_shift_norm=use_scale_shift_norm,
    )

同时,反向扩散生成图片时,由模型预测方差的代码在improved_diffusion/gaussian_diffusion.py的p_mean_variance方法中,如下:

        # 模型预测
        model_output = model(x, self._scale_timesteps(t), **model_kwargs)

        if self.model_var_type in [ModelVarType.LEARNED, ModelVarType.LEARNED_RANGE]:
            assert model_output.shape == (B, C * 2, *x.shape[2:])
            # 按列拆分模型输出,后半部分作为v
            model_output, model_var_values = th.split(model_output, C, dim=1)
            if self.model_var_type == ModelVarType.LEARNED:
                model_log_variance = model_var_values
                model_variance = th.exp(model_log_variance)
            else:
                # 按公式exp(v·log(β_t)+(1-v)·log(β_t))计算方差 
                min_log = _extract_into_tensor(
                    self.posterior_log_variance_clipped, t, x.shape
                )
                max_log = _extract_into_tensor(np.log(self.betas), t, x.shape)
                # The model_var_values is [-1, 1] for [min_var, max_var].
                frac = (model_var_values + 1) / 2
                model_log_variance = frac * max_log + (1 - frac) * min_log
                model_variance = th.exp(model_log_variance)

而代码如何调整损失函数的计算,在下一节介绍。

训练时采用Importance Sampling

图17 损失函数取值随着训练迭代变化的曲线

图18 L_vlb每项的损失

论文进一步发现,将损失函数替换为LvlbL_\text{vlb}LhybridL_\text{hybrid}后,损失函数取值随着训练迭代变化的曲线比较波动,不易收敛,如图17所示。LvlbL_\text{vlb}包含多项,每项对应一个步数tt,且每项的取值量纲差别较大,如图18所示,而原始DDPM算法训练时随机采样步数tt,因此论文认为每次训练迭代随机采样步数tt、并进而计算量纲差别大的LtL_t导致LvlbL_\text{vlb}LhybridL_\text{hybrid}的波动。论文通过训练时采用Importance Sampling来解决上述波动问题。Importance Sampling中,LvlbL_\text{vlb}可表示为以下公式:

Lvlb=Etpt[Ltpt],where ptE[Lt2] and pt=1L_\text{vlb}=E_{t\sim p_t}\left[\frac{L_t}{p_t}\right],\text{where }p_t\propto\sqrt{E\left[L_t^2\right]}\text{ and }\sum{p_t}=1

E[Lt2]E\left[L_t^2\right]无法提前求解,且在训练过程中会变化,因此,论文对LvlbL_\text{vlb}的每一项LtL_t保留最新的10个取值,并在训练过程中动态更新。训练初期,仍是随机采样步数tt,直至所有的LtL_t均有10个取值,再采用Importance Sampling。从图3可以看出,经过Importance Sampling后的LvlbL_\text{vlb}随着训练迭代变化的曲线比较平滑,且损失最小。

如果需要使用Importance Sampling后的LvlbL_\text{vlb}作为损失函数,可以将训练脚本参数DIFFUSION_FLAGS中的use_kl设置为True、TRAIN_FLAGS中的schedule_sampler设置为loss-second-moment。Importance Sampling的相关代码在improved_diffusion/resample.py的LossSecondMomentResampler类中,如下:

class LossSecondMomentResampler(LossAwareSampler):
    def __init__(self, diffusion, history_per_term=10, uniform_prob=0.001):
        self.diffusion = diffusion
        self.history_per_term = history_per_term
        self.uniform_prob = uniform_prob
        self._loss_history = np.zeros(
            [diffusion.num_timesteps, history_per_term], dtype=np.float64
        )
        self._loss_counts = np.zeros([diffusion.num_timesteps], dtype=np.int)

    def weights(self):
        # 训练初期,仍是随机采样步数
        if not self._warmed_up():
            return np.ones([self.diffusion.num_timesteps], dtype=np.float64)
        # 根据L_t的历史值计算p_t
        weights = np.sqrt(np.mean(self._loss_history ** 2, axis=-1))
        weights /= np.sum(weights)
        weights *= 1 - self.uniform_prob
        weights += self.uniform_prob / len(weights)
        return weights

    def update_with_all_losses(self, ts, losses):
        for t, loss in zip(ts, losses):
            if self._loss_counts[t] == self.history_per_term:
                # Shift out the oldest loss term.
                self._loss_history[t, :-1] = self._loss_history[t, 1:]
                self._loss_history[t, -1] = loss
            else:
                self._loss_history[t, self._loss_counts[t]] = loss
                self._loss_counts[t] += 1

    def _warmed_up(self):
        return (self._loss_counts == self.history_per_term).all()

使用LvlbL_\text{vlb}作为损失函数的相关代码在improved_diffusion/gaussian_diffusion.py的_vb_terms_bpd方法中,如下:

    def _vb_terms_bpd(
        self, model, x_start, x_t, t, clip_denoised=True, model_kwargs=None
    ):
        """
        Get a term for the variational lower-bound.

        The resulting units are bits (rather than nats, as one might expect).
        This allows for comparison to other papers.

        :return: a dict with the following keys:
                 - 'output': a shape [N] tensor of NLLs or KLs.
                 - 'pred_xstart': the x_0 predictions.
        """
        true_mean, _, true_log_variance_clipped = self.q_posterior_mean_variance(
            x_start=x_start, x_t=x_t, t=t
        )
        out = self.p_mean_variance(
            model, x_t, t, clip_denoised=clip_denoised, model_kwargs=model_kwargs
        )
        kl = normal_kl(
            true_mean, true_log_variance_clipped, out["mean"], out["log_variance"]
        )
        kl = mean_flat(kl) / np.log(2.0)

        decoder_nll = -discretized_gaussian_log_likelihood(
            x_start, means=out["mean"], log_scales=0.5 * out["log_variance"]
        )
        assert decoder_nll.shape == x_start.shape
        decoder_nll = mean_flat(decoder_nll) / np.log(2.0)

        # At the first timestep return the decoder NLL,
        # otherwise return KL(q(x_{t-1}|x_t,x_0) || p(x_{t-1}|x_t))
        output = th.where((t == 0), decoder_nll, kl)
        return {"output": output, "pred_xstart": out["pred_xstart"]}

效果

针对上述改进,论文在ImageNet 64×64和CIFAR-10这两个数据集上分别进行消融实验以验证各改进的有效性,如图19和图20所示。

图19 ImageNet 64×64数据集上的消融实验结果

图20 CIFAR-10数据集上的消融实验结果

其中,有效性指标使用了NLL和FID。NLL(Negative Log Likelihood)等价于损失函数LvlbL_\text{vlb},NLL越小,说明生成图像与真实图像的分布越接近。FID(Fréchet Inception Distance)是另一种用于图像生成质量评估的指标,它可以评估生成图像与真实图像之间的相似度。FID指标的计算方法是使用Inception-v3模型对生成图像和真实图像进行特征提取,并计算两个特征分布之间的Fréchet距离。FID越小,说明生成图像与真实图像越相似。从实验结果上看,噪声Schedule采用余弦函数、对方差进行学习并且训练时损失函数采用Importance Sampling后的LvlbL_\text{vlb},NLL最低,但FID较高,而噪声Schedule采用余弦函数、对方差进行学习并且训练时损失函数采用LhybridL_\text{hybrid},在NLL、FID上都能取得较小的值。

图21 和其他基于似然预估的模型进行对比实验

另外,论文还和其他基于似然预估的模型进行了对比实验,如图21所示。优化后的DDPM虽然在NLL和FID上还不是SOTA,但相对也是较优的效果,仅次于基于Transformer的网络结构。

加速

DDPM在生成图片时需要从完全噪声开始执行多步降噪操作,而每步操作均需要将当前步带噪声的图片作为输入由模型预测噪声,导致生成图片需要较多的步骤和计算量。论文也采用了《Denoising Diffusion Implicit Models》提出的采样方法——DDIM,减少步数。

Classifier Guidance

Improved DDPM虽然对DDPM进行了改进,但在一些大数据集上(如ImageNet 256×256)生成图片的实验效果(FID)仍是低于GAN。因此,OpenAI继续对DDPM进行改进,在2021年随后又发表了论文《Diffusion Models Beat Gans on Image Synthesis》在模型结构上进一步优化,同时引入Classifier Guidance技术,通过图片分类器的梯度调节反向扩散过程,在尽量保持图片生成多样性的前提下,提升准确性,从而在多个数据集的实验效果(FID)上超过了GAN,实现了SOTA。论文将改进后的模型称为ADM(Ablated Diffusion Model)

改进

网络结构

DDPM和Improved DDPM中的模型均使用U-Net,ADM在其网络结构的基础上,进一步增加以下数项改进:

  • 增加网络结构的宽度和深度;
  • 在注意力机制上,DDPM原先只在16×16这一层增加单头注意力层(缩放点积注意力),而ADM在32×32、16×16、8×8各层均增加了多头注意力层;
  • 在上下采样上,DDPM原先在下采样使用池化或卷积层、在上采样使用插值或卷积层,而ADM使用残差卷积块;

ADM的代码开源,代码地址是:github.com/openai/guid…,其是在Improved DDPM的代码基础上进行修改。网络结构定义的相关代码在guided-diffusion/unet.py的UNetModel类中,例如,ADM使用残差块卷积进行下采样的代码如下:

            if level != len(channel_mult) - 1:
                # 除编码器最后一层外的其他层,需要进行下采样输出到下一层
                out_ch = ch
                self.input_blocks.append(
                    TimestepEmbedSequential(
                        # 若标记resblock_updown为True,则使用残差卷积块进行下采样
                        ResBlock(
                            ch,
                            time_embed_dim,
                            dropout,
                            out_channels=out_ch,
                            dims=dims,
                            use_checkpoint=use_checkpoint,
                            use_scale_shift_norm=use_scale_shift_norm,
                            # 设置残差卷积块中需进行下采样
                            down=True,
                        )
                        if resblock_updown
                        else Downsample(
                            ch, conv_resample, dims=dims, out_channels=out_ch
                        )
                    )
                )
                ch = out_ch
                input_block_chans.append(ch)
                ds *= 2
                self._feature_size += ch

自适应组归一化

图22 各种归一化方式的示意

ADM还使用了自适应组归一化(Adaptive Group Normalization,AdaGN),组归一化如图22最右侧所示,即对一个图片样本的所有像素,按通道分组进行归一化,而自适应归一化可表示为以下公式:

AdaGN(h,y)=ysGroupNorm(h)+yb\text{AdaGN}(h,y)=y_s\text{GroupNorm(h)}+y_b

其中,hh是残差卷积块中第一个卷积层的输出,ysy_syby_b分别是步数和图片分类的Embedding向量经过线性层后的投影。自适应归一化的代码如下所示:

            # 经过第一个卷积层的输出
            h = self.in_layers(x)
            ......
            out_norm, out_rest = self.out_layers[0], self.out_layers[1:]
            # 取y_s和y_b
            scale, shift = th.chunk(emb_out, 2, dim=1)
            # 按y_s * GroupNorm(h) + y_b进行自适应组归一化
            h = out_norm(h) * (1 + scale) + shift
            # 经过第二个卷积层输出
            h = out_rest(h)

论文通过实验发现,使用自适应组归一化能够进一步优化FID。

Classifier Guidance

除了在网络结构上精心设计和优化外,GAN还在有条件(已知图片类别)的图片生成中大量使用了图片类别信息。基于此,ADM一方面在自适应组归一化中引入图片类别的Embedding向量作为模型输入,另一方面设计了Classifier Guidance机制,通过引入一个分类器指导反向扩散过程:预先使用带噪声的图片xtx_t训练分类器pϕ(yxt)p_{\phi}(y|x_t)实现对类别的预测;在逐步反向扩散生成图片时,DDPM在每一步基于扩散模型预测噪声ϵθ(xt)\epsilon_\theta(x_t)和方差Σθ(xt)\Sigma_\theta(x_t),并由公式μθ(xt)=1αt(xt1αt1αˉtϵθ(xt))\mu_\theta(x_t)=\frac{1}{\sqrt{\alpha_t}}\left(x_t-\frac{1-\alpha_t}{\sqrt{1-\bar{\alpha}_t}}\epsilon_\theta(x_t)\right)计算得到μθ(xt)\mu_\theta(x_t),即得到了xt1x_{t-1}高斯分布的均值和方差,在此基础上,ADM使用分类器pϕ(yxt)p_{\phi}(y|x_t)输出对数的梯度对均值进行调整,并使用调整均值后的高斯分布进行采样得到xt1x_{t-1},均值调整公式如下所示:

μ^θ(xty)=μθ(xty)+sΣθ(xty)xtlogpϕ(yxt)\hat{\mu}_\theta(x_t|y)=\mu_\theta(x_t|y)+s\cdot\Sigma_\theta(x_t|y)\nabla_{x_{t}}\log{p_{\phi}{(y|x_t)}}

其中,系数ss被称为Guidance Scale,论文通过实验发现随着ss的增加,生成图片的质量会提升,但多样性会减少。引入Classifier Guidance后的采样算法步骤如图23所示。

图23 引入Classifier Guidance后的采样算法

ADM中使用的分类器网络结构和扩散模型网络结构近似,均采用U-Net,但只使用编码器层和中间层,而没有解码器层,另外,由于分类器的目标是预测类别,因此类别没有作为输入。分类器网络结构定义的相关代码在guided-diffusion/unet.py的UNetModel类中。

在分类器和扩散模型训练完成后,便可使用其进行图片采样。采样时使用分类器输出梯度对均值进行调整的代码在guided-diffusion/gaussian_diffusion.py的condition_mean方法中,如下所示:

    def condition_mean(self, cond_fn, p_mean_var, x, t, model_kwargs=None):
        """
        Compute the mean for the previous step, given a function cond_fn that
        computes the gradient of a conditional log probability with respect to
        x. In particular, cond_fn computes grad(log(p(y|x))), and we want to
        condition on y.

        This uses the conditioning strategy from Sohl-Dickstein et al. (2015).
        """
        # 使用分类器输出分类器梯度
        gradient = cond_fn(x, self._scale_timesteps(t), **model_kwargs)
        # 根据分类器梯度、均值和方差计算新均值
        new_mean = (
            p_mean_var["mean"].float() + p_mean_var["variance"] * gradient.float()
        )
        return new_mean

而其中使用分类器输出分类器梯度的cond_fn方法代码如下所示:

    def cond_fn(x, t, y=None):
        assert y is not None
        with th.enable_grad():
            x_in = x.detach().requires_grad_(True)
            # 分类器输出分类结果
            logits = classifier(x_in, t)
            log_probs = F.log_softmax(logits, dim=-1)
            selected = log_probs[range(len(logits)), y.view(-1)]
            # 计算梯度并返回
            return th.autograd.grad(selected.sum(), x_in)[0] * args.classifier_scale

效果

图24 对比实验结果

论文在多个数据集上对ADM(仅做网络结构优化)、ADM-G(同时引入Classifier Guidance机制)和其他模型进行了对比实验,从FID指标上,ADM、特别是ADM-G超过了GAN,实现了SOTA。

Classifier-Free Guidance

在Classifier Guidance机制被提出后,紧接着Google于2021年发表了论文《Classifier-Free Diffusion Guidance》这篇论文指出Classifier Guidance仍存在以下不足:一是Classifier Guidance需要额外训练分类器,二是Classifier Guidance会导致基于梯度的对抗攻击,欺骗FID、IS这类基于分类器的评估指标。因此,这篇论文提出了一种不需要训练分类器、但仍可以基于类别信息指导反向扩散过程的机制——Classifier-Free Guidance

之前ADM等扩散模型可使用类别信息进行有条件的图片生成,由模型基于xtx_t和类别yy预测噪声ϵθ(xt,y)\epsilon_\theta(x_t,y),或是不使用类别信息进行无条件的图片生成,由模型仅基于xtx_t预测噪声ϵθ(xt)\epsilon_\theta(x_t),但这两种情况需要分别训练模型,而Classifier-Free Guidance的思想是在模型训练时,按一定比例丢弃类别信息,使得模型能够同时学习有条件的图片生成和无条件的图片生成,这样在采样生成图片时,由同一个模型预测ϵθ(xt,y)\epsilon_\theta(x_t,y)ϵθ(xt)\epsilon_\theta(x_t),并使用两者的差值等价替换分类器输出的梯度对ϵθ(xt,y)\epsilon_\theta(x_t,y)进行调整,调整公式如下:

ϵ^θ(xt,y)=ϵθ(xt)+s(ϵθ(xt,y)ϵθ(xt))\hat{\epsilon}_\theta(x_t,y)=\epsilon_\theta(x_t)+s\cdot(\epsilon_\theta(x_t,y)-\epsilon_\theta(x_t))

再基于调整后的ϵ^θ(xt,y)\hat{\epsilon}_\theta(x_t,y)计算均值,并从高斯分布中采样得到xt1x_{t-1}

CLIP Guidance

Classifier Guidance使用类别信息指导反向扩散过程,那是否可以使用除类别外的其他信息指导反向扩散过程?2022年发表的论文《More Control for Free! Image Synthesis with Semantic Diffusion Guidance》尝试使用了其他信息,其中包括在多模态领域应用比较广泛的CLIP模型

CLIP模型包括两部分,图片编码器f(x)f(x)和文本编码器g(c)g(c),其中xx为图片,cc为文本。训练阶段,采用对比学习,使得正确图片、文本对(x,c)(x,c)的点积f(x)g(c)f(x)\cdot g(c)尽可能大,错误图片、文本对的点积尽可能小。因此,在推理阶段,可以进行文本和图片相关性的比较。关于CLIP模型的详细介绍,可以阅读原论文《Learning Transferable Visual Models From Natural Language Supervision》《AIGC系列-CLIP论文阅读笔记》

在Classifier Guidance中可以使用CLIP模型替换分类器,对于xtx_t,使用f(xt)g(c)f(x_t)\cdot g(c)的梯度调整μθ(xtc)\mu_\theta(x_t|c),公式如下所示:

μ^θ(xtc)=μθ(xtc)+sΣθ(xtc)xt(f(xt)g(c))\hat{\mu}_\theta(x_t|c)=\mu_\theta(x_t|c)+s\cdot\Sigma_\theta(x_t|c)\nabla_{x_t}(f(x_t)\cdot g(c))

和Classifier Guidance中的分类器类似,需使用带噪声的图片和文本对训练CLIP模型以获得正确的梯度。

GLIDE

在上述工作的基础上,OpenAI于2022年发表了论文《GLIDE: Towards Photorealistic Image Generation and Editing with Text-Guided Diffusion Models》其中发布了GLIDE(Guided Language to Image Diffusion for Generation and Editing)模型,用于基于文本的图片生成。图25是使用GLIDE模型基于文本生成的图片。

图25 使用GLIDE模型基于文本生成的图片

基于文本的图片生成

一般的扩散模型从随机采样的高斯噪声开始,不能生成特定的图片,而GLIDE在已有扩散模型的基础上,使用文本信息指导扩散过程,对于带噪声的图片xtx_t和文本cc,能够通过模型预测pθ(xt1xt,c)p_\theta(x_{t-1}|x_t,c),从而逐步降噪,实现了基于文本的图片生成

具体实现上,GLIDE基于ADM模型,但模型参数和训练数据规模更大,模型参数达到35亿。GLIDE先将文本cc转化为长度为KK的token序列,再通过Transformer输出文本的Embedding向量,最后使用文本Embedding向量替换原ADM模型输入中的类别Embedding向量。另外,文本Embedding向量还会经过投影与注意力层中的KKVV拼接在一起,通过注意力机制指导扩散过程。

GLIDE的代码开源,代码地址是:github.com/openai/guid…。通过Transformer输出文本的Embedding向量作为模型输入的代码在glide_text2im/text2im_model.py中,如下所示:

    def forward(self, x, timesteps, tokens=None, mask=None):
        hs = []
        emb = self.time_embed(timestep_embedding(timesteps, self.model_channels))
        if self.xf_width:
            text_outputs = self.get_text_emb(tokens, mask)
            xf_proj, xf_out = text_outputs["xf_proj"], text_outputs["xf_out"]
            emb = emb + xf_proj.to(emb)
        else:
            xf_out = None
        h = x.type(self.dtype)
        for module in self.input_blocks:
            h = module(h, emb, xf_out)
            hs.append(h)
        h = self.middle_block(h, emb, xf_out)
        for module in self.output_blocks:
            h = th.cat([h, hs.pop()], dim=1)
            h = module(h, emb, xf_out)
        h = h.type(x.dtype)
        h = self.out(h)
        return h

注意力层拼接文本Embedding向量的代码在glide_text2im/unet.py中,如下所示:

class QKVAttention(nn.Module):
    """
    A module which performs QKV attention. Matches legacy QKVAttention + input/ouput heads shaping
    """

    def __init__(self, n_heads):
        super().__init__()
        self.n_heads = n_heads

    def forward(self, qkv, encoder_kv=None):
        """
        Apply QKV attention.

        :param qkv: an [N x (H * 3 * C) x T] tensor of Qs, Ks, and Vs.
        :return: an [N x (H * C) x T] tensor after attention.
        """
        bs, width, length = qkv.shape
        assert width % (3 * self.n_heads) == 0
        ch = width // (3 * self.n_heads)
        q, k, v = qkv.reshape(bs * self.n_heads, ch * 3, length).split(ch, dim=1)
        if encoder_kv is not None:
            assert encoder_kv.shape[1] == self.n_heads * ch * 2
            ek, ev = encoder_kv.reshape(bs * self.n_heads, ch * 2, -1).split(ch, dim=1)
            # 拼接文本Embedding向量到k
            k = th.cat([ek, k], dim=-1)
            # 拼接文本Embedding向量到v
            v = th.cat([ev, v], dim=-1)
        scale = 1 / math.sqrt(math.sqrt(ch))
        weight = th.einsum(
            "bct,bcs->bts", q * scale, k * scale
        )  # More stable with f16 than dividing afterwards
        weight = th.softmax(weight.float(), dim=-1).type(weight.dtype)
        a = th.einsum("bts,bcs->bct", weight, v)
        return a.reshape(bs, -1, length)

Classifier-Free Guidance

GLIDE也使用了Classifier-Free Guidance机制,只是将类别替换为文本,因此,对模型所预测噪声ϵθ(xt,y)\epsilon_\theta(x_t,y)进行调整的公式如下:

ϵ^θ(xt,c)=ϵθ(xt)+s(ϵθ(xt,c)ϵθ(xt))\hat{\epsilon}_\theta(x_t,c)=\epsilon_\theta(x_t)+s\cdot(\epsilon_\theta(x_t,c)-\epsilon_\theta(x_t))

GLIDE在上一节已训练得到基于文本的模型、可预测ϵθ(xt,c)\epsilon_\theta(x_t,c)的基础上,对模型进行微调,将20%的文本Token序列替换成空序列,从而使得模型在具备预测ϵθ(xt,c)\epsilon_\theta(x_t,c)的基础上,能够进一步预测ϵθ(xt)\epsilon_\theta(x_t),从而在采样时,能够基于调整后的噪声ϵ^θ(xt,c)\hat{\epsilon}_\theta(x_t,c)降噪生成图片。

CLIP Guidance

GLIDE也使用了CLIP Guidance机制,但从实验结果上,其效果不如Classifier-Free Guidance。

Stable Diffusion

Stable Diffusion是由Stability AI发布的开源文生图、图生图工具,其GitHub地址为:github.com/Stability-A…。Stable Diffusion基于LDM(Latent Diffusion Model),介绍LDM的论文《High-Resolution Image Synthesis with Latent Diffusion Models》于2021年发表。LDM相较于DDPM,在隐向量空间进行反向扩散过程,而不是在原始图片上,从而能够减少训练资源、提升推理速度《The Illustrated Stable Diffusion》是一篇介绍Stable Diffusion的经典博客,以下内容是对这篇博客的翻译和精简。

整体架构

图26 Stable Diffusion整体架构

Stable Diffusion的整体架构如图26所示,其先在隐向量空间中进行反向扩散过程得到图片的Embedding向量,再通过图片解码器得到最终的图片,该流程包含以下核心组件:

  • Text Encoder,将文本转化为Embedding向量,其中包含文本的语义信息;
  • Image Information Creator,Stable Diffusion的核心,使用扩散模型,基于文本Embedding向量通过逐步的反向扩散过程,生成相应的图片Embedding向量,反向扩散步数默认50到100;反向扩散过程在图片的隐向量空间进行,并不是在图片的原始像素空间,这样相比于原先的扩散模型可以减少计算量;
  • Image Decoder,基于图片Embedding向量生成图片。

Stable Diffusion各组件的模型细节:

  • Text Encoder,使用CLIP的文本编码器,输入文本,输出77个Token的Embedding向量,每个Token的Embedding向量的维度为768;
  • Image Infomation Creator,使用扩散模型,网络结构为UNet,输入文本Embedding向量、步数Embedding向量、带噪声的图片Embedding向量(64×64、4通道的张量),输出图片Embedding向量(64×64、4通道的张量);
  • Image Decoder,使用自编码器模型中的解码器部分,输入图片Embedding向量(64×64、4通道的张量),输出生成的图片(512×512、3通道的张量)。

Image Infomation Creator在反向扩散过程中逐步去噪的是图片Embedding向量(隐向量),如果将每步去噪后的图片Embedding向量直接输入Image Decoder并输出相应的图片,可以看到初始的Embedding向量对应完全噪声,中间的Embedding向量对应的图片逐渐清晰,直至得到最终生成的图片,如图27所示。

图27 Stable Diffusion反向扩散过程示例

模型训练

Stable Diffusion基于LDM(Latent Diffusion Model),LDM论文中的模型结构如图28所示,和上一节的Stable Diffusion整体架构图基本一致,其中xx表示原始图片,x~\tilde{x}表示生成的图片,E\mathcal{E}D\mathcal{D}分别表示自编码器的编码器和解码器(已预训练),通过E\mathcal{E}得到xx的Embedding向量zz,基于zz进行正向扩散过程,逐步增加噪声得到z1,...,zTz_1,...,z_T,将z1,...,zTz_1,...,z_T作为训练样本,进行LDM模型的训练,LDM模型的输出仍是预测噪声ϵθ(zt,t)\epsilon_\theta(z_t,t),损失函数为:

LLDM:=EE(x),ϵN(0,1),t[ϵϵθ(zt,t)22]L_{\text{LDM}}:=\mathbb{E}_{\mathcal{E}(x),\epsilon\sim\mathcal{N}(0,1),t}\left[\parallel\epsilon-\epsilon_\theta(z_t,t) \parallel_2^2\right]

扩散模型训练过程可以参见DDPM算法中的介绍。模型训练完成、进行采样生成时,基于完全噪声进行反向扩散,通过模型逐步预测噪声并降噪,得到最终的图片Embedding向量z~\tilde{z},再通过D\mathcal{D}得到x~\tilde{x}

图28 LDM模型结构

基于已预训练的自编码器,使用编码器生成带噪声的图片Embedding向量用于扩散模型训练,并使用解码器基于扩散模型采样、生成的图片Embedding向量解码得到图片的流程如图29所示。

图29 基于已预训练自编码器的扩散模型训练和采样流程

基于文本生成

Stable Diffusiion在图片生成过程中引入文本信息作控制,先使用CLIP的文本编码器将文本转化为Embedding向量,而扩散模型所使用网络结构UNet由多层ResNet组成,在每层ResNet输出的后面,增加一个Attention层,将前序ResNet的输出,和文本Embedding向量作为该Attention层的输入,通过注意力机制基于文本调整前序ResNet的输出,再输入到下一个ResNet中,整体流程如图30所示。

图30 通过注意力机制实现基于文本生成

用公式描述该过程,Attention层可用以下公式表示:

Attention(Q,K,V)=softmax(QKTd)V\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{QK^T}{\sqrt{d}}\right)\cdot V
Q=WQ(i)ψi(zt),K=WK(i)τθ(y),V=WV(i)τθ(y)Q=W_Q^{(i)}\cdot\psi_i(z_t),K=W_K^{(i)}\cdot\tau_\theta(y),V=W_V^{(i)}\cdot\tau_\theta(y)

其中ψi(zt)RN×dϵi\psi_i(z_t)\in\mathbb{R}^{N\times d_\epsilon^i}表示前序ResNet的输出,τθ(y)RM×dτ\tau_\theta(y)\in\mathbb{R}^{M\times d_{\tau}}表示文本经过CLIP文本编码器后的Embedding向量,WV(i)Rd×dϵiW_V^{(i)}\in\mathbb{R}^{d\times d_\epsilon^i}WQ(i)Rd×dτW_Q^{(i)}\in\mathbb{R}^{d\times d_{\tau}}WK(i)Rd×dτW_K^{(i)}\in\mathbb{R}^{d\times d_{\tau}}表示Attention层中用于线性投影的参数矩阵,即将文本的Embedding向量作为注意力机制中的KKVV,将前序ResNet的输出作为注意力机制中的QQ。若进一步了解注意力机制,可以参见Transformer部分的介绍。

Diffusion Transformer(DiT)

2022年由William Peebles和谢赛宁发表的论文《Scalable Diffusion Models with Transformers》提出了Diffusion Transformer(可简称为DiT)。原先的扩散模型中,噪声预测的模型主要使用U-Net,而DiT的创新点是在LDM的基础上,使用Transformer替换U-Net,并论证了DiT在图片生成领域也满足Scaling Law。目前,William Peebles是OpenAI Sora项目的负责人之一,而Sora的技术报告指出DiT是其核心技术之一。

方法

DiT整体基于LDM,在隐空间中进行逆向扩散过程。在逆向扩散过程的每一步,将带噪声的隐空间图片和条件信息(当前扩散步数和图片类别)作为模型输入,由模型预测噪声和方差(DiT参考了Improved DDPM算法,既预测噪声,也预测方差),然后基于预测的噪声和方差对带噪声的隐空间图片进行采样、去噪。和LDM不同之处是,LDM中的预测模型采用U-Net,DiT中的预测模型采用Transformer。以下主要介绍DiT如何在隐空间中使用Transformer预测噪声和方差。

按块切分图片

和ViT类似,DiT也将图片按块进行切分。因为LDM是在图片编码压缩(如原始图片长、宽为256,通道数为3,编码压缩后的图片长、宽为32,通道数为4)后的隐空间(Latent Space)进行扩散过程,所以DiT实际是在编码压缩后的隐空间图片上按块进行切分,如图31所示,令编码压缩后的图片长和宽均为II、每个像素的通道数为CC,块的长和宽均为pp,则块序列的长度T=(I/p)2T=(I/p)^2。和ViT类似,每个块通过一个线性映射被转化为dd维向量,且每个块的dd维向量也会叠加其相应的位置编码。

图31 DiT对隐空间图片按块切分

编码器设计

隐空间图片被按块切分、转化为块向量序列后,DiT再将块向量序列和条件信息(包括扩散过程当前的步数tt、图片类别cc以及图片的自然语言描述)的Embedding向量输入NN层Transfomer的编码器,每层编码器为一个Block,如图32所示。

图32 DiT Block的设计

DiT的Block在原Transformer编码器结构的基础上进行升级,共设计了4种结构:

  • In-Context conditioning,即对原Transformer编码器结构不做改动,只是将扩散过程当前的步数tt、图片类别cc的Embedding向量追加到块向量序列的后面;
  • Cross-attention block,扩散过程当前的步数tt、图片类别cc的Embedding向量不再追加到块向量序列的后面,而是将两者拼接在一起后,作为一个独立的长度为2的条件信息向量序列,并在原Transformer编码器的多头自注意力层和多层感知机之间,增加一个多头交叉注意力层,多头自注意力层中的QQKKVV均取自块向量序列,而多头交叉注意力层和LDM中的使用方式类似,其中的QQ取自块向量序列,KKVV取自条件信息向量序列;增加多头交叉注意力层会增加约15%的计算量;
  • Adaptive layer norm(adaLN) block,在原Transformer编码器的层归一化后增加自适应层归一化(adaptive layer norm,adaLN);将扩散过程当前的步数tt、图片类别cc的Embedding向量拼接在一起,作为一个新的多层感知机的输入,由其预测原Transformer编码器中多头自注意力层和多层感知机的自适应层归一化的缩放和偏移参数γ1\gamma_1β1\beta_1γ2\gamma_2β2\beta_2
  • adaLN-Zero block,将自适应层归一化的参数初始化为0,即初始时残差连接模块等价为Identity函数;另外,在原Transformer编码器中多头自注意力层和多层感知机的输出后再增加一层缩放,并由预测自适应层归一化参数的多层感知机预测上述缩放参数α1\alpha_1α2\alpha_2

论文通过实验论证adaLN-Zero block在上述4种结构中的效果最好,因此,后续如无特别说明,DiT的Block均采用adaLN-Zero block。DiT的代码地址是:github.com/facebookres…,其中,adaLN-Zero block的代码如下所示:

class DiTBlock(nn.Module):
    """
    A DiT block with adaptive layer norm zero (adaLN-Zero) conditioning.
    """
    def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, **block_kwargs):
        super().__init__()
        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
        self.attn = Attention(hidden_size, num_heads=num_heads, qkv_bias=True, **block_kwargs)
        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
        mlp_hidden_dim = int(hidden_size * mlp_ratio)
        approx_gelu = lambda: nn.GELU(approximate="tanh")
        self.mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, act_layer=approx_gelu, drop=0)
        self.adaLN_modulation = nn.Sequential(
            nn.SiLU(),
            nn.Linear(hidden_size, 6 * hidden_size, bias=True)
        )

    def forward(self, x, c):
        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=1)
        x = x + gate_msa.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift_msa, scale_msa))
        x = x + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
        return x

解码器设计

最后一层Block的输出也是各个块的dd维向量,通过线性映射将各个块的dd维向量转化为p×p×2Cp\times p\times 2C维向量,即转化后向量的每个元素对应到块中每个像素、每个通道的噪声或方差,然后合并各个块通过解码器输出的向量,得到输入模型的带噪声的隐空间图片每个像素、每个通道的噪声和方差,然后基于Improved DDPM算法中的采样方法,对带噪声的隐空间图片进行降噪。

效果

DiT通过使用不同的超参配置,构建了四种规模的模型,分别是DiT-S、DiT-B、DiT-L、DiT-XL,如图33所示。四个模型的规模逐渐增长,Block的层数从12增大到28,Block中向量的维度从384增大到1152,模型计算量从1.4Gflops增大到29.1Gflops。

图33 不同规模的DiT模型

论文评估模型规模(S、B、L、XL)和块的大小(8、4、2)对图片生成效果的影响,如图34所示。图34中的7张图表,横坐标为训练步数,纵坐标为FID值(值越小,图片生成效果越好)。图34中第一行3张图表为块的大小分别是8、4、2时,各规模模型的FID值,说明块的大小固定时,模型规模越大,图片生成效果越好,第二行4张图表为模型规模分别是S、B、L、XL时,各块大小的FID值,说明模型的规模固定时,块的大小越小,图片生成效果越好。

图34 模型规模(S、B、L、XL)和块的大小(8、4、2)对图片生成效果的影响

图35 不同模型规模和块的大小下生成图片的对比

不同模型规模和块的大小下生成图片的对比如图35所示。对于图中红框部分,模型规模从左到右逐渐增大(S、B、L、XL),块的尺寸从上到下逐渐减少(8、4、2,即块向量序列从上到下逐渐增大),从中可以直观地感受到,从左上到右下,模型越大,块越小,生成的图片效果越好,但计算量也越大。

图36 ImageNet 256×256数据集上基于类别的图片生成各模型效果对比

图37 ImageNet 512×512数据集上基于类别的图片生成各模型效果对比

基于类别的图片生成各模型对比如图36和图37所示,DiT-XL实现了SOTA。

Sora

2024年2月15日,OpenAI发布了Sora的技术报告《Video Generation Models as World Simulators》,其中对Sora的技术原理做了简要的介绍。Sora是一个生成式模型,可以基于文本描述、图片和视频生成新的视频或图片,其底层技术类似DiT,在隐空间进行逆向扩散过程,基于条件信息使用Transformer对带噪声的隐空间视频进行噪声预测,并去噪,最后通过解码器生成像素空间的清晰视频

方法

按块切分视频

图38 按块切分视频

视频压缩网络

在LDM、DiT中,通过编码器将像素空间的图片编码压缩至隐空间,在隐空间进行逆向扩散过程,生成隐空间图片,再通过解码器将隐空间图片解码还原成像素空间的图片。

视频可以认为是由多帧图片组成,因此和LDM、DiT类似,Sora也使用编码器将像素空间的视频编码压缩至隐空间,压缩至隐空间的视频同样具有原视频空间和时间的信息,对于视频的每帧图片,类似LDM、DiT,将其压缩至隐空间图片,构成隐空间的空间信息,而隐空间的各帧图片序列,构成隐空间的时间信息。

同样和LDM、DiT类似,Sora也使用解码器将生成的隐空间视频解码还原至像素空间。

按块切分视频

类似DiT对隐空间图片按块切分,将隐空间图片转化为块(Patch)向量序列,隐空间视频也是由多个隐空间图片构成,因此也可以将隐空间视频的各帧隐空间图片按块切分,得到各隐空间图片的块向量序列,再合并各隐空间图片的块向量序列,得到隐空间视频的块向量序列,该块向量序列同样具有原视频空间和时间的信息。

类似NLP中,通过将句子切分为词元向量序列后,可以在Transformer中处理最大词元向量序列长度内任意长度的句子,Sora中通过将隐空间视频切分为块向量序列后,也可以处理并生成任意分辨率和播放时长的的视频。

视频生成

Sora本质是一个扩散模型,报告中并没有详细介绍其细节,只是介绍其使用类似DiT的技术方案,并给出了如图39所示的示例。关于DiT的细节,可以参见上一章节的详细介绍。另外,和DiT类似,Sora也论证了其满足Scaling Law,即模型规模越大、计算量越大,其生成的视频效果越好。

图39 基于扩散模型的视频生成

整体架构

《复刻Sora有多难?一张图带你读懂Sora的技术路径》这篇文章中,通过图40描述Sora的整体架构,该架构与类似DiT,在LDM的基础上,将扩散过程中预测噪声的模型由U-Net替换为Transformer。

图40 Sora的整体架构

效果

可变的分辨率、播放时长和宽高比

得益于将视频按块切分为块向量序列,相比过去的方法只能生成固定分辨率和播放时长(且较短)的视频,Sora可以生成不同分辨率、不同播放时长和不同宽高比的视频。

图41 Sora可以生成1920x1080的宽屏视频或1080x1920的竖屏视频或介于两者之间的视频。

文生视频

为了训练Sora能够基于文本生成视频,需要大量带有文本描述的视频作为训练数据集。Sora采用类似OpenAI文生图模型——DALL-E3中的方法,先训练文本生成模型,为训练集中的视频生成文本描述。另外,Sora在文生视频时,先将用户输入的简短提示输入ChatGPT,由ChatGPT生成详细的符合SORA规范的文本描述,再将详细的文本描述输入Sora,从而能够生成符合用户提示的高质量的视频。

以下是文生视频的一个例子的截图。用户提示为“a woman wearing blue jeans and a white-shirt taking a pleasant stroll in Johanesburg, South Africa during a beautiful sunset”(一位身穿蓝色牛仔裤和白色衬衫的女人在南非约翰内斯堡美丽的日落中愉快地散步),而所生成的视频中均准确体现了“日落”、“蓝色牛仔裤”、“白色衬衫”、“散步”等描述。

图42 文生视频示例

图生视频、视频生视频

除了基于文本生成视频,Sora还可以基于图片、视频生成视频,基于视频生成视频时,Sora可以:

  • 对已有视频按时序进行前向和后向扩展,并可生成一段无限循环的视频(如绕圈的骑行视频);
  • 基于文本描述对视频进行编辑(如将汽车行驶的环境由山林变成丛林);
  • 将两个不同的视频拼接在一起,两个视频之间进行平滑的过渡。

文生图

将图片看作只有一帧的视频,Sora也支持生成2048x2048内任意尺寸的图片,以下是文生图的一个例子的截图。用户提示为“Vibrant coral reef teeming with colorful fish and sea creatures”(充满多彩鱼类和海洋生物的珊瑚礁)。

图43 文生图示例

世界模拟能力

Sora还具备模拟真实物理世界的能力,包括:

  • 当镜头移动时,视频中的物体变化符合三维透视原理;
  • 保持视频的连贯性,即使某个物体在视频中被遮挡,该物体在遮挡前后能够保持连贯性;
  • 模拟动作对环境的影响,例如当一个人咬汉堡时,会在汉堡上留下咬痕。

Sora还能够模拟数字世界,例如,若文本描述中包含“Minecraft”(游戏我的世界),则Sora能够生成该游戏玩家视角的视频。

局限

Sora的技术报告也说明了其存在的一些局限,比如Sora并不能总是准确模拟真实物理世界,一个例子是玻璃杯掉落时,没有准确模拟玻璃的破裂。

图44 未准确模拟真实物理世界的示例

参考文献