Diffusion 模型基本原理
本文介绍 Diffusion 模型的基本原理,帮助你理解训练框架是如何构建的。为了让读者更轻松地理解这些复杂的数学理论,我们重构了 Diffusion 模型的理论框架,抛弃了复杂的随机微分方程,用一种更简洁易懂的形式进行介绍。
引言
Diffusion 模型通过多步迭代式地去噪(denoise)生成清晰的图像或视频内容,我们从一个数据样本 \(x_0\) 的生成过程开始讲起。直观地,在完整的一轮 denoise 过程中,我们从随机高斯噪声 \(x_T\) 开始,通过迭代依次得到 \(x_{T-1}\)、\(x_{T-2}\)、\(x_{T-3}\)、\(\cdots\),在每一步中逐渐减少噪声含量,最终得到不含噪声的数据样本 \(x_0\)。
这个过程是很直观的,但如果要理解其中的细节,我们就需要回答这几个问题:
每一步的噪声含量是如何定义的?
迭代去噪的计算是如何进行的?
如何训练这样的 Diffusion 模型?
现代 Diffusion 模型的架构是什么样的?
本项目如何封装和实现模型训练?
每一步的噪声含量是如何定义的?
在 Diffusion 模型的理论体系中,噪声的含量是由一系列参数 \(\sigma_T\)、\(\sigma_{T-1}\)、\(\sigma_{T-2}\)、\(\cdots\)、\(\sigma_0\) 决定的。其中
\(\sigma_T=1\),对应的 \(x_T\) 为纯粹的高斯噪声
\(\sigma_T>\sigma_{T-1}>\sigma_{T-2}>\cdots>x_0\),在迭代过程中噪声含量逐渐减小
\(\sigma_0=0\),对应的 \(x_0\) 为不含任何噪声的数据样本
至于中间 \(\sigma_{T-1}\)、\(\sigma_{T-2}\)、\(\cdots\)、\(\sigma_1\) 的数值,则不是固定的,满足递减的条件即可。
那么在中间的某一步,我们可以直接合成含噪声的数据样本 \(x_t=(1-\sigma_t)x_0+\sigma_t x_T\)。
迭代去噪的计算是如何进行的?
在理解迭代去噪的计算前,我们要先搞清楚,去噪模型的输入和输出是什么。我们把模型抽象成一个符号 \(\hat \epsilon\),它的输入通常包含三部分
时间步 \(t\),模型需要理解当前处于去噪过程的哪个阶段
含噪声的数据样本 \(x_t\),模型需要理解要对什么数据进行去噪
引导条件 \(c\),模型需要理解要通过去噪生成什么样的数据样本
其中,引导条件 \(c\) 是新引入的参数,它是由用户输入的,可以是用于描述图像内容的文本,也可以是用于勾勒图像结构的线稿图。
而模型的输出 \(\hat \epsilon(x_t,c,t)\),则近似地等于 \(x_T-x_0\),也就是整个扩散过程(去噪过程的反向过程)的方向。
接下来我们分析一步迭代中发生的计算,在时间步 \(t\),模型通过计算得到近似的 \(x_T-x_0\) 后,我们计算下一步的 \(x_{t-1}\):
完美!与时间步 \(t-1\) 时的噪声含量定义完美契合。
(这部分可能有点难懂,请不必担心,首次阅读本文时建议跳过这部分,不影响后文的阅读。)
完成了这段有点复杂的公式推导后,我们思考一个问题,为什么模型的输出要近似地等于 \(x_T-x_0\) 呢?可以设定成其他值吗?
实际上,Diffusion 模型依赖两个定义形成完备的理论。在以上的公式中,我们可以提炼出这两个定义,并导出迭代公式:
数据定义:\(x_t=(1-\sigma_t)x_0+\sigma_t x_T\)
模型定义:\(\hat \epsilon(x_t,c,t)=x_T-x_0\)
导出迭代公式:\(x_{t-1}=x_t + (\sigma_{t-1} - \sigma_t) \cdot \hat \epsilon(x_t,c,t)\)
这三个数学公式是完备的,例如在刚才的推导中,我们把数据定义和模型定义代入迭代公式,可以得到与数据定义吻合的 \(x_{t-1}\)。
这是基于 Flow Matching 理论构建的两个定义,但 Diffusion 模型也可用其他的两个定义来实现,例如早期基于 DDPM(Denoising Diffusion Probabilistic Models)的模型,其两个定义及导出的迭代公式为:
数据定义:\(x_t=\sqrt{\alpha_t}x_0+\sqrt{1-\alpha_t}x_T\)
模型定义:\(\hat \epsilon(x_t,c,t)=x_T\)
导出迭代公式:\(x_{t-1}=\sqrt{\alpha_{t-1}}\left(\frac{x_t-\sqrt{1-\alpha_t}\hat \epsilon(x_t,c,t)}{\sqrt{\sigma_t}}\right)+\sqrt{1-\alpha_{t-1}}\hat \epsilon(x_t,c,t)\)
更一般地,我们用矩阵描述迭代公式的导出过程,对于任意数据定义和模型定义,有:
数据定义:\(x_t=C_T(x_0,x_T)^T\)
模型定义:\(\hat \epsilon(x_t,c,t)=C_T^{[\epsilon]}(x_0,x_T)^T\)
导出迭代公式:\(x_{t-1}=C_{t-1}(C_t,C_t^{[\epsilon]})^{-T}(x_t,\hat \epsilon(x_t,c,t))^T\)
其中,\(C_t\)、\(C_t^{[\epsilon]}\) 是 \(1\times 2\) 的系数矩阵,不难发现,在构造两个定义时,需保证矩阵 \((C_t,C_t^{[\epsilon]})^T\) 是可逆的。
尽管 Flow Matching 与 DDPM 已被大量预训练模型广泛验证过,但这并不代表这是最优的方案,我们鼓励开发者设计新的 Diffusion 模型理论实现更好的训练效果。
如何训练这样的 Diffusion 模型?
搞清楚迭代去噪的过程之后,接下来我们考虑如何训练这样的 Diffusion 模型。
训练过程不同于生成过程,如果我们在训练过程中保留多步迭代,那么梯度需经过多步回传,带来的时间和空间复杂度是灾难性的。为了提高计算效率,我们在训练中随机选择某一时间步 \(t\) 进行训练。
以下是训练过程的伪代码
从数据集获取数据样本 \(x_0\) 和引导条件 \(c\)
随机采样时间步 \(t\in(0,T]\)
随机采样高斯噪声 \(x_T\in \mathcal N(O,I)\)
\(x_t=(1-\sigma_t)x_0+\sigma_t x_T\)
\(\hat \epsilon(x_t,c,t)\)
损失函数 \(\mathcal L=||\hat \epsilon(x_t,c,t)-(x_T-x_0)||_2^2\)
梯度回传并更新模型参数
现代 Diffusion 模型的架构是什么样的?
从理论到实践,还需要填充更多细节。现代 Diffusion 模型架构已经发展成熟,主流的架构沿用了 Latent Diffusion 所提出的“三段式”架构,包括数据编解码器、引导条件编码器、去噪模型三部分。
数据编解码器
在前文中,我们一直将 \(x_0\) 称为“数据样本”,而不是图像或视频,这是因为现代 Diffusion 模型通常不会直接在图像或视频上进行处理,而是用编码器(Encoder)-解码器(Decoder)架构的模型,通常是 VAE(Variational Auto-Encoders)模型,将图像或视频编码为 Embedding 张量,得到 \(x_0\)。
数据经过编码器编码后,再经过解码器解码,重建后的内容与原来近似地一致,会有少量误差。那么,为什么要在编码后的 Embedding 张量上处理,而不是在图像或视频上直接处理呢?主要原因有亮点:
编码的同时对数据进行了压缩,编码后处理的计算量更小。
编码后的数据分布与高斯分布更相似,更容易用去噪模型对数据进行建模。
在生成过程中,编码器部分不参与计算,迭代完成后,用解码器部分解码 \(x_0\) 即可得到清晰的图像或视频。在训练过程中,解码器部分不参与计算,仅编码器用于计算 \(x_0\)。
引导条件编码器
用户输入的引导条件 \(c\) 可能是复杂多样的,需要由专门的编码器模型将其处理成 Embedding 张量。按照引导条件的类型,我们把引导条件编码器分为以下几类:
文本类型,例如 CLIP、Qwen-VL
图像类型,例如 ControlNet、IP-Adapter
视频类型,例如 VAE
前文中的模型 \(\hat \epsilon\) 指代此处的所有引导条件编码器和去噪模型这一整体,我们把引导条件编码器单独拆分列出,因为这类模型在 Diffusion 训练中通常是冻结的,且输出值与时间步 \(t\) 无关,因此引导条件编码器的计算可以离线进行。
去噪模型
去噪模型是 Diffusion 模型真正的本体,其模型结构多种多样,例如 UNet、DiT,模型开发者可在此结构上自由发挥。
本项目如何封装和实现模型训练?
请阅读下一文档:标准监督训练