Diffusion_Autoencoder：扩散自编码器

1 前言

​ 本期内容，我们讲Diffusion Autoencoder，这是一个基于自编码器和扩散模型组合而成的模型，主要为了通过他们，去学习到图像的语义表征。然后再使用这些表征去用于其他下游任务。

​ 视频讲解：Diffusion AutoEncoder：扩散自编码器

​ 参考论文：[Diffusion Autoencoders: Toward a Meaningful and Decodable Representation：

2 引入

​ 在一般的扩散模型中，我们已经可以和GAN、VAE等模型那般，去合成质量不错的图像了。然而，扩散模型与GAN、VAE还有一个非常显著的区别。在GAN和VAE中，往往可以将图像的语义信息提取出来（如StyleGAN中的风格编码w、VAE中的编码向量）；而对于Diffusion便没有诸如此类的东西，整个生成过程只是依赖于扩散时间步T。因此，该工作的目的，就是想要利用Diffusion，去学习到图像的语义信息。以用于其他下游任务。

3 DDIM

​ 在讲解之前，让我们先来回顾一下DDIM。在DDIM中，设定去噪采样公式为（$x_i\to x_{i-1}$）
$$x_{i-1}=\sqrt{{\bar{\alpha }}_{i-1}}\frac{x_i-\sqrt{1-{\bar{\alpha }}_i}{ϵ}_{\theta }(x_i,i)}{\sqrt{{\bar{\alpha }}_i}}+\sqrt{1-{\bar{\alpha }}_{i-1}-{\sigma }_i^2}{ϵ}_{\theta }(x_i,i)+{\sigma }_{i}z\text{\hspace{1em}(1)}$$
​ 其中，里面的${\sigma }_i$是由我们自己去选定的，当${\sigma }_i=0$时，即为确定性采样——见Eq.(2)
$$x_{i-1}=\sqrt{{\bar{\alpha }}_{i-1}}\frac{x_i-\sqrt{1-{\bar{\alpha }}_i}{ϵ}_{\theta }(x_i,i)}{\sqrt{{\bar{\alpha }}_i}}+\sqrt{1-{\bar{\alpha }}_{i-1}}{ϵ}_{\theta }(x_i,i)\text{\hspace{1em}(2)}$$
​ Eq.(2)和Eq.(1)是概率分布$q(x_{i-1}|x_i,x_0)$重参数化的结果，Eq.(2)对应的概率密度表示为
$$q(x_{i-1}|x_i,x_0)\sim N(x_{i-1}|\sqrt{{\bar{\alpha }}_{i-1}}\frac{x_i-\sqrt{1-{\bar{\alpha }}_i}{ϵ}_i}{\sqrt{{\bar{\alpha }}_i}}+\sqrt{1-{\bar{\alpha }}_{i-1}-{\sigma }_i^2}{ϵ}_i,0)\text{\hspace{1em}(3)}$$
​ 与DDPM一样，DDIM的优化目标为
$${\mathcal{L}}_{\text{sample}}=\Vert {ϵ}_i-{ϵ}_{\theta }(x_i,i){\Vert }_2^2\text{\hspace{1em}(4)}$$
​ 其中，$x_i=\sqrt{{\bar{\alpha }}_i}{x}_0+\sqrt{1-{\bar{\alpha }}_i}{ϵ}_i$

4 条件生成

​ 在扩散模型的条件生成当中，我们还曾提到过，可以在预测${ϵ}_{\theta }$的网络中，多输入一个条件c（比如文本），以让模型关注到条件信息，按条件生成对应的图像
$${\mathcal{L}}_{\text{c}}=\Vert {ϵ}_i-{ϵ}_{\theta }(x_i,c,i){\Vert }_2^2\text{\hspace{1em}(5)}$$
​ 使用Eq.(5)进行训练，当模型收敛，那么模型便可以学习到去运用条件信息。而条件在输入到模型训练的过程中，也会被编码成模型可以识别的条件信息。

5 Diffusion Autoencoder

​ 对于条件c，他可以是图像的标签或者其他的东西。那么，假如说条件c是图像的编码向量呢？即假设原始图像的编码向量为
$$z_{\text{sem}}=\text{Enc}(x_0)$$
​ 其中$x_0$是原图像，Enc是一个编码器，他的输入是图像，输出的编码向量，稍后我们再说这个编码器的结构。

​ 我们让条件c=$z_{\text{sem}}$，即Eq.(5)可以变换为
$${\mathcal{L}}_{\text{z}}=\Vert {ϵ}_i-{ϵ}_{\theta }(x_i,z_{\text{sem}},i){\Vert }_2^2\text{\hspace{1em}(6)}$$
​ 使用Eq.(6)去训练一个条件扩散模型。那么当模型收敛，便正如我们前面所说，这个编码向量$z_{\text{sem}}$就会被当作是条件信息，就会被模型使用预测去噪后的图像。同时，模型也会学习到这个编码向量。而这个编码向量，正是图像的语义信息，也正是我们所需要的。他的流程图可见下图

​ 里面的Semantic encoder就是我们上面提到的编码器，在训练的过程中，是编码器和${ϵ}_{\theta }$一起去训练。当模型收敛，Semantic encoder这个编码器就可以提取语义信息，我们只需要把图像输入进这个编码器，所得到的编码向量就是我们需要的语义表征。

6 条件输入Unet

​ 在Diffusion Autoencoder当中，预测${ϵ}_{\theta }$的网络使用的仍是Unet，但是这个Unet是一个修改版，具体请看参考【1】。那么如何将条件编码向量$z_{\text{sem}}$注入到Unet网络当中呢？论文使用的是自适应组归一化进行注入（AdaGN）
$$\text{AdaGN}(h,t,z_{\text{sem}})=z_s(t_s\text{GroupNorm}(h)+t_b)$$
​ 其中，$h$是Unet对应的特征图输出，GroupNorm是组归一化，$z_s=\text{Affine}(z_{sem})$，Affine是仿射变换。（$t_s$，$t_b$）是先对时刻t进行正弦编码（记作$\psi$）后再进行MLP映射得到的，即$(t_s,t_b)=\text{MLP}(\psi (t))$

7 Semantic encoder

​ Semantic encoder是语义编码器，他的架构使用Unet的解码器的前半部分

8 图像重建

​ 对于前向加噪过程，有
$$\begin{gathered} x_i=\sqrt{{\bar{\alpha }}_i}{x}_0+\sqrt{1-{\bar{\alpha }}_i}ϵ \\ {x}_{i+1}=\sqrt{{\bar{\alpha }}_{i+1}}{x}_0+\sqrt{1-{\bar{\alpha }}_{i+1}}ϵ\text{\hspace{1em}(7)} \end{gathered}$$
​ 把Eq.(7)的第一个等式移项，可得
$$x_0=\frac{x_i-\sqrt{1-{\bar{\alpha }}_i}ϵ}{\sqrt{{\bar{\alpha }}_i}}\text{\hspace{1em}(8)}$$
​ 将Eq.(8)代入至Eq.(7)的第二个等式，可得
$$x_{i+1}=\sqrt{{\bar{\alpha }}_{i+1}}\frac{x_i-\sqrt{1-{\bar{\alpha }}_i}ϵ}{\sqrt{{\bar{\alpha }}_i}}+\sqrt{1-{\bar{\alpha }}_{i+1}}ϵ\text{\hspace{1em}(9)}$$
​ 我们可以构造这样的一个形式，原图像–>加噪–>加噪图像–>去噪–>原图像

​ 也就是说，我们希望，把原始图像加噪之后，再去噪，能够得到原图像。想要达到这种效果，那么里面的随机噪声便不能够再去随机采样，而是使用网络去预测，那么Eq.(9)就变成Eq.(10)

​
$$x_{i+1}=\sqrt{{\bar{\alpha }}_{i+1}}\frac{x_i-\sqrt{1-{\bar{\alpha }}_i}ϵ(x_i,i,z_{\text{sem}})}{\sqrt{{\bar{\alpha }}_i}}+\sqrt{1-{\bar{\alpha }}_{i+1}}ϵ(x_i,i,,z_{\text{sem}})\text{\hspace{1em}(10)}$$
​ 记
$$f(x_i,i,,z_{\text{sem}})=\frac{x_i-\sqrt{1-{\bar{\alpha }}_i}{ϵ}_{\theta }(x_i,i,z_{\text{sem}})}{\sqrt{{\bar{\alpha }}_i}}$$
​ 最终
$$x_{i+1}=\sqrt{{\bar{\alpha }}_{i+1}}f(x_i,i,,z_{\text{sem}})+\sqrt{1-{\bar{\alpha }}_{i+1}}ϵ(x_i,i,,z_{\text{sem}})\text{\hspace{1em}(11)}$$
​ Eq.(11)这个过程在论文中被称为 Stochastic encoder 。

9 Diffusion Autoencoder 采样生成

​ 由于有$z_{\text{sem}}$是原始图像编码的，在Diffusion采样过程中，我们是没有$z_{\text{sem}}$的，故而我们没有办法进行随机采样生成图像，仅仅能够用于图像重建或者语义编码。

​ 因此，想要进行随机采样生成，我们就需要生成一个$z_{\text{sem}}$，如何生成呢？假如说$z_{\text{sem}}$这个值服从某个简单的分布就好了，比如服从标准高斯分布。但是若是强行干预$z_{\text{sem}}$让他服从标准高斯，又恐怕会难以学习到更加丰富的语义信息，毕竟语义信息总是复杂的，用一个高斯分布去囊括确实不妥。另一个办法就是用GAN去拟合$z_{\text{sem}}$所服从的分布。但话又说回来了，GAN是出了名的不稳定。于是，论文便提出了另一种方法——构造一个Latent DDIM去学习$z_{\text{sem}}$

​ 一般来说，Diffusion是把一张图像加噪成一个标准高斯分布，我们同样可以将$z_{\text{sem}}$当作是一张图像，然后进行加噪，然后进行训练。当模型收敛，只需要从标准高斯分布中采样一个噪声，然后慢慢去噪，就可以得到$z_{\text{sem}}$，即
$${\mathcal{L}}_{\text{latent}}=\Vert {ϵ}_i-{ϵ}_{\omega }(z_{\text{sem,i}},i){\Vert }_1\text{\hspace{1em}(12)}$$
​ 其中，$z_{\text{sem,i}}=\sqrt{{\bar{\alpha }}_i}{z}_{\text{sem}}+\sqrt{1-{\bar{\alpha }}_i}$.。我们注意到，Eq.(12)使用的是L1范数。论文发现这里使用L1范数是优于L2范数的。

10 结束

​ 好了，本篇文章到此为止，还有一些网络架构的内容，本文主要讲Diffusion Autoencoder的原理，相关的网络架构不在涉及范畴，感兴趣的可自行阅读论文。如果问题，还望指出，阿里嘎多！

参考文献

[1] Diffusion Models Beat GANs on Image Synthesis