MiOIR | 直面 “多合一”图像复原，港理工张磊团队提出MiOIR，融顺序学习与提示学习于一体！

本文首发: AIWalker

https://arxiv.org/abs/2401.03379
https://github.com/Xiangtaokong/MiOIR

尽管单任务图像复原已取得了极大成功，但单模型处理多复原任务仍极具挑战。本文对 “多合一”图像复原任务进行了深入探索并指出其所面临的两个关键挑战并提出了对应的应对方案：

• (1) 多目标优化问题：提出了顺序学习策略而非混合学习策略，通过增量学习机制引导网络学习；
• (2) 多任务自适应性：提出了提示学习策略解决不同复原任务的自适应性，辅助网络理解特定的任务并提升泛化能力。

研究结果表明，采用所提的顺序学习和提示学习策略可以显著提高CNN和Transformer模型在多合一复原任务上的性能。此外，这两种学习策略还能够相互促进，增强模型的鲁棒性。

“多合一”图像复原范式

“多合一”图像复原(Multiple-in-one Image Restoration, MiOIR)旨在通过一个模型处理多个图像复原问题，即不同任务的输入图像具有不同类型的画质退化问题。我们采用 { X t } t ∈ [ T ] \{X^t\}_{t \in [T]} {Xt}t∈[T]​表示MiOIR的输入，T表示任务数量， { X t } \{X^t\} {Xt}表示第t个任务；MiOIR对应的真值表示为 { Y } \{Y\} {Y}。数据样例可以表示为 { x n 1 , ⋯   , x n T , y n } n ∈ [ N ] \{x_n^1, \cdots, x_n^T, y_n\}_{n \in [N]} {xn1​,⋯,xnT​,yn​}n∈[N]​，N表示样本数量， { x n t } t ∈ [ T ] , y n \{x_n^t\}_{t\in [T]}, y_n {xnt​}t∈[T]​,yn​分别表示第t个任务第n个样例输入与输出。需要注意的是：不同任务的输入图像 { x n 1 ∼ T } \{x_n^{1 \sim T}\} {xn1∼T​}具有相同的高质量真值图像$y_n$。

MiOIR可以描述为: F ( { X t } ; θ ) : X t → Y F(\{X^t\}; \theta) : X^t \rightarrow Y F({Xt};θ):Xt→Y。它可以通过最小化如下损失进行学习，学习策略类似上图。

$$\begin{gathered} mi{n}_{\theta }\sum _{t=1}^T\frac{1}{T}{L}^t(\theta ) \\ {L}^t(\theta )=\frac{1}{N}\sum _{i=1}^N{L}^t(F(x_i^t;\theta ),y_i) \end{gathered}$$

但是，MiOIR模型训练存在两个关键性挑战：

• (1) 模型优化：不同的图像复原任务具有不同的退化问题，这无疑会导致严重的训练冲突。训练曲线会随不同的输入而剧烈波动，进而陷入较差的局部最小值。
• (2) 任务自适应性：MiOIR期盼模型能识别不同的退化类型并进行对应的复原能力。也就是说，MiOIR需要能以高精度适配不同的复原任务。

本文方案

针对于上述问题，本文提出了两种优化策略：顺序学习策略与提示学习策略。接下来，我们针对这两种策略进行逐步解析。

顺序学习

顺序学习策略是针对“多合一”复原任务的多样性目标优化难问题而设计。如前述所提，在MiOIR模型训练过程中，所有训练数据均是可获取的的，关键是如何找一个更优的学习策略。

最直接的策略就是将所有训练数据混合后直接进行训练。但是，已有预训练相关研究表明：非相关图像复原任务预训练都可以为其他图像复原任务提供一个足够好的起点。基于此，如果我们让模型现在某些任务上学习，这些前置任务可以视作预训练任务，因此可以为后续其他任务的训练提供一个好的起点。

有很多种任务拆分方式，作为探索，我们采用了最简单的方法，如左图所示。顺序学习策略通过增量方式(即逐个添加复原任务)对MiOIR进行训练学习。经实验论证，有很多学习顺序都可以导致性能改善。

关于顺序学习顺序，作者建议：高频重建类任务(如超分、去模糊等)在前，全局亮度调整类任务(如去雾、低光增强)在后。

提示学习策略

针对任务自适应挑战，作者提出了提示学习策略。合适的提示有助于网络更好的理解图像复原任务，进而针对性的进行图像复原处理。作者提出了两种典型的提示学习策略：显式提示学习(右图上)与自适应提示学习(右图下)。

作者采用三个CNN层(表示为$F_{ext}(\cdot )$)进行提示信息P的提取，然后通过全连接层将其转换为合适形状的调制尺度s与调试偏置b参数，最后通过如下方式对复原阶段的特征进行调制处理：

$$\begin{array}{rl}{s}_m,b_m& =F{C}_m(F_{ext}(P))\\ f_m^{prompt}& =f_m\ast s_m+b_m\end{array}$$

需要注意：两种提示学习方式的区别仅在于如何得到提示信息，后续复原阶段的提示注入方式是相同的。

• 显示提示学习：作者采用额外的输入作为提示以告知网络需要处理的任务类型。在实际执行阶段，固定提示参数并将其与不同复原任务相绑定，然后训练一个分类器根据对应复原任务选择提示信息(域内测试集准确率高达99.7%)。考虑到显示学习策略的难度并不高，该方案应当具有较好的域内性能。
• 自适应提示学习：在面对真实低质图像时，有时用户可能无法准确判断其类型，导致显示提示学习策略失效。针对这种情况，自适应提示学习策略可作为补救措施。它可以自动地从输入图像中提取任务相关的视觉提示信息，从而有效地指导用户完成任务。尽管自适应学习策略的训练难度较高，但由于其与输入图像的相关性，该方案在跨域泛化性方面具有很高的潜力。

本文实验

学习策略有效性验证

如上表所示，在域内测试集上，相比于混合学习策略

• 顺序学习平均可以提升SRResNet指标达0.29dB、提升SwinIR指标达0.85dB；
• 显示提示学习对两个骨干的指标提升平均达0.7dB；
• 自适应提示学习对于SRResNet仅提升0.08dB，甚至导致SwinIR指标降低0.6dB。这是因为该学习策略更难以训练，同样解释了为什么已有自适应提示学习需要额外的约束信息。
• 通过组合两种学习策略，模型性能取得进一步的巨大提升，即两者具有互补性。比如，SwinIR-S+EP与SwinIR-S+AP指标提升分别达1.21dB与0.95dB。
• 从右图可以看到：所提方案重建结果视觉效果更优！

作者进一步在域外与未知退化测试集上验证了所提方案的泛化性，结果见上表，可以看到：

• 由于域外具有与域内相近的退化类型，显示提示学习策略表现优于自适应提示学习；
• 在未知退化测试集上，自适应提示学习因自身更优的泛化性取得了更佳的PSNR指标。

前面已经验证了所提方案的有效性，作者进一步验证该方案是否可以提升当前SOTA方案PromptIR的性能，结果见上表。

• 在域内/域外测试集上，顺序学习可以提升PromptIR指标约0.75dB；
• 当同时使用顺序学习与显式提示学习时，域内/域外测试集提升超1dB，未知退化测试集指标提升0.17dB。
• 此外值得一提的是：PromptIR+显式提示仅需26.7M参数量，仅为原始方案的75%。

作者进一步验证了所提方案在其他骨干下的性能对比，见上表，可以看到：

• Restormer与Uformer具有与SRResNet和SwinIR相似的行为表现；
• 顺序学习与提示学习可以改善模型在几乎所有测试集上性能，且两者可以相互促进。

最后，作者从可解释性维度(DDR、CHI)对所提方案进行了分析，可参考左图示意。

对比左图a与b可以看到：

• SRResNet-M+EP与SRResNet-S+EP的聚类信息非常清晰，7个任务可以得到很好的分离；
• 从CHI角度看，顺序学习可以进一步改善混合学习的聚类性能。

对比左图c与d可以看到：

• 尽管顺序学习具有更高的CHI指标，但自适应提示学习的聚类性能明显弱于显示提示学习，这是以为自适应提示学习更难以训练。
• 尽管如此，自适应提示学习可以赋予模型更优的泛化性能，可参考Tab4的客观指标。

个人反思

从最近几年研究来看，底层视觉领域的研究方向已经开始朝向“工业应用”偏移。针对真实场景下的图像复原问题，我们需要关注如何有效解决复杂退化场景中的图像超分、人脸复原以及面向传感器特性的图像降噪等问题。在应用层面上，这些场景可以分为限定场景下的图像复原和通用场景下的图像复原，而“多合一”图像复原是通用场景图像复原在实际应用中必须克服的难点。

事实上，围绕提示学习(任务自适应性)的研究还是挺多的，比如SRMD、UDVD、DASR、MRDA，再比如Diffusion时代的各种提示学习下的图像复原/超分方案等等。但是，多任务学习在底层视觉领域受到的关注貌似并不多。结合上述信息，这篇论文的关键其实是多任务学习中的“多目标训练冲突”问题，在引入提示学习/条件学习也就是非常自然的事情了。

多任务学习最简单而直接方法肯定是混合学习，但受限于种种因素容易陷入局部最优，甚至偶尔会出现训练不收敛问题。增量渐进式学习策略可以极大程度上避免训练不收敛问题，但恐仍无法避免局部最优问题。此外，在实际应用中，可能还需要考虑不同任务的不平衡问题，而这好像在底层视觉领域尚未得到应有的重视~

从某种程度上讲，这篇论文为真实场景图像复原引入了一个新的观察视角，在迈向通用复原大模型的路上又前进了一步~