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就像一個即插即用的模塊，在不改變現有恢復模型的情況下優化訓練結果。