Faster Diffusion: Rethinking the Role of UNet Encoder in Diffusion Models

Faster Diffusion: Rethinking the Role of UNet Encoder in Diffusion Models

这也是一篇加速的文章，本质上是探索了不同时间步的编码器重用，以及先验噪声注入策略。来减少推理时间。

Abstract

扩散模型的关键组件之一：UNet 通常是作为噪声预测来看待的。之前的一些工作已经探索了UNet作为解码器的基本属性，但其作为编码器在很大程度上仍未被探索。

原因是因为作者发现：编码器特征变化缓慢，而解码器特征在不同的时间步长上表现出实质性的变化。
这一发现启发作者在一定的相邻时间内省略编码器，循环重用前面时间步的编码器特征来用于解码。

基于这一观察，本文引入了一种简单而有效的编码器传播方案来加速不同任务的扩散采样。此外，还引入了噪声先验来提高生产图像中的纹理细节。

在不使用任何知识蒸馏的方法下，本方法将稳定扩散(SD)和DeepFloyd-IF模型的采样速度分别提高了41%和24%。

---

Introduction

我们知道基于扩散模型的文生图模型表现高度依赖于 UNet 的去噪网络以及高质量的图像数据集。

UNet 结构包括 3 个部分：1）encoder $\mathcal{E}$     2）a bottleneck $\mathcal{B}$     3）decoder $\mathcal{D}$

经过 encoder 编码出来后的特征 are skipped to 到 decoder 中；
之前有一些工作将 UNet 解码后的特征用于各种各样的下游任务。例如

• PnP 用于文本引导的 Image-to-Image 翻译任务
• DIFT 发现了一种主要存在于解码特征的主要突发现象；
• ControlNet 微调一个额外的 encoder，用预训练的UNet的编码器初始化，并保留原始的解码器。
• T2I-Adapter 将条件信息与可训练编码器一起注入到UNet的编码器中，并保留初始解码器。这些工作表明了解码器在生成高质量图像中的关键作用。

本文聚焦于探索 UNet 编码器的关键作用，而关键问题就是：

• 来自不同层的特征在推理时间内是否有不同的贡献。
• 如果有，那么如何重用这些特征。

基于上述问题，本文提出：

• encoder propagation : 一种简单有效的扩散采样过程中基于编码器特征复用的传播方案。
• 虽然传播方案提高了扩散采样的效率，但我们观察到会导致生成结果中纹理信息的一些丢失。为了解决这个问题，我们引入了 先验噪声注入策略 来保留生成图像中的纹理细节

可能看下面这个图会更加直观，这个第二点感觉还是比较牵强……（因为我好像乍一看看不出啥）

对比

---

Methods

前面说过 UNet 包括 3 个部分：

• encoder $\mathcal{E}$    
• a bottleneck $\mathcal{B}$    
• decoder $\mathcal{D}$

从编码器提取出来的特征会通过跳跃连接传递到解码器中。

为了更好的进行表述，我们将 UNet 划分为几个块： E = { E ( ⋅ ) s } \mathcal{E} = \{\mathcal{E}(\cdot)_s\} E={E(⋅)s​} , B = { B ( ⋅ ) 8 } \mathcal{B} = \{\mathcal{B}(\cdot)_8\} B={B(⋅)8​} , D = { D ( ⋅ ) s } \mathcal{D} = \{\mathcal{D}(\cdot)_s\} D={D(⋅)s​}
其中， s ∈ { 8 , 16 , 32 , 64 } s \in \{8, 16, 32, 64\} s∈{8,16,32,64}。而 $\mathcal{E}(\cdot {)}_s$ 和 $\mathcal{D}(\cdot {)}_s$ 分别表示编码器和解码器中具有输入分辨率 $s$ 的 block layer。

Analyzing the UNet in Diffusion Model

特征在时间步长上的演化

>先说结论吧<

编码器特征在相邻的时间步长上表现出细微的变化，而解码器特征在不同的时间步长上表现出实质性的变化。

Analyzing the UNet in Diffusion Model. (Fig.2)

(a)通过MSE测量相邻时间步长的特征演化。(b)提取UNet各层在每个时间步长的分层特征输出，沿通道维数对其进行平均，得到二维分层特征，并计算其 f-范数。© UNet编码器的分层特征显示出较低的标准偏差，而解码器的分层特征显示出较高的标准偏差。

Visualising the hierarchical features. (Fig.3)

我们将PCA应用于 PnP 之后的分层特征，并使用前三个主要成分作为RGB图像进行可视化。编码器特征变化缓慢，在许多时间步长上具有相似性 (左图)，而解码器特征在不同时间步长上表现出 实质性变化(右图)。

>接着说说具体做法<

在扩散过程种，我们会进行迭代采样，因此也就会产生相应的 latent code ${\mathbf{z}}_{\mathbf{t}}$ , 包括相应的层级特征： { E ( z t , c , t ) s } \{\mathcal{E}(z_t,c,t)_s\} {E(zt​,c,t)s​} ， { D ( z t , c , t ) s } ( s ∈ { 8 , 16 , 32 , 64 } ) \{\mathcal{D}(z_t,c,t)_s\}(s\in\{8, 16, 32, 64\}) {D(zt​,c,t)s​}(s∈{8,16,32,64})

想要将层次特征变化量化出来，我们可以有:

$${\Delta }_{\mathcal{E}(\cdot {)}_s}=\frac{1}{d\times s^2}\Vert \mathcal{E}(z_t,c,t{)}_s,\mathcal{E}(z_{t-1},c,t-1{)}_s{\Vert }_2^2$$

其中，$d$ 表示编码器特征 $\mathcal{E}(z_t,c,t{)}_s$ 通道数。${\Delta }_{\mathcal{B}(\cdot {)}_s}$ 和 ${\Delta }_{\mathcal{D}(\cdot {)}_s}$ 同理。

特征在跨层上的演化

同样的，编码器变化的强度是轻微的，而解码器的变化是剧烈的。

具体来说，我们计算了所有时间步长的层级特征 { E ( z t , c , t ) s } \{\mathcal{E}(z_t,c,t)_s\} {E(zt​,c,t)s​} 的 Frobenius范数 ，表示为

F E ( ⋅ ) s = { F E ( z T , c , T ) s , . . . , F E ( z 1 , c , 1 ) s } . \mathcal{F}_{\mathcal{E}(\cdot)_s}=\{\mathcal{F}_{\mathcal{E}(z_T,c,T)_s},...,\mathcal{F}_{\mathcal{E}(z_1,c,1)_s}\}. FE(⋅)s​​={FE(zT​,c,T)s​​,...,FE(z1​,c,1)s​​}.

计算${\mathcal{F}}_{\mathcal{B}(\cdot {)}_8}$ 和 ${\mathcal{F}}_{\mathcal{D}(\cdot {)}_s}$ 同理。

能否在某些时间步骤忽略编码器？

从前面的分析可以知道，编码在大多数时间步骤上变化是不大的，所以就考虑是否可以通过重用编码器特征来加速扩散采样。

具体来说，我们删除了时间步长为 $t-1(t-1<t)$ 的编码器，相应的解码器(包括跳过连接)将前一个时间步长为 $t$ 的编码器 $\mathcal{E}$ 的层级输出作为其输入，而不是像标准SD采样那样从当前时间步长为t−1的输出作为输入。（这一点看下面的图很好理解）

>这里有几个现象<

Fig4

• (a) 和 （b）比较，如果忽略了编码器，那么可以生成与标准 SD 模型非常相似的图像；
• (a) 和 （d）比较，我们把同样的方法复用到解码器上，发现生成的图像经常无法覆盖文本提示中的某些特定对象。
  比如文本提示：“一个留着胡子、戴着眼镜、戴着无边帽的男人”，这是因为语义主要包含在解码器而不是编码器的特征中。

---

编码器传播方案

• 对于采样过程 t = { T , … , 1 } t = \{T,\dots,1\} t={T,…,1}, 将那些部署编码器传播的时间步定义为 non-key 时间步，表示为：$t^{non\text{-}key}=\left\{t_0^{non\text{-}key},...,t_{N-1}^{non\text{-}key}\right\}$ ；
• 除了非关键步，那就是关键步了,就表示为：$t^{key}=\left\{t_0^{key},t_1^{key},...,t_{T-N-1}^{key}\right\}$。这些步上是不采用编码器传播方案的；
• 初始时间步$(t_0^{key}=T)$。

因此，diffusion 的推理时间步表示为 $\left\{t^{key},t^{non\text{-}key}\right\}$, 同时 t k e y ∪ t n o n - k e y = { T , . . . , 1 } t^{key}\cup t^{non\text{-}key}=\{T,...,1\} tkey∪tnon-key={T,...,1} 以及 $t^{key}\cap t^{non\text{-}key}=\acute{\varnothing }.$

non-/uniform encoder propagation

如图Fig.2a所示，在整个推理过程中，初始推理阶段的编码器特征变化比后期阶段更大。

• 因此，我们在初始推理阶段选择了更多的 key time-steps;
• 在之后的阶段选择更少的 key time-steps

• 经过实验，定义为 key time-steps 为： t k e y = { 50 , 49 , 48 , 47 , 45 , , 40 , 35 , 25 , 15 } t^{key} = \left\{50,49,48,47,45,,40,35,25,15\right\} tkey={50,49,48,47,45,,40,35,25,15} (SD with DDIM Solver)
• 对于 DeepFloyd-IF：key time-steps 为： t k e y = { 100 , 99 , 98 , … t^{key}=\{100,99,98,\ldots tkey={100,99,98,…,$92,91,90,85,80,\dots ,25,20,15,14,13,\dots ,2,1\}$, ${ 50, 49, \ldots , 2, 1} $ and ${ 75, 73, 70, 66, 61, 55, 48, 40, 31, $ $21,10\}$ 3阶段

其余的时间步骤被归类为 non-key time-steps。我们将此策略定义为非均匀编码器传播(non-uniform encoder propagation, 见图4e)。

如图4c所示，固定步长的时间步长选择，称为均匀编码器传播。

表3报告了消融研究的结果，考虑了关键和非关键时间步长的各种组合。这些结果表明，关键时间步长集合在生成图像方面具有更好的性能。

Parallel non-uniform encoder propagation

我们知道编码器传播的时候，也就是在那些 time-step $t\in t^{non-key}$ ，解码器的输入不依赖于当前时间 $t$ 的编码器输出，而是依赖于前一个最近的 key time step $t-1$ 的编码器输出。

这也就意味着，我们其实可以进行并行操作，如图4f 所示。

我们从 $t-2$到$t-k+1$时间步 并行执行解码。该技术可以同时进行多个时间步长的译码向前，进一步提高了推理效率。我们把这些 non-key time-steps称为 parallel-batch non-key time-steps.

如下图所示，推理时间可以减少 41%。

Prior noise injection(先验噪声注入)

前面说过虽然编码器传播可以提高推理阶段的效率，但观察到它会导致生成结果中纹理信息的轻微丢失（图 6 所示）。因此作者提出了一种先验噪声注入策略。

说白了就是它将初始潜码 ${\mathbf{z}}_{\mathbf{t}}$ 合并到后续时间步(即${\mathbf{z}}_{\mathbf{T}}$)的生成过程中，

$${\mathbf{z}}_{\mathbf{t}}={\mathbf{z}}_{\mathbf{t}}+\alpha \cdot {\mathbf{z}}_{\mathbf{T}},\mathrm{if}t<\tau .$$

其中 $\alpha =0.003$ 是控制${\mathbf{z}}_{\mathbf{T}}$影响的尺度参数。我们从 $\tau =25$的步骤开始使用这个注入机制。这种策略的结合成功地改善了纹理信息。重要的是，它需要的额外计算资源几乎可以忽略不计。

我们观察到，纹理信息的丢失发生在频域的所有频率上(见图6(右红蓝曲线))。这种方法确保了SD和zT注入在频域生成的结果非常相似(见图6(右，红色和绿色曲线))，生成的图像保持所需的保真度(见图6(左，下))。

---

Experiments

有关视频任务的这里就不细说了

Dataset

MS-COCO2017 中随机选取 10K prompts，然后生成相应的 10K 个图像。

Metric

• FID、 CLIP Score；
• 计算负载的平均值量：GFLOPs/image
• 采样时间

文本生成图像

Figure 7. Generated images at different time-steps, with the prompt “A man with a beard wearing sunglasses and a beanie”.

Quantitative evaluation in both SD model and DeepFloydIF diffuison model on MS-COCO 2017 10K subset.

也可以泛化到其他任务上，像定制化的生成（Dreambooth、Custom Diffusion），或者是参考图像生成（ControlNet）

Quantitative evaluation on text-to-video, personalized generation and reference-guided generation tasks.

消融实验

不同策略的消融

• 可视化结果