Learning a Spiking Neural Network for Efficient Image Deraining

1 Introduction

Single image deraining aims to reconstruct high-quality rainfree images from rain-affected ones. This challenging problem has witnessed significant advances due to the development of various effective image priors [Luo et al., 2015; Li et al., 2016] and deep learning models [Chen et al., 2023b]. However, early prior-based methods typically necessitate numerous iterative optimization steps to ascertain the optimal solution, also constraining their practical applicability.

*Corresponding author.

In the deep learning era, numerous convolutional neural networks (CNNs) based models [Li et al., 2018; Fu et al., 2019; Ren et al., 2019; Chen et al., 2021] have been developed as better solutions for image deraining. The success of these methods greatly stems from their elaborate network architectures, such as multi-scale [Jiang et al., 2020] and multistage [Zamir et al., 2021]. However, these networks heavily rely on increasing the depth or complexity of models in order to expand the receptive fields and achieve better performance.

Recently, to facilitate global information aggregation, Transformers [Xiao et al., 2022; Chen et al., 2024b] have been applied to image deraining field and achieved significant advancements as they can model the non-local information for better image reconstruction. However, the computation of the scaled dot-product self-attention in Transformers results in quadratic space and time complexity as the number of tokens increases. To alleviate the computational and memory burden, some efforts have been proposed to reduce model parameters or complexity (FLOPs), such as pruning [Zou et al., 2022], feature reuse [Fu et al., 2019], pixel-wise dilation filtering [Guo et al., 2021], token sparsification [Chen et al., 2023a] and low-rank matrices [Xiong et al., 2021].

Although these approaches have made initial explorations, achieving a better performance-efficient trade-off remains a major challenge. Spiking neural networks (SNNs) [Zhou et al., 2022; Su et al., 2023; Qiu et al., 2023], as the third generation of neural networks, which aim to simulate the way signals are transmitted in the brain. Unlike the artificial neural networks (ANNs) that directly learn embedding features from the input image, information in SNNs is represented as binary spike sequences and transmitted by neurons. Due to strong reflections caused by rain, the values of rain pixels in a rainy image are typically higher than those of their neighboring non-rain pixels. As shown in Figure 2, we observe that higher pixel values may result in a more pronounced intensity of spike signals in SNNs because they cause more neurons to fire spikes. This naturally arise a question: whether we can harmonize the characteristics of rain perturbation and SNN for high-quality and low-energy-consuming rain removal?

In fact, two challenges still need to be addressed when applying the SNN to the image deraining task. Firstly, to effectively learn complex and diverse rain streaks, the network needs to possess powerful information representation capabilities. Unfortunately, most existing SNN-based methods suffer from information loss caused by discrete binary activation, which may significantly interfere with the subsequent clear image reconstruction. To enhance the representation capability of SNNs, researchers have explored effective network structures, such as MS-ResNet [Hu et al., 2021b] and SEWResNet [Fang et al., 2021]. Nevertheless, these networks are not elaborately designed for image deraining task. Secondly, the non-differentiability of spike signals poses a challenge in using traditional backpropagation training strategies, causing difficulties in training the network. To address this issue, previous studies have employed the ANN-SNN conversion strategy. However, this strategy has limitations, as it heavily relies on the performance of the original ANNs and the converted SNN models fail to surpass the performance of the original artificial neural network (ANNs). A promising approach is to directly train SNNs using surrogate gradients, which can gain better performance with shorter time steps.

To this end, we propose an efficient and effective SNNbased method, called ESDNet, to solve the image deraining problem. Specifically, to tackle the weak representation capabilities and scale problems of SNN, we design a spiking residual block (SRB), which consists of a spiking convolution unit (SCU) and a mixed attention unit (MAU). SRB first converts the input image into spike signals, then adaptively optimizes the membrane potential by introducing attention weights to adjust spike responses, alleviating information loss caused by discrete binary activation. Furthermore, to overcome the difficulty of training SNNs, we apply the gradient surrogate strategy to directly train the model rather than the ANN-SNN conversion strategy. Compared to other ANN-based deraining models, our method exhibits lower energy consumption and superior performance (see Figure 1).

The main contributions can be summarized as follows:

• We propose a deep spike neural network (ESDNet) for single image deraining, and provide a new perspective on developing efficient models for resource-limited devices in real-world applications.
• We design a spiking residual block, which can effectively alleviate the problem of information loss caused by discrete binary activation and learn multi-scale rain streaks, facilitating high-quality image reconstruction.
• Extensive experimental results on single image deraining datasets display that our proposed model can obtain comparable rain removal performance against ANNsbased methods while reducing energy cost by 54%.

2 RelatedWork

2.1 Single Image Deraining

Traditional methods for image deraining explored various handcrafted prior knowledge to provide additional constraints for this ill-posed inverse problem. These methods involved complex optimization problems and only learned the characteristics of rain streaks without considering the underlying clear information [Ding et al., 2016; Luo et al., 2015]. Recently, existing most methods used CNN-based frameworks, treating image deraining as a pixel regression task and achieving significant performance gains. [Jiang et al., 2020], [Zamir et al., 2021] and [Chen et al., 2024a] decomposed the deraining task into multiple subspaces, further acquiring multi-scale information to enhance image reconstruction quality.

With the successful application of Vision Transformer (ViT) in various advanced visual tasks, the ViT has also been applied to image rain removal, achieving better performance than CNN-based models. For example, Restormer [Zamir et al., 2022] enhanced representation ability by calculating the similarity of tokens at the channel dimension. [Xiao et al., 2022] used a hierarchical architecture of window-based attention to capture global correlations. Despite these methods achieving impressive reconstruction performance, their heavy computation burden made it difficult to deploy them in practical applications on resource-constrained devices.

To improve efficiency, many efforts have been made. [Fu et al., 2019] constructed a lightweight network based on the Laplace pyramid architecture, which significantly improves the model's efficiency. [Guo et al., 2021] proposed pixelwise dilated filtering to predict multi-scale information corresponding to each pixel and developed an efficient rain removal framework. [Ren et al., 2019] designed a deraining network using a recursive computing architecture to accelerate model inference. In contrast to CNN-based methods that aimed to improve efficiency through carefully designed architectures, Transformer-based approaches focused more on optimizing attention to enhance computational efficiency. Sparse representation [Chen et al., 2023a; Song et al., 2023b] and low-rank matrix approximation calculation [Xiong et al., 2021] were employed to compress and accelerate the inference speed of the model, improving efficiency. Although these methods have successfully reduced computational costs, they sacrifice the network's representational ability and cannot effectively explore spatial information, which may affect deraining performances. In contrast to these ANN-based rain removal methods, we propose using SNN with lower energy efficiency as a framework to achieve efficient image rain removal.

2.2 Spiking Neural Networks

Currently, there are two main routes to implementing deep SNNs. One approach is the ANN-SNN conversion strategy, which involves converting a well-trained and highperformance ANN into a SNN [Diehl et al., 2015]. In this conversion strategy, the basic idea is to replace the ReLU activation functions in the ANN with spiking neurons to achieve the conversion to an SNN. However, this strategy has some inherent limitations. Firstly, achieving a conversion close to the performance of the ANN typically requires longer time steps, which leads to higher latency and extra computational consumption [Yao et al., 2023b]. Secondly, the performance of the converted SNN is highly dependent on the original ANN and cannot surpass it. The other promising route is to utilize surrogate gradient functions to unfold the SNN over simulation time steps and directly train the SNN using backpropagation or spike timing dependent plasticity (STDP) [Wang et al., 2023c]. Compared to the ANN-SNN conversion strategy, directly trained SNNs can significantly reduce the simulation time step, making them more attractive in terms of energy efficiency [Kim et al., 2020]. For example, [Su et al., 2023] achieved performance equivalent to an ANN with the same architecture using only 4 time steps. [Zheng et al., 2021] proposed temporal delay batch normalization (tdBN), which greatly improved the depth of the SNN model. [Hu et al., 2021b; Yao et al., 2023a] applied SNNs to classification tasks and achieved good performance. Although some works have applied directly trained deep SNNs to regression tasks such as object detection, there has been little exploration in pixellevel regression tasks such as image deraining.

3 ProposedMethod

In this section, we first offer an overview of the proposed ESDNet. Then, we introduce the spiking residual block (SRB) and the spike convolution unit (SCU). Next, we describe the feature refinement block (FRB). The training strategy and model optimization are described in the final section.

3.1 Data Preprocessing

Encoding to generate spike signals globally is a typical approach in SNNs for simulating pixel intensity signals in images. Considering the spatiotemporal properties of SNNs, we first perform direct encoding on the input rainy image Xt to generate a sequence X = {Xt}Tt=1 , ı.e., copying the single degraded image Xt as the input for each time step t.

3.2 Network Architecture

The proposed network architecture is shown in Figure 3. This network takes a single rainy image as input and adopts an encoder-decoder framework to learn hierarchical representations for capturing information in the network and achieving high-quality deraining output. First, for the rainy input, the network employs the encoding layer to generate the rainy sequence X and extract shallow features from the X by 3 × 3 convolutional layers. Then, the captured shallow features will be fed into N stacked SRBs to conduct spike transformation, feature extraction, and fusion of rain spike information through SCU and MAU. For SCU, the Leaky Integrate-and-Fire neurons (LIF) [Fang et al., 2023; Deng et al., 2022] convert and weight the inputs to generate output spike sequences, which consist of 0 and 1. The value of the output spike from LIF is 1 when the membrane potential exceeds the threshold, otherwise 0. Next, MAU can adaptively optimize the membrane potential by applying attention weights, leading to adjusting the spike response in a datadriven manner and learning the distribution of rain streaks.

After the decoder stage, the network adopts a FRB to convert the features extracted from the spike sequence into accurate continuous value representations, while refining the potential representations which are used for restoring clear images [Zheng et al., 2023]. Finally, the network generates high-quality and clear image output from the obtained representations by applying a 3 × 3 convolutional layer. Furthermore, skip connections are employed to assist model training during the feature extraction and image reconstruction stages.

3.3 Spiking Residual Block

As previously mentioned, regions degraded with rain perturbation have higher pixel values than those of no rain regions, exhibiting noticeable fluctuations. Correspondingly, this may lead to pronounced firing spike rates of neurons in SNNs. Hence, we develop the spiking residual block (SRB) to learn the spike information from the converted rain streaks, which is displayed in Figure 3. We first adopt SCU to convert the inputs into a discrete spike sequence and mine the rain streak clues included in the spike sequence. Then, an additional branch comprising solely a convolutional layer and a tdBN layer [Zheng et al., 2021] is constructed to mitigate the significant decline in performance resulting from the inadequate information conveyed by SNNs. Furthermore, inspired by the previous work [Yao et al., 2023b], we introduce MAU to augment the spike activity rate of neurons and regulate spike responses, better guiding the process of eliminating rain streaks. Finally, residual learning is applied to SRB to alleviate the issues of gradient vanishing and explosion. Assuming that the output of the n-th SRB at the t-th time step is Xt,n , the calculated process of SRB can be depicted as:

tdBN refers to threshold-dependent batch normalization. It considers both temporal and spatial dimensions during the normalization process to relieve the gradient explosion. Profiting from the above designs, SRB can enable adaptively learning the spike fluctuations induced by rain streaks, thereby effectively simulating degradation in diverse regions.

3.4 Spike Convolution Unit

Compared to ANNs, which employ continuous decimal values for information propagation and calculation, SNNs exploit discrete binary spike values for communication. Spiking neurons can convert received continuous values into discrete spike sequences, encoding information with temporal precision. In this paper, we select LIF neurons to convert input signals into spike sequences, as they better balance biological characteristics and computational complexity [Guo et al., 2023b]. The explicit dynamic equation for LIF neurons can be expressed mathematically as:

where t and n represent the t-th time step and the n-th layer. Ht,n represents membrane potential, which is generated by a combination of time input U t-1,n and spatial input Xt,n . τ is the membrane time constant. Among them, Vthr represents the activation threshold, used to determine whether to output S or maintain it at 0. Θ(·) denotes the step function. When x ≥ 0, Θ(x) = 1, otherwise Θ(x) = 0. Vreset denotes the potential reset after S activates the output. β represents the decay factor and ⊙ indicates the element-wise multiplication.

In Eq. 2, the spatial representationXt,n+1 of rain streaks is captured from the output spike sequence U t,n by convolution operators. This calculation process can be defined as:

where St,n denotes a spike tensor, which only included 0 and 1. tdBN refers to threshold-dependent batch normalization. Due to space limitations, specific descriptions of MAU are included in the supplementary materials.

3.5 Feature Refinement Block

To transform discrete pulse sequences into continuous pixel values, a typical approach is to apply mean sampling in the time dimension and then scale it to the input size. Nevertheless, this sampling process may result in the loss of some crucial structural elements, which could affect the final image restoration quality. To address this issue, we introduce a FRB that can adaptively aggregate sampled features and mitigate information loss. Given the input spike sequence Y , the calculation process of FRB can be presented as:

where ˜F refers to the final output feature and GAP(·) denotes the global average pooling operation. σ(·) and φ(·) indicate the Sigmoid and ReLU functions. Subsequently, the output feature of the FRB will be fed into a 3×3 convolutional layer to map it back to the original input resolution. By incorporating FRB into our proposed model, the network can better learn and refine the latent representation of clear images for restoring fine structural and textural details in images.

3.6 Training Strategies

In recent years, backpropagation-based training schemes have achieved significant progress in facilitating the training of SNNs, showing considerable superiority over transformation-based strategies. Since the non-differentiable nature of binary activations in SNN neurons, most existing methods use gradient-proxy functions to implement backpropagation training. Inspired by the previous wisdom, the Sigmoid function serves as the gradient-surrogate function learning to train SNNs. It can effectively handle binary neuron outputs and enable gradient propagation during backpropagation. The gradient surrogate function can be defined as:

where α refers to a hyper-parameter, adopted to adjust the gradient of the surrogate function. The larger the value of α, the greater the gradient of the function.

3.7 Model Optimization

To better recover more similar structure details, we apply SSIM loss as the reconstruction loss and train the model by minimizing it. The loss function can be depicted as:

where S(·) is the proposed ESDNet, X and Y are the input rainy image and its corresponding ground-truth, respectively.

4 Experiments

In this section, we conduct comprehensive experiments on commonly used benchmark datasets to evaluate the effectiveness of the proposed method. The experiments involve rigorous testing and analysis, aiming to provide a thorough assessment of the performance. More results are included in the supplemental material. The training code and test models will be available to the public.

4.1 Experimental Setup

Datasets. Since the original comparison methods are trained on different datasets, we retrained all models on the four publicly available datasets (Rain12 [Li et al., 2016], Rain200L [Yang et al., 2017], Rain200H [Yang et al., 2017], Rain1200 [Zhang and Patel, 2018]) to ensure a fair comparison of all methods. The Rain12 dataset consists of 12 synthetic rain images that were used for testing. The Rain200L and Rain200H datasets contain 1800 synthetic rain images for training, along with 200 images designated for testing. The Rain1200 dataset includes 12000 training images and 1200 test images, which have varying rain directions and density levels. Furthermore, we exploit the RW-Data [Zhang et al., 2019], a real-world rain benchmark, consisting of 185 realworld rainy images.

Evaluation Metrics. In light of practical considerations, we select two common evaluation metrics, e.g., Peak Signal to Noise Ratio (PSNR) [Hore and Ziou, 2010] and Structural Similarity (SSIM) [Wang et al., 2004], to perform a quantitative comparison of the synthesized rainy dataset. Following the previous approaches [Chen et al., 2023a], we compute these metrics in the Y channel of the YCbCr space. For the real-world rain dataset, we adopt non-reference image quality evaluation metrics such as Perception-based Image Quality Evaluator (PIQE) [Venkatanath et al., 2015] and Metalearning-based Image Quality Assessment (MetaIQA) [Zhu et al., 2020] for assessing performance.

Comparison Methods. We compare our proposed method with various image rain removal baselines, including seven CNN-based methods (i.e., RESCAN [Li et al., 2018], PreNet [Ren et al., 2019], RCDNet [Wang et al., 2020], MPRNet [Zamir et al., 2021], Effderain [Guo et al., 2021], NAFNet [Chen et al., 2022], and SFNet [Cui et al., 2022]) and four transformer-based networks (i.e., DRT [Liang et al., 2022], ELFormer [Jiang et al., 2022], HPCNet [Wang et al., 2023a], and SmartAssign [Wang et al., 2023b]). For the recent representative method SmartAssign, we encounter a lack of available code. Hence, we recreate the corresponding network according to the provided information in their paper and retrain it. For other competitive methods, if the authors provide pre-trained models available, we utilize their online code to evaluate performance, otherwise, we will retrain them using the same settings as ESDNet to ensure a fair comparison.

Implementation Details. During the training process, we conduct the proposed network in the PyTorch framework with an Adam optimizer and a batch size of 12. We set the learning rate to 1×10-3 and apply the cosine annealing strategy [Song et al., 2023a] to steadily decrease the final learning rate to 1× 10-7 . For Rain200L, Rain200H, and Rain1200 datasets, we train the model by 1000 epochs. We set the stacking numbers of SRB to [4,4,8] in the encoder stage and [2,2] in the decoder stage. For the α of the gradient proxy function, it is set to 4 according to [Su et al., 2023]. All experiments are executed on an NVIDIA GeForce RTX 3080Ti GPU (12G). To ensure a fair comparison, for all retrained methods, we uniformly cut the data into 64 patch sizes for training, and apply the sliding window slicing strategy for testing.

4.2 Experimental Results

Synthetic Datasets. Table 1 provides a comprehensive comparison between our proposed method and 12 representative and competitive training methods. It is evident that the incorporation of spike convolution units in our ESDNet significantly improves the performance in terms of PSNR and SSIM values compared to all other baselines. Notably, our approach achieves more appealing results on the Rain200H benchmark, surpassing the recent CNN-based method SFNet by 0.26 dB in PSNR. The performance improvement, when compared to existing ANNs-based deraining methods, shows that our framework offers a fresh perspective for deep image training architectures. Furthermore, Figure 3 presents the qualitative evaluation results on the Rain200H benchmark. The visual comparison indicates that our method exhibits enhanced contrast and reduced color distortion in comparison to other approaches, corroborating the quantitative findings. Specifically, Effderain and DRT exhibit noticeable rain artifacts in Figure 3, while also tending to produce overly smooth results. In contrast, our model preserves more details and achieves superior perceptual quality.

Real-world Datasets. We conduct additional experiments using the RW-Data benchmark dataset. The obtained quantitative results are presented in Table 2. Analysis of these results reveals that our network achieves lower PIQE values compared to other comparative models. This indicates that in real rain scenarios, our network exhibits superior output quality with clearer content and enhanced perceptual quality. To further evaluate the qualitative performance, we implement a quality comparison analysis as shown in Figure 4. The results demonstrate that most models exhibit sensitivity to spatial long rain streaks, which often leads to noticeable rain effects in their outputs. In contrast, our network effectively eliminates a significant portion of rain disturbances, resulting in visually pleasing restoration effects. Overall, the results suggest that our network outperforms existing models in terms of output quality, perceptual clarity, and the ability to handle rain disturbances in real scenarios.

Model Effciency. In ANNs, each operation involves the multiply-accumulate operations (MAC) and floating-point operations (FLOPs) are used to estimate computational burden. SNNs exhibit energy-saving properties in neuromorphic hardware, as neurons only participate in accumulated operations (AC) during spikes and can complete AC through approximately the same number of synaptic operations (SOPs). In this work, we exploit 4 time steps to estimate energy costs. It is observed that the first convolutional encoding layer, RB block, and last convolutional encoding layer require floatingpoint operations due to they enjoy the matrix multiplication. The other layers are calculated through SOPs (synaptic operations). Following the previous works [Guo et al., 2023a], SOPs can be computed through s × T × A, where s represents the average sparsity, T is the time step, and A is the additive number in the artificial neural network. For the binary SNNs, the sparsity is 16.42%. It is noted that SOPs can be triggered once when the sign function output is 1 due to the membrane potential exceeding the threshold. One sign (using energy per spike to calculate) needs 3.7pJ. Furthermore, the sign function only exists in LIF, and the number of LIF is limited compared to convolution operations. According to [Hu et al., 2021a], we compute the energy as follows: a FLOP requires 12.5pJ, an SOP requires 77fJ, and a Sign (calculated using energy per spike) requires 3.7pJ.

The results are displayed in Table 1. It can be observed that our network has achieved performance comparable to the recent ANNs-based deraining approaches, and has lower energy costs. In addition, we compared the computational complexity of various deep derivation transformer-based methods, including the number of trainable parameters and FLOPs on 256 × 256 images, as reported in Figure 1.

4.3 Ablation Studies

We conduct the ablation experiments to better understand the impact of different residual blocks, time steps, and model components on performance. Unless otherwise specified, all ablation experiments are conducted on the Rain200H dataset using the experimental setup described in Section 4.1.

Effectiveness of Different Residual Blocks. To investigate the impact of different residual modules on model performance, Table 3 presents the PSNR/SSIM values of the corresponding models. From the experimental results, we note that our SRB obtains a more attractive performance compared to the other two spiking residual blocks, thanks to the additional branches in SRB and the information compensation brought by MAU. By introducing a lower computational burden alone, the model performance is significantly improved.

Effectiveness of Different Time Steps. The performance results of the model with different time steps are reported in Table 4. According to the experimental results, it can be found that the longer the time step, the higher the quality of image restoration. This indicates that higher time steps can enhance the model to better obtain feature representations. Similarly, this also brings about a significant increase in energy costs. To better balance performance and efficiency, we chose a time step of 4 as the default time step in this work.

Effectiveness of Different Designs. To evaluate the effectiveness of the model designs, we conduct experiments based on different model variables in Table 5. Compared to the baseline model (a), model (b) provides additional performance advantages due to the consideration of temporal correlation. In addition, we note a significant decrease in performance when MAU is not applied, indicating that MAU can effectively enhance the spike activity rate of neurons and regulate spike responses.

5 Concluding Remarks

This paper presents an image deraining approach using spiking neural networks. The proposed network learns rain streak information by leveraging the spike convolution unit. Initially, continuous pixel values are transformed into spike sequences by LIF neurons. Convolution operations are then applied to learn the spike information associated with rain streaks. To enhance the accuracy of rain streak removal, a mixed attention unit is introduced, which adjusts spike responses based on the input data, thereby facilitating the extraction of rain streak clues. In addition, feature refinement blocks are incorporated to further improve the model representation for gaining high-quality images. The efficacy of the proposed method is demonstrated through extensive experiments conducted on synthetic and real rain datasets.