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),这是一种基于扩散模型的图像生成算法。所谓的扩散,可以简单的理解为物体从有序到无序,最后到稳定的过程,比如一滴墨水滴入到水中逐渐变成一个均匀分布的水溶液。
那么以此类比,一滴墨水扩散到均匀无序的状态就像是一个图像从清晰变成噪声的过程,就像下图中石原里美姐姐笑容逐渐消失的过程。
假设扩散过程就是马尔可夫过程(可以简单理解为: 未来的状态或分布只与当前的状态有关,与过去的状态无关),扩散过程的逆过程符合高斯分布。从石原里美小姐姐笑容消失这个例子中理解就是, 左起第三张图的状态只与第二张图有关,与第一张图有关。
在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表示噪声的权重(正态分布)。用 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.4 模型分析
从训练过来看,就是预测一个图像中的噪声,然后把噪声恢复出来(或者恢复出清晰度更高的图);采样的过程就是一步步从带有噪声的图像中恢复原图。那么能不能直接从噪声中恢复出原图呢(这里是不是类似于VAE中的Decoder)?理论上应该可行。猜测通过这样一步步的恢复,效果可能会比较好,另外也降低了模型学习的难度。
二、崭露头角(LDM)
Diffusion 中提出使用UNet 预测噪声的方式来实现图像生成,在图像生成中需要迭代 T T T 个时间步。由于每次迭代都需要过一遍 UNet 网络,当 T T T比较大时,将会非常耗时;另外,Diffusion 无法控制生成的图像,只能随机生成。LDM (High-Resolution Image Synthesis with Latent Diffusion Models)从两个方面解决这些问题:
(1)通过编码器、解码器将图像降维到更低的尺度(下图中左边红色区域),然后再基于扩散模型进行训练(下图中中间绿色区域);
(2)增加条件控制模块(下图中右边模块),可以将图像、文本等特征注入到 UNet 中,指导图像生成。
2.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.2
未完待续~
参考文献(相关论文的引用直接在正文中贴上了论文名称及链接,可以直接点击跳转):
[1] https://zhuanlan.zhihu.com/p/563543020 (前面讲的关于GAN、VAE等模型的对比讲的好)
[2] https://zhuanlan.zhihu.com/p/663880249 (反向过程中问题说明讲的好)