Learning Dynamic Prompts for All-in-One Image Restoration

Abstract

引言

例如，AirNet [19] 采用自监督预训练来学习退化表示。PromptIR [20] 引入了一种自适应提示框架，其中提示直接进行参数化，并与恢复模型联合优化。这种端到端的方法利用模型自身的优化来捕捉数据和任务先验，提供了一种高效的解决方案。PromptIR 取得了显著的性能提升，并启发了众多后续研究 [22]、[23]。但是，这些提示学习缺乏明确的约束，可能导致表达和判别能力有限。端到端学习过程受数据集的影响很大，而提示学习缺乏约束可能会导致不同退化类型之间的混淆，无法实现任务感知的提示表示。
如图1所示，我们绘制了PromptIR中的提示集成权重。可以发现，在多个退化任务中分配的权重几乎是均匀的，这表明缺乏有区分性的、任务感知的表示。此外，在推理阶段，提示成为固定参数，并且在推理过程中这些固定提示的线性组合（例如，通过softmax）本质上限制了表示能力。这种限制显著地限制了目标提示表示的范围，阻碍了模型的能力。
后续的研究，如《MiOIR》[21]，进一步分析了提示词的影响，并将分类器特征表示用作提示词。通过利用从预训练的退化分类器中提取的特征，它们提供了更准确且与任务相关的表示，相较于自适应提示取得了显著改进。然而，这些方法依赖于使用有标记的退化数据进行有监督的预训练，这引入了额外的限制。
同时，基于语义提示的最新方法引入了文本信息，并利用了预训练的多模态模型[24]-[26]。这些方法通过利用丰富的先验知识实现了准确且高效的恢复。然而，由于使用了大规模的预训练模型，它们通常会带来巨大的计算开销。显然，具有判别性且丰富的退化表示能显著提升图像恢复模型的性能。在本文中，我们专注于基于自适应提示的方案，旨在在保持其简单性的同时增强其判别能力和表现力。一个自然的问题出现了：我们如何获得更具判别性且更丰富的自适应提示？
为应对这些挑战，我们提出了一种新颖的方法，该方法引入了具有判别性、可扩展性和动态性的任务提示，用于一体化图像修复。

我们的方法引入了两个关键组件：Degradation prototype assignment（DPA）和Prompt Distribution Learning （PDL）。首先，为了增强提示的判别能力，我们引入了退化原型分配。我们利用原型学习为不同的退化类型预先定义具有判别性的原型表示，如图2所示。通过将自适应提示学习转化为将提示表示与原型匹配的过程，我们获得了更具判别性的退化表示。这种分配过程有效地将退化表示与预定义的原型对齐，有助于在无需监督预训练的情况下更好地分离不同的退化类型。其次，为了克服静态提示表示空间有限的问题，我们提出了提示分布学习。我们不再学习固定的提示参数，而是将提示学习过程建模为在多个退化任务上学习一种分布，如图3所示。通过从这种学习到的提示分布中采样，我们可以生成丰富多样的提示表示，极大地丰富了提示表示空间。通过整合退化原型分配（DPA）和提示分布学习（PDL），我们的方法有效地解决了现有方法的局限性，为一体化图像恢复提供了一个更灵活、更强大的框架。此外，我们将提出的DPPD框架应用于重新训练PromptIR，在3任务和5任务设置下都取得了显著的改进（如图4所示）。具体而言，配备我们DPPD模块的PromptIR在三个和五个退化任务上分别获得了0.53 dB和0.94 dB的提升，证明了DPPD在一体化图像恢复中的有效性。

贡献：

• 该方法将退化表示分配给预定义的判别原型。这种方法无需监督预训练即可提供准确且可扩展的退化表示，有效增强了提示的判别能力。
• 我们引入了提示分布学习（PDL），它将提示建模为分布而非固定参数。
• 我们将退化原型分配（DPA）和提示分布学习（PDL）整合到一个统一的框架中，称为退化原型分配与提示分布学习（DPPD），这代表了一种全新的一体化图像恢复提示学习范式。DPPD作为一个即插即用的组件，与现有的图像恢复主干网络兼容。

方法

在本节中，我们详细介绍了用于一体化图像恢复的退化原型分配与提示分布学习（DPPD）框架。如图6所示，我们的方法基于Restormer [12]和PromptIR [20]等现有架构构建，集成了一个新颖的提示模块，这是我们方法的核心。它包括所提出的DPPD模块和一个带有残差连接的交叉融合模块。核心思想是通过有效地学习和利用有区别性的退化提示，增强模型处理多样且复杂退化情况的能力。接下来，我们详细介绍我们方法的各个组成部分，包括退化原型分配和提示分布学习。

Prototype Acquistion

$z_i\in R^d$代表退化图像的特征表示，其中d为潜在维度。我们引入了一组预定义退化原型 $p_j\in R^d$.我们的目标是给每个退化表征$z_i$分配对应的原型 $p_j$.
为了获得具有判别性的退化信息，遵循原型学习[50]，将原型设计为满足：

我们采用满足公式（1）的退化原型的简单初始化方法。通过将退化表示锚定到这些固定的原型上，我们将多任务学习问题转化为原型匹配任务，从而实现更具判别性的退化表示学习。

Degradation prototype assignment

给定一个输入的退化图像x，我们从编码器的中间层提取其特征表示F∈R H×W×C。我们执行全局平均池化，以获得一个紧凑的特征向量f∈R C。将f通过全连接层获得退化嵌入$z\in R^d$,接下来我们计算z和每个prototype $p_k$之间的余弦相似性：

我们使用这些相似性来计算分配权重${\alpha }_k$通过softmax函数：

任务权重反映了输入退化与每个原型的匹配程度，实际上是后续提示生成的一个条件。

Prompt Distribution Learning

基于DPA模块的分配权重，我们将提示建模为条件高斯混合模型。我们提议不再使用固定提示，而是将每个提示建模为一种分布，以实现动态且自适应的提示生成。具体来说，提示P的建模方式如下：

其中µk和Σk分别是第k个高斯分量的均值和协方差，代表与原型$p_k$相关的即时分布。这些参数是可学习的，并且能够捕捉每种退化类型中的变异性。遵循文献[51]，我们使用重参数化技巧对动态提示进行采样：

通过集成DPA和PDL模块，我们的方法能够根据每个输入图像的退化特征，有效地生成自适应提示。DPA模块通过将输入特征分配给原型来提供一种有区分性的退化表示，而PDL模块则利用这些信息，通过条件高斯混合模型生成动态提示。

Reamrk

从变分推断的角度来看，我们的提示分布学习可以表述为优化证据下界（ELBO）[51]。我们考虑提示上的先验分布p§，将其建模为标准高斯分布N(0, I)。条件混合分布p(P | zˆ)作为近似后验。根据变分推断的原理，我们最小化学习到的分布与先验之间的KL散度：

这个正则化项鼓励模型保持多样化的提示分布，同时防止它们偏离先验分布太远。