Restormer: Efficient Transformer for High-Resolution Image Restoration

Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Shahbaz Khan, Ming-Hsuan Yang

Introduction

Image restoration is the task of reconstructing a high-quality image by removing degradations (e.g., noise, blur, rain drops) from a degraded input. Due to the ill-posed nature, it is a highly challenging problem that usually requires strong image priors for effective restoration. Since convolutional neural networks (CNNs) perform well at learning generalizable priors from large-scale data, they have emerged as a preferable choice compared to conventional restoration approaches.

The basic operation in CNNs is the ‘convolution’ that provides local connectivity and translation equivariance. While these properties bring efficiency and generalization to CNNs, they also cause two main issues. (a) The convolution operator has a limited receptive field, thus preventing it from modeling long-range pixel dependencies. (b) The convolution filters have static weights at inference, and thereby cannot flexibly adapt to the input content. To deal with the above-mentioned shortcomings, a more powerful and dynamic alternative is the self-attention (SA) mechanism that calculates response at a given pixel by a weighted sum of all other positions.

Self-attention is a core component in Transformer models but with a unique implementation, i.e., multi-head SA that is optimized for parallelization and effective representation learning. Transformers have shown state-of-the-art performance on natural language tasks and on high-level vision problems . Although SA is highly effective in capturing long-range pixel interactions, its complexity grows quadratically with the spatial resolution, therefore making it infeasible to apply to high-resolution images (a frequent case in image restoration). Recently, few efforts have been made to tailor Transformers for image restoration tasks . To reduce the computational loads, these methods either apply SA on small spatial windows of size $8<math xmlns="http://www.w3.org/1998/Math/MathML" display="block"><semantics><mrow><mo>×</mo></mrow><annotation encoding="application/x-tex">\times</annotation></semantics></math>×8$ around each pixel , or divide the input image into non-overlapping patches of size $48<math xmlns="http://www.w3.org/1998/Math/MathML" display="block"><semantics><mrow><mo>×</mo></mrow><annotation encoding="application/x-tex">\times</annotation></semantics></math>×48$ and compute SA on each patch independently . However, restricting the spatial extent of SA is contradictory to the goal of capturing the true long-range pixel relationships, especially on high-resolution images.

In this paper, we propose an efficient Transformer for image restoration that is capable of modeling global connectivity and is still applicable to large images. Specifically, we introduce a multi-Dconv head ‘transposed’ attention (MDTA) block (Sec. 3.1) in place of vanilla multi-head SA , that has linear complexity. It applies SA across feature dimension rather than the spatial dimension, i.e., instead of explicitly modeling pairwise pixel interactions, MDTA computes cross-covariance across feature channels to obtain attention map from the (key and query projected) input features. An important feature of our MDTA block is the local context mixing before feature covariance computation. This is achieved via pixel-wise aggregation of cross-channel context using $1<math xmlns="http://www.w3.org/1998/Math/MathML" display="block"><semantics><mrow><mo>×</mo></mrow><annotation encoding="application/x-tex">\times</annotation></semantics></math>×1$ convolution and channel-wise aggregation of local context using efficient depth-wise convolutions. This strategy provides two key advantages. First, it emphasizes on the spatially local context and brings in the complimentary strength of convolution operation within our pipeline. Second, it ensures that the contextualized global relationships between pixels are implicitly modeled while computing covariance-based attention maps.

A feed-forward network (FN) is the other building block of the Transformer model , which consists of two fully connected layers with a non-linearity in between. In this work, we reformulate the first linear transformation layer of the regular FN with a gating mechanism to improve the information flow through the network. This gating layer is designed as the element-wise product of two linear projection layers, one of which is activated with the GELU non-linearity . Our gated-Dconv FN (GDFN) (Sec. 3.2) is also based on local content mixing similar to the MDTA module to equally emphasize on the spatial context. The gating mechanism in GDFN controls which complementary features should flow forward and allows subsequent layers in the network hierarchy to specifically focus on more refined image attributes, thus leading to high-quality outputs.

Apart from the above architectural novelties, we show the effectiveness of our progressive learning strategy for Restormer (LABEL:{Progressive_Learning}). In this process, the network is trained on small patches and large batches in early epochs, and on gradually large image patches and small batches in later epochs. This training strategy helps Restormer to learn context from large images, and subsequently provides quality performance improvements at test time. We conduct comprehensive experiments and demonstrate state-of-the-art performance of our Restormer on $16$ benchmark datasets for several image restoration tasks, including image deraining, single-image motion deblurring, defocus deblurring (on single-image and dual pixel data), and image denoising (on synthetic and real data); See Fig. 1. Furthermore, we provide extensive ablations to show the effectiveness of architectural designs and experimental choices.

The main contributions of this work are summarized below:

We propose Restormer, an encoder-decoder Transformer for multi-scale local-global representation learning on high-resolution images without disintegrating them into local windows, thereby exploiting distant image context.

We propose a multi-Dconv head transposed attention (MDTA) module that is capable of aggregating local and non-local pixel interactions, and is efficient enough to process high-resolution images.

A new gated-Dconv feed-forward network (GDFN) that performs controlled feature transformation, i.e., suppressing less informative features, and allowing only the useful information to pass further through the network hierarchy.

Background

Image Restoration. In recent years, data-driven CNN architectures have been shown to outperform conventional restoration approaches . Among convolutional designs, encoder-decoder based U-Net architectures have been predominantly studied for restoration due to their hierarchical multi-scale representation while remaining computationally efficient. Similarly, skip connection based approaches have been shown to be effective for restoration due to specific focus on learning residual signals . Spatial and channel attention modules have also been incorporated to selectively attend to relevant information . We refer the reader to NTIRE challenge reports and recent literature reviews , which summarize major design choices for image restoration.

Vision Transformers. The Transformer model is first developed for sequence processing in natural language tasks . It has been adapted in numerous vision tasks such as image recognition , segmentation , object detection . The Vision Transformers decompose an image into a sequence of patches (local windows) and learn their mutual relationships. The distinguishing feature of these models is the strong capability to learn long-range dependencies between image patch sequences and adaptability to given input content . Due to these characteristics, Transformer models have also been studied for the low-level vision problems such as super-resolution , image colorization , denoising , and deraining . However, the computational complexity of SA in Transformers can increase quadratically with the number of image patches, therby prohibiting its application to high-resolution images. Therefore, in low-level image processing applications, where high-resolution outputs need to be generated, recent methods generally employ different strategies to reduce complexity. One potential remedy is to apply self-attention within local image regions using the Swin Transformer design . However, this design choice restricts the context aggregation within local neighbourhood, defying the main motivation of using self-attention over convolutions, thus not ideally suited for image-restoration tasks. In contrast, we present a Transformer model that can learn long-range dependencies while remaining computationally efficient.

Method

Our main goal is to develop an efficient Transformer model that can handle high-resolution images for restoration tasks. To alleviate the computational bottleneck, we introduce key designs to the multi-head SA layer and a multi-scale hierarchical module that has lesser computing requirements than a single-scale network . We first present the overall pipeline of our Restormer architecture (see Fig. 2). Then we describe the core components of the proposed Transformer block: (a) multi-Dconv head transposed attention (MDTA) and (b) gated-Dconv feed-forward network (GDFN). Finally, we provide details on the progressive training scheme for effectively learning image statistics.

The major computational overhead in Transformers comes from the self-attention layer. In conventional SA , the time and memory complexity of the key-query dot-product interaction grows quadratically with the spatial resolution of input, i.e., $\mathcal{O}(W^{2}H^{2})$ for images of $W<math xmlns="http://www.w3.org/1998/Math/MathML" display="block"><semantics><mrow><mo>×</mo></mrow><annotation encoding="application/x-tex">\times</annotation></semantics></math>×H$ pixels. Therefore, it is infeasible to apply SA on most image restoration tasks that often involve high-resolution images. To alleviate this issue, we propose MDTA, shown in Fig. 2(a), that has linear complexity. The key ingredient is to apply SA across channels rather than the spatial dimension, i.e., to compute cross-covariance across channels to generate an attention map encoding the global context implicitly. As another essential component in MDTA, we introduce depth-wise convolutions to emphasize on the local context before computing feature covariance to produce the global attention map.

2 Gated-Dconv Feed-Forward Network

where $\odot$ denotes element-wise multiplication, $\phi$ represents the GELU non-linearity, and LN is the layer normalization . Overall, the GDFN controls the information flow through the respective hierarchical levels in our pipeline, thereby allowing each level to focus on the fine details complimentary to the other levels. That is, GDFN offers a distinct role compared to MDTA (focused on enriching features with contextual information). Since the proposed GDFN performs more operations as compared to the regular FN , we reduce the expansion ratio $\gamma$ so as to have similar parameters and compute burden.

3 Progressive Learning

CNN-based restoration models are usually trained on fixed-size image patches. However, training a Transformer model on small cropped patches may not encode the global image statistics, thereby providing suboptimal performance on full-resolution images at test time. To this end, we perform progressive learning where the network is trained on smaller image patches in the early epochs and on gradually larger patches in the later training epochs. The model trained on mixed-size patches via progressive learning shows enhanced performance at test time where images can be of different resolutions (a common case in image restoration). The progressive learning strategy behaves in a similar fashion to the curriculum learning process where the network starts with a simpler task and gradually moves to learning a more complex one (where the preservation of fine image structure/textures is required). Since training on large patches comes at the cost of longer time, we reduce the batch size as the patch size increases to maintain a similar time per optimization step as of the fixed patch training.

Experiments and Analysis

We evaluate the proposed Restormer on benchmark datasets and experimental settings for four image processing tasks: (a) image deraining, (b) single-image motion deblurring, (c) defocus deblurring (on single-image, and dual-pixel data), and (d) image denoising (on synthetic and real data). More details on datasets, training protocols, and additional visual results are presented in the supplementary material. In tables, the best and second-best quality scores of the evaluated methods are highlighted and underlined.

Implementation Details. We train separate models for different image restoration tasks. In all experiments, we use the following training parameters, unless mentioned otherwise. Our Restormer employs a 4-level encoder-decoder. From level-1 to level-4, the number of Transformer blocks are [ $4,6,6,8$ ], attention heads in MDTA are [ $1,2,4,8$ ], and number of channels are $ $. The refinement stage contains$ 4 $blocks. The channel expansion factor in GDFN is$ \gamma{=}2.66 $. We train models with AdamW optimizer ($ \beta_{1}{=}0.9 $,$ \beta_{2}{=}0.999 $, weight decay$ 1e^{-4} $) and L1 loss for$ 300 $K iterations with the initial learning rate$ 3e^{-4} $gradually reduced to$ 1e^{-6} $with the cosine annealing . For progressive learning, we start training with patch size$ 128 $\times$ 128 $and batch size$ 64 $. The patch size and batch size pairs are updated to [($ 160^{2} $,$ 40 $), ($ 192^{2} $,$ 32 $), ($ 256^{2} $,$ 16 $), ($ 320^{2} $,$ 8 $), ($ 384^{2} $,$ 8 $)] at iterations [$ 92 $K,$ 156 $K,$ 204 $K,$ 240 $K,$ 276$K]. For data augmentation, we use horizontal and vertical flips.

We compute PSNR/SSIM scores using the Y channel in YCbCr color space in a way similar to existing methods. Table 1 shows that our Restormer achieves consistent and significant performance gains over existing approaches on all five datasets. Compared to the recent best method SPAIR , Restormer achieves $1.05$ dB improvement when averaged across all datasets. On individual datasets, the gain can be as large as $2.06$ dB, e.g., Rain100L. Figure 3 shows a challenging visual example. Our Restormer reproduces a raindrop-free image while effectively preserving the structural content.

2 Single-image Motion Deblurring Results

We evaluate deblurring methods both on the synthetic datasets (GoPro , HIDE ) and the real-world datasets (RealBlur-R , RealBlur-J ). Table 2 shows that our Restormer outperforms other approaches on all four benchmark datasets. When averaged across all datasets, our method obtains a performance boost of $0.47$ dB over the recent algorithm MIMO-UNet+ and $0.26$ dB over the previous best method MPRNet . Compared to MPRNet , Restormer has $81\%$ fewer FLOPs (See Fig. 1). Moreover, our method shows $0.4$ dB improvement over the Transformer model IPT , while having $4.4\times$ fewer parameters and runs $29\times$ faster. Notably, our Restormer is trained only on the GoPro dataset, yet it demonstrates strong generalization to other datasets by setting new state-of-the-art. Fig. 4 shows that the image produced by our method is more sharper and visually closer to the ground-truth than those of the other algorithms.

3 Defocus Deblurring Results

Table 3 shows image fidelity scores of the conventional defocus deblurring methods (EBDB and JNB ) as well as learning based approaches on the DPDD dataset . Our Restormer significantly outperforms the state-of-the-art schemes for the single-image and dual-pixel defocus deblurring tasks on all scene categories. Particularly on the combined scene category, Restormer yields ${\sim 0.6}$ dB improvements over the previous best method IFAN . Compared to the Transformer model Uformer , our method provides a substantial gain of $1.01$ dB PSNR. Figure 5 illustrates that our method is more effective in removing spatially varying defocus blur than other approaches.

4 Image Denoising Results

We perform denoising experiments on synthetic benchmark datasets generated with additive white Gaussian noise (Set12 , BSD68 , Urban100 , Kodak24 and McMaster ) as well as on real-world datasets (SIDD and DND ). Following , we use bias-free Restormer for denoising.

Gaussian denoising. Table 5 and Table 5 show PSNR scores of different approaches on several benchmark datasets for grayscale and color image denoising, respectively. Consistent with existing methods , we include noise levels $15$ , $25$ and $50$ in testing. The evaluated methods are divided into two experimental categories: (1) learning a single model to handle various noise levels, and (2) learning a separate model for each noise level. Our Restormer achieves state-of-the-art performance under both experimental settings on different datasets and noise levels. Specifically, for the challenging noise level $50$ on high-resolution Urban100 dataset , Restormer achieves $0.37$ dB gain over the previous best CNN-based method DRUNet , and $0.31$ dB boost over the recent transformer-based network SwinIR , as shown in Table 5. Similar performance gains can be observed for the Gaussian color denoising in Table 5. It is worth mentioning that DRUNet requires the noise level map as an additional input, whereas our method only takes the noisy image. Furthermore, compared to SwinIR , our Restormer has $3.14\times$ fewer FLOPs and runs $13\times$ faster. Figure 6 presents denoised results by different methods for grayscale denoising (top row) and color denoising (middle row). Our Restormer restores clean and crisp images.

Real image denoising. Table 6 shows that our method is the only one surpassing $40$ dB PSNR on both datasets. Notably, on the SIDD dataset our Restormer obtains PSNR gains of $0.3$ dB and $0.25$ dB over the previous best CNN method MIRNet and Transformer model Uformer , respectively. Fig. 6 (bottom row) shows that our Restormer generates clean image without compromising fine texture.

5 Ablation Studies

For ablation experiments, we train Gaussian color denoising models on image patches of size $128<math xmlns="http://www.w3.org/1998/Math/MathML" display="block"><semantics><mrow><mo>×</mo></mrow><annotation encoding="application/x-tex">\times</annotation></semantics></math>×128$ for $100$ K iterations only. Testing is performed on Urban100 , and analysis is provided for a challenging noise level $\sigma{=}50$ . FLOPs and inference time are computed on image size $256<math xmlns="http://www.w3.org/1998/Math/MathML" display="block"><semantics><mrow><mo>×</mo></mrow><annotation encoding="application/x-tex">\times</annotation></semantics></math>×256$ . Table 7-10 show that our contributions yield quality performance improvements. Next, we describe the influence of each component individually.

Improvements in multi-head attention. Table 7c demonstrates that our MDTA provides favorable gain of $0.32$ dB over the baseline (Table 7a). Furthermore, bringing locality to MDTA via depth-wise convolution improves robustness as removing it results in PSNR drop (see Table 7b).

Improvements in feed-forward network (FN). Table 7d shows that the gating mechanism in FN to control information flow yields $0.12$ dB gain over the conventional FN . As in multi-head attention, introducing local mechanism to FN also brings performance advantages (see Table 7e). We further strengthen the FN by incorporating gated depth-wise convolutions. Our GDFN (Table 7f) achieves PSNR gain of $0.26$ dB over the standard FN for the noise level $50$ . Overall, our Transformer block contributions lead to a significant gain of $0.51$ dB over the baseline.

Design choices for decoder at level-1. To aggregate encoder features with the decoder at level-1, we do not employ $1<math xmlns="http://www.w3.org/1998/Math/MathML" display="block"><semantics><mrow><mo>×</mo></mrow><annotation encoding="application/x-tex">\times</annotation></semantics></math>×1$ convolution (that reduces channels by half) after concatenation operation. It is helpful in preserving fine textural details coming from the encoder, as shown in Table 8. These results further demonstrate the effectiveness of adding Transformer blocks in the refinement stage.

Impact of progressive learning. Table 9 shows that the progressive learning provides better results than the fixed patch training, while having similar training time.

Deeper or wider Restormer? Table 10 shows that, under similar parameters/FLOPs budget, a deep-narrow model performs more accurately than its wide-shallow counterpart. However, the wider model runs faster due to parallelization. In this paper we use deep-narrow Restormer.

Conclusion

We present an image restoration Transformer model, Restormer, that is computationally efficient to handle high-resolution images. We introduce key designs to the core components of the Transformer block for improved feature aggregation and transformation. Specifically, our multi-Dconv head transposed attention (MDTA) module implicitly models global context by applying self-attention across channels rather than the spatial dimension, thus having linear complexity rather than quadratic. Furthermore, the proposed gated-Dconv feed-forward network (GDFN) introduces a gating mechanism to perform controlled feature transformation. To incorporate the strength of CNNs into the Transformer model, both MDTA and GDFN modules include depth-wise convolutions for encoding spatially local context. Extensive experiments on 16 benchmark datasets demonstrate that Restormer achieves the state-of-the-art performance for numerous image restoration tasks.

Acknowledgements. Ming-Hsuan Yang is supported by the NSF CAREER grant 1149783. Munawar Hayat is supported by the ARC DECRA Fellowship DE200101100. Special thanks to Abdullah Abuolaim and Zhendong Wang for providing the results.