家用扩散模型 (2)：Latent Diffusion Model

2023-08-13 16:22

前文导航

在之前的三篇中，我们大致介绍了扩散模型背后的数学原理以及如何从一个训完的扩散模型中采样。

但是，还有一个最重要的问题没有解答：那我怎么估算 \(\boldsymbol{\epsilon}_\theta(x, t)\) 呢？

有请本文的主角，generative model第一次火出圈的LDM [1] （aka Stable Diffusion 1.0）

TL; DR

这篇文章的思路非常简单：

首先train一个VAE把high-res的图片压到一个low-res的latent space（注意被压成feature map了还是2D，和许多VAE直接压成1D latent code不同）。backbone使用了VQGAN。

虽然作者提了一下loss可以用传统VAE的KL loss，或者用VQGAN的quantization layer，但不管选哪个regularization，backbone结构都是一样的——因为VQGAN虽然名义上是GAN，但其实还是有encoder-decoder，不要像鸽子一样望文生义。

然后就可以在latent space上做diffusion了，因此我们可以通过cross attention把各种condition fuse进来，从而让各种task都复用同一个backbone。例如，如果是从text prompt出图，就把encoder从图的encoder换成文字的encoder。具体可以看同属SD系列的img2img源码和text2img源码。

从image space到latent space：VQGAN系列

首先，回忆一下没有V的AE（autoencoder）是img \(x\) --> encoder \(e\) --> latent \(z_e(x)\) --> decoder --> output。

传统VAE可以看成，对posterior \(p(z|x) = \mathcal{N}(z; \mu, \sigma^2)\) 优化ELBO。

Vector Quantization (VQ) 做了什么呢？它把\(z_e(x)\)给最近邻到codebook里的nearest neighbour code \(z_q\)。

因此VQVAE可以看成，对posterior \(p(z=z_k|x) :=\boldsymbol{1}_{k=\arg \min_j\|z_e(x)-e_j\|}\) 优化ELBO。Posterior变成了一个冲击——因此他就不variational了。

回忆一下我们的ELBO，它由reconstruction项和KL项组成，reconstruction项是一致的，但代替Gaussian的KL的是VQ regularization \(\|z_q-z\|_2^2\) 。

注意这个loss本身不可导（因为argmin没有导数），所以需要一个trick叫stop-gradient (字面意思，就是让\(sg(x)=x\), \(\frac{\partial}{\partial x}sg(x)=0\))，因此真正的loss term是

\|z_q-\text{sg}(z)\|_2^2+\beta\|\text{sg}(z_q)-z\|_2^2

其中参数 \(\beta\) 通常取0.25。详见这里的介绍。

除了直接用这个MSE loss来update codebook之外，也可以用Exponential Moving Averages（EMA）来update codebook（相当于加了个momentum），详见VQVAE Appendix A。

那这么绕了一圈，换了一个discrete posterior，好处是什么呢？——好处是我们tokenize了latent space feature map! 这样你就可以做next token prediction了。这就是我们要介绍的VQGAN[2]：

这里就是train了一个autoregressive model，然后从第一个token开始不断做next token prediction，最后从所有token decoder回图。

当然，token prediction的顺序当然会影响结果，但实际上只要从左上开始按顺序一路采样就行。作者们在附录里作了一番实验，结果惊奇的发现按顺序效果最好（而不是那些特意设计的采样顺序）

另外还有一个补丁：transformer的长度是有限制的（不能给太大的context window），但CNN Encoder-Decoder不能层数太多，不然VQVAE无法recon。因此作者们使用一个sliding attention window让某一个patch的token概率只和邻近的patch有关。

Google此后又做了ViT-VQGAN[3]。相对VQGAN，首先backbone encoder/decoder换成ViT，其次是对vq层codebook的改进，包括两条：

normalized code，这样code到code的l2 distance就可以被解释为cosine similarity。
传统VQGAN一直受到codebook usage不足的困扰（大多数code根本没人用），作者提出只要简单地把高维的code project到低维（比如32维），然后在低维做最近邻，就能有效解决usage的问题。

这也非常合理，想象一下一个1024维的球面，球面上随便点两个vector几乎肯定是正交的，但在32d上就好多了。

U-Net Denoiser

一个更好的介绍在这里，是SD源码的逐行解析：U-Net for Stable Diffusion

Unconditioned的情况： \(\boldsymbol{\epsilon}_\theta(x, t)\)

U-Net大家都十分熟悉了，四舍五入约等于带skip connection的encoder-decoder。

但是t是怎么混进U-Net的呢——答案是把t看成一个position信息。

在SD中使用的是最传统的sinusoidal encoding（也叫Absolute Position Encoding），来自17年的万恶之源Attention is all you need[4]。它把 t 编码成一个向量：前一半维度是 \(\cos(\omega_k\cdot t)\)，后一半是 \(\sin(\omega_k\cdot t)\)，其中频率 \(\omega_k=\theta^{-2k/d}\)（\(d\) 是 embedding 的总长度，\(k\) 从 0 取到 \(d/2\)，max period \(\theta\) 一般取 10000，这也是 SD 中的取值）。

使用三角函数的目的是使得t平移dt对encoding的变换（aka "relative position"）是线性且独立于t的：

\begin{bmatrix} \cos(\omega_k \cdot \phi) & \sin(\omega_k \cdot \phi) \\ -\sin(\omega_k \cdot \phi) & \cos(\omega_k \cdot \phi) \end{bmatrix} \begin{bmatrix} \sin(\omega_k \cdot t) \\ \cos(\omega_k \cdot t) \end{bmatrix} = \begin{bmatrix} \sin(\omega_k \cdot (t + \phi)) \\ \cos(\omega_k \cdot (t + \phi)) \end{bmatrix}

原文的说法是We chose this function because we hypothesized it would allow the model to easily learn to attend by relative positions（section 3.5）

注：虽然在NLP里的本意是为了相对位置，但对于diffusion其实也是很有道理的——回忆一下 \(\boldsymbol{\epsilon}_\theta(x, t)\) 其实是尝试approximate噪声 \(\boldsymbol{\epsilon}\) ，而噪声是Wiener过程，它也是Markov的，两个噪声的差只和时间差有关。

而使用geometric progression的\(\omega_k\)，乍一看不太自然——因为大家熟悉的Fourier变换需要线性的 \(\omega_k\) ，搞成geometric就不能碰瓷Fourier了。但其实这个选择也有些道理：类比一下，如果采样int的话，一个自然的想法是采样n进制上的各位，位的底是指数增长的，而位上的值也是一个周期函数。那如果采样float呢？我们这个 \(\omega_k\) 实际上可以看成 \(\theta^{2/d}\) 进制的编码（

虽然这个思路最早是在reddit上发现的，但后来我发现苏老师也想到了一样的解释（

得到position encoding之后，过两层MLP就得到t的embedding。

注意embedding和encoding的区别，encoding是人为选择的，而embedding是mlp学出来的（

然后就把t的embedding插入回去。插入方法非常的暴力，unet每一层input先conv2d(3,1,1)到hidden layer（dim和t_emb一样），然后加上t_emb，最后conv2d(3,1,1)回去。其他的部分就和传统unet一样了。

为什么是加呢？这个也是NLP那边传过来的操作，主要也是为了方便model理解relative position（注意到对linear layer来说concat和加其实是一样的）。

那么可不可以换个别的方法呢？有的有的，时代已经变了，现在人人都用RoPE了。

RoPE[5]

之前我们提到了relative position。

这个性质其实可以更进一步的formulate为：我们希望对于 \(\boldsymbol{QK}^T\) 中的每一项， \((\boldsymbol{QK}^T)_{mn}=\boldsymbol{q}_m^T\boldsymbol{k}_n\) 的值只和原信息x1/x2，以及位置差 \(m-n\) 有关（注意[5]中的符号和上一张图mn的顺序是反的）。

我们记给x1/x2添加位置信息的过程为 \(\boldsymbol{q}_m=\boldsymbol{f}_q(\boldsymbol{x}_m, m),\boldsymbol{k}_n=\boldsymbol{f}_k(\boldsymbol{x}_n, n)\) ，其中m和n是x1和x2对应位置的值，那么我们想要的其实是