A Mountain-Shaped Single-Stage Network for Accurate Image Restoration

In recent years, the majority of image restoration research has focused on single-stage architecture. Among these architectures, the encoder-decoder based U-Net [12, 13, 18, 33, 25, 8, 34, 35] and dual network structure [36, 37, 38, 39, 40] are mainly included.

Encoder-Decoder Approaches. In recent years, encoder-decoder have gained great attention from researchers in the field of image restoration thanks to its ability to capture multi-scale information. To construct an effective and efficient Transformer-based architecture for image restoration, [18] introduce a novel locally-enhanced window and multi-scale restoration modulator to create a hierarchical encoder-decoder network. [30] utilize selective kernel feature fusion to realize the information exchange of different scales and information aggregation based on attention. [22] develops a simple yet effective boosted decoder to progressively restore the haze-free image by incorporating the strengthen-operate-subtract boosting strategy in the decoder. By eliminating or substituting the nonlinear activation function, [12] establishes a simple baseline that yields measurable outcomes while requiring fewer computing resources.

Dual Network Approaches. The Dual Networks architecture is designed with two parallel branches that separately estimate the structure and detail components of the target signals from the input. These components are then combined to reconstruct the final results according to the specific task formulation module. This architecture was first proposed by [41] and has since inspired a lot of subsequent work, including in the areas of image dehazing [36, 37, 38], image deraining [42], image denoising [39], and image super-resolution/deblurring [40]. The Dual Networks approach has proven to be effective in addressing various low-level vision problems, by enabling a better separation of the structure and detail information, leading to improved performance in terms of both accuracy and computational efficiency. Furthermore, the flexibility of this architecture makes it adaptable to different types of data, making it a popular choice in the field of image restoration.

Despite the significant achievements made by these networks, it remains a challenge to effectively balance these competing goals of preserving spatial details and contextualized information while recovering images.

The multi-stage networks are shown to be more effective than their single-stage counterparts in high level vision problems  [43, 44, 45, 46]. In recent years, there have been some attempts [47, 48, 11, 10, 8, 49, 50] to apply multi-stage networks to image restoration. They aims to break down the image restoration process into several manageable stages, enabling the use of lightweight subnetworks to progressively restore clear images. This approach facilitates the capture of both spatial details and contextualized information by individual subnetworks at each stage. To prevent the production of suboptimal results that may arise from using the same subnetwork at each stage, a supervisory attention mechanism was proposed along with the adoption of distinct subnetwork structures [8]. Additionally, [49] present a novel self-supervised event-guided deep hierarchical Multi-patch Network to handle blurry images and videos through fine-to-coarse hierarchical localized representations. Nevertheless, this approach elevates the complexity of the system as refining the previous stage’s results is required in subsequent stages.

By incorporating an encoder and decoder network in the initial stage, followed by a network operating at the original image input resolution in the final stage, the multi-stage network can produce high-quality images with accurate spatial details and reliable contextual information. However, the latter stage of this process requires revising the results of the previous stage, which adds a slight level of complexity to the system. While a single-stage network has relatively less complexity, it may struggle to balance spatial details and context information effectively. Therefore, we are exploring a middleware mechanism for feature fusion that enables a simple single-stage U-Net architecture to achieve the same functionality as a multi-stage architecture, without sacrificing smoothness of expression or compromising on the original meaning.

As shown in Fig. 4(a), the feature fusion middleware(FFM) is a nonlinear activation-free block (NAFBlock) with upsample and feature fusion. We have introduced the FFM between the encoder-decoder architectural levels to integrate upper-layer information into adjacent lower layers. This integration takes place sequentially, from the highest layer down to the lowest, until all information is fused into the original image resolution manipulation level. This approach enhances the network’s capacity to capture and fuse multi-scale features, ranging from simple patterns at low levels, such as corner or edge/color connections, to more complex higher-level features, such as significant variations and specific objects. As a result, this structure preserves spatial details while integrating contextual information, ultimately ensuring high-quality image restoration. Formally, let $\mathbf{FE_{i}}\in\mathbb{R}^{\frac{H}{i^{2}}\times\frac{W}{i^{2}}\times{i^{2}}C}$ be the output in the $i$-th level encoder $(i=1,2,3,4)$. At each level, the feature fusion information $FFM_{s,i}$ is given as:

where $\oplus$ denote the element-wise addition, $UP(\cdot)$ represents the up-sampling operation and $H_{Naf_{s,i}}(\cdot)$ is the $s$-th FFM in the $i$-th level.

This design offers two benefits. Firstly, the feature fusion middleware integrates multi-scale information, allowing the network model to capture abundant context information. Secondly, the feature fusion middleware in the $1_{th}$ layer operates on the original image resolution, without employing any subsampling operation, thereby enabling the network model to acquire detailed spatial information of high-resolution features.

Finally, the depth features $\mathbf{F_{DF}}$ are obtained through this single-stage architecture, as demonstrated below:

where $FD_{i}$ is the output in the $i$-th level decoder, and -$1$ indicates that this is the last feature fusion middleware at this level.

The transformer model [18, 16, 17] has gained popularity in image restoration tasks due to its capability to capture global information, as evidenced by its increasing usage in recent years. The computation on a global scale results in a quadratic complexity in relation to the number of tokens as shown in Eq. 15, rendering it inadequate for the representation of high-resolution images.

To alleviate this issue, [17, 16, 32] etc., proposed various methods to reduce complexity. In this paper, we propose a multi-head attention middle block (MHAMB) as the bridge of the encoder-decoder, shown in Fig. 4(d).MHAMB utilizes global self-attention to process and integrate the feature map information that is generated by the last layer of the encoder. This approach is particularly efficient in handling large images because convolution downsamples space, while attention can effectively process smaller resolutions for better performance. From the last layer of the encoder output $FE_{4}$, our MHAMB first apply a $1\times 1$ convolution and a $3\times 3$ depth-wise convolution to generate the $query,key$ and $value$ matrices $\mathbf{Q}\in\mathbb{R}^{H\times W\times C},\mathbf{K}\in\mathbb{R}^{H\times W\times C}$ and $\mathbf{V}\in\mathbb{R}^{H\times W\times C}$ as follows:

Then, we reshape $\mathbf{Q}$, $\mathbf{K}$, $\mathbf{V}$ to $\mathbf{\hat{Q}}\in\mathbb{R}^{(HW)\times\frac{C}{h}\times h}$, $\mathbf{\hat{K}}\in\mathbb{R}^{(HW)\times\frac{C}{h}\times h}$, $\mathbf{\hat{V}}\in\mathbb{R}^{(HW)\times\frac{C}{h}\times h}$, where $h$ is the number of head. Next, the attention matrix is thus computed by the self-attention mechanism as :

where $\beta$ is a learning scaling parameter used to adjust the magnitude of the dot product of $\hat{Q}$ and $\hat{K}$ prior to the application of the softmax function. Finally, we reshape the attention matrix back to its original dimensions of $\mathbb{R}^{H\times W\times C}$ and apply a $1\times 1$ convolution. The resulting output is then added to $FE_{4}$ and passed to the decoder. This allows us to capture more information and enhance the overall performance of the model as shown in the following experiment.

We use PSNR and SSIM as quality assessment metrics. To report the reduction in error for each method relative to the best-performing method, we convert PSNR to RMSE ($RMSE\propto\sqrt{10^{-PSNR/10}}$) and SSIM to DSSIM (DSSIM = (1 - SSIM)/2).

We train the proposed M3SNet without any pre-training and separate models for different image restoration tasks. We utilize the following block configurations in our network for each level: $[1,1,1,28]$ blocks for the encoder, $[1,1,1,1]$ blocks for the decoder, $[2,2,1,0]$ blocks for the FFM. And one MHAMB for the bridge of the encoder and decoder. We train models with Adam [58] optimizer($\beta_{1}=0.9,\beta_{2}=0.999$) and PSNR loss for $5\times 10^{5}$ iterations with the initial learning rate $1\times 10^{-3}$ gradually reduced to $1\times 10^{-7}$ with the cosine annealing [59]. We extract patches of size $256\times 256$ from training images, and the batch size is set to $32$. We adopt TLC [60] to solve the issue of performance degradation caused by training on patched images and testing on the full image. For data augmentation, we perform horizontal and vertical flips.

In our image deraining task, we compute the PSNR/SSIM scores using the Y channel in the YCbCr color space, which is consistent with previous works such as[8, 61, 62]. Our method has been demonstrated to outperform existing approaches significantly and consistently, as presented in Table 3. Specifically, our method achieves a remarkable improvement of 0.93 dB and a $10.2\%$ error reduction on average across all datasets when compared to the best CNN-based method, SPAIR [62]. Moreover, we can achieve up to 2.76 dB improvement over HINet [9] on individual datasets, such as Rain100L.

In addition to quantitative evaluations, Fig. 6 presents qualitative results that demonstrate the effectiveness of our M3SNet in removing rain streaks of various orientations and magnitudes while preserving the structural content of the images.

The performance evaluation of image deblurring approaches on the GoPro [55] and HIDE [57] datasets is presented in Table 4. Our M3SNet outperformed other methods, with a performance gain of 0.09dB when averaging across all datasets [55, 57]. Compared to our baseline network NAFNet [12], we improve 0.04 dB and 0.05 dB at 32 widths and 64 widths, respectively. It is worth noting that even though our network is trained solely on the GoPro Dataset, it still achieves state-of-the-art results (31.49 dB in PSNR) on the HIDE dataset. This demonstrates its impressive generalization capability.

Figure 7 displays some of the deblurred images produced by our method. Our model recovered clearer images that were closer to the ground truth than those by others.

The ablation studies are conducted on image deblurring (GoPro [55]) to analyze the impact of each of our model components. It is worth noting that we use NAFNet-32 [12] as the baseline network to demonstrate the effectiveness by adding our proposed component. Table. 5 shows the effectiveness of our M3SNet. Next, we describe the impact of each component.

FFM We demonstrate the effectiveness of the proposed feature fusion middleware by adding it to NAFNet-32 [12]. As the Table. 5 shows that the FFM added in each level of encoder-decoder leads to improved performance (32.90 dB) as compared to the baseline. This indicates that our proposed FFM can help the model capture more multi-scale information while also preserving spatial detail information.

MHAMB As the Table. 5 shows a tiny improvement in PSNR from 30.87 dB to 30.88 dB after we add a multi-head attention middle block (MHAMB) to capture more global information in the topmost level.

Adding both of these components improves the performance by a large margin from 32.87 dB to 32.91 dB in PSNR.

Deep learning models have become increasingly complex in order to achieve higher accuracy. However, larger models require more resources and may not be practical in certain contexts. Therefore, there is a need to design lightweight image restoration models that can achieve high accuracy. In our work, we design a mountain-shaped single-stage network. This architecture optimizes the balance between spatial details and contextual information while minimizing the computational resources required to restore images.

Our M3SNet has been shown to outperform other models, as demonstrated in Table 6. Despite having 0.6M higher parameters than MIMO-unet++ [33], our proposed M3SNet-32 still achieves better performance, while using significantly fewer computational resources, with MACs approximately 40 times smaller than that of MIMO-unet++. Although the MACs of our model are higher than Uformer [18] and NAFNet [12], our model parameters are smaller, yet still perform better. Considering all factors, including model parameters, MACs, and performance, our model is the optimal choice.