Strip Attention for Image Restoration

1 Introduction

Image restoration aims to reconstruct a high-quality image from the observation suffering from various degradations (e.g., blur, snowflake, haze), playing an essential role in many fields, such as surveillance, medical imaging, and remote sensing. It is an inverse problem and has an ill-posed nature. To resolve this challenging problem, a multitude of conventional algorithms have been developed based on handcrafted features, which are impractical in more complicated real-world scenarios [Zhang et al., 2022].

In recent years, convolutional neural networks (CNNs) have witnessed a significant development of image restoration and achieved remarkable performance compared to traditional approaches by virtue of the powerful mapping capability. A great number of CNN-based methods have been proposed for varied image restoration tasks by designing or borrowing advanced units, including U-shaped backbone [Lee et al., 2021], residual connection [Cho et al., 2021], dilated convolution [Son et al., 2021], and attention modules [Qin et al., 2020; Cui et al., 2023b]. Nevertheless, CNN has two defects that are not beneficial for image restoration: (a) The convolution operator has static filters that are not applicable to the dynamic and non-uniform blur. (b) The convolution filter has a limited receptive field that is not capable of modeling longrange pixels interactions for large-size blur. Despite many efforts to enlarge the receptive field by stacking deep layers or using dilated convolution [Son et al., 2021], these remedies entail heavy computation burden and still struggle to obtain the global receptive field.

More recently, Transformer model borrowed from natural language processing has shown state-of-the-art performance on high-level vision tasks. The core element, self-attention mechanism, is capable of modeling long-range dependencies effectively. However, its quadratic complexity with respect to the spatial resolution makes it infeasible for image restoration, which always involves high-resolution images, e.g.×1120 image size for defocus deblurring in

DPDD [Abuolaim and Brown, 2020]. To alleviate this issue, many measures have been taken to improve efficiency in the realm of image restoration. For instance, a few methods restrict the operation region of self-attention to reduce complexity [Liang et al., 2021; Wang et al., 2022]. Restormer [Zamir et al., 2022] applies self-attention among the channel dimension rather than the spatial dimension. Stripformer [Tsai et al., 2022] develops strip-type self-attention for image deblurring. Though these methods realize the goal of reducing complexity to some extent, they do not break the nature of selfattention, i.e., they still have quadratic complexity to the size of windows, channels, or strips.

In this paper, we exploit a strip attention mechanism for image restoration to harvest contextual information and meanwhile maintain high efficiency. Concretely, for each pixel, we perform information aggregation from its adjacent pixels in the same horizontal or vertical direction. This process is guided by the weights generated by convolutional layers. With joint horizontal and vertical aggregation, each location can implicitly receive information from a large region centered at itself. Furthermore, to enhance feature representation learning, we empirically adopt distinct receptive fields in feature groups to deal with degradation blurs of different sizes.

Our strip attention module has several key advantages. Firstly, by disintegrating attention into two directions, it significantly improves efficiency and can achieve large-scale receptive fields with negligible introduced computational complexity. Secondly, compared to the static filters of convolution operators, it is content-aware to adapt to the different input and blur. Thirdly, it is capable of capturing multi-scale contextual information. Our design is distinguished from other strip-type attention approaches. Specifically, CCNet [Huang et al., 2019] utilizes recurrent criss-cross attention to capture full-image dependencies for semantic segmentation. CSWin transformer [Dong et al., 2022] and Stripformer [Tsai et al., 2022] execute self-attention within the strip-shaped regions in different directions. The attention weights of these methods are produced by matrix multiplication or affinity operation, which entails quadratic complexity. Differently, we generate weights from a simple bypass network and conduct integration in a cheap manner. Moreover, we exploit multi-scale receptive fields to boost performance.

Equipped with the proposed strip attention module, our SANet performs favorably against state-of-the-art algorithms on several image restoration tasks. For dehazing, as shown in Figure 1, SANet outperforms PMNet [Ye et al., 2022] by 2.99 dB on the SOTS-Indoor [Li et al., 2018] benchmark with 54% fewer MACs. For the defocus blur removal, SANet obtains 26.29 dB PSNR on DPDD [Abuolaim and Brown, 2020], an improvement of 0.31 dB over the strong Transformer model Restormer [Zamir et al., 2022]. Our model also displays the potential on the desnowing task, surpassing NAFNet [Chen et al., 2022] by 1.26 dB on CSD [Chen et al., 2021].

The main contributions of the paper are as follows:

• We propose a strip attention module for image restoration that integrates multi-scale contextual information efficiently by performing horizontal and vertical local attention successively.
• Based on the proposed strip attention module, we establish SANet that performs favorably against state-of-theart algorithms on several image restoration tasks.

2 RelatedWork

2.1 Image Restoration

Since image restoration plays an important role in photography, self-driving techniques, and medical imaging, it has drawn substantial attention from the industrial community and academia. This inverse problem has an ill-posed nature. To constrain the solution space, a flurry of conventional methods have been developed based on various assumptions and hand-crafted features [Zhang et al., 2022]. Lately, the datadriven CNN-based frameworks have significantly advanced the performance of image restoration [Ren et al., 2016; Ren et al., 2018; Cui et al., 2023a]. Among these networks, the U-shaped architecture [Ronneberger et al., 2015] is a popular solution for hierarchical feature representation learning. Besides, numerous advanced modules have been created or borrowed from high-level tasks, including dilated convolution [Son et al., 2021], skip connection [Liu et al., 2019b], and multifarious attention mechanisms [Qin et al., 2020]. More recently, Transformer models have been introduced into low-level tasks to help model long-range dependencies [Liang et al., 2021].

2.2 AttentionMechanism

Attention mechanisms have been widely used in the computer vision community. In the context of image restoration, a great number of attention modules have been developed to capture inter-dependencies along channels [Liu et al., 2019a; Zamir et al., 2022], spatial coordinates [Zamir et al., 2021], or both [Chen et al., 2023]. For instance, FFA-Net [Qin et al., 2020] leverages channel attention and pixel attention to deal with different types of information flexibly. GridDehazeNet [Liu et al., 2019a] utilizes channel-wise attention to adjust the contributions of different streams for feature fusion. MPRNet [Zamir et al., 2021] leverages the supervised attention module for feature filtering. These attention modules have boosted the performance of image restoration tasks.

Another line of this topic is to devise efficient self-attention for image restoration. Specifically, resembling Swin Transformer [Liu et al., 2021], Uformer [Wang et al., 2022] and SwinIR [Liang et al., 2021] apply self-attention within local regions. Restormer [Zamir et al., 2022] switches its focus from spatial dimensionality to channel self-attention. Stripformer [Tsai et al., 2022] develops interlaced intra-strip and inter-strip attention layers for motion blur removal. However, these remedies still have quadratic complexity to the size of the region, channel, or strip.

In this paper, we present an ingenious strip attention module that performs efficient information integration in horizontal and vertical directions successively. Compared to the convolution operator, our paradigm not only inherits its high efficiency but also produces dynamic aggregation weights and an enlarged receptive field.

3 Methodology

In this section, we first describe the strip attention operation and then present the strip attention module. Next, we delineate the architecture of SANet for image restoration. Finally, we introduce the loss functions used for training.

Skip Connection

3.1 Strip Attention

Our main goal is to exploit a unit that can perform information integration efficiently and effectively. Before describing the formulation of the proposed strip attention, we first provide complexity analyses of self-attention.

Self-Attention

Self-attention has achieved successful stories in high-level vision tasks. However, due to its quadratic complexity, it is infeasible for image restoration tasks that always involve high-resolution images. Formally, give an input tensor X ∈ RH×W×C , where H ×W denotes spatial coordinates and C is the number of channels, self-attention can be expressed as,

where Q, K, V ∈ RHW×C , which are generated by using corresponding projection matrices (WQ,WK , and WV ) and reshaping. We omit the normalization term for simplicity.

From Eq. 1, we can observe that the complexity of selfattention comes from three aspects: (a) the production of query (Q), key (K), and value (V) with the complexity of 3HWC 2 ; (b) generation of the attention map based on key-query dot-product with the complexity of (HW)2C; (c) the weighted summation process with the complexity of (HW)2C. We can see that in the last two terms, the complexity is quadratic to the spatial size.

Strip Attention

We aim to devise an efficient operator for information aggregation from the perspective of reducing the complexities of the above-mentioned three steps. Here, we take the horizontal strip attention as an example. Concretely, given any input feature X∈ RH×W×C , we remove the procedure of producing Q, K, and V, and instead directly yield the attention weights via an extremely lightweight branch that consists of global average pooling (GAP) followed by 1 × 1 convolution layer and Sigmoid function. This process can be formally expressed as,

where W1×1 is a 1 × 1 convolution layer and σ denotes the Sigmoid function. A∈ RK , where K specifies the length of the strip for integration. Note that we share the resulting attention weights across both spatial and channel dimensions for further efficiency.

Regarding the weighted sum operation, instead of operating on the whole image like self-attention or on the strip of size n ×W (n < H) like Stripformer [Tsai et al., 2022] and CSWin Transformer [Dong et al., 2022], we execute our information integration within the strip of size 1×K (K < W) based on the obtained attention weights, which can be formally expressed as,

Rather than generating the attention weights with a similar shape to that of self-attention and then performing integration via matrix multiplication, inspired by [Zhou et al., 2021], we adopt a more reasonable convolution-type integration as shown in Figure 3 (c), where each pixel receives information from the region centered at itself.

To summarize, our strip attention operator can be formally expressed as:

3.2 Strip AttentionModule

It has been illustrated in prior works [Zamir et al., 2022; Wang et al., 2022] that enlarging the receptive field of the network is beneficial to image restoration. Motivated by this fact, we present an efficient manner to expand the receptive field of each pixel by exploiting the above-mentioned strip attention operator. Specifically, we develop a strip attention module that carries out strip attention operations in both vertical and horizontal directions to harvest long-range contexts, as shown in the bottom part of Figure 2. Furthermore, we combine different K within each attention operation to pursue multi-scale receptive fields. More concretely, with the input tensor X∈ RH×W×C , we first divide it into two parts by splitting the channel dimension evenly and then impose the horizontal strip attention on each part separately with different strip lengths K. Next, we perform the multi-scale strip attention in the vertical direction. The final output is produced by adding the original input X. The entire process of the proposed strip attention module can be formally expressed as:

where S Ho and S V e denote the horizontal attention and vertical attention, respectively; [·,·] is concatenation; X1 and X2 are obtained by splitting the feature on channel dimension evenly. Our strip attention module implicitly enlarges the receptive field of the network. As shown in Figure 3, the horizontal and vertical strip attention perform information integration in two directions, respectively. For convenience, we only pick a few representative pixels for illustration. The horizontal one gives B = wABA + wBBB + wCBC, where w denotes the attention weight. By using two-directional strip attention successively, the value of pixel D in Figure 3 (b) is computed by:

As a consequence, the pixel in the center receives contexts from the whole region determined by K.

3.3 Overall Architecture

The overall pipeline of the proposed SANet is illustrated in Figure 2 (Top). SANet adopts the popular encoder-decoder architecture to learn hierarchical representations efficiently and consists of six scales in total. Specifically, given a degraded image with the shape of R3×H×W , a single convolution layer is utilized to generate the shallow feature map of size RC×H×W . Then, the resulting feature is fed into the encoder layers (Scale 1-3). In this process, the number of channels is expanded, while the spatial resolution is reduced gradually from RC×H×W to R4C× H4 × W4 . Each stage contains a stack of residual blocks, and the last one involves the proposed strip attention module. The downsampling operation is accomplished by the strided convolution. Next, the feature with the lowest resolution passes through the decoder layers (Scale 4-6) to recover the high-resolution representations progressively. For feature upsampling, we adopt the transposed convolution. To alleviate the issue of information loss caused by downsampling, we apply the featurelevel skip connections as previous works [Zamir et al., 2022;

Wang et al., 2022]. Concretely, the encoder features are concatenated with the corresponding decoder features, followed by a convolution layer to adjust the channel dimension. The final sharp image is produced by adding the original input image, which forces the network to focus only on the residual information learning. Besides, to ease the training difficulty, multi-input and multi-output strategies are adopted following recent methods [Cho et al., 2021; Mao et al., 2021].

3.4 Loss Functions

To facilitate feature refinement in spatial and frequency domains simultaneously, we use the dual-domain L1 loss [Cho et al., 2021] to train our network. For each output, the loss function is given by:

where ˆI, I are the predicted image and ground-truth, respectively; S depicts the total elements for normalization; and F is the fast Fourier transform (FFT). λ is set to 0.1.

4 Experiments

To verify the effectiveness of our SANet, we conduct extensive experiments on several image restoration tasks, including single-image defocus deblurring (DPDD [Abuolaim and Brown, 2020]), image dehazing (RESIDE [Li et al., 2018]), and image desnowing (CSD [Chen et al., 2021]). In the following, we first introduce the training settings, and then we report our results on the above datasets. Finally, we carry out a series of ablation experiments.

4.1 Implementation Details

We train the proposed network via Adam optimizer with β1 = 0.9, β2 = 0.999. The initial learning rate is set to 1e-4 and reduced to 1e-6 gradually with the cosine annealing. The batch size is set as 8 for the RESIDE-Outdoor [Li et al., 2018] dataset and 4 for others. Models are trained on the patch size of 256×256. We adopt only horizontal flips for data augmentation. We choose k1 = 7 and k2 = 11 in Eq. 5. According to the task complexity, we deploy varying numbers of residual blocks N in each scale for different tasks, i.e., N = 4 for image dehazing and desnowing, and N = 16 for image defocus deblurring.

PSNR

Reference

18.91 dB Input

27.07 dB FFA-Net

24.64 dB MAXIM

27.03 dB DeHamer

31.24 dB SANet

Figure 5: Image dehazing comparisons on the SOTS-Outdoor [Li et al., 2018] dataset among FFA-Net [Qin et al., 2020], MAXIM [Tu et al., 2022], DeHamer [Guo et al., 2022], and our SANet.

4.2 Main Results

Image dehazing. We train the network on the RESIDE [Li et al., 2018] dataset and test on the SOTS [Li et al., 2018] dataset. The results are reported in Table 1. Our SANet achieves better performance with lower complexity than most approaches. Particularly on the SOTS-Outdoor dataset, SANet yields a 2.83 dB performance gain over the expensive Transformer model DeHamer [Guo et al., 2022] with only 76%MACs and 3% parameters. Compared to the recent algorithm DehazeFormer-L [Song et al., 2022], our model surpasses it by 0.35 dB in terms of PSNR on SOTS-Indoor, while having 6.68× fewer parameters and 7.5× fewer MACs. The qualitative comparisons on the SOTS-Indoor and SOTSOutdoor datasets are exhibited in Figure 4 and Figure 5, respectively. We can see that SANet is more effective in removing haze blur, and the images produced by our model are visually closer to the target ones than other algorithms.

Single-image defocus deblurring. We compare image fidelity scores of our method with both learning-based singleimage defocus deblurring methods and conventional ones, e.g., JNB [Shi et al., 2015] and EBDB [Karaali and Jung, 2018], on the DPDD [Abuolaim and Brown, 2020] dataset. The comparison results in Table 2 show that our model outperforms the strong Transformer model Restormer [Zamir et al., 2022] in most cases. Particularly in the indoor scene category, SANet produces a substantial gain of 0.43 dB over Restormer. Furthermore, our method outperforms DRBNet [Ruan et al., 2022] by 0.56 dB PSNR on the combined category. The visual results in Figure 6 illustrate that the proposed network recovers more faithful details than other competitive frameworks.

Image desnowing. The desnowing comparisons on the CSD [Chen et al., 2021] dataset are provided in Table 3. We can see that our method obtains higher scores than other approaches. Compared to the recent algorithm NAFNet [Chen et al., 2022], SANet provides a performance boost of 1.26 dB PSNR. Furthermore, our model shows a 2.64 dB improvement over the Transformer model MSP-Former [Chen et al., 2023]. Visual results presented in Figure 7 show that our SANet generates a cleaner image than other algorithms.

4.3 Ablation Studies

For ablation experiments, we study diverse design choices for the strip attention module, including the combination pattern of strip attention, different strip lengths, and activation functions. Furthermore, we compare our module with other attention units and depth-wise convolution to demonstrate the effectiveness of our method. To this end, we train SANet on the dehazing task with the RESIDE-Indoor [Li et al., 2018] dataset. Unless specified otherwise, the hyperbolic tangent function serves as the activation function in Eq. 2, and we only adopt the single receptive field for the strip attention module with K = 5. The training configurations are consistent with the main experiment except that N is set to 1.

MACs are computed on the size of 256 × 256. The baseline model is obtained by removing the proposed attention module from our model.

Improvements of strip attention module. Table 4 shows that two-directional strip attention units both produce favorable gains over the baseline model with negligible introduced parameters and complexity. Using two strip attention operators in different directions leads to further accuracy improvement. Especially for the horizontal-vertical version, our model achieves a gain of 3.39 dB over the baseline, while only consuming additional 0.04 M parameters and 0.07 G

Table 6: Comparisons with other attention modules.

MACs, illustrating the effectiveness of our design.

Different receptive fields. We further exploit the impact of the receptive field by changing the strip size in the strip attention module. The results are shown in Figure 8. As the increase of the strip length, we can observe a consistent improvement in terms of PSNR. Our model receives a remarkable gain of 1.36 dB when the receptive field is enlarged from 3 to 11, while only introducing 0.12 G MACs. Furthermore, to deal with blurs of various sizes, we adopt the multi-scale receptive fields, i.e., 7 and 11, as we elaborate in Sec. 3.2. This strategy leads to 35.45 dB PSNR, 0.05 dB and 0.14 dB higher than a single kernel 11 and 7, receptively.

Design choices for activation function. Instead of inheriting the Softmax function from the canonical self-attention, we explore more choices in Table 5 based on the verticalhorizontal variant in Table 4. The Sigmoid version obtains higher accuracy than Softmax by breaking the sum-to-one property and is 0.32 dB higher than that of the Tanh version.

Comparisons with alternatives. As our strip attention module implicitly receives the same receptive field as the depth-wise convolution when K is equal to the kernel size of the latter, we compare our module with depth-wise convolution in Figure 8. With the same receptive field, our model consistently outperforms the depth-wise convolution version with fewer extra complexities. To further verify the superiority of our method, we provide comparisons between other self-attention units and our module. We can see from Table 6 that, with the best choices of receptive field, activation, and combination order of strip attention units, our final design is superior to the global self-attention and the window-based variant in terms of accuracy and computation overhead. Due to the large complexities of global self-attention, we only insert it into scale 3-4, which have the lowest resolution.

5 Conclusion

In this paper, we develop a novel image restoration model that is computationally efficient in integrating contexts for feature representation enhancement. Specifically, our strip attention unit realizes efficient information aggregation by modifying the three steps of self-attention, while maintaining the content-aware property based on the learned attention weights. Furthermore, the proposed strip attention module enlarges the receptive field by combining two-directional strip attention units and adopts multi-scale kernels to well handle blurs with various sizes. Comprehensive experiments on several image restoration tasks demonstrate that SANet performs favorably against state-of-the-art algorithms.