Diff-Plugin: Revitalizing Details for Diffusion-based Low-level Tasks (Paper reading)

Yuhao Liu, City University of Hong Kong, CVPR2024, Papar, Code

1. 前言

我们提出了一种新的 Diff-Pluggin 框架，使单个预训练的扩散模型能够在各种低级任务中生成高保真结果。具体来说，我们首先提出了一个轻量级的任务-Pluggin模块，该模块具有双分支设计，以提供特定于任务的先验，指导扩散过程保留图像内容。然后，我们提出了一种插件选择器，它可以根据文本指令自动选择不同的 Task-Plugins，允许用户通过指示具有自然语言的多个低级任务来编辑图像。我们对 8 个低级视觉任务进行了广泛的实验。结果表明，Diff-Pluggin 优于现有方法，尤其是在现实场景中。我们的消融进一步验证了 Diff-Pluggin 是稳定的、可调度的，并支持跨不同数据集大小的稳健训练。

2. 整体思想

写作很好，故事讲的挺不错，把很简单的一个东西讲的很高级。就是在预训练的扩散模型中插入个模块，然后训练这些模型的参数，使这个模块适应你的任务和扩散模型的先验。模型很大很大，用了4个A100 80G。

3. 方法

在这项工作中，我们观察到描述任务目标的准确文本提示已经可以指示预训练的扩散模型来解决许多低级任务，但通常会导致明显的内容失真，如图 2 所示。首先，我们使用具有全文描述的 SD 生成左图像。然后，我们删除不需要的属性，可选地添加新属性，并使用 SD 中的 img2img 函数，使用左侧图像作为条件来生成右侧的编辑图像。我们对这个问题的理解是，包含任务的引导信息和输入图像的空间信息的特定任务先验可以充分引导预先训练的扩散模型处理低级别任务，同时保持高保真度的内容一致性。为了利用这一潜力，我们提出了Diff-Plugin，这是第一个实现预先训练的扩散模型的框架，以适应各种低级任务，而不损害其原始生成能力。

Task-Plugin：任务插件模块由两个分支组成：任务提示分支（TPB）和空间互补分支（SCB）。TPB包括三个具有Layer Normalization和LeakyReLU的MLP层（除了最后一层），确保仅保留最特定于任务的属性。SCB中Res和Att块表示来自扩散模型的标准ResNet和Cross Attention Transformer块。我们使用通过预先训练的CLIP视觉编码器提取的视觉提示，将模型的注意力引向与任务相关的模式（例如，用于减噪的雨带和用于除雪的雪花）。具体而言，对于输入图像$I$，编码器$En{c}_I$首先提取一般视觉特征，然后由TPB提取这些特征，以产生有区别的视觉引导先验$F^p$。我们利用预先训练的VAE编码器$En{c}_V$来捕获输入图像$I$的全部内容。当与来自$F^p$的语义指导相结合时，该综合图像细节由我们的SCB进行处理，以提取空间特征$F^s$。Res的输出作为att块的Query，$F^p$作为Key和Value。

具体来说，$F^p$被插在扩散模型中的cross attenrion层，$F^s$被直接融合在Decoder的最后一层。我的理解，简单来说，经过TPB处理的$F^p$代替了原来SD中的img_condition。如果用降质图像作为condition，那么噪声也会被带入生成的图像中。这里同样是CLIP的编码，而CLIP是高质量图片训练的，那么TPB的作用其实就是对噪声图像的编码与CLIP的清晰先验得到的编码进行一个对齐，或者说这一步就是对噪声编码进行去噪。但是你condition的强度完全不够，模型更多从空间$x_t$中进行内容的生成，那么是不是我们应该把稍微复原的图像空间编码$F^s$给到$x_t$中。这里的方式就是直接把去噪的编码加到了扩散模型解码器中。把这个Task-Plugin一顿折腾到预训练的SD中，就可以训练了$||ϵ-{ϵ}_{\theta }(z_t,t,F^p,F^s)|{|}_2^2$。这就相当于重新训练扩散模型嘛，只不过训练的是Task-Plugin的参数而已，也就是让Task-Plugin学会适应SD并能够图像复原。我对这个时间复杂度和空间复杂度担忧，这里额外增加了那么多参数（包括$Enc$的参数），且影响SD的推理时间。而且作者的实验都是在很轻微的降质上进行的，如果降质很严重，估计Task-Plugin的复原能力就不够了。

Plugin-Selector：这里其实就是文本和图像的对齐，我觉得这里有点多余，感觉没什么说的，有兴趣建议看原文。

4. 实验

不同任务用不同low-level的数据集（差评），图像大小512方块，用批次为64的4个80G A100训练Task-Plugin，批次8训练Task-Selector（有钱）。对比实验包括特定任务的回归模型(ALL第一行)，回归多任务模型，基于扩散的模型。但有一个问题，你这算力这么猛，对比他们公平吗？而且只有FID和KID很玄幻的指标，在合成数据集上的一些其他指标作者并未给出，不知道会怎么样。此外，这些测试集的降质程度都很低。

---