2024-04-10

Stable Diffusion Paper 精读

随着 GPT-3 和 Stable Diffusion 的成功，AI技术进入大模型时代，Large Language Model (LLM) 和 Diffusiom Model (DM) 成为大模型的核心技术。作为为 Stable Diffusion 提供最重要技术突破的论文 High-Resolution Image Synthesis with Latent Diffusion Models，迄今引用已经达到6k+，具有重要的学习价值。本文为该工作的阅读笔记，主要记录笔者对工作的整体理解和一些细节展开。

1. 主要内容和贡献
2. 重要依赖组件
3. 实现细节和结果
- 3.1 Encoder的预训练
- 3.2 Conditional Diffusion Model的结构
4. 未来发展
Reference

1. 主要内容和贡献

基于 Diffusion Model 的生成式模型，因为是直接学习分布，因此比依赖对抗学习的GAN更稳定。除图像生成外，DM还实现了图像超分辨¹、图像修复²等功能。但现有DM通常是在像素空间进行操作的，使得训练和推理需要的计算量都很大。该工作主要贡献有两点：

提出两阶段模型，即LDM：第一阶段通过 Auto Encoder 学习 latent space representation，将图像压缩至隐空间；第二阶段在隐空间训练DM，从而大大降低运算量。两阶段模型不仅降低了训练和推理阶段的运算量，还避免了联合训练需要平衡AE和DM训练目标的问题³。
设计了一种通用的添加 condition 的结构：在U-net内部feature中加入QKV Attention层，以cross attention的形式提取condition信息。该结构实现了class-conditional, text-to-image和layout-to-image等任务，都达到了SOTA，并且可以多模态地训练。

笔者个人认为，第一个贡献点是比较容易想到的。因为在Vahdat et. al.³的工作中，已经有了在latent space训练DM的方法，区别是在该工作中AE和DM是同时训练，因此训练时运算开销很大，还需要平衡两者的训练目标。于是很自然地想到可以将联合训练拆分为两阶段训练，这里的难点就变成了如何得到一个较好的encoder。作者们精心设计的训练方法使得encoder达到最小perceptual distortion，取得了足以服众的压缩效果，也是很亮眼。

第二个贡献点是很重大的，具备广泛的可控性让LDM一下子有了成为image领域大统一模型的潜力。如果只有第一个贡献点，本文将只是一个改进DM训练推理效率的一般文章，只能做image synthesis这一个任务。但作者们通过提出一个通用的添加condition的结构，实现了各种任务的统一：image synthesis, class-conditional image synthesis, text-to-image, layout-to-image, super-resolution, semantic synthesis, inpainting 等等。因此这篇读书笔记也会主要去关注第二个贡献点的实现细节。

2. 重要依赖组件

2.1 VQ-VAE

第一阶段的 encoder 和用以解码LDM输出的 decoder 的训练采用了VQ-VAE，即使用 Vector Quantilization（向量量化）进行图像的压缩表征，主要参考文献 Neural Discrete Representation Learning。原始图像 $W \times H \times 3$ 在经过若干层CNN后，成为 $w \times h \times c$，然后通过查询code book $K \times c$ 中最近元素得到每个位置feature的编码，并以code book中feature替换之，最后经过decoder还原图像。

训练过程中主要用到一个trick叫梯度直传，用以解决量化code不能反传梯度的问题，其简洁实现如下

1
2
3

def quantile(x, code_book):
    z = nearest(x, code_book)
    return x + tf.stop_gradient(z - x)

笔者推测选择VQ-VAE是因为量化离散能够提供较为鲁棒的decoder。该过程将LDM的输出重新投影到量化空间，对微小变动是鲁棒的，并且可以处理数值超出表示范围的情况。所以相比于连续空间，离散的量化空间能提供更高容错。

2.2 U-net

LDM的 diffusion model 中函数$\epsilon_\theta(x, t)$的结构主体为U-net。U-net首先由Ronneberger等人⁴在医疗图像分割中使用。在最新使用DM的图像生成¹和图像超分辨²工作中$\epsilon_\theta(x, t)$正是使用U-net结构来表达。其中图像生成任务中U-net已经引入了attention结构，而超分辨任务中作为condition的低分辨率图像通过concatenation $X_t$进入$\epsilon_\theta(x, t)$。可以明显看到LDM的结构除引入cross attention外，其他部分和这两篇工作几乎一模一样。可以合理推测LDM作者们是在此基础上对U-net结构进行改进，引入了对 semantic feature 的 cross attention，从而得到一个更普适的 conditional U-net。Conditional U-net 的通用结构如下，区别在于如何结合context：

def unet_call(self, x, time_emb, context):
    h = x
    for module in self.down_blocks:
        h = module(h, time_emb, context)
        hs.append(h)
    h = self.middle_block(h, time_emb, context)
    for module in self.up_blocks:
        h = tf.cat([h, hs.pop()], axis=-1)
        h = module(h, time_emb, context)
    return h

2.3 Diffusion Model

在diffusion model⁵出现之前，对于分布的学习主要依赖于GAN或者Norm Flow等方法。Diffusion model将任意的分布通过扩散过程（Diffusion Process）与高斯分布联系在一起，然后以Deep Learning的方法反解对应的随机微分方程，得到扩散过程的逆过程，从而可以从高斯分布还原原始分布。DM具有很高的稳定性和精确性，在提出之后迅速得到学术界的重视。

对于一个 sample $X_0$，$1 > \bar{\alpha}_0 > \bar{\alpha}_1 … > \bar{\alpha}_T \approx 0$描述了扩散过程中不同时刻的 $X_t \sim N(\sqrt{\bar{\alpha}_t} X_0, (1-\bar{\alpha}_t) I)$，训练阶段的优化目标为

$\mathcal{L}_\theta = \|\epsilon_t - \epsilon_\theta(X_t, t)\|_2^2，X_t = \sqrt{\bar\alpha_t} X_0 + \sqrt{1 - \bar\alpha_t} \epsilon_t,$

推理阶段则有

$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_t z.$

其中$\alpha_t = \frac{1-\bar\alpha_t}{1-\bar\alpha_{t-1}}, z \sim N(0, I).$ OpenAI员工Lilian有关diffusion model的博客非常适合用于深入了解diffusion model.

2.4 Transformer

Transformer 结构最早由 Attention is All You Need⁶提出，通过注意力机制解决传统NLP解决方案中RNN网络结构不能处理长程关联的问题。不久后OpenAI和DeepMind各自基于该工作提出了General Pre-trained Transformer （GPT）和 BERT，刷新了NLP task的得分。不仅如此，OpenAI的后续研究⁷揭示Transformer具有相当好的规模化效应（scaling law），拥有可随着模型增大、训练数据变多、训练算力增长而变强的能力。随着2022年Chat-GPT的横空出世，证明了 Transformer 结构是一条可能通往通用人工智能（AGI）的新道路，对学术界外乃至整个社会都产生了颠覆性的冲击。LDM这篇文章就是和Chat-GPT同时代的工作，作者们敏锐地察觉到 Transformer 结构和LLM所蕴含的通用能力，便将其引入图像生成领域，促成了 Stable Diffusion 的诞生。

3. 实现细节和结果

3.1 Encoder的预训练

LDM的Encoder-Decoder训练是基于同一作者的早期工作⁸。不同于使用基于像素的$L_2$或者$L_1$损失，作者使用了一种结合perceptual distortion和patch-based adversarial的训练方法。这种方法能够保证重建图像落在真实数据流形上，且避免了基于像素的损失会产生的模糊结果。

$E, G, Z, D$ 分别表示Encoder, Generator(Decoder), Code Book和Discriminator。$\hat{x} = G(z) = G(VQ(E(x), Z))$ 为经编码后重建的图像，则VQ-VAE的优化目标为

$\mathcal{L}_{VQ}(E, G, Z) = \mathcal{L}_{rec}(x, \hat{x}) + \|sg[E(x)] - z\|^2_2 + \|sg[z] - E(x)\|^2_2,$

此处$\mathcal{L}_{rec}(x, \hat{x})$为重建误差，在VQ-VAE论文中为pixle-wise $L_2$ norm，但此处被替换为perceptual loss⁹，即将图像输入VGG-net等深度神经网络后观察各feature layer的加权$L_2$距离。Patch-based adversaial 训练loss为

$\mathcal{L}_{GAN}(D) = \log D(x) + \log(1-D(\hat x)) .$

则总优化问题定义为

$\operatorname*{argmin}_{E,G,Z} \max_{D} E_{x\sim p(x)} [\mathcal L_{VQ} + \lambda \mathcal L_{GAN}(D)],$

其中$\lambda = \frac{\nabla_{G_L}[\mathcal L_{rec}]}{\nabla_{G_L}[\mathcal L_{GAN}] + \delta}$，$\nabla_{G_L}$表示Generator最后一层的梯度，用以平衡两个优化目标。

从作者在Github公开的repo CompVis/latent-diffusion 可知最终用于大多数任务（text-to-image, layout-to-image, inpainting, …）的设定为VQ-VAE。code book大小为 $8192 \times 3$；降采样倍率为4；Encoder内部结构为3层down-sample layer，Res-Attn-Res中间层，还有kernel=3的conv输出层；每层 down-sample layer 中有两个 residual 结构；每降采样一层，channel数就涨一倍，从而保证每一层总feature数基本不变。

3.2 Conditional Diffusion Model的结构

构成$\epsilon_\theta(x, t)$的U-net基础组成模块除ResnetBlock外，就是AttentionBlock。AttentionBlock有两种结构，一种是self-attention，另一种是cross-attention。AttentionBlock通常置于U-net每一层的最后，输入为$H\times W \times C$，在attention运算时，将每个像素位上维度为$C$的vector作为一个token。因此 self-attention 时是 $Q=K=V=(H*W)\times C$ ，而 cross-attention 时 $Q=(H*W)\times C,K=V= M \times C$ 。

U-net 的 condition 输入分为两种：图像形式（sematic map, low-resolution image, …），语义形式（text, layout, class）。

图像形式的condition因为和输入的latent code map具有相同的长宽，因此可以直接在channel层concatenate即可。
语义形式的condition首先经过一个LLM生成 $M \times C$ 的token，然后通过AttentionBlock和U-net产生联系。

U-net还有一个输入是time embedding，其实现为在每个ResnetBlock中通过MLP层将time映射为一个embedding，然后concatenate到feature的每一个像素位后。

笔者比较奇怪的一点是，其他的使用Transformer进行图像处理的模型比如ViT系列，都是会给每个图像patch一个positional embedding，而奇怪的是LDM和其前序工作中使用attention时都没有给token以position位置信息，不知道为何这么做，也不知道这样一来attention如何区别不同位置的信息。

4. 未来发展

未来的发展方向有两个：

寻找更加scalable的$\epsilon_\theta(x)$结构，能够吸收更多的数据、随算力的增长而成长。
从2D走向2D+time，即对视频进行类似的操作。

幸运的是这两个发展方向已经成为了现实：2024年2月，SORA¹⁰的发布标志着LDM生成长视频能力的成功。在诸多分析报告中，大家一致推测SORA通过使用Diffusion Transformers¹¹获得更强的scale能力，并将视频tokenize化为三维tensor进行了处理，但即便进行了latent representation的压缩，SORA的训练和推理运算量也十分庞大。

Reference

1.Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In NeurIPS, 2020. ↩
2.Chitwan Saharia, Jonathan Ho, William Chan, Tim Salimans, David J. Fleet, and Mohammad Norouzi. Image super-resolution via iterative refinement. CoRR, abs/2104.07636, 2021. ↩
3.Arash Vahdat, Karsten Kreis, and Jan Kautz. Score-based generative modeling in latent space. CoRR, abs/2106.05931, 2021 ↩
4.Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmen- tation. In MICCAI (3), volume 9351 of Lecture Notes in Computer Science, pages 234–241. Springer, 2015. ↩
5.Jascha Sohl-Dickstein, Eric A. Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. CoRR, abs/1503.03585, 2015. ↩
6.Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In NIPS, pages 5998–6008, 2017. ↩
7.Jared Kaplan, Sam McCandlish, Tom Henighan, et al. Scaling laws for neural language models. CoRR, abs/2001.08361, 2020. ↩
8.Patrick Esser, Robin Rombach, and Bj¨orn Ommer. Taming transformers for high-resolution image synthesis. CoRR, abs/2012.09841, 2020. ↩
9.Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang. The Unreasonable Effectiveness of Deep Features as a Perceptual Metric. In CVPR, 2018. ↩
10.https://openai.com/sora ↩
11.William Peebles, and Saining Xie. Scalable Diffusion Models with Transformers. Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023: 4195-4205. ↩