Abstract

一体化图像恢复是一项基础的底层视觉任务，在现实世界中有重要应用。主要挑战在于在单个模型中处理多种退化情况。虽然当前方法主要利用任务先验信息来指导恢复模型，但它们通常采用统一的多任务学习，忽略了不同退化任务在模型优化中的异质性。为消除偏差，我们提出一种任务感知优化策略，为多任务图像恢复学习引入自适应的特定任务正则化。具体而言，我们的方法在训练过程中动态加权并平衡不同恢复任务的损失，促使实施最合理的优化路径。通过这种方式，我们能够实现更稳健且有效的模型训练。值得注意的是，我们的方法可作为即插即用策略，在不改变推理过程的情况下增强现有模型。在各种一体化恢复设置中的大量实验证明了我们方法的优越性和通用性。例如，使用TUR重新训练的AirNet在三个不同任务上平均提升1.16分贝，在五个不同的一体化任务上平均提升1.81分贝。这些结果突出了TUR在推动一体化图像恢复领域当前最优技术方面的有效性，为更稳健和通用的图像恢复铺平了道路。

Introduction

最近的工作MioIR（Kong、Dong和Zhang，2024年）通过提出一种顺序训练策略，在这个方向上取得了进展。该策略根据观察到的特征对退化任务进行分组。然而，它缺乏应对新退化任务的灵活性。关键挑战仍然存在：我们如何开发一种灵活有效的方法，在多任务优化过程中增强对任务先验的解释，以实现一体化图像恢复？
为应对这一挑战，我们从贝叶斯视角重新审视一体化图像恢复问题。我们针对多种退化任务提出了一种任务分组正则化方法。
它结合了一种与任务相关的正则化方法，该方法能够捕捉多种退化任务的独特特征，超越了现有混合多任务训练方法中使用的统一先验。更重要的是，这种依赖于任务的正则化提供了一种灵活且动态的机制，用于在多个退化任务中组合目标损失。
受神经网络不确定性的贝叶斯建模（阿卜达尔等人，2021年）启发，我们引入了任务不确定性正则化（TUR）。
这种任务分组正则化的实际实现提供了一种有原则的方法，可同时平衡和优化多个恢复任务，以适应每种退化类型的独特特征。此外，TUR通过再训练提供了一种即插即用的解决方案，以增强现有的一体化模型。如图1所示，我们在五种不同的一体化图像恢复设置和多个模型上验证了我们的方法，结果表明其性能始终优于其他方法。值得注意的是，使用TUR重新训练的AirNet在三个不同的任务上平均提升了1.16 dB，在五个不同的一体化任务上平均提升了1.81 dB。

Method

Preliminaries

传统的图像恢复方法通常使用最大似然估计来优化θ。给定退化的输入y，观察到干净图像x的似然被建模为：

Challenges of Multiple Degradation Types

对于涵盖K种退化任务的一体化图像恢复，我们将多种单一退化情况结合起来。设T={t1,t2,…,tK}表示K个恢复任务的集合。整体损失函数变为

虽然现有方法侧重于改进恢复模型中的任务先验，但它们往往忽略了优化目标。当前方法中多个任务损失的统一组合未能捕捉到每个任务的独特分布，而这是多任务学习中的一个关键考量因素。

Task Uncertainty Regularization

为应对这一挑战，我们将不确定性估计纳入一体化图像恢复框架。从贝叶斯观点出发，我们将恢复过程建模为：

基于此，我们引入任务不确定性正则化来捕获特定任务的不确定性，
$p(x\mid y,\theta ,{\sigma }_k)=\mathcal{N}(x\mid {\mathcal{M}}_{\theta }(y),{\sigma }_k^{2}I)$
其中${\sigma }_k$代表任务不确定性。对于任务k，我们的损失函数定义如下：

为了在所有任务中进行联合优化，我们提出以下多任务损失函数：

Remarks

自适应多任务组合：
为清晰起见，我们将公式（6）重新表述如下：

其中$w_k=1/2{\sigma }^2$，提供灵活且动态的多任务重新加权，$R={\sum }_k\log {\sigma }_k$。这种公式化方法会根据不同任务的估计不确定性自动学习平衡它们的贡献，并提供一个自适应、无偏差的优化目标。

实现细节

我们的TUR方法为一体化图像恢复模型提供了灵活的增强。在训练过程中，如图2所示，我们引入一个由三个堆叠的卷积层和激活层组成的辅助投影头，作为不确定性估计模块。这可以在不进行重大修改的情况下轻松集成到现有模型中。重要的是，在推理过程中不需要UEM，从而保持原始模型结构。因此，TUR就像一个即插即用的模块，在不改变现有恢复模型的情况下优化训练结果。