A Decoder-free Transformer-like Architecture for High-efficiency Single Image Deraining

1 Introduction

Image deraining is a classical problem in computer vision, which is highly desired in consumer photography and image processing. However, it is a challenging task since distant rain streaks usually combine with water droplets to form a veil over the backdrop, substantially reducing the visibility of the image. Recently, with the emergence of convolutional neural networks (CNNs), the state-of-the-art (SOTA) methods for image deraining are dominated by CNNs. However, the receptive field of convolution operations is limited, and CNNs usually have bulky structures to boost performance.

Transformer, first applied in natural language processing, is increasingly being used to replace CNNs in various vision tasks. Pioneering works in Transformer [Vaswani et al., 2017], [Liu et al., 2021] widely involve self-attention mechanisms as basic blocks to realize a strong feature representation ability. Most Transformer-based methods consist of three parts: a) patch generator; b) MLP; c) multi-head selfattention (MSA). In general, Transformer can be written in the following form:

* Corresponding authors.

where X ∈ RN×C denotes the input features, N is the number of patches and C is channels. WQ, WK, WV are weights of query (Q), key(K) and value (V) in the linear projections (LP). FFN(·) is the feed-forward network. The complete Transformer is achieved by stacking Eq. 1 to build Transformer encoder block (TEB) and Transformer decoder block (TDB). Therefore, the computational complexity of vanilla Transformer is O(N 2 + NC) that is more cumbersome than CNNs. Meanwhile, we could build cross-covariance attention (XCA) by matrix multiplication VT KQT , which is O(NC + C 2 ). The computational complexity is quadratic with the channels on the input images.

By analyzing previous works such as IPT [Chen et al., 2021], Uformer [Wang et al., 2021b], and Restormer [Zamir et al., 2021] (please refer to Fig. 1), we find redundancy that they just simply apply TEB-to-TEB Transformer as the encoder-decoder architecture and decoders perform feature interaction across inputs of encoders rather than encoders and TDB, as shown in Fig. 2 (a). In other words, these decoders actually result in a slight or no improvement in their performance, but with high costs in terms of the parameter number, FLOPs, and training/inference time. It motivates us to build a decoder-free Transformer-like (DFTL) framework.

In this paper, we first adopt a cheap linear projection (LP) to generate intrinsic feature maps with lower complexity. Then, we introduce the Transformer encoder block (TEB) to represent desired features. Finally, Progressive Patch Merging (PPM) is proposed to restore rich features with different spatial resolution representations. Meanwhile, the architecture is still under the self-attention mechanism, which attains the powerful contextual representation ability of the vanilla Transformer. The advances of DFTL can be summarized as 1) Efficient Transformer: propose cost-effective operations for better results. 2) Flexible structure: a scalable framework for the restoration of rainy images.

The contributions of this work are as follows:

1. We propose a new DFTL architecture that could achieve competitive performance with less GPU memory and hold satisfying complexity/performance trade-offs compared to other self-attention-based techniques.
2. We adopt a Cheap LP and MSA/XCA to capture the multi-scale contextual details. Based on these modules optimized for Transformer, we develop two variants named DFTL-W and DFTL-X, which consistently achieve comparable performance to the prior arts.
3. A new hybrid loss is designed for more effective training, which could favor the potential convergence and improve the final testing performance expectedly.

2 RelatedWorks

In general, existing single image deraining methods can be roughly divided into two categories, i.e., optimization-based and deep learning-based methods.

2.1 Optimization-based DerainingMethods

Optimization-based deraining methods view the rainy image as components assembled with the background image B and the rain streaks R. The whole process can be expressed as the following formula:

To remove R and obtain a clean image Y, several works are proposed for single image deraining by founding effective optimization models based on image priors, e.g., directional sparsity prior [Jiang et al., 2019; Jiang et al., 2017; Deng et al., 2018]. With the development of deep learning, optimization-based techniques may be insufficient. They are only capable of dealing with particular situations.

2.2 Deep Learning-based DerainingMethods

CNN-based techniques. With the powerful representation and extraction ability of CNN, diverse CNN-based structures are designed to improve the performance of image deraining. In [Li et al., 2018], the authors utilize dilated CNN and squeeze-and-excitation blocks to deal with the task of image deraining, obtaining superior outcomes compared to traditional optimization-based methods. PReNet given by [Ren et al., 2019] is to deepen shallow ResBlock via recurrent operations progressively, achieving competitive results. Wang et al. introduce a convolutional dictionary learning mechanism to remove rain streaks [Wang et al., 2020] effectively. In addition, Fu et al. firstly attempt to use the so-called graph convolution network to extract relation-aware features and exploit pixel-level global spatial relationship [Fu et al., 2021], getting the SOTA deraining results. However, all the techniques mentioned above are based on the characteristics of CNN, which requires stacking several convolution blocks to enlarge the receptive field and improve performance, lacking long-range interactions and ignoring geometric details extraction.

Vision Transformer techniques. In [Dosovitskiy et al., 2021], it introduces Vision Transformer (ViT) to treat input image as 16 × 16 words and attains excellent results on image recognition. Besides, IPT [Chen et al., 2021] pretrains the model on ImageNet dataset and achieves SOTA on several low-level visual tasks. SwinIR [Liang et al., 2021] builds a single way structure based on Swin Transformer to achieve image restoration. More recently, Restormer [Zamir et al., 2021] designs Transformer with convolution projections and cross-covariance across feature channels, further improving the structures of Transformer.

3 ProposedMethod

3.1 Pipeline

An overview of the DFTL structure is presented in Fig. 2 (b). Given an image I ∈ RH×W×3 , we first progressively divide it into 4-level patches by Patchify module. These multi-scale patches are {1, 1/2, 1/4, 1/8} relative to H,W, termed as Hl ,W l , abbreviated as P l standing for patch size P at the specific scale l by the following Eq. 3:

where N is the number of patches, k is patch/kernel size, pad means padding, and s denotes stride. In what follows, we describe proposed modules: (a) Cheap Linear Projection (LP); (b) Transformer Encoder Block (TEB); (c) Progressive Patch Merging (PPM).

3.2 Cheap Linear Projection

Linear projection (LP) could represent patch-level features into long range dependencies. However, according to Eq. 1, we can find one of complexity of Transformer is LP. Thus, we sort out several classic LP variants and show Cheap LP in Fig. 3. Convolution LP embeds spatial relationships into feature maps based on strided convolution operations. ViT

LP projects the flattened patches and window based LP represents window based self-attention (W-MSA) along channel direction. Compared to depthwise convolution, they would occupy more memory and FLOPs. Hence, given the input X ∈ RC×H×W , X could be partitioned into patches {Xp ∈ Rk×k×C }Pp=1 by Eq. 3, where Xp is pth patch (p = 1, · · · , P). Cheap LP can be expressed as,

3.3 Transformer Encoder Block

Self-attention (SA) is the core of Transformer but is infeasible to GPU overheads. Thus, we analyze two SA mechanisms in this section, including 1) W-MSA through introducing one local window to consistent with the efficiency and performance, 2) XCA by channel correlation with a lower spatial cost.

Window Based Self-Attention (W-MSA). To reduce cost in spatial resolution, W-MSA encodes local pixel similarity in a window M × M, whose computational complexity is O(M 2 + MC), which reduces the computational/memory overhead. With W-MSA followed by FFN module, TEB can be computed as follows,

where Xp and Yp denote the input and output of TEB. Zp indicates the intermediate feature maps, and d is a scalar. DFTL-W is built by the sequential form of LP and W-MSA.

Cross-Covariance Attention (XCA). Considering the feature may have more channels, DFTL-W may produce limited performance because W-MSA only encodes feature vectors in each pixel. We introduce XCA to compute channel correlation along channel dimension:

where KT Q is also termed as the attention matrix with RC×C . The model with XCA is named as DFTL-X, which is the sequential form of TEB and LP.

3.4 Progressively PatchMerging (PPM)

It is redundant that decoders perform feature interaction across the Query, Key and Value from encoders. Considering generations of patch embeddings, we develop PPM to replace decoders of Transformer. Specifically, we use bilinear interpolation to upsample the patch-level features. Then we restore detailed information and shrink the channels to obtain the outputs by two convolution layers. In this process, we use skip-connection to achieve concatenation of patches progressively. The outputs of PPM are then added with the original inputs to remove rain streaks and produce clean images. The whole process is presented as follows:

where Fl represents l th-level feature maps, I and O are the input and output of our DFTL, respectively.

3.5 Hybrid Loss Function

Our work proposes a novel gradient-based hybrid loss function (GBHL) to achieve better results. We empirically present the Eq. 8, in which the significant numerical value is provided by SSIM [Wang et al., 2004] or mean square error (MSE), etc. Mathematically, we impose a regularizer as a constraint for SSIM or MAE as follows:

Here ⊙ represents hadamard product. ∥Lk∥F is set to requires grad = False. ∥ · ∥F is Frobenius norm, and k = [ssim, mae, mse].

4 Experiments

In this section, we demonstrate the advantages of proposed method via comprehensive experiments on both synthetic and real datasets. In particular, we also compare our DFTL with other Transformer architectures to prove the efficiency of our methods. Please refer to the supplementary materials for more details, e.g., datasets, implementation details and discussion of model structures and complexity.

4.1 Comparison on Synthetic and Real Datasets

We evaluate our model on five synthetic datasets to compare quantitative results. Besides, we list the maximum batch (MaxBatch) of each model that can be trained simultaneously on single GPU. Also, we perform tests on two real scenarios. Synthetic cases. The performance of all compared methods on five synthetic datasets is reported in Table 1. It can be observed that our models could exceed the most advanced techniques. Besides, we show visual comparisons on Rain200H, DDN, and DID datasets, see Fig. 4. For rainy streaks in different directions, our methods could preserve more details and attain better disentanglement from complex scenes. We also visualize the trade-off analysis between latency and performance among these deraining models in Fig. 5.

Real cases. To demonstrate the generalization ability, we perform the visual evaluation on the real-world rainy images in Fig. 6. It can be observed that DFTL-W/-X could remove more rainy streaks than other methods and have better visual qualities on the whole image.

4.2 Comparison with other Transformers

In this section, we perform a comparison with existing vision Transformers to verify the effectiveness of DFTL targeting the specific image deraining task. We select seven classic Transformer architectures, i.e., ViT256, DETR [Carion et al., Latency (s/it)

Figure 5: Comparison between latency and performance.

2020], DeformDETR [Zhu et al., 2021], Uformer [Wang et al., 2021b], IPT [Chen et al., 2021] and Restormer [Zamir et al., 2021]. All models are trained in the same framework with default settings as their original codes. In Table 2, we summarize the PSNR and SSIM of all outcomes produced by different architectures on Rain200L dataset. Besides, we report parameter number (Param) and FLOPs, demonstrating that our DFTL can significantly reduce computational costs. Furthermore, we show the comparisons on PSNR and latency in Fig. 5. Our methods could obtain better PSNR in an efficient manner.

4.3 Ablation Study

Ablation on Cheap LP. In Table 3, we compare different locations replacing LP with Cheap LP in DFTL-X over Rain200H dataset. Cheap LP can achieve competitive performance and high efficiency. On the one hand, benefiting from depthwise operation, Patchify with Cheap LP could generate more feature maps to improve performance with lower complexity. On the other hand, the network with Cheap LP leads to weak performance because of poor feature extraction.

Ablation on DFTL. We compare the effects of different selfattention in DFTL-X on Rain200L dataset. In Table 4, SA and XCA significantly surpass other compared counterparts. However, SA has higher computational loads and better results than XCA. W-MSA has the same parameter number as SA and larger memory than XCA. Compared to it, DFTLW avoids the problem by employing the sequential form of LP and W-MSA (see more details in the supplementary material). Thus, it attains competitive results and handles larger MaxBatch with limited resources in Table 1.

Ablation on PPM. This part investigates the effects of conventional decoders and proposed PPM on two representative Transformers, i.e., IPT and Restormer. Since IPT is a single way architecture to restore images, PPM is used for multiscale architecture and can't be directly integrated into IPT. Thus we only ablate the decoders of IPT. Restormer's decoders are replaced with PPM to validate the effectiveness. Through the comparisons in Fig. 1, decoders of Transformer have a slight performance gain but lead to colossal GPU memory occupations (2x∼3x) and costs of latency (∼1.5x). Thus, it is significant for an efficient Transformer to build a decoder-free Transformer-like architecture.

Ablation on Hybrid Loss. We compare the results between different weighted combinations and proposed GBHL. Due to the numerical differences between LMAE and LSSIM are about 10 times in the initial training phase, α is set to 10. Two groups shown in Table 5 are adopted to conduct a comparison on Rain200H. Experimental results show that our GBHL makes the highest value for both PSNR and SSIM.

5 Conclusion

This paper presents a decoder-free Transformer-like architecture (DFTL) for image deraining to analyze popular Transformer architectures from a new perspective. It reveals that decoders are redundant to Transformer. Proposed modules are more computationally efficient compared with standard Transformer modules. The comparisons with several competitive Transformers show our methods have good feature representation ability at low computational costs. Moreover, we propose a novel gradient-based hybrid loss function to produce more reasonable results. Extensive experiments demonstrate DFTL can achieve the comparable to SOTA methods.

Acknowledgments

This research is supported by NSFC (12171072, 61876203, 61702083), Key Projects of Applied Basic Research in Sichuan Province (Grant No. 2020YJ0216), and National Key Research and Development Program of China (Grant No. 2020YFA0714001)