Source-linked AI summary
Learning A Sparse Transformer Network for Effective Image Deraining
Xiang Chen, Hao Li, Mingqiang Li, Jinshan Pan
TL;DR
Image deraining is challenged by complex rain patterns, limited CNN receptive fields, and irrelevant global interactions in Transformer attention. DRSformer addresses these issues with sparse top-k attention, mixed-scale feature processing, and expert-based refinement, outperforming existing derainers on benchmark datasets, including a 0.4 dB average PSNR gain over IDT.
Problem
Image deraining methods remain limited by handcrafted priors, CNNs’ local receptive fields, and Transformers’ insufficient modeling of local image details.
Method
DRSformer combines learnable top-k sparse attention, mixed-scale feed-forward processing, and mixture-of-experts feature compensation for image deraining.
Results
0.4 dB average PSNR gain over IDT accompanies DRSformer’s higher performance than other derainers across reported benchmark datasets.
Takeaways & Limitations
DRSformer can handle diverse spatially varying rain streaks while preserving finer reconstructed image details.
Abstract
from arXiv · showhide
Transformers-based methods have achieved significant performance in image deraining as they can model the non-local information which is vital for high-quality image reconstruction. In this paper, we find that most existing Transformers usually use all similarities of the tokens from the query-key pairs for the feature aggregation. However, if the tokens from the query are different from those of the key, the self-attention values estimated from these tokens also involve in feature aggregation, which accordingly interferes with the clear image restoration. To overcome this problem, we propose an effective DeRaining network, Sparse Transformer (DRSformer) that can adaptively keep the most useful self-attention values for feature aggregation so that the aggregated features better facilitate high-quality image reconstruction. Specifically, we develop a learnable top-k selection operator to adaptively retain the most crucial attention scores from the keys for each query for better feature aggregation. Simultaneously, as the naive feed-forward network in Transformers does not model the multi-scale information that is important for latent clear image restoration, we develop an effective mixed-scale feed-forward network to generate better features for image deraining. To learn an enriched set of hybrid features, which combines local context from CNN operators, we equip our model with mixture of experts feature compensator to present a cooperation refinement deraining scheme. Extensive experimental results on the commonly used benchmarks demonstrate that the proposed method achieves favorable performance against state-of-the-art approaches. The source code and trained models are available at https://github.com/cschenxiang/DRSformer.
1. Introduction
DRSformer addresses image deraining as an ill-posed inverse problem where handcrafted priors and CNN locality limit restoration under complex rain. Its sparse Transformer architecture uses learnable top-k attention and mixed-scale feed-forward processing to improve feature aggregation and multi-scale representations, with favorable benchmark performance against SOTA methods.
- Problem: Image deraining is an ill-posed inverse problem because both the clean image and rain streaks are unknown.Handcrafted priors based on statistical properties are not robust to complex and varying rainy scenarios, limiting deraining performance.
- Motivation: CNNs’ local receptive fields and input-content independence hinder removal of long-range rain degradation perturbations.Transformers have therefore been applied to image deraining to address these limitations.
- Motivation: Dense self-attention aggregates relations from irrelevant key tokens, allowing smaller similarities to interfere with latent clear image restoration.This deficiency arises from the native dense calculation pattern of standard Transformers.
- Method: DRSformer introduces sparse attention that adaptively retains the most useful self-attention values for feature aggregation.Its sparse Transformer block contains top-k sparse attention and a mixed-scale feed-forward network for multi-scale features.
- Method and Results: A mixed-scale fusion feed-forward network explores multi-scale representations, while experiments on various benchmarks show favorable performance against SOTA approaches.The architecture is intended to generate high-quality deraining results with more accurate detail and texture recovery.
2. Related Work
Related work progresses from handcrafted-prior and CNN-based image deraining toward sparse attention and top-k selection mechanisms in vision Transformers. The paper positions its approach as an efficient top-k useful channel selection operator within this line of research.
- Single image deraining: Traditional deraining methods develop handcrafted image priors, but their empirical basis limits modeling of clear-image properties.These methods address the ill-posed nature of image deraining by providing additional constraints.
- Single image deraining: CNN-based frameworks were developed to overcome handcrafted-prior limitations and achieved decent restoration performance.
- Sparse attention: Sparse attention in vision can be categorized as data-based fixed sparsity or content-based sparsity, with data-based methods often attending only to local windows.
- Sparse attention: Prior work introduced sparsity into Transformer backbones and developed attention-retractable Transformer designs.
- Top-k selection: Top-k selection began as an explicit mechanism in NLP, later supporting k-NN attention in vision Transformers; this work instead selects useful channels efficiently.
3. Proposed Method
DRSformer uses a hierarchical encoder-decoder built from sparse Transformer blocks for image deraining. Its key components—top-k sparse attention, mixed-scale feed-forward processing, and mixture-of-experts feature compensation—target irrelevant interactions, multi-scale rain streaks, and diverse rainy appearances.
- Sparse Transformer Block: Each sparse Transformer block contains top-k sparse attention and a mixed-scale feed-forward network as its two principal elements.These components form the fundamental feature-extraction unit of DRSformer.
- Overall Architecture: The hierarchical encoder-decoder embeds overlapping 3 × 3 image patches and stacks Ni∈[1,2,3,4] sparse Transformer blocks across levels with distinct resolutions and channel dimensions.The network extracts features for spatially varying rain distributions and multi-scale rain degeneration.
- Top-k Sparse Attention: Top-k sparse attention masks lower-weight query-key interactions and retains only the most useful scores, reducing dense attention’s involvement of irrelevant information.It computes channel-wise self-attention, selects top-k values for softmax normalization, and assigns zero probability to smaller scores.
- Mixed-scale Feed-forward Network: The mixed-scale feed-forward network inserts two depth-wise convolution paths to model correlations among multi-scale rain streaks.The described paths use 3 × 3 and 5 × 5 depth-wise convolutions with channel-wise concatenation.
- Mixture-of-Experts Feature Compensator: MEFC combines parallel sparse CNN experts and uses self-attention rather than an external gating network to adaptively select representations for diverse rainy appearances.Its experts include pooling, separable convolutions with multiple kernel sizes, and dilated convolutions; the module remains linked to the main sparse Transformer blocks.
4. Experiments and Analysis
Experiments evaluate DRSformer across synthetic and real-world deraining benchmarks against prior-based, CNN-based, and Transformer-based methods. Results show stronger deraining performance, while ablations attribute gains to top-k attention selection, mixed-scale feed-forward processing, and multi-expert feature compensation.
- Datasets: DRSformer is evaluated on Rain200L/H, DID-Data, DDN-Data, and the real-world SPA-Data benchmark.The datasets include synthetic rainy images with varying directions and densities, plus 638,492 training pairs and 1,000 testing images in SPA-Data.
- Comparison methods: DRSformer is compared with prior-based, CNN-based, and recent Transformer-based deraining methods, including IDT, Restormer, and Uformer.The comparison also includes DSC, GMM, DDN, RESCAN, PReNet, MSPFN, RCDNet, MPRNet, DualGCN, and SPDNet.
- Synthetic datasets: 0.4 dB: DRSformer surpasses the concurrent approach IDT by this average PSNR margin on the synthetic benchmarks.The passage states that DRSformer outperforms all other derainers, achieves the highest PSNR/SSIM value, and generalizes across datasets.
- Effectiveness of Top-k selection: Top-k selection improves PSNR and preserves finer high-frequency details by reducing irrelevant long-range context and avoiding oversmoothed reconstructions.The ablation reports lower PSNR without top-k selection, while high-pass-filter visualizations show better fine-detail reconstruction with the proposed strategy.
- Effect of the number of k: 32.18 dB: the best result is achieved when the TKSA sparsity interval [∆1, ∆2] is assigned to [1, 2, 4, 5].Performance declines when k is too small because global information aggregation is insufficient, and decreases again as irrelevant features enter with larger k.
- Effectiveness of MEFC: Multi-expert MEFC outperforms a single-expert variant, with diverse receptive fields and CNN structures contributing to restoration gains.MEFC provides additional benefits over the baseline, and its placement at different network stages has specific effects on restoration performance.
5. Concluding Remarks
The paper presents DRSformer, a sparse Transformer for image deraining that uses top-k sparse attention and mixed-scale feed-forward modeling to improve feature aggregation and multi-scale representation. Its efficiency remains limited by 33.7 Million parameters and 242.9G FLOPs for a 256×256 image, motivating future pruning or distillation.
- Contributions: DRSformer addresses irrelevant global interactions in vanilla self-attention through top-k sparse attention that retains useful self-attention values for feature aggregation.The method is designed to facilitate aggregated features for rain removal.
- Contributions: DRSformer develops a mixed-scale feed-forward network to better explore multi-scale representation for image deraining.
- Limitations: 33.7 Million parameters and 242.9G FLOPs are required for one 256×256 image, limiting model efficiency.The authors propose applying pruning or distillation to preserve deraining performance while achieving model compression.