Boosting Image Restoration via Priors from Pre-trained Models

Fig. 2 depicts the restoration pipeline using our method. Given an input image $I_{d}$, we have an initial restoration result as $\hat{I}_{c}=\mathcal{F}(I_{d})$. We aim to refine the result using the proposed pre-training guided refinement module (PTG-RM) $\mathcal{R}$, resulting in $\bar{I}_{c}=\mathcal{R}(\hat{I}_{c},I_{d},g)$. The key of this approach is to distill restoration-related information from the prior $g$.

$\mathcal{R}$ is a simple encoder-decoder structure. The encoder and decoder of $\mathcal{R}$ are denoted as $\mathcal{R}_{e}$ and $\mathcal{R}_{d}$, respectively. To ensure lightweight implementation, distillation occurs in the latent space, avoiding the need to align $g$ with restoration-related features. The latent feature $f$ is derived through a comparison between the initial enhanced results and the original input images, given as $f=\mathcal{R}_{e}(\hat{I}_{c}\oplus I_{d})$, where $\oplus$ denotes the concatenation operation. The resulting $f$ is in $\mathbb{R}^{h\times w\times c}$, with $h$, $w$, and $c$ representing feature height, width, and channel number, respectively. The priors are used in further learning the latent feature as $\bar{f}=\mathcal{C}(\mathcal{A}(f,g),g)$, where $\mathcal{A}$ and $\mathcal{C}$ represent the Pre-Train-Guided Spatial-Varying Enhancement (PTG-SVE) and Pre-Train-Guided Channel-Spatial Attention (PTG-CSA) modules, respectively. The final enhancement is obtained from the decoder as $[I_{m},I_{r}]=\mathcal{R}_{d}(\hat{f})$, comprising two components. The first component, $I_{m}$, represents the correction mask used to mitigate errors in the initial enhancement results. The second component, $I_{r}$, is the residual refinement that addresses artifacts and adds additional details. The final result is denoted as

In PTG-SVE, we argue that $g=\mathcal{G}(I_{d})$ may contain information reflecting the pixel-level image quality of $I_{d}$. For areas with poor quality, long-range operations are used to capture non-local features, while regions with relatively good quality prioritize local features for accurate restoration.

In Fig. 3, the primary objective is to predict the optimal neural operation range for each location of the feature map $f$, which we refer to as the “range score map”, denoted as $M$. To ensure a general $\mathcal{R}$ with unified 1D priors $g$ from various models, we propose adding location-aware embeddings for the priors, thereby adaptively discovering quality information for different pixels. Let $S=\{(x,y)|x\in[1,w],y\in[1,h]\}$ represent the 2D coordinate map with dimensions $h\times w\times 2$. We use a position embedding module $\mathcal{P}$ to generate spatial representation, denoted as $\mathcal{S}_{m}=\mathcal{P}(S)$, where $\mathcal{S}\in\mathbb{R}^{h\times w\times c}$. Furthermore, to determine the admired neural operation range for each location of $f$, we use a learnable mapping function $\mathcal{T}_{m}$ to transform the priors to another space that can more effectively decide the optimal range. To obtain $M$, we use a range-learning module $\mathcal{R}_{m}$, which takes the encoder’s feature $f$, the pre-trained prior $g$, and the spatial representation $\mathcal{S}_{m}$ as inputs. The procedure is denoted as

Nach [42], we use CNN for the short-range operation, denoted as $\mathcal{R}_{s}$, and transformer for the long-range operation, represented as $\mathcal{R}_{l}$. Specifically, we employ the Restormer backbone for $\mathcal{R}_{l}$ and ResNet for $\mathcal{R}_{s}$. Suppose the features after the short- and long-range operation are $f_{s}$ and $f_{l}$, respectively. We can obtain the refined feature $\hat{f}$ as

The previous approach [42] relies on pre-computed SNR values, which may not always be accurate and can fail to enhance results, especially when the initial results from $\mathcal{F}$ have reached their upper bound. In contrast, our score range map is learned online based on the input image, restoration-related priors, and explicit spatial features that are learnable. This flexibility allows us to handle various situations, resulting in better performance and generalization (as demonstrated in the ablation study).

As shown in Fig. 3, we further introduce a lightweight component that utilizes pre-trained priors $g$ to create an effective attention mechanism in $\mathcal{R}$. Optimizing the feature attention in $\mathcal{R}$ is crucial for effectively identifying helpful features to enhance the initial results $\hat{I}_{c}$. This involves both spatial-level and channel-level attentions. The hidden restoration-related information in $g$ can be discovered by using $g$ to improve the restoration features in $\mathcal{R}$ conditioned on them.

We begin by formulating the attention computation at the channel level. We introduce a mapping function $\mathcal{T}_{c}$ to transform $g$ into the attention-prediction space, and utilize the channel attention computation module $\mathcal{R}_{c}$. The formulation of the channel attention is

where $\mathcal{O}$ is the pooling operation, and $\mathcal{M}_{c}\in\mathbb{R}^{c}$.

As for the spatial-attention computation, we utilize the 1D pre-trained prior $g$ to predict location-wise attention based on the feature distribution of each location in $\hat{f}$. Simply using the spatial location information, as shown in Eq. 2, results in each pixel’s feature considering a similar condition for neighboring features, limiting the elimination of spatial artifacts. Therefore, we propose an alternative strategy by predicting the neural operation parameters for each location, optimizing the spatial attention based on the varying location-wise feature distribution. We denote the spatial attention computation module as $\mathcal{R}_{p}$, and first formulate the location-wise convolution map, as

where the obtained convolution map $\mathcal{C}_{p}\in\mathbb{R}^{h\times w\times(k_{h}\times k_{w}\times c)}$, $k_{h}$ and $k_{w}$ are the convolution kernel size, and $\mathcal{S}_{c}$ is another learnable position embedding here. The obtained convolution maps can be utilized to optimize the feature, and spatial attention can be obtained as

where $*$ is the convolution operation for each location, and $\mathcal{R}_{o}$ is another learnable operation which mapps the feature channel $c$ to 1, eliminating the influence from the channel-level dependency. Further, the feature after spatial attention can be described as $\hat{f}_{s}=\hat{f}\times M_{s}$.

The features after spatial and channel attentions can be merged via a fusion module as

where $\mathcal{R}_{f}$ denotes the fusin module. The obtained feature $\bar{f}$ can be processed via a decoder $\mathcal{R}_{d}$ to obtain the residual refinement $I_{r}$ and the mask $I_{m}$ as indicated in Eq. 1.

Our designed $\mathcal{R}$ can be jointly trained with the model $\mathcal{F}$. Suppose the paired ground truth for the input image $I_{d}$ is $\mathcal{I}_{c}$, and the loss function for the model $\mathcal{F}$ is denoted as $\mathcal{L}_{g}(\hat{I}_{c},\mathcal{I}_{c})$ (is usually the reconstruction loss in the pixel level or perceptual loss, and can also be the unsupervised loss), then the loss function for the refinement module can be written as $\mathcal{L}_{g}(\bar{I}_{c},\mathcal{I}_{c})$. In summary, the overall loss is

where $\lambda_{1}$ is the loss weight and remains robust across various tasks and networks (in our experiments, $\lambda 1$ is always set as 1).

For low-light enhancement, we use the SID [5] and LOL-real [49] datasets. For deraining, we use the Rain13K [52] dataset for training and test on Rain100H [48], Rain100L [48], Test100 [55], Test1200 [54], and Test2800 [8] datasets. For gaussian denoising, we use two settings: synthetic noise on Set12 [56], BSD68 [23], CBSD68 [23], Kodak [7], McMaster [58], and Urban100 [10]; and real-world denoising on SIDD [2]. For single-image motion deblurring, we use the GoPro [24] dataset for training and evaluate on synthetic datasets (GoPro [24], HIDE [30]) and real-world datasets (RealBlur-R [28], RealBlur-J [28]). For defocus deblurring, we use the DPDD [3] training data and test on the EBDB [13] and JNB [31] datasets.

Moreover, we conducted an experiment by adding the refinement module to the intermediate layer of $\mathcal{F}$, refining features of the target model. The refinement module is added to the deepest feature layer for efficiency, producing the residual feature map and the mask information for refinement. These results are denoted as $-f$. The improvement achieved by this operation is also evident as displayed in Table 1.

We conduct experiments without SP, CA, or SA to analyze whether additional parameters or priors take a prominent role. The short-range and long-range results are fused via a simple sum, and the spatial-channel attention is conducted using only the features themselves. Additionally, we increase the feature channel number fourfold to add more parameters. The results, denoted as “Large $\mathcal{R}$ w/o SP/CA/SA” in Table 4, are still lower than our full setting, indicating the effectiveness of our proposed approach over simply increasing parameters.

In addition, we perform an experiment by removing the learnable position embeddings $\mathcal{S}_{m}$ and $\mathcal{S}_{c}$, denoted as “w/o Position Embedding for Priors” in Table 4. This comparison highlights the importance of using spatial-aware representations for the pre-trained features.

In this section, we conduct experiments using CLIP as the pre-trained model ($-c$). CLIP is chosen for its efficiency and convenience compared to other pre-trained models.