基于Difussion图像、视频生成综述
2024年大年初七(02.16)OpenAI 发布视频生成模型在各大平台转疯了,和2022年发布ChatGPT3.5时一样的疯狂。在开工第一天,我就去官网上看了,遗憾的是,在这份技术报告中只披露了一些模型思想和效果展示,对于算法细节介绍的不多。一些大佬对此做了一些解读,提到最多的是DiT这篇论文。直接讲 DiT 没有相关的知识铺垫难免有点像论文解读,为此本篇博文从Diffusion说起,梳理近几年
2024年大年初七(02.16)OpenAI 发布视频生成模型 Sora 在各大平台转疯了,和2022年发布ChatGPT3.5时一样的疯狂。在开工第一天,我就去官网上看了 Sora 的技术报告,遗憾的是,在这份技术报告中只披露了一些模型思想和效果展示,对于算法细节介绍的不多。一些大佬对此做了一些解读,提到最多的是 DiT (Scalable Diffusion Models with Transformers)这篇论文。直接讲 DiT 没有相关的知识铺垫难免有点像论文解读,为此本篇博文从Diffusion说起,梳理近几年来图像、视频生成相关的进展。对于理解不准确或者有误的地方,还需要多多交流讨论。
一、开山鼻祖(Diffusion)
随着深度学习的发展,图像生成一直是AI相关领域的研究热点,例如 GAN、VAE 等模型。2020年Jonathan Ho等人提出 DDPM 算法(Denoising Diffusion Probabilistic Models,2020,NeurIPS),这是一种基于扩散模型的图像生成算法。所谓的扩散,可以简单的理解为物体从有序到无序,最后到稳定的过程,比如一滴墨水滴入到水中逐渐变成一个均匀分布的水溶液。
那么以此类比,一滴墨水扩散到均匀无序的状态就像是一个图像从清晰变成噪声的过程,就像下图中石原里美姐姐笑容逐渐消失的过程。
假设扩散过程就是马尔可夫过程(可以简单理解为: 未来的状态或分布只与当前的状态有关,与过去的状态无关),扩散过程的逆过程符合高斯分布。从石原里美小姐姐笑容消失这个例子中理解就是, 左起第三张图的状态只与第二张图有关,与第一张图有关。
在DDPM这篇论文中,将扩散过程分成两个部分,一个是前向过程,一个是反向过程。 前向过程就是图像加噪的过程(下图中从右往左), 反向过程就是去噪的过程(下图中从左到右)。
1.1 加噪过程
对于加噪过程的马尔可夫过程如公式(1)所示:
q
(
x
t
∣
x
t
−
1
)
=
N
(
x
t
;
1
−
β
t
x
t
−
1
,
β
t
I
)
(1-1)
q(x_t|x_{t-1})=N(x_t;\sqrt{1-β_t}x_{t-1},β_tI) \tag{1-1}
q(xt∣xt−1)=N(xt;1−βtxt−1,βtI)(1-1)
其中
1
−
β
t
\sqrt{1-β_t}
1−βt表示 t 时刻下的样本权重,
β
t
β_t
βt表示噪声的权重(正态分布)。【注意,这里的
β
t
\beta_t
βt是越来越大的,比如
1
0
−
4
10^{-4}
10−4到
1
0
−
2
10^{-2}
10−2,为什么呢?可以直观的理解就是清晰的图像加上一点噪声就有变化,当噪声很多时,需要加更多的噪声才能看到变化】用
z
z
z表示噪声,服从正态分布,将
α
t
=
1
−
β
t
α_t=1-β_t
αt=1−βt,
α
t
ˉ
=
∏
i
=
1
T
α
i
\bar{α_t}=\prod_{i=1}^Tα_i
αtˉ=∏i=1Tαi带入计算,可得:
x
t
=
α
t
x
t
−
1
+
1
−
α
t
z
t
−
1
=
α
t
α
t
−
1
x
t
−
2
+
1
−
α
t
α
t
−
1
z
ˉ
t
−
2
=
.
.
.
=
a
ˉ
t
x
0
+
1
−
α
ˉ
t
z
\begin{align} x_t &=\sqrt{α_t}x_{t-1}+\sqrt{1-α_t}z_{t-1} \tag{1-2} \\ &=\sqrt{α_tα_{t-1}}x_{t-2}+\sqrt{1-α_tα_{t-1}}\bar z_{t-2} \tag{1-3} \\ &=...=\sqrt{\bar a_t}x_0+\sqrt{1-\bar α_t}z \tag{1-4} \end{align}
xt=αtxt−1+1−αtzt−1=αtαt−1xt−2+1−αtαt−1zˉt−2=...=aˉtx0+1−αˉtz(1-2)(1-3)(1-4)
从这样的推理来看,
x
t
x_t
xt可以由
x
0
x_0
x0和噪声得到,即:
q
(
x
t
∣
x
0
)
=
N
(
x
t
;
a
ˉ
t
x
0
,
(
1
−
α
ˉ
t
)
I
)
(1-5)
q(x_t|x_0)=N(x_t;\sqrt{\bar a_t}x_0,(1-\bar α_t)I) \tag{1-5}
q(xt∣x0)=N(xt;aˉtx0,(1−αˉt)I)(1-5)
其中上式(3)中的
z
ˉ
t
−
2
\bar z_{t-2}
zˉt−2 是什么呢?其实是
z
t
−
1
,
z
t
−
2
z_{t-1},z_{t-2}
zt−1,zt−2的加权求和,两个独立的正态分布相加之后还是一个正态分布。公式里面为了简化表示成一个正态分布了。
那么这一大推公式推导的目的是什么? 加噪的过程可以看成是马尔可夫过程,当前状态都取决于上一步的状态,现在推导发现,从
x
0
x_0
x0可以直接计算出来
x
t
x_t
xt,这样就简单很多了呀!想象一下,石原姐姐的微笑,你一次就能让她消失,是不是会惊喜! 不用按照马尔可夫过程训练一个样本要迭代
T
T
T 次。
1.2 去噪过程
去噪过程也就是从噪声中恢复数据,这个过程要比加噪声复杂很多。在加噪声的时候,我们已知t-1时刻的图像,已知噪声分布,那么加上一个噪声的权重就可以得到加噪之后 t 时刻的图像; 反过来的话,我们只知道 t 时刻加噪之后的图像,不知道
t
−
1
t-1
t−1 时刻的原图和
t
−
1
t-1
t−1 时刻增加的噪声,这样就比较难搞了。 或许可以学习一个网络模型
f
f
f 用来预测增加的噪声,即
z
t
−
1
=
f
(
x
t
)
z_{t-1}=f(x_t)
zt−1=f(xt) 【默认
x
t
x_t
xt=
α
t
x
t
−
1
+
1
−
α
t
z
t
−
1
\sqrt{\alpha_t}x_{t-1} + \sqrt{1-\alpha_t}z_{t-1}
αtxt−1+1−αtzt−1】,但是这样训练过程会非常复杂。想象一下:对于一个样本,需要迭代 t 个时间步,每迭代一个时间步都需要把 t-1 时刻的图像、 t-1 时刻添加的噪声以及 t 时刻的结果记录下来,然后拿过来去训练模型预测噪声
。这样模型训练时需要先扩散到 t 时间步,然后训练一次 t-1 时间步,复杂度会非常高,如果 T 比较大更是增加了训练的难度。有人问了,上一节中推导可以直接从
x
0
x_0
x0 得到
x
t
x_t
xt, 那么有没有什么方法可以从
x
t
x_t
xt 得到
x
t
−
1
x_{t-1}
xt−1 呢? 也就是求解
x
t
−
1
=
h
(
x
t
,
x
0
)
x_{t-1}=h(x_t,x_0)
xt−1=h(xt,x0), 即求
q
(
x
t
−
1
∣
x
t
,
x
0
)
q(x_{t-1}|x_t,x_0)
q(xt−1∣xt,x0),其实是可以的。注意:这里把
x
0
x_0
x0 当做已知变量,为什么已知呢?因为它就是原图呀。
t-1时刻的数据可以表示为与上面的推导类似 [2],即:
q
(
x
t
−
1
∣
x
t
,
x
0
)
=
q
(
x
t
−
1
,
x
t
,
x
0
)
q
(
x
t
,
x
0
)
=
q
(
x
t
∣
x
t
−
1
,
x
0
)
q
(
x
t
−
1
,
x
0
)
q
(
x
t
,
x
0
)
=
q
(
x
t
∣
x
t
−
1
,
x
0
)
q
(
x
t
−
1
∣
x
0
)
q
(
x
t
∣
x
0
)
=
q
(
x
t
∣
x
t
−
1
)
q
(
x
t
−
1
∣
x
0
)
q
(
x
t
∣
x
0
)
\begin{align} q(x_{t-1}|x_t,x_0)&=\frac{q(x_{t-1},x_t,x_0)}{q(x_t,x_0)} \tag{1-6} \\ &=\frac{q(x_t|x_{t-1},x_0)q(x_{t-1},x_0)}{q(x_t,x_0)} \tag{1-7} \\ &=\frac{q(x_t|x_{t-1},x_0)q(x_{t-1}|x_0)}{q(x_t|x_0)} \tag{1-8} \\ &=\frac{q(x_t|x_{t-1})q(x_{t-1}|x_0)}{q(x_t|x_0)} \tag{1-9} \end{align}
q(xt−1∣xt,x0)=q(xt,x0)q(xt−1,xt,x0)=q(xt,x0)q(xt∣xt−1,x0)q(xt−1,x0)=q(xt∣x0)q(xt∣xt−1,x0)q(xt−1∣x0)=q(xt∣x0)q(xt∣xt−1)q(xt−1∣x0)(1-6)(1-7)(1-8)(1-9)
其中
q
(
x
t
−
1
∣
x
0
)
q(x_{t-1}|x_0)
q(xt−1∣x0)就是公式(1)的扩散过程,
q
(
x
t
∣
x
0
)
q(x_t|x_0)
q(xt∣x0)就是公式(5), 带入一通操作(我也没有算明白,哈哈哈 ),对于 t 时刻:
q
(
x
t
−
1
∣
x
t
,
x
0
)
=
N
(
x
t
−
1
;
μ
~
(
x
t
,
x
0
)
,
β
~
t
I
)
(1-10)
q(x_{t-1}|x_t, x_0)=N(x_{t-1};\tilde \mu(x_t,x_0),\tilde \beta_{t}I)\tag{1-10}
q(xt−1∣xt,x0)=N(xt−1;μ~(xt,x0),β~tI)(1-10)
其中:
β
~
t
=
1
−
α
ˉ
t
−
1
1
−
α
ˉ
t
β
t
,
μ
~
t
(
x
t
,
x
0
)
=
α
ˉ
t
−
1
β
t
1
−
α
ˉ
t
x
0
+
α
t
(
1
−
α
ˉ
t
−
1
)
1
−
α
ˉ
t
x
t
(1-11)
\tilde \beta_t=\frac{1-\bar \alpha_{t-1}}{1-\bar \alpha_t} \beta_t ,\tilde \mu_t(x_t,x_0)=\frac{\sqrt{\bar \alpha_{t-1}\beta_t}}{1-\bar\alpha_t}x_0+\frac{\sqrt{\alpha_t}(1-\bar\alpha_{t-1})}{1-\bar\alpha_t}x_t\tag{1-11}
β~t=1−αˉt1−αˉt−1βt,μ~t(xt,x0)=1−αˉtαˉt−1βtx0+1−αˉtαt(1−αˉt−1)xt(1-11)
从公式(4)可以得到:
x
0
=
1
α
ˉ
t
(
x
t
−
1
−
α
ˉ
t
z
)
(1-12)
x_0=\frac{1}{\sqrt{\bar\alpha_t}}(x_t-\sqrt{1-\bar\alpha_t}z)\tag{1-12}
x0=αˉt1(xt−1−αˉtz)(1-12)
μ
~
t
(
x
t
,
t
)
=
1
α
t
(
x
t
−
β
t
1
−
α
ˉ
t
z
ˉ
t
)
(1-13)
\tilde\mu_t(x_t,t)=\frac{1}{\sqrt{\alpha_t}}(x_t-\frac{\beta_t}{\sqrt{1-\bar\alpha_t}}\bar z_t)\tag{1-13}
μ~t(xt,t)=αt1(xt−1−αˉtβtzˉt)(1-13)
到了这里局势明朗了,
q
(
x
t
−
1
∣
x
t
,
x
0
)
q(x_{t-1}|x_t,x_0)
q(xt−1∣xt,x0)可以直接由
x
t
x_t
xt 和 后面的噪声
z
z
z 决定了。那么反向扩散的推理就有了:
x
t
−
1
=
u
~
t
+
β
~
t
z
t
=
1
α
t
(
x
t
−
1
−
α
t
1
−
α
ˉ
t
z
ˉ
t
)
+
1
−
α
ˉ
t
−
1
1
−
α
ˉ
t
β
t
z
t
(
x
t
,
t
)
\begin{align} x_{t-1}&=\tilde u_t+\tilde \beta_tz_t \tag{1-14} \\ &=\frac{1}{\sqrt{\alpha_t}}(x_t-\frac{1-\alpha_t}{\sqrt{1-\bar\alpha_t}}\bar z_t) + \frac{1-\bar \alpha_{t-1}}{1-\bar \alpha_t}\beta_tz_t(x_t,t) \tag{1-15} \end{align}
xt−1=u~t+β~tzt=αt1(xt−1−αˉt1−αtzˉt)+1−αˉt1−αˉt−1βtzt(xt,t)(1-14)(1-15)
好了,说到这里优雅的理论推理就结束了,从这里可以看到当有了
x
t
x_t
xt通过
z
ˉ
t
\bar z_t
zˉt就能得到
x
t
−
1
x_{t-1}
xt−1了。那么怎么去计算或者估计
z
ˉ
t
\bar z_t
zˉt呢,直接用神经网络来预测就好(既然算不出来,就学出来,嗯哼!!)。【换个直观的理解就是,重建过程中,前面的均值可以认为是确定的方向,后面加了一个扰动,产生多样性。】
1.3 模型的训练与推理
对于模型的训练,简单来说就是对图像进行加噪声,然后通过 UNet去预测的噪声。由公式(4-5)中推理证明了,时间 t 时刻的样本可以通过一次加噪实现,然后模型去预测这个噪声(如下图左图)。采样过程就是输入一个正态分布的噪声,然后迭代
T
T
T次,逐步预测
x
t
,
x
t
−
1
,
x
t
−
2
,
.
.
.
,
x
1
,
x
0
x_t,x_{t-1},x_{t-2},...,x_1,x_0
xt,xt−1,xt−2,...,x1,x0 (如下图右图)。
1.3.1 预测噪声?预测 x 0 x_0 x0?
论文中提出的方法是预测噪声,那么有人问了,预测噪声这么抽象,能不能直接预测
x
0
x_0
x0 呢?其实也是可以的,论文的源码里面实现是预测的噪声(DDPM源码中通过噪声预测原图代码)。pytorch版本里面实现可以直接预测
x
0
x_0
x0 (链接),还有预测
v
v
v (链接)。
对于预测噪声或者
x
0
x_0
x0有什么区别呢?
(1)如果预测噪声,对于一个batch中采样不同的时间步,损失优化时会平等对待;
(2)如果预测
x
0
x_0
x0,由于时间步的关系,不同的时间步下噪声的权重不同,相当于时间步越大,权重越大。
1.3.2 图像过渡(插值)
实验部分提到此方法可以实现图像的插值,直接对图像进行线性叠加(两个图像以不同的权重叠加会有伪影),编码之后恢复的图像可以消除伪影。实验中发现,对眼镜、闭嘴、发型等可以进行变换,但是对于眼镜这种装饰是不能变换的(可能是为了说明眼镜这些穿戴设备本身就是一个整体?)。这种插值也说明了,diffusion 中编码的特征空间的连续的。
1.4 模型分析
从训练过来看,就是预测一个图像中的噪声,然后把噪声恢复出来(或者恢复出清晰度更高的图);采样的过程就是一步步从带有噪声的图像中恢复原图。那么能不能直接从噪声中恢复出原图呢(这里是不是类似于VAE中的Decoder)?理论上应该可行。猜测通过这样一步步的恢复,效果可能会比较好,另外也降低了模型学习的难度。
1.5 如何降低采样次数?(DDIM ,Denoising Diffusion Implicit Models)
上一节中(1.4)中提到「能不能直接从噪声中恢复出来图像呢?」其实这对模型来说很难,但是能不能减少采样次数呢?DDIM(Denoising Diffusion Implicit Models,2021,ICML)中提不改变 DDPM 的训练过程,仅需要改变采样过程就可以将采样次数降低20倍、50倍,对效果影响也不大。DDIM 不使用马尔可夫链对降噪过程进行推理,可以由
x
t
x_t
xt 得到前面间隔多步的
x
p
r
e
v
x_{prev}
xprev, 具体推导过程可以参考这个博客。那么对于 DDPM 中降噪过程的时间步由
T
=
1000
,
999
,
998
,
.
.
.
,
1
,
0
T={1000,999,998,...,1,0}
T=1000,999,998,...,1,0 变成了
T
′
=
1000
,
950
,
900
,
850
,
.
.
.
,
50
,
0
T'={1000,950,900,850,...,50,0}
T′=1000,950,900,850,...,50,0 大大降低了推理时间。
其实换一种思路来看,加噪过程可以从
x
0
x_0
x0 到
x
t
x_t
xt 一步完成,那么降噪过程应该也可以从
x
t
x_t
xt 到
x
0
x_0
x0。但是实验表明,采样部署太少的话效果还是比较差的,但是20、50步是完全够的。
二、崭露头角(LDM,DiT)
2.1 Latent Diffusion Models(LDM)
Diffusion 中提出使用UNet 预测噪声的方式来实现图像生成,在图像生成中需要迭代
T
T
T 个时间步。由于每次迭代都需要过一遍 UNet 网络,当
T
T
T比较大时,将会非常耗时;另外,Diffusion 无法控制生成的图像,只能随机生成。LDM 是一个两阶段的图像生成模型,第一个阶段先训练一个图像编解码器,第二阶段进行图像生成。(High-Resolution Image Synthesis with Latent Diffusion Models,2022,CVPR),主要从两个方面解决这些问题:
(1)通过编码器、解码器将图像降维到更低的尺度(下图中左边红色区域),然后再基于扩散模型进行训练(下图中中间绿色区域);
(2)增加条件控制模块(下图中右边模块),可以将图像、文本等特征注入到 UNet 中,指导图像生成。
2.1.1 图像编码-解码器
对于图像编码部分,将使用
E
\Epsilon
E作为编码器,
D
D
D 作为解码器,有:
x
~
=
D
(
z
)
=
D
(
E
(
x
)
)
(2-1)
\tilde x=D(z)=D(\Epsilon(x)) \tag{2-1}
x~=D(z)=D(E(x))(2-1)
其中
x
∈
R
H
∗
W
∗
3
,
z
∈
R
h
∗
w
∗
c
x∈R^{H*W*3},z∈R^{h*w*c}
x∈RH∗W∗3,z∈Rh∗w∗c, 下采样的倍数为 2 的次方。论文中提到两种视频编码模块,一种的VQ-reg,一种是 KL-reg,看代码本质都是采用对抗的方式进行训练的。下图是 LDM 中使用的Encoder-Decoder的结构(注意为什么没有残差连接呢?因为要做特征编码,如果有残差连接岂不是无法编码了),代码在这里。(相关VQ-VAE、VQ-GAN等内容讲解参考,感兴趣可以移步了解)。
对于VQ-reg, 类似于一个VQ-GAN的损失如下,其中GEN 和 DIS 的损失是间隔一个 iter 进行优化的。
L
V
Q
−
G
E
N
=
L
p
i
x
e
l
+
L
p
e
r
c
e
p
t
u
a
l
+
L
g
e
n
+
L
c
o
d
e
b
o
o
k
L
V
Q
−
D
I
S
=
L
d
i
s
\begin{align} L_{VQ-GEN}&=L_{pixel}+L_{perceptual}+L_{gen}+L_{codebook}\tag{2-2} \\ L_{VQ-DIS}&=L_{dis}\tag{2-3} \end{align}
LVQ−GENLVQ−DIS=Lpixel+Lperceptual+Lgen+Lcodebook=Ldis(2-2)(2-3)
对于KL-reg, 类似 VAE,增加了 KL 损失,总得损失如下:
L
A
E
K
L
−
G
E
N
=
L
p
i
x
e
l
+
L
p
e
r
c
e
p
t
u
a
l
+
L
g
e
n
+
L
k
l
L
A
E
K
L
−
D
I
S
=
L
d
i
s
\begin{align} L_{AEKL-GEN}&=L_{pixel}+L_{perceptual}+L_{gen}+L_{kl}\tag{2-4} \\ L_{AEKL-DIS}&=L_{dis}\tag{2-5} \end{align}
LAEKL−GENLAEKL−DIS=Lpixel+Lperceptual+Lgen+Lkl=Ldis(2-4)(2-5)
2.1.2 LDM (Latent Diffusion Models)
LDM中直接使用 DDPM 的模型,代码中的 LDM 类直接继承 pytorch 版本的 DDPM,里面给出了预测噪声和
x
0
x_0
x0两个版本,默认预测噪声。UNet 的输入为 Encoder 对图像的编码,输出为噪声。由于经过 Encoder 编码之后,输入到Unet 的信息降低很多,论文中对比了
1
,
1
/
2
,
1
/
4
,
1
/
8
,
1
/
16
,
1
/
32
1,1/2,1/4,1/8,1/16,1/32
1,1/2,1/4,1/8,1/16,1/32这几种降采样的参数。当下采样倍数较小时(1、2)模型收敛比较慢,需要更长的训练时间;当下采样倍数较高时(32),由于损失了较多的洗洗脑,模型训练困难;4、8 能达到效果和速度的平衡。损失函数为:
L
L
D
M
=
L
p
i
x
e
l
+
L
v
l
b
(2-6)
L_{LDM}=L_{pixel}+L_{vlb}\tag{2-6}
LLDM=Lpixel+Lvlb(2-6)
其中 pixel 损失为 L1 或者 L2 损失,vlb(变分下界) 损失为时间步对应的权重乘以 pixel 损失。
2.1.3 条件控制机制
DDPM、LDM 其实只能随机的生成图像,无法进行内容控制。因此作者提出在图像生成阶段加上控制信息,如图2-1中右边模块。比如,对于文本类型的数据,先经过 BERTEmbedder 获得文本向量,再通过 Attention 融合到 UNet 中去。计算方式如下:
A
t
t
e
n
t
i
o
n
(
Q
,
K
,
V
)
=
s
o
f
t
m
a
x
(
Q
K
T
d
)
V
(2-7)
Attention(Q,K,V)=softmax(\frac {QK^T}{\sqrt d})V \tag{2-7}
Attention(Q,K,V)=softmax(dQKT)V(2-7)
Q
=
W
q
ϕ
(
z
t
)
,
K
=
W
k
τ
(
y
)
,
V
=
W
v
τ
(
y
)
(2-8)
Q=W_q \phi(z_t),K=W_k\tau(y),V=W_v\tau(y)\tag{2-8}
Q=Wqϕ(zt),K=Wkτ(y),V=Wvτ(y)(2-8)
其中
ϕ
\phi
ϕ 表示某一层Unet网络中的输出,
τ
\tau
τ 表示文本处理模型(比如BERT),
y
y
y 表示文本 prompt,这一部分的代码在这里。
对于语义合成(semantic synthesis)、指定类别图像生成等,使用不同的 embeding 模块,关于图像相关的控制都是通过 concat 来实现的。关于不同条件的控制机制可以看代码
(图2-1中的 switch 机制),可以使用单一的图像控制、单一的文本控制,也可以两者混合。
2.2 Scalable Diffusion Models with Transformers (DiT)
DiT 是 Sora 出来之后被多次提起的论文之一,论文来源于 Scalable Diffusion Models with Transformers,2023,ICCV,代码在这里。在基于 diffusion 的模型中,都是采用 UNet 的结构用于噪声或者
x
0
x_0
x0 的生成,DiT 在 LDM 的基础上,提出使用 Transformer 来代替 UNet 结构。DiT 的模型结构如下:
2.2.1 DiT Block
如图2-3 所示,论文中提出了三种类型的 DiT Block,三种类型分别是 In-Context Condition、Cross-Attention、AdaLN-zero(AdaLN).
- In-Context. 将时间步 t t t 和条件控制 c c c 作为两个额外的 token 拼接到输入 Atten 模块中,类似于 ViT 中的 c l s cls cls token,实现方式与 ViT block 一样。在最后一个 block 中把这个 token 移除掉。这样做相比之前的模型几乎不增加额外的计算量。
- Cross-Attention. 将时间步和条件控制独立输入到 Multi-Head Cross-Attention 中,并接在 Multi-Head self-Attention 后面。这样做相比之前增加了约 15% 的计算量。
- AdaLN(Adaptive Layer Norm ). 自适应归一化层(adaptive normalization layers) 在GAN,图像风格迁移中广泛应用,其思想在于将一种特征风格迁移到另一种特征风格。Transformer 常用 Layer Norm,GAN 常用 Instance Norm,图像风格迁移常用 Adaptive Instance Norm。这里将从时间步和条件控制中学习缩放、位移、权重项,然后经过 Attention 模块融入到特征中。这里可以理解为,输入的时间步、条件控制经过 AdaLN 之后,作为 Attention 模块中的 key、value,原始特征作为 query。
- AdaLN-zero. 零初始化可以相当于进行恒等映射,这里将 AdaLN 进行零初始化可以加速训练。这一块的代码在这里。
2.2.2 Patch 模块
在 UNet 中采用卷积的方式处理图像尺度的噪声,在这里将图像处理成 patch 然后展开成 token 输入到 transformer 模块中,如图2-4所示。另外,将位置编码相加到 patch 之后的特征上。
三、百家争鸣(文生图)
3.1 Imagen
2022年Google在NeurIPS发表 Photorealistic Text-to-Image Diffusion Models
with Deep Language Understanding,2022, NeurIPS 论文,也就是 Imagen,项目主页,代码(非官方实现,官方未开源)。论文的创新点主要有以下几点:
- 文本预训练大模型输出的特征用于 text2img 效果要更好,文本模型编码器的大小相比图像扩散模型的大小,对结果的影响更大。
- 噪声生成中,动态采样相比静态采样效果更好(静态采样时[-1,1]的区间可能会导致边界溢出,造成生成图像受限)。
- 轻量化的 UNet模型。
- 在 COCO 数据集上达到 SOTA、提出新的基准(DrawBench)。
图3-1 Imagen 中 网络结构和 Efficient UNet (图源论文补充材料)
3.1.1 文本编码器
实验对比了不同大小的文本编码器和 UNet 对实验结果的影响。如下图(a)(b)所示,实验发现文本编码器的容量大小对于实验结果的影响比较大,UNet 的大小相对对于实验结果影响较小,也就是文本的特征其实对于图像生成来说更加重要(毕竟是指导图像生成的信息来源)。
3.1.2 动态阈值
之前的 DDPM 模型默认都将输出的结果通过 clip 截断到 [-1, 1] 的区间,这样做的目的是将结果维持在一定的范围,便于模型学习。论文中提到增加文本的引导权重可以使得生成的图像和文本对齐的更好,但是这样操作会导致生成的图像饱和度较高(简单来说饱和度越高表示单一色彩的纯度越高,看起来越越鲜艳,色调更硬)。这里需要插播一下相关知识,在训练条件模型的时候,会让条件和无条件联合训练,如下面公式:
ϵ
~
θ
(
z
y
,
c
)
=
w
ϵ
θ
(
z
y
,
c
)
+
(
1
−
w
)
ϵ
(
z
t
)
(3-1)
\tilde \epsilon_{\theta}(z_y,c)=w\epsilon_{\theta}(z_y,c)+(1-w)\epsilon(z_t)\tag{3-1}
ϵ~θ(zy,c)=wϵθ(zy,c)+(1−w)ϵ(zt)(3-1)
其中
w
ϵ
θ
(
z
y
,
c
)
w\epsilon_{\theta}(z_y,c)
wϵθ(zy,c)就是条件噪声预测结果,
(
1
−
w
)
ϵ
(
z
t
)
(1-w)\epsilon(z_t)
(1−w)ϵ(zt)就是无条件预测。当
w
w
w为 1 时,后一项为0,表示禁用无分类引导。由于训练的时候结果都为 [-1, 1] 区间,但是测试的时候发现使得
w
>
1
w>1
w>1 可以增强生成图像和文本的一致性,但是这样就导致了结果区间不在 [-1, 1] 了,就出现了训练-测试不一致的问题。
下图所示为阈值对比实验,可以看到:
- 无阈值时,生成的图像色彩更单一(白、黑),随着引导权重的变大,图像饱和度增大;
- 增加了[-1,1]的静态阈值之后,生成的图像饱和度降低,色彩更丰富,但是增加了引导权重饱和度还是比较高(背景大面积黑色);
- 设置为动态阈值时,色彩更加丰富,内容更丰富,看起来更真实,增加引导权重饱和度仍然较低。
论文中动态阈值的设置根据当前特征的分布进行计算的,看下面的代码片段,是缩放值其实是取一定百分比的像素值(比如取 95% 作为阈值),然后再进行缩放,这样做可以比较好的适应当前色彩空间,让色差分布更加均匀。
p=0.95
s = tprch.quantile(jnp.abs(x0_t), p)
s = torch.max(s, 1.0)
x0_t = jnp.clip(x0_t, -s, s) / s
3.1.3 层级扩散模型(Cascaded Diffusion Model)
如图 3-1 所示,模型生成的图像为 64 像素,然后超分到 256,再超分到 1024,这种级联的超分策略可以有效的提高图像的清晰度和生成图像的质量。在训练时,往低分辨率的图像上面加上一些噪声;在采样时,按照时间步迭代采样,类似于 DDPM 的过程。伪代码如下图所示:
3.2 RePaint (图像重绘/图像补全)
RePaint 是用于图像重绘、补全的一个论文,来源于 RePaint: Inpainting using Denoising Diffusion Probabilistic Mode,2022,CVPR,代码这这里。改论文的核心点就是基于 DDPM 模型在不需训练的情况下进行图像补全。方法也比较简单,就是把
x
0
′
x_0'
x0′ (部分缺失的图像)加噪,然后乘以 mask,然后再采样噪声,补上缺少部分的噪声, 这样就得到了
x
t
x_t
xt,然后用
x
t
x_t
xt 生成图像
x
0
x_0
x0。然后问题来了,这样做之后生成的图像和其他区域容易有明显的差异,局部比较奇怪, 如下图中 n=1 所示。
从上图的看:
(1)上一半看,就是输入图像
x
t
x_t
xt,加上噪声,得到
x
t
−
1
k
n
o
w
n
x_{t-1}^{known}
xt−1known,乘以
m
a
s
k
mask
mask;
(2)下一半看,输入噪声去噪,得到
x
t
−
1
u
n
k
n
o
w
n
x_{t-1}^{unknown}
xt−1unknown, 乘以
1
−
m
a
s
k
1-mask
1−mask;
(3)最后,两个合成一起,得到完整的
x
t
−
1
x_{t-1}
xt−1。
乍一看下面一半输入中,怎么输入的是 x t x_t xt 加噪的感觉呢?其实不是这样的,在迭代第一次的时候,噪声是随机生成,到了第二次之后才会循环沿用 x t − 1 x_{t-1} xt−1。由于噪声生成是随机的,因此采样不同的噪声补全的效果也不同。
3.3 DALLE2
DALLE2 是 openai 在2022年4月挂在 arxiv 上的一篇论文 《Hierarchical Text-Conditional Image Generation with CLIP Latents》,代码在这里。整个算法的结构图如下:
- 自回归(Autoregressive,AR)Prior:将CLIP的图像编码器编码的embedding 转换为离散编码序列作为监督目标,并根据文本标题进行自回归预测。
- Diffusion Prior:直接使用CLIP的图像编码器编码的embedding 作为监督目标,并以文本为条件的 Diffusion 模型。
以10%的概率丢弃文本条件,以 cfg 的方式进行训练。为了更高效的训练 AR Prior 模型,用 PCA 对 1000 纬度的特征进行降维到 319 纬度。其他参数细节比较多,这里不再一一赘述,感兴趣可以看下论文和源码。
3.3.3 实验
实验做的非常充分,包括图像多样化、图像插值、文本差异,Prior 模型的重要性等。
3.4 SDXL
SDXL来源于 SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis 这篇论文,github,是一个两阶段的级联扩散模型,由 Base + Refiner 两个模型构成。可以移步大佬详细说明。
3.4.1 模型总览
SDXL模型总览如下,其由 SD Base 和 SD Refiner 构成,其中:
(1) Base 模型由VAE、UNet、Text Encoder (CLIP ViT-L + CLIP ViT-bigG) 三个模块组成;
(2)Refiner 模型由 VAE、UNet、Text Encoder(CLIP ViT-bigG) 三个模块构成,其中 VAE 和 Text Encoder 与 Base 模型共用;
3.4.1 VAE
VAE 使用的是 LDM 中相同的结构,只不过是使用了加大的 batch size 重新训练的(SDXL 256 vs LDM 9),论文中说可以改善生成图像的高频细节。同时,使用 ema 权重更新策略。
3.4.2 UNet
相比 LDM 中的UNet结构,由4阶段的改成 3 阶段,并且只进行了两次下采样(将 8 倍以下的下采样舍弃掉了)。在第一个阶段中不使用 Transformer 结构(resnet block + downsample),在第二和第三阶段使用(resnet block + transformer + downsample ),可以有效的降低显存和计算量。整个UNet的参数量相比之前达到了之前的 3 倍左右(2.6B)。
3.4.3 Text Encoder
SDXL的base 模型使用两个 Text Encoder (OpenCLIP ViT-bigG 和 CLIP ViT-L), 特征使用模型的最后一层,所以特征纬度为 1280(ViT-bigG) + 768 (ViT-L) = 2048。在 Imagen 中也介绍了增加文本编码的容量相比增加 UNet 的容量更重要。
3.4.4 Refiner
SDXL的Refiner模块采用四个阶段的UNet,并且在第一个阶段未使用 Transformer 结构。文本编码只用了 OpenCLIP ViT-bigG 的特征。因为 base 模型已经训练出来比较好的效果了,因此只在前面 200 个时间步上面进行训练。
3.4.5 训练策略
(1)图像分辨率编码
之前 SD 的模型训练是两阶段的,先训练 256 尺度的,然后再训练 512 尺度的高分辨的图像。这在第二阶段会导致很多低分辨率的图像没有用上。如果使用超分策略会导致生成的图像会有一些糊影。在 SDXL 中将输入图像的高宽度经过傅里叶特征编码,concat 在 time embeding 后面 作为额外的条件嵌入U-Net模型中。这样做使得模型在训练的时候将尺度作为一定的参考信息,可以提高训练质量,推理时可以适应不同尺度的要求。训练之后发现,生成的图像尺度小的较糊,尺度大的清晰度较高。如下图为不同分辨率生成效果展示。
(2)图像分组与裁剪编码
在 LDM 的模型训练时,为了进行多 batch 训练会统一输入尺度的大小,比如短边尺度对齐之后进行中心裁剪、随机裁剪、缩放等,这样做会导致图像中一些部分裁剪掉,不符合真实输入的场景(比如把边缘、头、脚等区域裁剪掉)。而生成的图像也会有部分区域丢失的情况。SDXL 为了解决这个问题,采用了两个方法。
- 基于尺度进行分组。把图像按照长宽进行分组,每个分组内为尺度接近的图像,训练时先选中某个分组,然后从这个分组内选取图像;
- 在进行裁剪时,将裁剪的左上角坐标进行傅里叶编码后,与原始图像尺寸一起作为监督条件加在Time Embedding上,这样训练之后模型对裁剪的尺度有了一定的感知。
下图为不同裁剪参数生成的效果,可见模型已经学习到了裁剪坐标的含义。
(3)多尺度训练
模型在 256、512 分辨率上进行训练之后,在使用 1024 的分辨率进行微调时,调整不同的长宽比例,分成多个尺度组,每个组内的尺度都接近 1024 x 1024 个像素数,长宽比例从0.25 到 4 范围区间。
(4)额外噪声
由于 LDM 这些模型训练之后生成的图像亮度比较统一,对于不同光照强度的图像生成较弱,这里在采样的高斯噪声中加入额外的噪声破坏(noise offset),增强不同亮度的生成能力。
未完待续,当前章节安排比较乱,后续待调整~
参考文献(相关论文的引用直接在正文中贴上了论文名称及链接,可以直接点击跳转):
[1] https://zhuanlan.zhihu.com/p/563543020 (前面讲的关于GAN、VAE等模型的对比讲的好)
[2] https://zhuanlan.zhihu.com/p/663880249 (反向过程中问题说明讲的好)
开放原子开发者工作坊旨在鼓励更多人参与开源活动,与志同道合的开发者们相互交流开发经验、分享开发心得、获取前沿技术趋势。工作坊有多种形式的开发者活动,如meetup、训练营等,主打技术交流,干货满满,真诚地邀请各位开发者共同参与!
更多推荐
所有评论(0)