DFMDA-Net: Dense Fusion and Multi-dimension Aggregation Network for Image Restoration

1 Introduction

Bad weather or physical limitations of the camera can degrade the quality of captured images and further negatively impact the robustness of downstream tasks. Image restoration aims to eliminate those annoying degradation (such as noise, rain, and blur), and therefore plays an important role in surveillance, autonomous driving, remote sensing, etc. Due to the ill-posed nature, earlier model-based methods restrict the solution space by relying on handcrafted statistical priors. However, in complex real-world scenarios, it is difficult for these methods to recover faithful results.

With the fast development of deep learning, convolutional neural networks (CNN)- and transformer-based methods have achieved remarkable success by learning generalizable priors from large-scale datasets. In these methods, ingenious architecture or blocks are designed or borrowed from the highlevel computation vision tasks. For example, the U-shape (encoder-decoder) architecture [Ronneberger et al., 2015] is widely adopted for the effective hierarchy representation. The blocks are commonly designed based on residual learning and different attention mechanisms. Representative blocks include the residual block [Nah et al., 2017] and transformer block [Zamir et al., 2022].

∗Corresponding author. The authors are also with the Guangdong Province Engineering Laboratory for Digital Creative Technology and the Shenzhen Key Laboratory of Digital Creative Technology.

This paper aims to address the restoration task from both the U-shape architecture and the block design. In the encoder-decoder architecture, to assist the recovery, the features from the encoder are passed through to the decoder at the same level. For the feature fusion problem, current methods directly conduct the concatenation operation [Wang et al., 2022], the summation operation [Chen et al., 2022a], and first concatenation and then compression [Zamir et al., 2022]. These operations either cause information loss or make the complexity higher 1 . To alleviate this challenge, this paper proposes an Integration-Compression-Integration (ICI) mechanism for feature fusion. Our ICI first conducts integration to avoid information loss, then applies information compression to maintain efficiency, and finally conducts integration again for better information utilization. In this way, our ICI implements effective feature fusion while keeping efficient.

For block design, we revisit existing blocks from the perspective of information aggregation direction. Existing blocks either aggregate features from the spatial dimension such as the residual block [Nah et al., 2017] and windowbased attention [Liu et al., 2021] or from the channel dimension such as transposed attention [Zamir et al., 2022], and have achieved remarkable success. Nevertheless, how to effectively aggregate the information from both the channel and spatial dimensions, which can further strengthen the performance, deserves attention 2 . Since the window-based attention shares weights across channels, it is compatible with the transposed attention, as the shared operation limits the diversity of channel features. Moreover, the residual block Restormer SFNet ShuffleFormer Ours and window-based attention can only aggregate the local spatial information, which limits the receptive field. To address the issue, we resort to Fourier frequency learning since the Fourier features embrace global properties and can apply different weights to adjust the frequency points at different channels. To this end, we propose a multi-dimension aggregation (MDA) mechanism, capable of effectively aggregating features from both the channel and spatial dimensions.

1 More details will be illustrated in Section 3.3.

2 Note that aggregating information from channels in this paper means applying self-attention across channels.

Figure 1: Our method achieves significant performance gain on different image restoration tasks (Tables 2, 3, 4, 7).

By introducing ICI and MDA into a U-shape backbone, we construct the dense fusion and multi-dimension aggregation network (DFMDA-Net). The contributions of this paper can be summarized in four aspects:

• We propose dense fusion and multi-dimension aggregation network for image restoration from the perspective of both the U-shape architecture and block design.
• We propose a dense mechanism for feature fusion, termed Integration-Compression-Integration (ICI). This process is designed to prevent information loss by integration, maintain efficiency through compression in the middle, and thus enhance overall information utilization.
• We propose a multi-dimension aggregation (MDA) mechanism, capable of effectively aggregating features from both the channel and spatial dimension, thus enhancing the restoration performance.
• We conduct extensive experiments on 16 benchmark datasets for numerous image restoration tasks including image deraining, image real denoising, Gaussian image denoising, image motion deblurring, and image defocus deblurring. The results demonstrate that the proposed method achieves state-of-the-art performance, as illustrated in Fig. 1.

2 RelatedWork

Since image restoration is an ill-posed problem, earlier model-based methods restrict the solution space by exploiting handcrafted statistical priors, such as discriminative prior, channel prior, and gradient prior. These methods are heuristic and usually involve complex optimization problems. Limited by approximate models and the above factors, these methods are usually less effective, especially in complex real-world scenarios.

2.1 CNN-basedMethods

With the fast development of deep learning, deep neural networks, especially convolutional neural networks (CNN) have motivated numerous advanced image restoration methods [Cui et al., 2023d; Cui et al., 2023c] by learning dense priors from large-scale datasets. In these methods, there are several excellent convolutional designs which are commonly employed. For example, based on residual learning [He et al., 2016], Resblock [Nah et al., 2017] is proposed and has become the basic module for many image restoration methods. The underlying reason for the effectiveness can be attributed to the fact that the restoration task can be seen as the process of learning the residual between the degraded image and the ground-truth. Spatial and channel attention, which can enhance the desirable features and suppress the detrimental signals, have also achieved remarkable success [Zamir et al., 2021; Cui et al., 2023a]. Moreover, based on U-Net [Ronneberger et al., 2015], the U-shape (encoder-decoder) architecture has become the backbone of most successful image restoration models.

Albeit achieving significant performance gain over modelbased methods, CNN-based methods suffer limited receptive field and have poor content adaptation which originates from the basic convolution operation, which is crucial for effective image restoration.

2.2 Transformer-basedMethods

Motivated by the great success of transformers [Vaswani et al., 2017], researchers have proposed several excellent restoration models. The self-attention mechanism, which is a core component in transformers [Vaswani et al., 2017], can deal well with the shortcomings of the convolution operation. IPT [Chen et al., 2021a] applies the vanilla transformer on local patches, and has achieved promising performance gain. However, IPT depends on costly pretraining, and cannot be applied to high-resolution images due to the quadratic complexity of self-attention.

To address the quadratic complexity of self-attention, a line of methods applies self-attention on windows [Liu et al., 2021]. For example, SwinIR [Liang et al., 2021] and Uformer [Wang et al., 2022] adopt window-based selfattention (WSA). Although WSA reduces the complexity from quadratic to linear compared to self-attention, it dissevers the global pixel dependencies. Another line of methods applies self-attention across channels instead of spatial dimensions. Notably, Restormer [Zamir et al., 2022] proposes transposed attention (TA) with linear complexity, and has achieved remarkable success in multiple image restoration tasks. In essence, TA is a kind of channel self-attention, which aggregates the channel features with the similarity matrix of channels. Therefore, it still has limited receptive field.

2.3 Discussions

In this part, we discuss the significance of our work. First, most of the research focuses on exploring effective backbones while ignoring the architectural effects of the U-shape backbone. This paper aims to explore an effective and efficient feature fusion paradigm for the features from the encoder and the decoder at the same level. Second, this paper aims to aggregate desirable features (information) from both the spatial and channel dimensions for effective image restoration. Although Fourier domain learning is not new in image restoration [Mao et al., 2021; Zhou et al., 2023], these methods only achieve the aggregation of spatial information, and fail to adaptively aggregate the features from the channel dimension. The excellent performance of our work demonstrates that effectively aggregating information from both channel and spatial dimensions contributes to image restoration, which is expected to motivate more future research.

3 Methodology

3.1 Overall Pipeline

The over pipeline of the proposed dense fusion and multidimension aggregation (DFMDA-Net) is given in Fig. 2. Following most image restoration methods [Wang et al., 2022; Zamir et al., 2022], our DFMDA-Net is constructed based on the hierarchical encoder-decoder architecture. The detailed process is described as follows.

Given an input, our DFMDA-Net first applies a 3 × 3 convolution to obtain the embeddings. Then these embeddings are sequentially passed through a 3-level encoder, a bottleneck, and a 3-level decoder. Each level consists of several multi-dimension aggregation (MDA) blocks. After each level of the encoder, the output features are halved in spatial size and doubled in channel dimension with a 3 × 3 convolution and the pixel-unshuffle operation. Before each level of the decoder, the input features from the last level are doubled in spatial size and halved in channel dimension with a 3×3 convolution and the pixel-shuffle operation. In addition, to assist the recovery process, the output features from the same level of the encoder and the upsampled features are fused with the proposed Integration-Compression-Integration (ICI) module. Finally, the output features from the last level of the decoder are refined with a 3 × 3 convolution, and then added to the degraded image to obtain the restored image.

In the following, we introduce the proposed MDA and ICI mechanism.

3.2 Multi-dimension Aggregation

Analysis and Motivation

For block design, we revisit existing blocks from the perspective of information aggregation direction. Existing blocks can be divided into two groups: 1) aggregation from the spatial dimension and 2) aggregation from the channel dimension. CNN-based blocks such as the residual block [Nah et al., 2017], and transformer-based blocks such as recently window-based attention [Liu et al., 2021], lie in category 1. Transposed attention, i.e., applying self-attention across channels, lies in category 2. The effectiveness of these blocks demonstrates that aggregating information from either the spatial or channel direction is vital for effective image restoration. Therefore, joint spatial and channel information aggregation is an intuitive idea to strengthen the performance.

However, directly combining the attention in categories 1 and 2 cannot well accomplish the purpose. Specifically, since the window-based attention shares weights across channels, it is compatible with the transposed attention, as the shared operation limits the diversity of channel features. Moreover, the residual block and window-based attention can only aggregate the local spatial information. Therefore, to effectively integrate information from different dimensions, the spatially aggregated module should be equipped with the global properties and the weights on channels cannot be shared.

To this end, we resort to Fourier frequency learning with two utilizable properties. i) According to the spectral convolution theorem, the Fourier features embrace global properties. Updating a single value in the frequency domain globally affects all original spatial features. ii) Fourier transform can transfer the features into different frequency components, which can be seen as a frequency disentanglement and rearrangement. This property allows us to selectively enhance the necessary frequencies and suppress the detrimental frequencies in different channels.

Solution

Based on the above analysis, we propose a multi-dimension aggregation (MDA) block, capable of effectively aggregating features from both the channel and spatial dimension, as illustrated in Fig. 3. Specifically, we first transfer the input tensor X into Fourier domain with fast Fourier transform. Then we conduct two group 1 × 1 convolutions to conduct feature transformation. A GELU activation function is followed with each convolution. Finally, we transfer the obtained features back into the original space with inverse Fourier transform, and add the output features to those of the channel aggregation branch [Zamir et al., 2022] to get the result of MDA. Mathematically, the overall process is defined as:

where C 1 g and C 2 g are group 1 × 1 convolutions. σ denotes the GELU activation function. FFT and IFFT denote fast Fourier transform and inverse fast Fourier transform, respectively. XA is the result of the channel aggregation branch.

3.3 Integration-Compression-Integration Mechanism

In most of the existing U-shape-based image restoration methods, the features from the encoder are passed through to the decoder at the same level. For the feature fusion problem, current methods directly conduct the concatenation operation [Wang et al., 2022], the summation operation [Chen et al., 2022a], and first concatenation and then compression [Zamir et al., 2022], as illustrated in Fig. 4a-c, respectively. Since the parameters and computational cost of these transformer blocks are proportional to C 2 , where C is the number of channels, the concatenation manner (Fig. 4a) is parameter and computation costly. For the summation operation (Fig. 4b), it is ineffective since there exists the misaligned problem between the features of the encoder and the decoder. For the third case (Fig. 4c), it is too rough to conduct the direct compression with a 1×1 convolution to reduce channels, resulting in information loss.

To alleviate this challenge, this paper revisits the existing feature fusion mechanisms in Fig. 4a-c. By deeming the Di which consists of stacked blocks as integration, we can abstract them into two categories, i.e., Integration (Fig. 4a), and Compression 3 -Integration (Fig. 4b-c). As demonstrated above, the former is parameter and computation costly, and the latter suffers information loss. This paper proposes an Integration-Compression-Integration (ICI) mechanism for feature fusion, as illustrated in Fig. 4d. Our ICI first conducts integration to avoid information loss, then applies information compression to maintain efficiency, and finally conducts integration again for better information utilization. In this way, our ICI implements effective feature fusion while keeping efficient.

4 Relationship with Existing Feature Fusion Methods

Assume there are m blocks used for integration. The parameters and MACs of each block is O(C 2 ) and O(C 2HW). Fig. 4a means that the m blocks are directly applied to integrate the features with 2C channels. Fig. 4c means that the m blocks are used to integrate the features with C channels. We omit the analysis of the manner in Fig. 4b, since Fig. 4b can be seen as a particular case of Fig. 4c. ICI applies m1 ∈ [0,m] blocks for the first integration, and m-m1 blocks for the second integration (Di). In other words, m1 blocks are used to integrate the features with 2C channels, and m - m1 blocks are used to integrate the features with C channels. Therefore, the proposed ICI is a dense feature fusion paradigm compared to existing feature fusion manners shown in Fig. 4a-c. We also list the parameters and MACs of different feature fusion methods in Table 1 for better comparison.

5 Experiments

5.1 Setups

Evaluation Metrics

Following most of image restoration literature [Li et al., 2023; Zamir et al., 2022], we adopt PSNR and SSIM as the evaluation metrics. For defocus deblurring, extra metrics MAE and LPIPS are employed. It is worth noting that the PSNR and SSIM values are only calculated on the Y channel in YCbCr space for image deraining. We introduce the datasets in Appendix A.

Implementation Details

For different image restoration tasks, we train separate models. Unless otherwise specified, all the experiments are with the following settings. The number of blocks [N1, N2, N3, N4, N5, N6, N7] is set to [4, 6, 6, 8, 6, 6, 4]. The number of attention heads and groups of 1 × 1 convolution in MDA are both set to [1, 2, 4, 8, 4, 2, 1]. The initial channel dimension is 48. The channel expansion ratio in FFN is set to 2.66. m1 is set to 1. We adopt the progressive learning strategy [Zamir et al., 2022] to train our models for 300K iterations with L1 loss. The optimizer is AdamW (β1 = 0.9, β2 = 0.999, weight decay 1×10-4 ). The initial learning rate is 3×10-4 , and declines to 1 × 10-6 with the cosine annealing. For data augmentation, vertical and horizontal flips are employed. The visual results are given in Appendix B.

3 Summation can be regarded as a kind of compression.

5.2 Results on Image Deraining

The image deraining results on five datastets (Test100 [Zhang et al., 2020a], Rain100H [Yang et al., 2017], Rain100L [Yang et al., 2017], Test2800 [Fu et al., 2017], and Test1200 [Zhang and Patel, 2018]) are given in Table 2. On all five datasets, our DFMDA-Net achieves the best results on PSNR. Compared with the best transformer-based method, Restormer, our DFMDA-Net brings 0.27 dB PSNR improvement on average. When paying attention to the individual dataset Rain100H, our DFMDA-Net achieves 0.55 dB PSNR improvement. Compared with the frequency method Fourmer, our DFMDA-Net achieves a performance gain of more than 1 dB on three datasets including Test100, Rain100H, and Rain100L. Moreover, compared with the recent method FSNet, our DFMDA-Net brings 0.72 dB PSNR improvement on average. When paying to the individual datasets Test100 and Rain100L, the improvements are up to 1.27 dB and 1.15 dB, respectively.

5.3 Results onMotion Deblurring

We gives the image motion deblurring results on datastets (GoPro [Nah et al., 2017], HIDE [Shen et al., 2019], and Realbur [Rim et al., 2020] (Realbur-R and Realbur-J)) in Table 3. Compared with Restormer, our DFMDA-Net brings the PSNR improvement as large as 0.62 dB on average. When averaged across all the datasets, our DFMDA-Net achieves 0.40 dB PSNR improvement over the global ShuffleFormer. Moreover, compared with the global modeling method, GRL, which simultaneously models the pixel, region, and global level interactions, our DFMDA-Net achieves 0.25 dB and 0.48 dB improvements on GoPro and HIDE, respectively.

5.4 Results on Defocus Deblurring

Table 4 presents the results of image defocus deblurring on the DPDD [Abuolaim and Brown, 2020] dataset. Our DFMDA-Net achieves basically consistent performance gain on both scene categories. Compared with Restormer, our DFMDA-Net brings 0.38 dB PSNR improvement on the combined category. When paying attention to the indoor scene category, our DFMDA-Net achieves 0.48 dB PSNR improvement. Compared with GRL, our DFMDA-Net achieves a performance gain of 0.18 dB on the combined category. Moreover, compared with the recent FSNet, our DFMDA-Net brings 0.047 LPIPS improvement on average. When paying to the outdoor scene category, the improvement is up to 0.057.

5.5 Results on Image Denoising

For image denoising, we conduct the experiments on both synthetic and real datasets.

Gaussian Image Denoising

For Gaussian image denoising, we generate the synthetic benchmark datasets by adding additive white Gaussian noise on DIV2K [Agustsson and Timofte, 2017], Flickr2K, BSD500 [Arbelaez et al., 2010] and WED [Ma et al.,

2016]. The testing datasets are Set12 [Zhang et al., 2017a], BSD68 [Martin et al., 2001], Urban100 [Huang et al., 2015], Kodak24, and McMaster [Zhang et al., 2011]. We list the results in Table 5 for grayscale images and Table 6 for color images. Following DRUNet [Zhang et al., 2021], three noise levels (15, 25, and 50) are tested. Overall, our DFMDANet has consistent performance gain on different datasets and noise levels. Compared with Restormer, our DFMDA-Net brings 0.29 dB and 0.27 dB PSNR improvement on the challenging 50 noise level on grayscale and color images of Urban100, respectively.

Real Image Denoising

We conduct the real image denoising on the SIDD [Abdelhamed et al., 2018] dataset, and the result is shown in Table 7. Compared with the recent best method, Restormer, our DFMDA-Net achieves 0.30 dB PSNR improvement. Compared with the recent global model, ShuffleFormer, our DFMDA-Net brings 0.32 dB PSNR gain.

5.6 Ablation Study

In this section, we first ablate the proposed two modules, ICI and MDA. Then we investigate the design choices of these modules. The models are trained on the GoPro [Nah et al., 2017] dataset for 150K iterations. The initial channel dimension is set to 32. Unless specified otherwise, other model configurations and training settings are the same as those in Section 5.1.

Effectiveness of Each Module

The results are given in Table 8. The baseline model is Restormer [Zamir et al., 2022], and achieves 33.16 dB PSNR. When equipped with MDA and ICI, the models bring 0.18 dB and 0.10 dB performance gains over the baseline model. With the introduction of both MDA and ICI, the model achieves 33.45 dB PSNR. The above results demonstrate the effectiveness of the proposed two modules.

Number of Blocks (m1) in ICI

We investigate the effects of the number of blocks (m1) used for the first integration in ICI. Table 9 lists the results. As m1 increases, the performance boosts rapidly when m1=1, and then tends to saturate when m1=4. This phenomenon demonstrates that without the first integration, the performance drops significantly. However, applying all the blocks (m1=6) is parameter- and complexity-costly. These results demonstrate the effectiveness and necessity of the proposed ICI. Considering the performance and efficiency, we empirically choose m1=1.

Number of Groups in 1 × 1 Convolution in MDA

We investigate the effects of the number of groups in 1 × 1 convolution. The results are given in Table 10. 1 represents that the regular 1×1 convolution is employed. H denotes that the number of groups is the same as the number of heads in self-attention. 2H denotes that the number of groups is twice the number of heads in self-attention, and the same meaning for 4H. As can be seen, the H case achieves the best result.

Alternatives to ICI

We also give the results of different alternatives of ICI in Table 11. As can be seen, our ICI achieves a better trade-off between performance and efficiency.

6 Conclusions

This paper presents a dense fusion and multi-dimension aggregation network (DFMDA-Net) for image restoration. For the fusion of features between the encoder and the decoder at the same level, we propose an Integration-CompressionIntegration (ICI) mechanism. ICI effectively conducts dense fusion, and achieves a better trade-off between effectiveness and efficiency. In addition, we also propose a multidimension aggregation (MDA) mechanism, capable of effectively aggregating features from both the channel and spatial dimensions. Extensive experiments on 16 benchmark datasets demonstrate that DFMDA-Net achieves state-of-theart performance on image restoration tasks including image deraining, image motion deblurring, image defocus deblurring, real image denoising, and Gaussian image denoising.

Acknowledgments

This work was supported in part by the National Natural Science Foundation of China under Grants 62101336 and in part by the Guangdong Basic and Applied Basic Research Foundation under Grants 2022A1515011301.