扩散模型（三）

Latent Variable Space

潜在扩散模型（LDM；龙巴赫、布拉特曼等人，2022 年）在潜在空间而非像素空间中运行扩散过程，这使得训练成本更低，推理速度更快。其背后的动机源于这样一个观察结果：图像的大部分信息仅关乎感知细节，而在经过大幅压缩后，语义和概念层面的构图依然保留。LDM 通过生成式建模学习，大致将感知压缩和语义压缩分离开来。首先，利用自动编码器去除像素级冗余，然后在学习到的潜在空间上，通过扩散过程操控 / 生成语义概念。

感知压缩过程依赖于一个自动编码器模型。编码器 $\mathcal{E}$ 用于将输入图像 $\mathbf{x} \in \mathbb{R}^{H \times W \times 3}$ 压缩为一个更小的二维潜在向量： $\mathbf{z} = \mathcal{E}(\mathbf{x}) \in \mathbb{R}^{h \times w \times c}$ ，其中，下采样率 $f=H/h=W/w=2^m, m \in \mathbb{N}$ 。然后，解码器 $\mathcal{D}$ 从潜在向量中重建图像，即 $\tilde{\mathbf{x}} = \mathcal{D}(\mathbf{z})$ 。该论文在自动编码器训练中探索了两种正则化方法，以避免潜在空间中出现过高的方差。

KL 正则化：对学习到的潜在变量施加一个朝向标准正态分布的小 KL 惩罚项，这与变分自编码器（VAE）类似。
VQ 正则化：在解码器中使用向量量化层，类似于向量量化变分自编码器（VQVAE），但量化层被解码器吸收。

扩散和去噪过程发生在潜在向量 $\mathbf{z}$ 上。去噪模型是一个时间条件 U 型网络，通过增加交叉注意力机制来处理图像生成中灵活的条件信息（例如类别标签、语义图、图像的模糊变体）。这种设计等同于通过交叉注意力机制将不同模态的表示融合到模型中。每种类型的条件信息都与一个特定领域的编码器 $\tau_\theta$ 配对，以将条件输入 $y$ 投影到一个中间表示，该中间表示可以映射到交叉注意力组件中，即 $\tau_\theta(y) \in \mathbb{R}^{M \times d_\tau}$ :

$\begin{aligned} &\text{Attention}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) = \text{softmax}\Big(\frac{\mathbf{Q}\mathbf{K}^\top}{\sqrt{d}}\Big) \cdot \mathbf{V} \\ &\text{where }\mathbf{Q} = \mathbf{W}^{(i)}_Q \cdot \varphi_i(\mathbf{z}_i),\; \mathbf{K} = \mathbf{W}^{(i)}_K \cdot \tau_\theta(y),\; \mathbf{V} = \mathbf{W}^{(i)}_V \cdot \tau_\theta(y) \\ &\text{and } \mathbf{W}^{(i)}_Q \in \mathbb{R}^{d \times d^i_\epsilon},\; \mathbf{W}^{(i)}_K, \mathbf{W}^{(i)}_V \in \mathbb{R}^{d \times d_\tau},\; \varphi_i(\mathbf{z}_i) \in \mathbb{R}^{N \times d^i_\epsilon},\; \tau_\theta(y) \in \mathbb{R}^{M \times d_\tau} \end{aligned}$

在这里插入图片描述
图 14：The architecture of the latent diffusion model (LDM).

Scale up Generation Resolution and Quality

为了生成高分辨率的高质量图像，霍等人（2021）提议采用由多个扩散模型组成的流水线，这些模型的分辨率逐步提高。流水线模型之间的噪声条件增强对于最终图像质量至关重要，具体做法是对每个超分辨率模型 $p_\theta(\mathbf{x} \vert \mathbf{z})$ 的条件输入 $\mathbf{z}$ 应用强大的数据增强。条件噪声有助于减少流水线设置中的复合误差。在用于高分辨率图像生成的扩散建模中，U 型网络是模型架构的常见选择。

在这里插入图片描述
图 15：A cascaded pipeline of multiple diffusion models at increasing resolutions.

他们发现，最有效的噪声应用方式是在低分辨率时使用高斯噪声，在高分辨率时使用高斯模糊。此外，他们还探索了两种条件增强形式，这些形式只需对训练过程进行小幅修改。需注意的是，条件噪声仅在训练时应用，推理时不使用。

截断式条件增强会在低分辨率时，于步骤 $t > 0$ 提前终止扩散过程。
非截断式条件增强会在低分辨率时运行完整的反向扩散过程直至步骤 0，然后通过 $\mathbf{z}_t \sim q(\mathbf{x}_t \vert \mathbf{x}_0)$ 对其进行破坏，再将被破坏的 $\mathbf{z}_t$ 输入到超分辨率模型中。

两阶段扩散模型 unCLIP（拉梅什等人，2022 年）大量利用 CLIP 文本编码器来高质量地生成文本引导的图像。给定一个预训练的 CLIP 模型 $c$ 以及扩散模型的配对训练数据 $(\mathbf{x}, y)$ ，其中 $x$ 是一幅图像， $y$ 是相应的图像说明，我们可以分别计算 CLIP 文本嵌入和图像嵌入，即 $\mathbf{c}^t(y)$ 和 $\mathbf{c}^i(\mathbf{x})$ 。unCLIP 并行学习两个模型：

一个先验模型 $P(\mathbf{c}^i \vert y)$ ：给定文本 $y$ 时，输出 CLIP 图像嵌入 $\mathbf{c}^i$ 。
一个解码器 $P(\mathbf{x} \vert \mathbf{c}^i, [y])$ ：给定 CLIP 图像嵌入 $\mathbf{c}^i$ 以及（可选的）原始文本 $y$ 时，生成图像 $\mathbf{x}$ 。
这两个模型能够实现条件生成，原因是

$\underbrace{P(\mathbf{x} \vert y) = P(\mathbf{x}, \mathbf{c}^i \vert y)}_{\mathbf{c}^i\text{ is deterministic given }\mathbf{x}} = P(\mathbf{x} \vert \mathbf{c}^i, y)P(\mathbf{c}^i \vert y)$

在这里插入图片描述

图：The architecture of unCLIP

unCLIP 遵循两阶段图像生成过程：

1. 给定一段文本 $y$ ，首先使用 CLIP 模型生成文本嵌入 $\mathbf{c}^t(y)$ 。利用 CLIP 潜在空间能够通过文本实现零样本图像操控。
1. 一个扩散或自回归先验 $P(\mathbf{c}^i \vert y)$ 对这个 CLIP 文本嵌入进行处理，构建一个图像先验，然后一个扩散解码器 $P(\mathbf{x} \vert \mathbf{c}^i, y)$ 根据该先验生成一幅图像。这个解码器还能以一幅图像作为输入条件，生成图像变体，同时保留其风格和语义。

与 unCLIP 使用 CLIP 模型不同，Imagen（萨哈里亚等人，2022 年）使用预训练的大型语言模型（即固定的 T5-XXL 文本编码器）对用于图像生成的文本进行编码。总体趋势是，模型规模越大，图像质量和文本与图像的对齐度越高。他们发现，T5 - XXL 和 CLIP 文本编码器在 MS-COCO 数据集上表现相近，但在 DrawBench（涵盖 11 个类别的提示词集合）上，人工评估更倾向于 T5-XXL。
在应用无分类器引导时，增加某个参数（此处你未提及具体参数）可能会使文本与图像的对齐度提高，但图像保真度下降。他们发现这是由于训练 - 测试不匹配导致的，也就是说，因为训练数据 $\mathbf{x}$ 处于某个范围 $[- 1, 1]$ 内，测试数据也应该在这个范围内。为此引入了两种阈值策略：

静态阈值处理：将预测值 $\mathbf{x}$ 裁剪到 $[- 1, 1]$ 。
动态阈值处理：在每个采样步骤中，将计算 $s$ 为某个百分比的绝对像素值；如果 $s > 1$ ，将预测值裁剪到 $[- s, s]$ 并除以 $s$ 。

Imagen 对 U-Net 的一些设计进行了修改，以构建高效的 U-Net。

参数调整：通过为较低分辨率添加更多残差连接，将模型参数从高分辨率模块转移到低分辨率模块。
跳跃连接缩放：按比例 $1/\sqrt{2}$ 缩放跳跃连接。
操作顺序调整：颠倒下采样（将其移到卷积之前）和上采样操作（将其移到卷积之后）的顺序，以提高前向传播的速度。

他们发现，噪声条件增强、动态阈值处理和高效 U-Net 对图像质量至关重要，但扩展文本编码器的大小比扩展 U-Net 的大小更为重要。

Model Architecture

扩散模型有两种常见的骨干网络架构选择：U-Net 和 Transformer。
U-Net（罗内贝格等人，2015 年）由一个下采样堆栈和一个上采样堆栈组成。

下采样：每一步都重复应用两个 3×3 卷积（无填充卷积），每个卷积之后接一个 ReLU 激活函数和一个步长为 2 的 2×2 最大池化层。在每个下采样步骤中，特征通道的数量翻倍。
上采样：每一步都先对特征图进行上采样，然后接一个 2×2 卷积，每一步都会使特征通道的数量减半。
捷径连接（Shortcuts）：捷径连接将下采样堆栈中相应层的特征进行拼接，为上采样过程提供必要的高分辨率特征。

为了实现基于诸如 Canny 边缘、霍夫线、用户涂鸦、人体骨骼、分割图、深度和法线等额外图像信息进行图像合成，ControlNet（张等人，2023 年）对架构进行了改动。具体做法是在 U-Net 的每个编码器层中添加 “夹心” 式零卷积层，这些零卷积层使用原始模型权重的可训练副本。确切地说，对于一个神经网络模块 $\mathcal{F}_\theta(.)$ ，ControlNet 执行以下操作：

首先，冻结原始模块的原始参数 $\theta$ 。
将其克隆为一个具有可训练参数 $\theta_c$ 以及一个额外条件向量 $\mathbf{c}$ 的副本。
使用两个零卷积层，分别记为 $\mathcal{Z}_{\theta_{z1}}(.;.)$ 和 $\mathcal{Z}_{\theta_{z2}}(.;.)$ ，这两个层是 1×1 卷积层，权重和偏置都初始化为零，用于连接这两个模块。零卷积层在初始训练步骤中通过消除作为梯度的随机噪声来保护主干网络。

在这里插入图片描述
图：The ControlNet architecture.

最终输出为： $\mathbf{y}_c = \mathcal{F}_\theta(\mathbf{x}) + \mathcal{Z}_{\theta_{z2}}(\mathcal{F}_{\theta_c}(\mathbf{x} + \mathcal{Z}_{\theta_{z1}}(\mathbf{c})))$

用于扩散建模的扩散变换器（DiT；皮布尔斯和谢，2023 年）在潜在图像块上进行操作，采用与潜在扩散模型（LDM）相同的设计空间。DiT 具有以下设置：

1. 将输入 $\mathbf{z}$ 的潜在表示作为 DiT 的输入。
1. 将大小为 $\times I \times C$ 的噪声潜在表示分割成大小为 $p$ 的图像块，并将其转换为大小为 $I/p)^2$ 的图像块序列。
1. 然后，这个标记序列通过变换器模块。他们探索了三种不同的设计，以确定如何根据诸如时间步 $t$ 或类别标签 $c$ 等上下文信息进行生成。在这三种设计中，adaLN（自适应层归一化）-Zero 的效果最佳，优于上下文条件和交叉注意力模块。缩放和平移参数 $\gamma$ 和 $\beta$ 由 $t$ 和 $c$ 的嵌入向量之和回归得到。维度缩放参数 $\alpha$ 同样通过回归得到，并在 DiT 模块内的任何残差连接之前立即应用。
1. 变换器解码器输出噪声预测和输出对角协方差预测。

变换器架构易于扩展，这一点广为人知。根据实验，这是 DiT 的最大优势之一，因为其性能随着计算量的增加而提升，并且更大的 DiT 模型计算效率更高。

在这里插入图片描述

图：The Diffusion Transformer (DiT) architecture.

参考：
What are Diffusion Models?
Weng, Lilian. (Jul 2021). What are diffusion models? Lil’Log. https://lilianweng.github.io/posts/2021-07-11-diffusion-models/.

References

[1] Jascha Sohl-Dickstein et al. “Deep Unsupervised Learning using Nonequilibrium Thermodynamics.” ICML 2015.

[2] Max Welling & Yee Whye Teh. “Bayesian learning via stochastic gradient langevin dynamics.” ICML 2011.

[3] Yang Song & Stefano Ermon. “Generative modeling by estimating gradients of the data distribution.” NeurIPS 2019.

[4] Yang Song & Stefano Ermon. “Improved techniques for training score-based generative models.” NeuriPS 2020.

[5] Jonathan Ho et al. “Denoising diffusion probabilistic models.” arxiv Preprint arxiv:2006.11239 (2020). [code]

[6] Jiaming Song et al. “Denoising diffusion implicit models.” arxiv Preprint arxiv:2010.02502 (2020). [code]

[7] Alex Nichol & Prafulla Dhariwal. “Improved denoising diffusion probabilistic models” arxiv Preprint arxiv:2102.09672 (2021). [code]

[8] Prafula Dhariwal & Alex Nichol. “Diffusion Models Beat GANs on Image Synthesis.” arxiv Preprint arxiv:2105.05233 (2021). [code]

[9] Jonathan Ho & Tim Salimans. “Classifier-Free Diffusion Guidance.” NeurIPS 2021 Workshop on Deep Generative Models and Downstream Applications.

[10] Yang Song, et al. “Score-Based Generative Modeling through Stochastic Differential Equations.” ICLR 2021.

[11] Alex Nichol, Prafulla Dhariwal & Aditya Ramesh, et al. “GLIDE: Towards Photorealistic Image Generation and Editing with Text-Guided Diffusion Models.” ICML 2022.

[12] Jonathan Ho, et al. “Cascaded diffusion models for high fidelity image generation.” J. Mach. Learn. Res. 23 (2022): 47-1.

[13] Aditya Ramesh et al. “Hierarchical Text-Conditional Image Generation with CLIP Latents.” arxiv Preprint arxiv:2204.06125 (2022).

[14] Chitwan Saharia & William Chan, et al. “Photorealistic Text-to-Image Diffusion Models with Deep Language Understanding.” arxiv Preprint arxiv:2205.11487 (2022).

[15] Rombach & Blattmann, et al. “High-Resolution Image Synthesis with Latent Diffusion Models.” CVPR 2022.code

[16] Song et al. “Consistency Models” arxiv Preprint arxiv:2303.01469 (2023)

[17] Salimans & Ho. “Progressive Distillation for Fast Sampling of Diffusion Models” ICLR 2022.

[18] Ronneberger, et al. “U-Net: Convolutional Networks for Biomedical Image Segmentation” MICCAI 2015.

[19] Peebles & Xie. “Scalable diffusion models with transformers.” ICCV 2023.

[20] Zhang et al. “Adding Conditional Control to Text-to-Image Diffusion Models.” arxiv Preprint arxiv:2302.05543 (2023).