【论文阅读】开集图像恢复的测试时间退化自适应Test-Time Degradation Adaptation for Open-Set Image Restoration

Test-Time Degradation Adaptation for Open-Set Image Restoration

---

原文链接：https://arxiv.org/abs/2312.02197
项目链接：https://github.com/XLearning-SCU/2024-ICML-TAO.

摘要

与在预训练阶段就已定义好的退化集（close-set scenarios）中恢复图像不同，开放集图像恢复（open-set image restoration）旨在处理在预训练阶段未预见的未知退化。据我们所知，这方面的研究还相对较少。本研究探讨了这一挑战性问题，并揭示了其本质是测试和训练数据之间未识别的分布偏移。最近，测试时适应（test-time adaptation）已成为解决这种固有差异的基本方法。受此启发，我们提出了一个用于开放集图像恢复的测试时退化适应框架，该框架包括三个组成部分：i) 一个预训练且不依赖于特定退化类型的扩散模型，用于生成干净图像；ii) 一个测试时退化适配器（test-time degradation adapter），它在测试阶段基于输入图像适应未知退化；iii) 适配器引导的图像恢复（adapter-guided image restoration），通过适配器指导模型产生相应的干净图像。通过在多种退化上的实验，我们展示了我们的方法达到了与特定任务方法相当的甚至更好的性能。

---

首先文中指出了图像恢复（IR）任务的差异。

具体而言，(a)经典IR在训练和测试退化相同且已知的近集场景中工作，并为每个场景定制专门的模型。(b)All-in-one也适用于训练和测试退化相同但未知的近距离场景，并通过统一模型解决这些问题。(c )Zero-shot侧重于从单个退化图像中恢复，该图像不受训练退化的影响，但通常需要事先了解测试退化。(d) OIR在测试退化未知且与预训练不同的开放集场景中工作。这类似于自然语言处理中的挑战，它将预训练好的大型语言模型应用于预训练阶段未预定义的各种下游任务。

介绍

大多数现有方法的一个共同限制是它们依赖于闭集场景（close-set scenario）中的假设，即测试退化与预训练阶段遇到的退化非常相似。这种假设限制了这些方法应用于特定类型的退化，阻碍了它们适应更广泛的现实世界场景，其中可能出现多样化和不可预见的退化。

为了打破上述假设，本研究深入探讨了开放集图像恢复（OIR）。OIR要求模型处理训练数据中不存在的未知退化。这种转变旨在推动图像恢复的界限，促进开发出在面对多样化和意外的退化时表现出适应性和弹性的模型。通过解决OIR问题，本研究旨在为现实世界场景中更复杂退化的解决方案铺平道路。

为了应对这一挑战，测试时适应（TTA）最近已成为解决测试和训练数据之间固有差异的有效方法（Wang et al., 2021）。具体来说，它根据测试数据的特定特征在测试阶段调整预训练模型，使模型能够在更广泛的输入场景上表现更好。这在OIR中特别有用，因为测试数据可能与训练数据不同，预训练模型应该适应未知和未见过的退化，以获得最佳性能。

测试时适应（Test-Time Adaptation, TTA）是一种在机器学习领域中用于提高模型泛化能力的技术。它的核心思想是在测试阶段，利用无标签的测试数据对预训练模型进行调整，以适应目标域的数据分布。这种方法特别适用于那些训练数据与测试数据分布不一致的情况，也就是所谓的分布偏移（distribution shift）问题。

TTA的主要优势包括：

1. 实时适应性：可以在测试阶段快速适应新的数据分布。
2. 无需重训练：不需要重新训练整个模型，节省了计算资源。
3. 隐私保护：不需要访问原始训练数据，保护了数据隐私。
4. 灵活性强：可以应用于各种机器学习任务和模型架构。

TTA的挑战与未来方向包括提高计算效率、保证适应的稳定性、保护隐私、提高可解释性以及建立更完善的理论基础。未来的研究可能会探索结合神经架构搜索、自监督学习、联邦学习等其他领域的技术来提高TTA的效果。

总的来说，TTA作为一种新兴的技术，为解决实际应用中的分布偏移问题提供了一种有效的解决方案，通过在测试阶段实时适应新的数据分布，显著提高了模型的泛化能力和鲁棒性。随着研究的不断深入，TTA有望在更多领域发挥重要作用。

基于此启发，我们提出了一个用于开放集图像恢复的测试时退化适应框架，简称为TAO，它利用TTA的思想为OIR提供了一个健壮的解决方案。具体来说，TAO由三个组成部分构成，即预训练的图像扩散模型（PDM）、测试时退化适配器（TDA）和适配器引导的图像恢复（AIR）。在这个框架中，由于以下考虑，PDM被采用作为OIR的基础模型。首先，PDM捕获了生成各种高质量视觉场景的丰富知识，可以被认为是针对产生干净图像的OIR的通用预训练。其次，PDM不依赖于特定退化，测试数据中的任何退化都被视为未预见的。请注意，PDM的测试阶段指的是图像生成过程，即逐步去噪过程。

在每个去噪步骤之后，TDA和AIR依次进行，分别用于适应开放集场景和指导图像恢复。具体来说，TDA在测试阶段采用了一个可学习的适配器来适应PDM在测试退化图像上的表现。这个适配器旨在进行域对齐，将PDM的生成域对齐到测试图像的退化域。通过这种方式，生成的干净图像可以被翻译成相应的退化图像，这可以进一步被监督以适应测试退化图像，并反过来更新生成的干净图像。AIR被详细设计来进行这种监督更新以实现图像恢复。简而言之，AIR的灵感来自于观察到PDM在逐步去噪过程中展现出的有趣时间动态（Choi et al., 2022）。因此，AIR在不同的去噪步骤中动态并相应地调整监督策略，以实现更好的性能。

总的来说，我们的贡献可以总结如下：

• 据我们所知，这可能是第一项明确研究OIR并发现其本质的工作，即测试和训练数据之间未识别的分布偏移。
• 这项工作揭示了TTA是OIR的有效方法，通过在测试阶段对测试数据上的预训练模型进行适应，以解决测试和训练数据之间的固有差异。
• 这项工作提出了一个用于OIR的测试时退化适应框架，其中两个组成部分被设计用于适应开放集场景和指导图像恢复。实验表明了该框架的有效性。

基本方法

开放集图像恢复（Open-set image restoration）旨在处理在预训练阶段未曾预见的未知退化。对于OIR，应遵循两个基本原则：(i) 未知：OIR在测试阶段对退化一无所知，任何关于特定退化的先验知识都是禁止的。(ii) 未预见：OIR独立于特定的预训练任务，只要求预训练中未出现测试退化。

在测试时适应（TTA）中，我们有一个在分布$D_{\text{train}}$上预训练的模型 $M$。TTA 的目标是在测试样本$y\sim D_{\text{test}}$上优化模型$M$的一些参数$\theta$，使得目标函数${\min }_{\theta }L(y;\theta )$最小化，其中$D_{\text{test}}$与$D_{\text{train}}$是不同的。这个过程遵循两个基本原则：

1. TTA 从一个已经学习了目标任务通用知识的预训练模型开始。
2. TTA 在测试数据上工作，训练对是不可用的。

在这项工作中，我们采用了扩散模型作为预训练模型，生成和退化的图像作为训练对来实现开放集图像恢复（OIR）任务的 TTA。

框架概述

为了通过预训练的扩散模型（PDM）实现开放集图像恢复（OIR），关键在于实现$p_{\varphi }(\cdot )$来生成适当的梯度指导$G$，引导生成的图像${\hat{x}}_0$朝向测试退化图像$y$的干净图像$x_0$。在本工作中，我们通过测试时退化适配器（TDA）和适配器引导的图像恢复（AIR）来实现$p_{\varphi }(\cdot )$，它们分别针对开放集场景和图像恢复进行了设计。为了克服开放集场景带来的挑战，TDA在测试阶段（即PDM的去噪过程）引入了一个不依赖于特定退化的适配器，用于适配PDM以适应未知和未预见的退化。该适配器由一个简单的神经网络$\varphi$实现，并且在每个去噪步骤后针对对齐${\hat{x}}_0$的域到$y$的域进行优化。通过适配器，${\hat{x}}_0$可以被翻译成相应的退化图像$\varphi ({\hat{x}}_0)$，这可以进一步被监督以适应$y$，进而产生$G$来更新${\hat{x}}_0$朝向期望的$x_0$。AIR被设计来执行这个监督更新过程。

现有研究表明，去噪过程展现出有趣的时间动态，即${\hat{x}}_0$从不可识别的内容生成到感知丰富的内容，最后是细节不可见的内容。因此，AIR根据去噪过程中的不同阶段动态调整监督策略，从高层次内容逐步转移到低层次内容。

TAO的框架如图所示，其中PDM被固定为逐步去噪纯噪声$x_T$到干净图像$x_0$。对于测试退化图像$y$，TAO首先在$t=T$时采样两幅随机噪声图，因此每个去噪步骤后将生成两幅图像。我们将更新朝向$y$的干净图像的图像称为引导图像$x_t^g$，另一个称为辅助图像$x_t^a$。我们引入$x_t^a$以防止TDA陷入平凡解，并且协助AIR引导$x_t^g$朝向期望的干净图像$x_0^g$。在每个步骤$t$，图像$x_t^g$、$x_t^a$和$y$首先被送入TDA，通过域对齐的目标优化适配器$\varphi$，然后送入AIR，通过方程（4）中的条件指导更新$x_t^g$朝向$x_0^g$，这最好地匹配了经过$\varphi$的$y$。

测试时退化适配器

测试时退化适配器（Test-time Degradation Adapter, TDA）的核心是解决开放集图像恢复（OIR）中未识别的分布偏移问题，而测试时适应（Test-time Adaptation, TTA）最近作为一种有效的方法来解决这种固有的差异。基于这些动机，我们提出了TDA，用于在测试阶段使预训练的扩散模型（Pretrained Diffusion Model, PDM）适应下游的退化分布。

TDA采用了一个四层卷积神经网络作为适配器$\varphi$，该适配器被优化以对齐生成的图像${\hat{x}}_g^{(0)}$和${\hat{x}}_a^{(0)}$的域到测试退化图像$y$的域。具体来说，我们引入了对抗训练，这已被证明在对齐源和目标域方面是有效的，并且翻译后的样本通常捕获了目标域固有的丰富视觉特征。因此，TDA通过以下方式对齐$\varphi ([{\hat{x}}_g^{(0)},{\hat{x}}_a^{(0)}])$到$y$，其中$y$包含显著的视觉退化：

$L_{\text{adv}}=\log (1-D(\varphi ([{\hat{x}}_g^{(0)},{\hat{x}}_a^{(0)}])))$

其中$[\cdot ]$表示沿批量大小维度的连接，(D$是一个优化的判别器，通过以下方式优化：

$L_{\text{dis}}=-\log D([y,y])-\log (1-D(\varphi ([{\hat{x}}_g^{(0)},{\hat{x}}_a^{(0)}])))$

同时，TDA采用了重建损失来保持域对齐期间的内容一致性，即：

$L_{\text{rec}}={\Vert [{\hat{x}}_g^{(0)},{\hat{x}}_a^{(0)}]-\varphi ([{\hat{x}}_g^{(0)},{\hat{x}}_a^{(0)}])\Vert }_2^2$

此外，引入了感知损失以保持域对齐期间的语义一致性，即：

$L_{\text{perc}}={\Vert V({\hat{x}}_g^{(0)})-V(\varphi ({\hat{x}}_g^{(0)}))\Vert }_2^2$

其中$V(\cdot )$表示从预训练的 VGG 网络提取的特征图。总的来说，优化适配器$\varphi$的目标是：

$L_{\varphi }={\lambda }_1\cdot L_{\text{rec}}+{\lambda }_2\cdot L_{\text{perc}}+{\lambda }_3\cdot L_{\text{adv}}$

其中${\lambda }_1$，${\lambda }_2$，和${\lambda }_3$是损失权重。通过在去噪过程中的迭代优化，$\varphi$逐渐将${\hat{x}}_g^{(0)}$和${\hat{x}}_a^{(0)}$翻译到$y$的域中。注意，${\hat{x}}_a^{(0)}$被引入以防止$\varphi$陷入恒等映射的平凡解。

适配器引导的图像恢复

配备了适配器$\varphi$后，${\hat{x}}_g^{(0)}$可以通过$\varphi ({\hat{x}}_g^{(0)})$被翻译成退化域中的图像，这可以进一步被监督以适应$y$，进而产生指导$G$来更新生成的图像${\hat{x}}_g^{(0)}$至期望的干净输出$x_g^{(0)}$。适配器引导的图像恢复（Adapter-guided Image Restoration, AIR）被设计来执行这一监督更新过程。

现有研究表明，去噪过程展现出有趣的时间动态，即${\hat{x}}_g^{(t)}$从不可识别的内容逐渐生成到感知丰富的内容，最终达到细节不可见的内容。因此，AIR根据去噪过程中的不同阶段动态调整监督策略，从高层次内容逐步转移到低层次内容。

AIR将去噪步骤分为三个阶段，并采用不同的策略来引导图像恢复。在第一阶段，当$t\in [T,T/2)$时，生成的图像信噪比较低，内容不可识别，仅包含一些全局信息，如颜色和布局。因此，AIR在空间像素和颜色通道上应用均方误差损失来学习全局内容，即：

$G={\gamma }_1\cdot {\nabla }_{{\hat{x}}_g^{(0)}}{\Vert y-\varphi ({\hat{x}}_g^{(0)})\Vert }_2^2$

其中${\gamma }_1$是损失权重， x ^ g ( 0 ) ∈ { x ^ g ( t ) , x ^ a ( t ) } \hat{x}_g^{(0)} \in \{\hat{x}_g^{(t)}, \hat{x}_a^{(t)}\} x^g(0)​∈{x^g(t)​,x^a(t)​}。注意，AIR也在辅助图像$x_a^{(t)}$上使用这种指导，使其包含与$y$类似的全局属性，因为这有助于适配器在大多数退化上进行优化。此外，AIR在${\gamma }_1$上使用线性预热，逐渐引入这种指导。

在第二阶段，当$t\in [T/2,0)$时，PDM更倾向于生成感知丰富的内容，因此AIR更专注于从$y$中恢复感知内容。为此，AIR采用感知损失，即：

$G_{\text{perc}}={\nabla }_{{\hat{x}}_g^{(0)}}{\Vert V(y)-V(\varphi ({\hat{x}}_g^{(0)}))\Vert }_2^2$

同时，引入对抗损失来提高感知质量，即：

$G_{\text{adv}}={\nabla }_{{\hat{x}}_g^{(0)}}\log (1-D({\hat{x}}_g^{(0)}-y))$
$L_{\text{dis}}=-\log (1-D({\hat{x}}_g^{(0)}-y))-\log D([{\hat{x}}_g^{(0)},{\hat{x}}_a^{(0)}]-\varphi ([{\hat{x}}_g^{(0)},{\hat{x}}_a^{(0)}]))$

通过这种损失，${\hat{x}}_g^{(0)}$将被进一步更新到干净图像，其退化版本与$y$匹配，通过调整退化引起的内容损失域。此外，均方误差损失也被用来保持全局内容，即：

$G_{\text{mse}}={\nabla }_{{\hat{x}}_g^{(0)}}{\Vert y-\varphi ({\hat{x}}_g^{(0)})\Vert }_2^2$

总的来说，第二阶段的指导损失是：

$G={\gamma }_2\cdot G_{\text{mse}}+{\gamma }_3\cdot G_{\text{perc}}+{\gamma }_4\cdot G_{\text{adv}}$

其中${\gamma }_2-{\gamma }_4$是损失权重。

在第三阶段，当$t\in [0]$时，感知丰富的内容已经在${\hat{x}}_g^{(0)}$中呈现。AIR只需要恢复不可见的细节。换句话说，${\hat{x}}_g^{(0)}$非常接近自然图像，因此可以引入一些额外的图像先验来调节最终输出。考虑到通用性，我们只引入了总变分损失，即：

$G={\gamma }_5\cdot \text{TV Loss}({\hat{x}}_g^{(0)})$

其中${\gamma }_5$是损失权重。注意，尽管第三阶段主要影响不可见的细节，像素值仍然波动，因此在第二阶段的指导损失被用来获得更好的数值保真度。

第四章节主要介绍了该研究的实验部分，包括实验设置、定量和定性结果，以及对提出的测试时退化适配框架（TAO）的分析实验。以下是该章节的主要内容概述：

实验

本章节首先介绍了实验设置，包括用于评估TAO方法的图像恢复任务、数据集、评估指标以及与其它方法的比较。实验的目标是验证TAO在处理开放集图像恢复（OIR）问题时的有效性，特别是在面对未知和未预见的退化时。
作者详细描述了实验中使用的数据集，包括用于图像去雾、低光照图像增强和图像去噪的任务特定数据集。同时，也介绍了实验中使用的性能评估指标，如峰值信噪比（PSNR）和结构相似性指数（SSIM）。
展示了TAO在不同图像恢复任务上的定量结果，包括与现有方法的比较。结果表明，TAO在多个退化类型上均取得了与特定任务方法相当甚至更好的性能。此外，还提供了定性结果，如图像视觉效果对比，以展示TAO在恢复图像质量方面的优势。
这一部分对TAO的两个核心组件——测试时退化适配器（TDA）和适配器引导的图像恢复（AIR）的有效性进行了深入分析。通过对比实验，展示了TDA在适应未知和未预见退化方面的能力，以及AIR在不同去噪阶段采用不同策略进行图像恢复的有效性。此外，还包括了对不同域对齐策略的比较，以及对模型在不同测试条件下的鲁棒性的分析。
总结来说，第四章节通过一系列实验验证了TAO方法在开放集图像恢复任务中的有效性，并通过分析实验深入探讨了该方法的工作原理和优势。

结论

我们首先回顾了本研究的核心目标，即探索开放集图像恢复（OIR）问题，并提出了一种新的测试时退化适应框架（TAO）来解决这一挑战。我们明确了OIR问题的本质，即测试数据和训练数据之间存在的未识别的分布偏移，并展示了测试时适应（TTA）是解决这种差异的有效方法。
我们详细介绍了TAO框架，它包括三个关键组件：预训练的扩散模型（PDM）、测试时退化适配器（TDA）和适配器引导的图像恢复（AIR）。PDM作为通用的预训练模型，用于生成高质量的图像；TDA在测试阶段适配未知退化，通过引入一个适配器来对齐域；AIR则根据去噪过程中的不同阶段动态调整监督策略，以优化图像恢复结果。
通过在多种退化任务上的实验，我们证明了TAO框架的有效性。与特定任务的方法相比，TAO不仅在性能上与之相当，甚至在某些情况下表现得更好。此外，我们还通过一系列分析实验验证了TDA和AIR组件的有效性，展示了TAO在适应未知退化和图像恢复方面的优越性能。
最后，我们讨论了本研究的社会影响，虽然存在潜在的负面应用，如被用于恢复被故意损坏的图像或不真实的图像，但我们也强调了TAO在图像恢复领域的积极贡献，例如在提高图像质量、增强图像可读性等方面的应用。

更多详细内容见原文