Dual Prompt Learning for Continual Rain Removal from Single Images

1 Introduction

As one of the most common weather degradation, rain streaks heavily reduce visibility and corrupt the information captured by images, impacting both human visual experience and computer vision algorithms like detection [Carion et al., 2020], segmentation [Chen et al., 2018], and depth estimation [Wang et al., 2020], which is closely related to many practical applications, e.g., autonomous navigation and surveillance systems.

In recent years, remarkable progress has been achieved in single image deraining, especially for deep learning based methods. Many methods built on convolutional networks [Fu Model N-1

∗Corresponding Author.

Transfer

Model N

Dataset N

Train

(a) PIGWM framework

(b) PIGWM result

(d) DPL result

Figure 1: The comparison between the traditional continual learning method and our method. (a) Parameter regularization framework, e.g., PIGWM. [Zhou et al., 2021] (b) Result of PIGWM (PSNR/SSIM: 21.47 dB/0.6792). (c) Overview of our Dual Prompt Learning (DPL) framework. The DPL framework only utilizes a single model and prompt pools to store task-specific knowledge as prompts, i.e., learnable parameters, as the past memory. DPL selectively updates prompts in an instance-wise manner, which are optimized to guide model prediction and explicitly manage taskinvariant and task-specific knowledge while preserving model plasticity. (d) Result of DPL (PSNR/SSIM: 27.79 dB/0.8240).

et al., 2017; Ren et al., 2019; Zamir et al., 2021] continue to break new ground in performance. Recently, transformerbased methods [Valanarasu et al., 2022] have achieved excellent results. However, rain degradation is complex and diverse, and most existing models only learn fixed mappings between paired rainy and clean images. As a result, deep neural networks may lose previously acquired knowledge and experience performance decline when faced with changing data distributions, limiting their real-world applicability.

To alleviate this issue, continual learning methods are developed to overcome catastrophic forgetting. They can be mainly divided into three categories: parameter isolationbased methods [Zhang et al., 2020; Mallya et al., 2018; Xu and Zhu, 2018], replay-based mechanisms [Zhao et al., 2021; Kirkpatrick et al., 2017; Pascanu and Bengio, 2013; Aljundi et al., 2018], and regularization-based methods [Li and Hoiem, 2017; Zenke et al., 2017]. However, parameter isolation and replay-based methods are computationally expensive, while regularization-based methods become inefficient and the parameters increase linearly as tasks increase.

In addition, all three kinds of methods fail to automatically select relevant knowledge components for arbitrary samples without knowing their task identity, which is a tremendous limitation in real applications.

Lately, researchers have drawn inspiration from recent advances in prompt-based learning (prompting) [Liu et al., 2021], which reformulates learning tasks by designing prompts [Raffel et al., 2020] instead of directly adapting model weights. Prompts encode the knowledge specifically related to a particular task and allow for more efficient use of a pre-trained frozen model than fine-tuning, which provides an ideal mechanism [Wang et al., 2022] for continual learning.

In this paper, we propose a novel continual learning scheme for single image deraining called Dual Prompt Learning (DPL), which is orthogonal to existing replay-based and regularization-based methods. Fig. 1 demonstrates its effectiveness. It does not need to know the task identity or boundaries, therefore more applicable to real scenarios. Its objective is to learn to select and update compact prompts, i.e., small learnable parameters, to instruct the transformer-based rain removal model to maintain both task-invariant and taskspecific knowledge. In detail, a subset of prompts are selected from the image/feature prompt pools based on our proposed instance-wise query mechanism and are concatenated with input image/embedded tokens for further processing. These prompts applied at both image and feature levels effectively leverage transferred knowledge of these two levels jointly. Furthermore, a regularization technique is further adopted to penalize the intense changes of important parameters jointly with the learned prompts.

In summary, our work has the following contributions:

• We propose a novel prompt learning-based continual learning (CL) scheme to handle different types of rain streaks with a single model. To the best of our knowledge, it is the first time to apply this new kind of CL method for low-level vision, which leads to superior performance on various benchmarks.
• We develop a dual prompt learning method for deraining, where prompts are applied at both image and feature levels to leverage effectively transferred knowledge of images and features jointly.
• Our DPL is further augmented by parameter regularization. The joint regularization of model parameters and learnable prompts obviously further improves the performance.
• Even though our method is designed for the continual learning scenario, it achieves competitive results against state-of-the-art methods on the stationary distributed data.

2 RelatedWork

2.1 Single Image Rain Removal

Recently, there have been several methods that have achieved significant progress in single image rain removal. Yang et al. [Yang et al., 2017], and Fu et al. [Fu et al., 2017] proposed the first framework on the rain streak removal. In [Yang et al., 2017], the rain streak detection and removal are modeled in the multi-task manner. In [Fu et al., 2017], the residual learning is applied for rain removal. Later works are proposed with more complex architectures to improve rain removal performance, including joint rain density estimation and deraining [Ren et al., 2019], non-local operation-based encoder-decoder network [Li et al., 2018a], multi-stage network [Zhang and Patel, 2018], conditional generative adversarial network [Zhang et al., 2019], and deep convolutional and recurrent neural network [Li et al., 2018b] that removes rain streaks stage by stage, etc. Recently, transformer-based approaches [Valanarasu et al., 2022] have been explored for weather removal tasks, demonstrating their superior performance compared to convolutional networks. However, all the mentioned deep learning-based methods suffer from catastrophic forgetting issue, failing to maintain their effectiveness when applied to different types of rainy datasets/tasks in real applications. Comparatively, our work explores the continual learning approach for image rain removal to adapt the model to a series of rain streaks.

2.2 Continual Learning

A large amount of research in continual learning follows a learning paradigm that involves continuously adapting the model weights, either partially or fully, as the data distribution changes [De Lange et al., 2021; Mai et al., 2022]. These approaches focus on preserving previous knowledge while also adapting to shifting data. The methods of overcoming catastrophic forgetting can be mainly divided into three categories: replay-based mechanisms [Zhao et al., 2021; Kirkpatrick et al., 2017; Pascanu and Bengio, 2013; Aljundi et al., 2018], regularization-based methods [Li and Hoiem, 2017; Zenke et al., 2017], and parameter isolationbased methods [Zhang et al., 2020; Mallya et al., 2018; Xu and Zhu, 2018]. In detail, replay-based and parameter isolation-based methods are computationally expensive since they require recording the old tasks' targets and computing old tasks' forward pass process for each novel data sample. Regularization-based methods are cost-effective. The representative form of the classic regularization-based method is Elastic Weight Consolidation (EWC) [Kirkpatrick et al., 2017], which quantifies how essential each parameter is for a task with the diagonal of the Fisher information matrix[Pascanu and Bengio, 2013] and protects critical weights with an additional regularization to restrict their movement when updating for the new job. Further, memory-aware synapses (MAS) [Aljundi et al., 2018] compute the parameter importance based on how sensitive the predicted output function is to a change in this parameter, and penalize changed essential parameters. Zhou et al. [Zhou et al., 2021] proposed a first and second-order parameter importance to jointly estimate the status of one parameter. In this paper, we propose a novel continual learning scheme for single image deraining, which is orthogonal to existing replay-based and regularization-based methods.

2.3 Prompting for Transfer Learning

Prompting is applying a function to modify the input so that a model gets additional information about the task.

However, the design of a prompting function is challenging and requires heuristics. Recent works in NLP, including prompt tuning [Lester et al., 2021] and prefix tuning [Li and Liang, 2021], seek to address this issue by applying learnable prompts in a continuous space, achieving excellent performance for transfer learning. Prompts capture datasetspecific knowledge with much smaller additional parameters than its competitors, such as Adapter [Pfeiffer et al., 2020] and LoRA [Hu et al., 2021], which achieve remarkable performance in many downstream tasks. Prompt techniques have acted out fantastic value for transfer learning. Recently, Wang et al. [Wang et al., 2022] revealed its significance to continual learning problems. In our work, prompt learning is introduced to construct a continual rain removal method.

3 Dual Prompt Learning for Continual Rain Removal

Prompt learning is a recent popular technique for transfer learning and model adaptation. It uses task-specific prompt functions, which improves sequential modeling capacity for continuous feature learning. In this paper, we propose to utilize two-level prompt pools that lead to superior performance even without knowing task identities, as shown in Fig. 2. Instead of using the naive way of prompt learning that relies on task identities and lacks the flexibility to distinguish taskindependent/relevant knowledge, the prompt pools enable an online search for the prompts to share the knowledge when tasks are similar and maintain independent knowledge otherwise. This prompt learning based on prompt pools is applied at both image and feature levels, to jointly leverage transferred knowledge of two stages.

3.1 Definition and Formulation

We define a sequence of T de-raining tasks as D = {D1, · · · ,DT} where the t-th task Dt = {(xti , yt i )}nt i=1 contains nt pairs of samples, where each pair consists of a rainy image xti ∈ X and its clean corresponding background yt i ∈ Y. Our goal is to train a model f (·|θmodel) : X → Y parameterized by θmodel, where X is the rain images domain and Y is the corresponding clean images domain. However, the data of D1,D2, · · · ,Dn is not available when training Dn+1.

Given the input image x ∈ RC×H×W where C is the channel number, and H ×W is the size of the image. We build a transformer-based deraining backbone f = fr ◦ fe, where fe is the input embedding layer, and fr represents a stack of selfattention layers and the subsequent rain removal network.

3.2 Image-Level Prompt-Based Learning

To alleviate the deficiency of linear growth of space expenditure caused by the replay mechanism, and transfer the imagelevel prior knowledge among different tasks effectively, we integrate image-level prompt learning into our proposed continual learning model. An image prompt pool is defined as P = {P1, · · · , PM} containing a certain amount of trainable prompts Pi = (Ki, Vi) where value Vi ∈ RC×H×W is a representation with the same size as the input image, and key Ki = fe (Vi) is the embedded value to match the key of input.

Following the notations in Sec. 3.1, x and xe = fe (x) are the input and its corresponding embedding feature, respectively. Note that we omit the task index t of x in our notation as our method is general enough to be applied even without the task identity. Ideally, the input x instance itself decides which prompts to be chosen through matching. To this end, we utilize a distance-measuring function γ to measure the similarity of a prompt and the input image. We directly take γ as cosine distance, which is proven as a good choice empirically [Wang et al., 2022]. During the training of task n, we maintain a prompt frequency table Qn = [qn 1 , qn 2 , · · · , qnM ], where each entry represents the normalized frequency of prompt Pi being selected in task n. As the feature includes richer semantic information, the similarity is calculated at the feature level to better reflect more intrinsic information instead of the pixel level. During the training of task n, given an input x, we lookup the top-N prompts by simply solving the following objective:

where Px represents the subset of top-N prompts selected specifically for x from P, and the prompt frequency table actually encourages the choice of diversified prompts, which will be removed during the testing process.

After selection, we concatenate the top-N prompts with the input along the channel dimension and put them into the embedding layer together. These prompts have the same size as the input and thus they can jointly encode knowledge with the input in model training.

where [; ] represents concatenation along the channel dimension.

3.3 Feature-Level Prompt-Based Learning

Given the input x, and the transformer-based model f = fr ◦ fe. The embedding layer fe : RC×H×W → RL×D projects the patched image to the embedding feature xe = fe (x) ∈ RL×D where L is the length of a token and D is the embedding dimension. At the feature level, we also maintain a set of prompts P ′ = {P ′1 , P ′2 , · · · , P ′M } containing a certain amount of prompts P ′ i = (K′ i , V ′ i ) where K′ i = V ′ i ∈ RL×D , for i = 1, 2, · · · ,M, and its corresponding prompt frequency table for task n, Q′n = [qn 1 ′ , qn 2 ′ , · · · , qnM ′ ], where each entry represents the normalized frequency of prompt P ′ i being selected in task n. Similar to Sec. 3.2, we denote γ as a cosine distance function to score the match between xe = fe (x) and prompts P ′1 , P ′2 , · · · , P ′M . In other words, during task n, we lookup the top-N prompts by simply solving the following objective:

where {s′1 , s′2 , · · · , s′N } is a subset of N indices from [1,M], and we can then adapt the input embedding as follows:

where [; ] represents concatenation along the token length dimension. In Eqn. (4), x1 represents the synthesized input in Sec. 3.2, and note that when selecting these top-N prompts, we use x, but during training, we concatenate these prompts with x1 = [Vs1 ; Vs2 ; · · · ; VsN ; x]. After concatenation, x2 is put into the subsequent rain-removal network for further training. As in Sec. 3.2, these prompts are also trainable.

3.4 Joint Regularization of Parameters and Prompts

When the neural network is trained on Task n and Task n + 1 sequentially, we hope the network can still maintain the performance of Task n. Thus, the regularization is imposed on the model parameters and prompts jointly. On Task n, the rain removal model's parameters are denoted as θnmodel = {θn 1 , θn 2 , · · · , θn r } where r is the depth of the network, and the parameter set of the overall prompts is denoted as θn prompt = {θ′n 1 , θ′n 2 , · · · , θ′n s }, where s is the total quantity of all parameters pertained to these prompts. We signify θn = θnmodel ∪ θn prompt . For the sake of convenience, we imply θn = {θn 1 , θn 2 , · · · , θn r+s}. Xn and Y n indicate the rain images set and the clean images set on Task n, respectively. Suppose (x, y) is a rainy/clean image pair. When it is fed into the network, the degradation of performance on Task n introduced by the training of network on Task n + 1 can be evaluated as:

where |·| denotes the absolute value operator, l represents the weighted average of the perceptual loss, L1 loss, and prompt distance loss used for training the de-raining network. For the k-th depth parameter on Task n, θn k , we denote

To evaluate ∆f (θn+1 , θn , x, y), we take the Taylor expansion of l (f (x, θ) , y) at point θn+1 k , which is an infinite sum of terms that are expressed in the form of target functions derivatives at a single point:

where H = ∇2θn k l (θn k , x) denotes Hessian matrix for θn k . In the actual calculation, we abandon the items more than three times. Therefore we have:

A small I (θn+1 ) leads to more preservation of the knowledge of Task n and can be adopted as an effective regularization when training on Task n + 1.

3.5 Optimization Objective Function

At every training step, after selecting N image prompts following the aforementioned objective. The joint input x1 is fed into the embedding layer. After that, N feature prompts are selected and concatenated. The adapted embedding feature x2 is fed into the subsequent layers. Additionally, the prompt distance penalty and the penalty of parameters and prompts are also put into the optimization objective. Overall, on Task n, we seek to minimize the end-to-end training loss function:

where L1 and Lp represents smooth L1 loss and perceptual loss, respectively. The third and fourth items are surrogate losses to pull selected prompts closer to the corresponding query features and the fifth item is the parameter penalty aforementioned in Eqn. (8). α, β, ζ, η, ω are parameters that control the importance of each term.

3.6 Learning of Prompt Pools

For each task, P and P ′ are initialized randomly. In detail, for a Pi = (Ki, Vi) , 1 ≤ i ≤ M, Vi is randomly initialized and has the same size as the input image, and Ki is calculated by Ki = fe (Vi), where fe is the embedding layer of a particular model. Similarly, for a P ′ i = (K′ i , V ′ i ) , 1 ≤ i ≤ M, V ′ i is also randomly initialized and has the same size as the token, and K′ i = V ′ i is the same as V ′ i . During each training epoch, after we draw a mini-batch B = {(xti , yt i )}li=1 , where l is the batch size, two sets of chosen prompts PB,P ′B are respectively maintained for the dual prompt pools. Initially, they are both empty sets. For each pair of rainy/clean images, (x, y) in B, top-N image prompts Px and top-N feature prompts P ′x are obtained by solving Eqn. (1) and Eqn. (3). We update the sets of chosen prompts by PB = PB ∪ Px, P ′B = P ′B ∪ P ′x , and update Qt and Q′t by adding 1 to the frequency-items corresponding to the selected prompts. Then, we calculate the sample loss Lx of every input in the mini-batch by Eqn. (9). We calculate per batch loss LB by accumulating Lx, and we update the parameters in P and P ′ by back-propagation immediately. In other words, P and P ′ are updated once for each epoch. The whole process of our proposed DPL for continual rain removal is summarized in Algorithm 1.

4 Experiment Results

In order to fully evaluate the capabilities of our proposed continual learning scheme, we integrate it with a state-of-the-art rain removal baseline TransWeather [Valanarasu et al., 2022]. Through extensive experimentation on several widely-used rain removal datasets, our DPL method consistently demonstrates superior performance in comparison to other existing continual learning methods. In particular, the results of all experiments indicate that our method performs excellently in both the capacities to adapt to new tasks and maintain high performance on previous ones.

Algorithm 1 Dual Prompt Learning

Require: DataSet D = {D1,D2, · · · ,DT},where Di includes an amount of pairs of rainy/clean images{(xi, yi)}Li=1 and T is the number of tasks.

Parameter: Network f and overall prompts P,P ′ parameterized by θ, image-level prompt pool P = {P1, · · · , PM}, feature-level prompt pool P ′ = {P ′1, P ′2, · · · , P ′M}, where Pi = (Ki, Vi) , P ′ i = (K′ i, V ′ i ), for 1 ≤ i ≤ M. Maintained frequency set Qt = [qt1, qt2, · · · , qtM], F ′ t = [qt1 ′ , qt2 ′ , · · · , qtM ′ ] for the t-th task, number of training epochs of the t-th task Et, 1 ≤ t ≤ T, learning rate r, hyperparameters of loss function α, β, ζ, η, ω.

Begin

• 1: Load the pre-trained VGG-16 model, and get D = {D1,D2, · · · ,DT}.
• 2: for t = 1, 2, · · · , T do
• 3: Initialize Qt and Q′t with real number 1. Initialize P and P ′ with random prompt-sized items.
• 4: for e = 1, 2, · · · ,Et do
• 5: Draw a mini-batch B = {(xti, yt i )}li=1 . Initialize the sets of chosen prompts for this batch: PB = {}, P ′B = {}.
• 6: for (x, y) in B do
• 7: Calculate the embedding input xe = fe (x).
• 8: Lookup top-N image prompts by solving Eqn. (1).
• 9: Prepending x1 with corresponding top-N prompts by x1 = [Vs1 ; Vs2 ; · · · ; VsN ; x].
• 10: Calculate the embedding input fe (x1).
• 11: Lookup top-N feature prompts by solving Eqn. (3).
• 12: Prepending x2 with corresponding top-N prompts by x2 = [V ′ s′1 ; V ′ s′2 ; · · · ; V ′ s′N ; fe (x1)].
• 13: Calculate the rain-free result f(x) = fr(x2).
• 14: Calculate per sample loss by solving Eqn. (9).
• 15: Update sets of chosen prompts: PB = PB ∪ Px.
• 16: Update sets of chosen prompts: P ′B = P ′B ∪ P ′x.
• 17: Update Qt and Q′t by adding 1 to the frequency-items corresponding to the selected prompts.
• 18: end for
• 19: Calculate per batch loss LB by accumulating Lx.
• 20: Update P and P ′ .
• 21: end for
• 22: Update θ.
• 23: end for=0

End

4.1 Dataset and Performance Metrics

We evaluate our proposed continual learning scheme on three widely-used rain removal datasets, including Rain100H [Yang et al., 2019], Rain100L [Yang et al., 2019], and Rain800 [Zhang et al., 2019]. In detail, the model is trained on Rain800 (Task 1) and Rain100H (Task 2) sequentially, denoted as Rain800-Rain100H. In addition to the continual task sequence Rain800-Rain100H, we further experiment with continual task sequences Rain800-Rain100L. Both Rain100H and Rain100L consist of 1,800 rainy/clean image pairs for training and 100 pairs for testing while Rain800 possesses 600 training samples and 200 testing images. Peak-Signal-to-Noise Ratio (PSNR) and Structure SIMilarity (SSIM) [Wang et al., 2004] are employed for evaluating the model performance.

4.2 Training Details

For a fair comparison, all the parameters setting and training techniques of the baseline model keep consistent with experiments in the original papers. Furthermore, we design distance query function γ as cosine distance. For both prompt pools, we assign M to 100, and N to 18, since the mini-batch size is 18. In terns of the hyperparameters, we assign α to 1, β to 0.04, ω to 0.95, and ζ, η are both assigned to 1e-5. For each task, we train for 50 epochs.

4.3 Results on Benchmark Datasets

To demonstrate the effectiveness of our proposed continual learning algorithm, we conduct both qualitative and quantitative experiments on the above datasets and performance measures. Table 2 and Table 3 report the comprehensive comparison among the baselines, classic continual learning methods, and our method, which indicates that our method can better mitigate catastrophic forgetting on multiple datasets. Surprisingly, Table 2 shows that DPL also achieves competitive results against the baseline method even on a single dataset without applying it to continual learning (PSNR/SSIM: 27.52 dB/0.8667), which further demonstrates the rationality of our method design. Table 4 clearly demonstrates the effectiveness of our method in not only maintaining consistently excellent performance in the previous task but also in achieving satisfactory results in a new task. That is, our method effectively balances the capacity to adapt to new tasks while simultaneously preserving its performance on established ones. As shown in Table 1, dual prompts offer lower parameter complexity compared to traditional continual learning methods, adding only 5.2% extra parameters. By applying parameter regularization to prompts, our method provides flexible solutions for different scenarios: adding dual prompt pool yields comparable performance to parameter regularization at lower complexity; incorporating parameter regularization can further enhance performance if higher complexity is acceptable.

4.4 Ablation Study

In this section, we conduct ablation studies to verify the importance of each item in Eqn. (9). It can be seen clearly in Table 5 and Fig. 4 that the dual prompt pools are the key to overcoming catastrophic forgetting. Moreover, feature-level prompts play a more critical role, and the parameter regularization technique is able to further improve our model's performance.

4.5 Extension toMultiple Datasets

In this section, we prove that our method shows superior performance not only for two tasks but also for multiple tasks. Namely, for n task sequences, we continually train task n after the model has been trained for the first n-1 tasks. Taking three tasks as an example, the model is trained for the first two tasks with our proposed scheme. Then, we continue to train the model for Task 3 based on this model. Table 6 shows the result of a model trained on Rain800-Rain100H-Rain100L sequentially, to verify the effectiveness of our method. The results clearly demonstrate the superiority of our method.

5 Conclusion

In this paper, we propose a novel continual learning scheme for single image deraining, Dual Prompt Learning (DPL). It learns to select and update compact prompts, i.e., small learnable parameters, to instruct the transformer-based rain removal model at both image and feature levels for continual rain removal. As regularization technique is further adopted to penalize the intense changes of important parameters jointly with the learned prompts. Extensive experimentation on various rain streak benchmarks demonstrates the effectiveness of our proposed scheme. Additionally, this approach can be seamlessly integrated into the training of lower-level task models, resulting in improved adaptability and functionality in challenging environments.

Acknowledgements

This work is supported by the Fundamental Research Funds for the Central Universities and the National Natural Science Foundation of China under Contract No.61772043 and a research achievement of Key Laboratory of Science, Technology and Standard in Press Industry (Key Laboratory of Intelligent Press Media Technology). This research work is also partially supported by the Basic and Frontier Research Project of PCL and the Major Key Project of PCL.