摘要
在大型数据集上训练的扩散模型在图像合成方面取得了显著进展。然而,由于扩散过程中的随机性,它们在处理需要保留细节的各种低级任务时常常遇到困难。为了克服这一局限性,我们提出了一种新的Diff-Plugin框架,使单个预训练的扩散模型能够跨多种低级任务生成高保真度的结果。具体来说,我们首先提出了一个具有双分支设计的轻量级Task-Plugin模块,以提供任务特定的先验信息,指导扩散过程保留图像内容。然后,我们提出了一个Plugin-Selector,它可以根据文本指令自动选择不同的Task-Plugins,允许用户通过自然语言指示多个低级任务来编辑图像。我们在8个低级视觉任务上进行了大量实验。结果表明,Diff-Plugin相较于现有方法具有优越性,特别是在真实场景中。我们的消融实验进一步验证了Diff-Plugin的稳定性、可调度性,以及在不同数据集大小上支持鲁棒训练的能力。项目页面:https://yuhaoliu7456.github.io/Diff-Plugin
1. 引言
在过去两年里,扩散模型[9, 21, 22, 61]在图像生成方面取得了前所未有的成功,并展现出成为视觉基础模型的潜力。最近,许多工作[4,25,28,31,46,91,96]表明,在大规模文本到图像数据集上训练的扩散模型已经能够理解各种视觉属性,并为下游任务提供多功能的视觉表示,例如图像分类[31]、分割[25, 96]、翻译[46, 91]和编辑[4, 28]。
然而,由于扩散过程固有的随机性,现有的扩散模型无法保持与输入图像一致的内容,因此在处理低级视觉任务时失败。为此,一些方法[46,63]提出在编辑图像时通过DDIM反转[61]策略将输入图像作为先验,但在场景复杂时这些方法不稳定。其他方法[16,52,56,71,83]尝试从头开始在任务特定的数据集上训练新的扩散模型,这限制了它们只能解决单一任务。
在本文工作中,我们观察到,一个准确描述任务目标的文本提示已经可以指导预训练的扩散模型处理许多低级任务,但通常会导致明显的内容失真,如图2所示。我们对这个问题的见解是,包含任务指导信息和输入图像空间信息的任务特定先验可以充分指导预训练的扩散模型处理低级任务,同时保持高保真度的内容一致性。为了利用这一潜力,我们提出了Diff-Plugin,这是第一个框架,它使预训练的扩散模型(如稳定扩散[54])能够容纳各种低级任务,而不损害其原有的生成能力。
Diff-Plugin由两个主要组件组成。首先,它包括一个轻量级的Task-Plugin模块,用于帮助提取任务特定的先验信息。Task-Plugin分为任务提示分支(TPB)和空间补充分支(SCB)。TPB提炼任务指导先验,使扩散模型朝向指定的视觉任务,并最大限度地减少对复杂文本描述的依赖,而SCB利用来自TPB的任务特定视觉指导来辅助空间细节捕获和补充,提高生成内容的保真度。其次,为了方便使用多个不同的Task-Plugins,Diff-Plugin包括一个Plugin-Selector,允许用户通过文本输入选择所需的Task-Plugins(视觉说明如图1所示)。为了训练Plugin-Selector,我们采用多任务对比学习[49],使用任务特定的视觉指导作为伪标签。这使Plugin-Selector能够将不同的视觉嵌入与任务特定的文本输入对齐,从而增强Plugin-Selector的鲁棒性和用户友好性。
为了全面评估我们的方法,我们在八个不同的低级视觉任务上进行了大量实验。我们的结果证实,Diff-Plugin不仅在不同任务中表现稳定,而且展现出卓越的可调度性,促进了文本驱动的多任务应用。此外,Diff-Plugin还展示了其可扩展性,能够适应不同大小数据集上的各种任务,样本量从少于500到超过50,000个,且不影响已训练的插件。最后,我们的结果还表明,所提出的框架在视觉上和定量上都优于现有的基于扩散的方法,并且与基于回归的方法相比也取得了具有竞争力的性能。
我们的主要贡献总结如下:
- 我们提出了Diff-Plugin,这是第一个使预训练的扩散模型能够执行各种低级任务同时保持其原始生成能力的框架。
- 我们提出了Task-Plugin,这是一个轻量级的双分支模块,旨在将任务特定的先验知识注入到扩散过程中,以提高结果的保真度。
- 我们提出了Plugin-Selector,用于根据用户提供的文本选择适当的Task-Plugin。这扩展到了一个新的应用,允许用户通过文本指令对低级视觉任务进行图像编辑。
- 我们在八个任务上进行了大量实验,证明了Diff-Plugin相较于现有的基于扩散和回归的方法具有竞争性的性能。
2. 相关工作
扩散模型[60, 62]已被应用于图像合成[9,21,22,61]并取得了显著成功。随着大量文本-图像数据[59]和大规模语言模型[49, 50]的出现,基于扩散的文本引导图像合成[2, 42, 51, 54, 57]变得更加引人注目。利用文本引导的合成扩散模型,几种方法利用生成能力实现了文本驱动的编辑。零样本方法[19, 46, 63]依赖于正确的初始噪声[61]并操作注意力图以在精确位置编辑指定内容。基于调优的策略通过优化的DDIM反演[65]、注意力调优[29]、文本-图像耦合[28,55, 93]和提示调优[10,14,39]来努力平衡图像保真度和生成多样性。相反,InstructP2P[4,89]通过潜在扩散[54]和提示到提示[19]生成配对数据,用于训练和编辑。然而,扩散过程中的随机性和缺乏任务特定的先验知识使得它们不适用于需要保留细节的低级视觉任务。条件生成模型使用各种外部输入来确保输出与条件一致。无训练方法[8,76]通过操作注意力层可以在指定位置生成新内容,但条件类型有限。基于微调的方法通过训练一个新的扩散分支[40,90,94]或整个模型[1],向预训练的扩散模型注入额外的指导。尽管具有全局结构一致性,但这些方法由于随机性和生成性质,无法保证输出和输入图像细节之间的高保真度。
基于扩散的低级方法可以分为零样本和基于训练的方法。前者可以从预训练的基于去噪扩散的生成模型[22]中借用生成先验来解决线性[27, 70]和/或非线性[7, 12]图像恢复任务,但在真实世界数据上往往产生较差的结果。后者通常通过任务依赖的设计为不同任务训练或微调单个模型,如超分辨率[58, 74]、JPEG压缩[56]、去模糊[52, 73]、人脸恢复[71, 95]、低光增强[24, 83, 92]和阴影去除[16]。并发工作StableSR[66]和DiffBIR[34]使用可学习的条件扩散分支与退化或恢复的图像来专门训练用于盲人脸恢复的扩散模型。相比之下,我们的框架通过配备轻量级的任务特定插件,使一个预训练的扩散模型能够处理各种低级任务。
多任务模型可以学习不同任务之间的互补信息,例如目标检测和分割[18]、雨滴检测和去除[80]、恶劣天气恢复[45, 82, 98]和盲图像恢复[33, 47]。然而,这些方法只能在训练后处理预定义的任务。相反,我们的Diff-Plugin具有灵活性,可以通过任务特定的插件整合新任务,因为我们的Task-Plugins是单独训练的。因此,当向Diff-Plugin添加新的低级任务时,我们只需将预训练的Task-Plugins添加到框架中,而无需重新训练现有的插件。
3. 方法论
在本节中,我们首先回顾扩散模型公式(第3.1节)。然后,我们介绍我们的Diff-Plugin框架(第3.2节),该框架是从我们新提出的Task-Plugin(第3.3节)和Plugin-Selector(第3.4节)发展而来的。
3.1. 预备知识
扩散模型包括一个前向过程和一个反向过程。在前向过程中,给定一个清晰的输入图像 x 0 \mathbf{x}_{0} x0,扩散模型逐步向其添加高斯噪声,以在时间步 t ∈ { 0 , 1 , … , T } t \in\{0,1, \ldots, T\} t∈{0,1,…,T}得到噪声图像 x t \mathbf{x}_{t} xt,表示为 x t = α ˉ t x 0 + 1 − α ˉ t ϵ t \mathbf{x}_{t}=\sqrt{\bar{\alpha}_{t}} \mathbf{x}_{0}+\sqrt{1-\bar{\alpha}_{t}} \epsilon_{t} xt=αˉtx0+1−αˉtϵt,其中 α ˉ t \bar{\alpha}_{t} αˉt是预定义的调度变量, ϵ t ∼ N ( 0 , I ) \epsilon_{t} \sim \mathcal{N}(\mathbf{0}, \mathbf{I}) ϵt∼N(0,I)是添加的噪声。在反向过程中,扩散模型从标准高斯噪声 x T \mathbf{x}_{T} xT中迭代去除噪声,并最终估计出清晰的图像 x 0 \mathbf{x}_{0} x0。这通常用于训练噪声预测网络 ϵ θ \epsilon_{\theta} ϵθ,其监督信息来自噪声 ϵ t \epsilon_{t} ϵt,损失函数为 L = E x 0 , t , ϵ ∼ N ( 0 , 1 ) [ ∥ ϵ − ϵ θ ( x t , t ) ∥ 2 2 ] \mathcal{L}=\mathbb{E}_{\mathbf{x}_{0}, t, \epsilon \sim \mathcal{N}(0,1)}\left[\left\|\boldsymbol{\epsilon}-\boldsymbol{\epsilon}_{\theta}\left(\mathbf{x}_{t}, t\right)\right\|_{2}^{2}\right] L=Ex0,t,ϵ∼N(0,1)[∥ϵ−ϵθ(xt,t)∥22]。
3.2. Diff-Plugin
我们的关键观察是,预训练的扩散模型在执行低级视觉任务时具有固有的零样本能力,使它们能够生成多样化的视觉内容,而无需明确的特定任务训练。然而,这种能力在更细致的特定任务编辑中面临局限性。例如,在去雪任务中,模型理想情况下应仅去除雪而保持其他内容不变,如图2所示,但扩散过程的固有随机性往往导致场景中除去雪之外的意外改变。这种不一致性源于模型缺乏特定任务的先验知识,这对于低级视觉任务中精确细节保留至关重要。
受自然语言处理(NLP)中模块化扩展[75, 77]和GPT-4[43]的启发,这些扩展利用即插即用工具来增强大型语言模型在下游任务中的能力,同时不损害其核心竞争力,我们基于类似的想法引入了一种新颖框架——Diff-Plugin。该框架将几个轻量级插件模块(称为Task-Plugin)集成到预训练的扩散模型中,用于各种低级任务。Task-Plugin旨在提供必要的特定任务先验知识,引导模型生成高保真度和任务一致的内容。此外,虽然扩散模型可以根据文本指令为特定场景生成内容,但它们缺乏为不同低级任务调度Task-Plugin的能力。即使现有的条件生成方法[48,90]也只能通过输入条件图像来指定不同的生成任务。因此,为了促进平滑的文本驱动任务调度,并能够在复杂工作流程中切换不同的Task-Plugin,Diff-Plugin包括一个Plugin-Selector,允许用户通过文本命令选择和调度适当的Task-Plugin。
图3描绘了Diff-Plugin框架。给定一张图像,用户通过单个或多个文本提示指定任务,Plugin-Selector为其识别适当的Task-Plugin。然后,Task-Plugin处理图像以提取特定任务的先验知识,引导预训练的扩散模型产生用户期望的结果。对于单个插件范围之外的更复杂任务,Diff-Plugin通过预定义的映射表将其分解为子任务。每个子任务由指定的Task-Plugin处理,展示了该框架处理多样化和复杂用户需求的能力。
3.3. 任务插件
如图4所示,我们的任务插件模块由两个分支组成:任务提示分支(Task-Prompt Branch,TPB)和空间补充分支(Spatial Complement Branch,SCB)。任务提示分支对于为预训练的扩散模型提供任务特定的指导至关重要,类似于在文本条件图像合成中使用文本提示[54]。我们利用通过预训练的CLIP视觉编码器[49]提取的视觉提示,来引导模型关注与任务相关的模式(例如,去雨任务中的雨痕和去雪任务中的雪花)。具体来说,对于输入图像 I I I,编码器 E n c I ( ⋅ ) E n c_{I}(\cdot) EncI(⋅)首先提取一般视觉特征,然后由任务提示分支提炼这些特征,以产生判别性视觉引导先验 F p \mathbf{F}^{p} Fp:
F p = TPB ( Enc I ( I ) ) \mathbf{F}^{p}=\operatorname{TPB}\left(\operatorname{Enc}_{I}(\mathbf{I})\right) Fp=TPB(EncI(I))
其中,任务提示分支包含三个带层归一化和LeakyReLU激活(除最后一层外)的多层感知机(MLP)层,确保仅保留最任务特定的属性。这种方法使 F p \mathbf{F}^{p} Fp与扩散模型在文本驱动生成过程中通常使用的文本特征对齐,从而为插件选择器更好地进行任务对齐。此外,使用视觉提示简化了用户角色,无需进行复杂的文本提示工程,这对于特定视觉任务来说通常是具有挑战性的,并且对文本的小幅变化敏感[78]。
然而,尽管任务特定的视觉引导先验 F p \mathbf{F}^{p} Fp对于提示全局语义属性至关重要,但它不足以保留细粒度细节。在这方面,DDIM反演通过提供包含图像内容信息的初始噪声发挥着关键作用。没有这一步,推理将依赖于不包含图像内容的随机噪声,导致扩散过程中的结果可控性较差。然而,反演过程是不稳定的且耗时。为了缓解这个问题,我们引入空间补充分支来有效提取和增强空间细节保留。我们利用预训练的VAE编码器[11] E n c V ( ⋅ ) E n c_{V}(\cdot) EncV(⋅)来捕获输入图像 I \mathbf{I} I的全部内容,表示为 F \mathbf{F} F。这个全面的图像细节,与来自 F p \mathbf{F}^{p} Fp的语义引导相结合,然后由我们的空间补充分支处理,以提炼空间特征 F s \mathbf{F}^{s} Fs:
F s = SCB ( F , F t , F p ) = Att ( Res ( F , F t ) , F t , F p ) \mathbf{F}^{s}=\operatorname{SCB}\left(\mathbf{F}, \mathbf{F}^{t}, \mathbf{F}^{p}\right)=\operatorname{Att}\left(\operatorname{Res}\left(\mathbf{F}, \mathbf{F}^{t}\right), \mathbf{F}^{t}, \mathbf{F}^{p}\right) Fs=SCB(F,Ft,Fp)=Att(Res(F,Ft),Ft,Fp)
其中, F t \mathbf{F}^{t} Ft是时间嵌入,用于表示扩散过程中变化的时间步。Res和Att块分别表示来自扩散模型[54]的标准ResNet和交叉注意力Transformer块。Res的输出用作查询特征,而 F p \mathbf{F}^{p} Fp在交叉注意力层中既作为键特征也作为值特征。
然后,我们将任务特定的视觉引导先验 F p \mathbf{F}^{p} Fp引入扩散模型的交叉注意力层,以指导模型的生成过程符合低级视觉任务的具体要求。之后,我们直接将提炼的空间先验 F s \mathbf{F}^{s} Fs作为残差融入解码器的最后阶段。这一位置的选择基于我们在表4中的实验观察,该观察表明,在稳定扩散[54]中,空间细节的保真度往往从浅层到深层降低。通过在这一特定阶段添加 F s \mathbf{F}^{s} Fs,我们有效地抵消了这一趋势,从而增强了细粒度空间细节的保留。
为了训练任务插件模块,我们采用[54]中定义的降噪损失,将任务特定的先验引入扩散降噪训练过程:
L = E z 0 , t , F p , F s , ϵ ∼ N ( 0 , 1 ) [ ∥ ϵ − ϵ θ ( z t , t , F p , F s ) ∥ 2 2 ] \mathcal{L}=\mathbb{E}_{\boldsymbol{z}_{0}, t, \mathbf{F}^{p}, \mathbf{F}^{s}, \epsilon \sim \mathcal{N}(0,1)}\left[\left\|\epsilon-\epsilon_{\theta}\left(\boldsymbol{z}_{t}, t, \mathbf{F}^{p}, \mathbf{F}^{s}\right)\right\|_{2}^{2}\right] L=Ez0,t,Fp,Fs,ϵ∼N(0,1)[∥ϵ−ϵθ(zt,t,Fp,Fs)∥22]
其中, z t = α ˉ t z 0 + 1 − α ˉ t ϵ t \mathbf{z}_{t}=\sqrt{\bar{\alpha}_{t}} \mathbf{z}_{0}+\sqrt{1-\bar{\alpha}_{t}} \epsilon_{t} zt=αˉtz0+1−αˉtϵt表示时间 t t t时潜在空间图像的噪声版本,而 z 0 \mathbf{z}_{0} z0是真实图像 I ^ \hat{\mathbf{I}} I^的潜在空间表示,通过 z 0 = E n c V ( I ^ ) \mathbf{z}_{0}=E n c_{V}(\hat{\mathbf{I}}) z0=EncV(I^)获得。该损失函数确保任务插件得到有效训练,以在引导扩散过程中融入任务特定的先验。
3.4. 插件选择器
我们提出了插件选择器(Plugin-Selector),使用户能够通过文本输入选择所需的任务插件(Task-Plugin)。对于输入图像 I I I和文本提示 T \mathbf{T} T,我们将任务插件集合定义为 P = { P 1 , P 2 , ⋯ , P m } \mathcal{P}=\left\{\mathcal{P}_{1}, \mathcal{P}_{2}, \cdots, \mathcal{P}_{m}\right\} P={P1,P2,⋯,Pm},其中每个 P i \mathcal{P}_{i} Pi对应一个特定的视觉任务,将 I I I转换为任务特定的先验 ( F i p , F i s ) \left(\mathbf{F}_{i}^{p}, \mathbf{F}_{i}^{s}\right) (Fip,Fis)。然后,每个任务插件的视觉引导 F i p \mathbf{F}_{i}^{p} Fip通过共享的视觉投影头 V P ( ⋅ ) VP(\cdot) VP(⋅)映射到一个新的文本-视觉对齐的多模态潜在空间,并表示为 V = { v 1 , v 2 , ⋯ , v m } \mathcal{V}=\left\{\boldsymbol{v}_{1}, \boldsymbol{v}_{2}, \cdots, \boldsymbol{v}_{m}\right\} V={v1,v2,⋯,vm}。同时, T \mathbf{T} T通过 E n c T ( ⋅ ) [ 49 ] Enc_{T}(\cdot)[49] EncT(⋅)[49]编码为文本嵌入,然后通过文本投影头 T P ( ⋅ ) TP(\cdot) TP(⋅)投影到 q q q,使文本嵌入和视觉嵌入对齐。该过程表示为:
v i = V P ( F i p ) \boldsymbol{v}_{i}=VP\left(\mathbf{F}_{i}^{p}\right) vi=VP(Fip); q = T P ( E n c T ( T ) ) \quad \boldsymbol{q}=TP\left(Enc_{T}(\mathbf{T})\right) q=TP(EncT(T))
然后,我们使用余弦相似度函数比较文本嵌入 q \boldsymbol{q} q与每个视觉嵌入 v i ∈ V \boldsymbol{v}_{i} \in \mathcal{V} vi∈V,使得 s i = s i m ( v i , q ) s_{i}=sim\left(\boldsymbol{v}_{i}, \boldsymbol{q}\right) si=sim(vi,q),得到一组相似度分数 S = { s 1 , s 2 , ⋯ , s m } \mathcal{S}=\left\{\boldsymbol{s}_{1}, \boldsymbol{s}_{2}, \cdots, \boldsymbol{s}_{m}\right\} S={s1,s2,⋯,sm}。我们选择满足指定相似度阈值 θ \theta θ的任务插件 P s e l e c t e d \mathcal{P}_{selected} Pselected:
P s e l e c t e d = { P i ∣ s i ≥ θ , P i ∈ P } \mathcal{P}_{selected}=\left\{\mathcal{P}_{i} \mid s_{i} \geq \theta, \mathcal{P}_{i} \in \mathcal{P}\right\} Pselected={Pi∣si≥θ,Pi∈P}
我们采用 F i p \mathbf{F}_{i}^{p} Fip作为伪标签,并将其与任务特定的文本配对以构建训练数据。我们采用对比损失[5,49]来优化视觉和文本投影头,提高它们处理多任务场景的能力。这涉及最小化锚图像和正文本之间的距离,同时增加与负文本之间的距离。对于每个图像 I I I,采样一个与其任务相关的正文本(例如,对于去雨任务,“我想去除雨水”)和来自其他任务的 N N N个负文本(例如,对于人脸修复,“增强人脸”)。对于正例对 ( i , j ) (i, j) (i,j),损失函数如下:
ℓ i , j = − log exp ( s i m ( v i , q j ) / τ ) ∑ k = 1 N + 1 1 [ k c ≠ i c ] exp ( s i m ( v i , q k ) / τ ) \ell_{i, j}=-\log \frac{\exp \left(sim\left(\boldsymbol{v}_{i}, \boldsymbol{q}_{j}\right) / \tau\right)}{\sum_{k=1}^{N+1} \mathbb{1}_{\left[k_{c} \neq i_{c}\right]} \exp \left(sim\left(\boldsymbol{v}_{i}, \boldsymbol{q}_{k}\right) / \tau\right)} ℓi,j=−log∑k=1N+11[kc=ic]exp(sim(vi,qk)/τ)exp(sim(vi,qj)/τ)
其中, c c c表示每个样本的任务类型, 1 [ k c ≠ i c ] ∈ { 0 , 1 } \mathbb{1}_{\left[k_{c} \neq i_{c}\right]} \in\{0,1\} 1[kc=ic]∈{0,1}是指示函数,当且仅当 k c ≠ i c k_{c} \neq i_{c} kc=ic时取值为1。 τ \tau τ表示温度参数。
4. 实验
在本节中,我们首先介绍实验设置,包括数据集、实现方式和评估指标。然后,在4.1节中,我们将Diff-Plugin与当前基于扩散和回归的方法进行比较,并在4.2节中通过消融研究对Diff-Plugin进行组件分析。
数据集。为了训练Task-Plugins,我们为每个低级任务使用特定的数据集:去雪:Snow100K [36],去雾:Reside [32],去模糊:Gopro [41],去雨:合并的训练集 [86],人脸修复:FFHQ [26],低光增强:LOL [72],去摩尔纹:LCDMoire [85],以及去除高光:SHIQ [13]。在测试时,我们在真实世界的基准数据集上进行评估:去雪:realistic test [36],去雾:RTTS [32],去模糊:RealBlur-J [53],去雨:real test [68],人脸修复:LFW [23, 69],低光增强:合并的低光数据集 [17, 30, 38, 64, 67, 72],去摩尔纹:LCDMoire [85],以及去除高光:SHIQ [13]。为了训练Plugin-Selector,我们使用GPT [44]为每个任务生成文本提示。
实现方式。在训练和测试期间,我们将图像调整为 512 × 512 512 \times 512 512×512大小以进行公平比较。我们采用AdamW优化器 [37]及其默认参数(例如,betas、权重衰减)。我们的Task-Plugins在四个A100 GPU上进行训练,每个GPU具有80G内存,使用恒定学习率 1 e − 5 1e^{-5} 1e−5和批量大小64。为了训练Plugin-Selector,我们从每个任务中随机采样5000张图像,并通过随机组合来自不同任务的文本输入来增加文本多样性。我们将批量大小设置为8,并采用与Task-Plugins相同的学习率。对于负文本,我们默认设置 N = 7 N=7 N=7。在推理期间,我们将指定的相似度阈值 θ \theta θ设置为0。
评估指标。我们遵循[54],采用广泛使用的无参考感知评估指标FID [20]和KID [3],在真实数据上评估我们的Diff-Plugin,因为真实数据的GT(真实标签)并不总是可用的。至于Plugin-Selector,我们遵循多标签目标分类[6]来报告平均精度均值(mAP)、每类平均精度(CP)、F1(CF1),以及整体平均精度(OP)、召回率(OR)和F1(OF1)。对于每个类别(即任务类型),如果其置信度分数大于 θ \theta θ,则将其标签预测为正。我们进一步提出了一个严格的零容忍评估指标(ZTA),从用户优先的角度严格评估句子级别的分类结果,进行二分类以确保最高准确性:
Z T A = 1 Q ∑ i = 1 Q ( ( min j ∈ Y i S i j > θ ) ∧ ( max k ∈ H i S i k ≤ θ ) ) ZTA = \frac{1}{Q} \sum_{i=1}^{Q}\left(\left(\min _{j \in Y_{i}} S_{i j}>\theta\right) \wedge\left(\max _{k \in H_{i}} S_{i k} \leq \theta\right)\right) ZTA=Q1∑i=1Q((minj∈YiSij>θ)∧(maxk∈HiSik≤θ))
其中, Q Q Q是测试样本的总数, S i S_{i} Si是样本 i i i的预测相似度分数集合, Y i Y_{i} Yi是正类别(即用户感兴趣的任务)的索引集合, H i H_{i} Hi是负类别(即不相关的任务)的索引集合。
4.1. 与最先进方法的比较
我们将提出的Diff-Plugin与来自不同低级视觉任务的最先进方法进行了比较,包括基于回归的专用模型:DDMSNet [88]、PMNet [81]、Restormer [87]、NeRCO [79]、VQFR [15]、UHDM [84]、SHIQ [13],多任务模型:AirNet [33]、WGWS-Net [98]和PromptIR [47],以及基于扩散的模型:SD [54]、PNP [63]、P2P [46]、InstructP2P [4]、NullText [39]和ControlNet [90]。我们在真实世界数据集上进行实验,以比较泛化能力。
定性结果。图5展示了我们的Diff-Plugin在八个低级视觉任务上具有挑战性的自然图像上的优越性能。首先,使用SD的img2img [54]功能无法保证内容准确性,经常导致主要场景变化(第8列)。InstructP2P [4]由于缺乏任务特定的先验,也表现不佳,在去雾和低光增强等任务中产生了较差的结果(第7列)。由于缺乏任务特定的先验,P2P [46]和Null-Text [39]尽管使用了来自DDIM Inversion [61]的初始噪声,但生成的内容也不一致(第5列和第6列)。ControlNet [90]通过扩散分支提供条件信息来处理一些任务(第4列),但其强烈的颜色失真降低了在这些任务中的有效性。最新的多任务方法PromptIR [47](第3列)受到模型规模的限制,只能处理少数任务。相比之下,我们的方法为每个任务使用轻量级的任务特定插件,在所有任务中提供灵活性和稳定性能(第2列)。
定量结果。我们还在表1中提供了定量比较。与基于扩散的方法相比,我们的Diff-Plugin总体上取得了最优结果。虽然PNP [63]和InstructP2P [4]能够生成具有低FID和KID的高质量图像,但它们经常产生显著的内容更改(参考图5)。与基于回归的多任务方法相比,我们的方法在大多数任务中表现出竞争力,尽管在去摩尔纹和去除高光等稀疏退化任务中效果稍差。尽管专用模型在其各自领域可能优于我们的方法,但其任务依赖的设计限制了它们在其他任务中的适用性。请注意,本文的主要目标不是在所有任务中取得最佳性能,而是为未来的进步奠定基础。此外,Diff-Plugin还实现了文本驱动的低级任务处理,这是基于回归的模型所不具备的能力。
用户研究。我们进行了一项有46名参与者参与的用户研究,通过主观评估来评估各种方法。每位参与者审查了来自测试集的5组图像,每组包括一张输入图像和10张预测图像,共涉及8个任务。根据内容一致性、退化去除(例如,雨、雪、高光)和整体质量对图像进行排名。我们分析了1840个排名(46名参与者×40组),计算了每种方法的平均排名(AR)。表2显示了结果。显然,用户在很大程度上倾向于我们的方法。
4.2. 消融研究
任务插件。我们首先通过探索各种消融设计并比较它们在去雪和去雾任务上的性能来评估任务插件的有效性。除非另有说明,否则在推理过程中使用随机噪声。我们有五个消融模型。(1)反演+编辑:使用任务特定描述(例如,“一张下雪天的照片”)的DDIM反演将输入图像反演为初始噪声,同时保留内容。然后使用目标描述(例如,“一张晴天的照片”)进行编辑。(2)TPB:移除SCB,仅专注于TPB训练。(3)TPB+反演:仅训练TPB,但在推理过程中使用DDIM反演作为初始噪声。(4)SCB:移除TPB,仅训练SCB。(5)TPB+SCB(重建):训练首先从使用自重建去噪损失的SCB开始,然后使用固定的SCB进行TPB训练。性能结果和比较如图6和表3所示。
我们有以下观察结果。(1)反演+编辑能够捕获输入图像的全局结构,但会丢失详细内容。(2)TPB提供任务特定的视觉指导,但由于其仅关注高级特征,因此缺乏空间内容约束。(3)使用反演初始噪声的TPB在结构化场景(例如,大型建筑)中表现优异,但往往会加深颜色并为较小物体创建随机内容。(4)SCB保留内容细节,但由于缺乏任务特定的视觉指导,难以有效去除退化(例如,雪或雾)。(5)当TPB与基于重建的SCB结合时,它通过重建保留图像内容,同时仅依赖TPB来处理退化。然而,由于SCB在每个扩散迭代中重新引入所有图像特征,包括原始退化(例如,图6第2行的雾),因此它无意中损害了期望的结果。最后,在我们的任务插件中结合来自TPB和SCB的任务特定先验,能够实现高保真度的低级任务处理。
我们还确认了去雪任务中SPB在预训练SD模型中的位置,并在表4中展示了结果。显然,我们可以观察到,对于预训练SD[54]的编码器和解码器,保真度降低,性能从较浅阶段到较深阶段(例如,阶段1到4)逐渐下降。因此,我们将空间特征注入到解码器的最后阶段,以平衡性能和参数。值得注意的是,任务插件模块的参数仅为SD的 1.67 % 1.67\% 1.67%。
插件选择器。如表5所示,我们首先评估插件选择器在单任务和多任务场景(第1行和第2行)中的准确性,并观察到其准确性一直很高。此外,在包含120,000个样本的显著扩展测试(标记为“多任务*”)中,它实现了 0.936 0.936 0.936的mAP准确率,证明了其有效性。此外,在结合任务特定和任务无关文本的鲁棒性测试(标记为“单任务+非任务”)中,它仍然实现了显著的零容忍准确率 0.779 0.779 0.779。
我们还对插件选择器进行了消融研究,以评估每个组件的重要性,结果详见表6。(1)我们分别移除了视觉和文本投影头。(2)我们评估了对比训练中负样本数量变化的影响。结果首先表明,视觉和文本投影头都至关重要。省略视觉头会导致训练崩溃并输出NaN,而移除文本头会使ZTA指标降低 15.4 % 15.4\% 15.4%。它还表明,增加负样本的数量(例如,从 N = 1 N=1 N=1到15)可以一致地提高选择准确性。
多样应用。图7展示了Diff-Plugin的多功能性。第1行通过子任务集成(例如,旧照片修复可以大致分为修复和着色)展示了复杂低级任务的执行。第2行突出了其反转低级任务的能力,从而能够生成如雨和雪等特效。
- 结论
在本文中,我们提出了一种名为Diff-Plugin的新型框架,该框架专为增强预训练扩散模型在处理各种需要严格细节保留的低级视觉任务方面的能力而设计。我们的Task-Plugin模块采用双分支设计,有效地将任务特定的先验知识融入到扩散过程中,从而能够在不针对每个任务重新训练基础模型的情况下,生成高保真且细节保留丰富的视觉结果。Plugin-Selector通过文本输入进一步增加了直观的用户交互,实现了文本驱动的低级任务,并增强了框架的实用性。在各种视觉任务上的广泛实验表明,我们的框架相较于现有方法具有优越性,尤其是在真实场景中。
我们当前Diff-Plugin框架的一个局限性是无法进行局部编辑。例如,在图1中,我们的方法可能无法仅去除河面上的雪而保留天空中的雪。解决这一问题的一个可能方案是整合大型语言模型(LLMs)[35, 97]来指示执行任务的区域。