Data-free Distillation with Degradation-prompt Diffusion for Multi-weather Image Restoration
Pei Wang Xiaotong Luo Yuan Xie Yanyun Qu


Purification of Gaussian States by Photon Subtraction
    Mattia Walschaers
    September 9, 2024
=====================================================



§ ABSTRACT 

Multi-weather image restoration has witnessed incredible
progress, while the increasing model capacity and expensive data acquisition impair its applications in memory-limited devices.
Data-free distillation provides an alternative for allowing to learn a lightweight student model from a pre-trained teacher model without relying on the original training data. 
The existing data-free learning methods mainly optimize the models with the pseudo data generated by GANs or the real data collected from the Internet.
However, they inevitably suffer from the problems of unstable training or domain shifts with the original data. 
In this paper, we propose a novel Data-free Distillation with Degradation-prompt Diffusion framework for multi-weather Image Restoration (D4IR). 

It replaces GANs with pre-trained diffusion models to avoid model collapse and incorporates a degradation-aware prompt adapter to facilitate content-driven conditional diffusion for generating domain-related images. 
Specifically, a contrast-based degradation prompt adapter is firstly designed to capture degradation-aware prompts from web-collected degraded images. 
Then, the collected unpaired clean images are perturbed to latent features of stable diffusion, and conditioned with the degradation-aware prompts to synthesize new domain-related degraded images for knowledge distillation. 

Experiments illustrate that our proposal achieves comparable performance to the model distilled with original training data, and is even superior to other mainstream unsupervised methods. 












§ INTRODUCTION


Multi-weather image restoration (MWIR) aims to recover a high-quality image from a degraded input (e.g., haze, rain), which can be used in autonomous driving, security monitoring, etc. 

Nowadays, MWIR <cit.> has made significant progress relying on 



the rapid development of computing hardware and the availability of massive data. 

In actual scenarios, the increasing model complexity may impair its application on resource-constrained mobile vehicular devices. 

As a widely used technique, Knowledge Distillation (KD) <cit.> is often adopted for model compression.


However, the original training data is unavailable for some reasons, e.g., transmission constraints or 
privacy protection. 
Meanwhile, due to the variability of weather conditions, access to large-scale and high-quality datasets containing all weather conditions can be both difficult and expensive.
Therefore, it is necessary to develop data-free learning methods to compress existing IR models for adapting to different edge devices and more robust to various adverse weather conditions. 



Data-free knowledge distillation <cit.> paves such a way to obtain lightweight models without relying on the original training data. 
Its core concern is how to acquire data similar to the training data. 
The existing methods mainly achieve knowledge transfer by generating pseudo-data based on generative adversarial networks (GANs) <cit.> or collecting trust-worth data from the Internet <cit.>. 
However, these methods mainly focus on high-level tasks, lacking sufficient exploration in low-level image restoration for pixel-wise dense prediction. 







Recently, a few studies <cit.> have explored data-free learning for image restoration. 
However, there are still two underlying limitations. 

Firstly, they all adopt the GAN-based framework, which often faces unstable training and complex regularization hyperparameter tuning. 
Secondly, they use pure noise as input to generate pseudo-data that generally lack clear semantic and texture information. It is crucial for low-level vision tasks.
Although collecting data from the Internet can avoid the problem, it would inevitably face domain shift from the original data, which is difficult to solve for MWIR unlike simple perturbations based on class data statistics <cit.> in image classification. 




In order to mitigate the above issues, we advocate replacing GANs with a pre-trained conditional diffusion model and equipping it with degradation-aware prompts to generate domain-related images from content-related features. 
On the one hand, the diffusion models can avoid mode collapse or training instability of GANs and are superior in covering the modes of distribution <cit.>. 
On the other hand, by training on large-scale datasets, many conditional diffusion models (e.g., Stable Diffusion (SD) <cit.> ) demonstrate exceptional ability in creating images that closely resemble the content described in the prompts. 
Especially, some methods <cit.> resort to the powerful prior of these pre-trained models and introduce trainable adapters to align the internal learned knowledge with external control signals for task-specific image generation. 



In this paper, we propose a novel Data-free Distillation with Degradation-prompt Diffusion for multi-weather Image Restoration (D4IR). 


As shown in Fig. <ref>, unlike previous GAN-based data-free learning methods <cit.> for MWIR, 
our D4IR separately extracts degradation-aware and content-related feature representations from the unpaired web-collected images with conditional diffusion to better approach the source distribution. 

It aims to shrink the domain shift between the web-collected data and the original training data.


Specifically, our D4IR includes three main components: degradation-aware prompt adapter (DPA), content-driven conditional diffusion (CCD), and pixel-wise knowledge distillation (PKD). 
DPA and CCD are jointly utilized to generate degraded images close to the source data.
For DPA, a lightweight adapter is employed to extract degradation-aware prompts from web-collected low-quality images, which employs contrastive learning to effectively learn diverse degradation representations across different images. 

For CCD, the encoded features of web-collected clean images are perturbed to latent samples by forward diffusion, and then conditioned with the degradation-aware prompts for synthesizing data near the source distribution under the degradation reversal of the teacher model. 
With the newly generated images, the student network could be optimized to mimic the output of the teacher network through PKD. 
Experiments illustrate that our proposal achieves comparable performance to distill with the original training data, and is even superior to other mainstream unsupervised methods. 


In summary, the main contributions are four-fold: 

    
  * We propose a novel data-free distillation method for MWIR, which aims to break the restrictions on expensive model complexity and data availability.
    
    
  * We design a contrast-based adapter to encode degradation-aware prompts from various degraded images, and then embed them into stable diffusion.
    
    
    
    
    
  * We utilize the diffusion model to capture the latent content-aware representation from clean images, which combines the degradation-aware prompts to generate data that is more consistent with the source domain.
    
    
  * Extensive experiments demonstrate that our method can achieve comparable performance to the results distilled with the original data and other unsupervised methods. 





§ RELATED WORKS



 §.§ Multi-weather Image Restoration
 



MWIR can be divided into single-task specific models for deraining <cit.>, dehazing <cit.>, desnowing <cit.>, and multi-task all-in-one IR models <cit.>. 




Based on the physical and mathematical models, many MWIR methods <cit.> attempt to decouple degradation and content information from the training data. 
For example, DA-CLIP <cit.> adapts the controller and fixed CLIP image encoder to predict high-quality feature embeddings for content and degradation information. 

Recently, transformer-based models <cit.> have been introduced into low-level tasks to model long-range dependencies, significantly improving performance. 

Restormer <cit.> designs a efficient multi-head attention and feed-forward network to capture global pixel interactions. 

Though these methods have made powerful performance, the substantial storage space and computational resources make them challenging to deploy on resource-constrained edge devices.

Moreover, due to the difficulty in obtaining large-scale paired degraded-clean images, many methods use unpaired data to achieve unsupervised IR based on techniques like GANs <cit.>, contrastive learning <cit.>, etc. 
Unlike these methods, our proposal combines disentanglement learning and stable diffusion to generate data closer to the source domain for KD. 





 §.§ Data-free Knowledge Distillation 


Existing data-free distillation methods can be roughly classified into three types.
Firstly, the methods <cit.> reconstruct training samples in the distillation process with the “metadata" preserved during training. However, they are less feasible when only the pre-trained teacher model is accessible due to the necessity of “metadata". 

Secondly, the methods <cit.> optimize GANs to generate data similar to the distribution of original training data by a series of task-specific losses. 

DAFL <cit.> distills the student network by customizing one-hot loss, information entropy loss, and activation loss based on classification features. 
DFSR <cit.> introduces data-free distillation to image SR and designs the reconstruction loss with bicubic downsampling to achieve performance comparable to the student network trained with the original data. 

DFMC <cit.> adopts a contrastive regularization constraint to further improve model representation based on DFSR for MWIR. 

The last methods <cit.> optimize with web-collected data and try to address the distribution shift between collected data and original training data. KD3 <cit.> selects trustworthy instances based on classification predictions and learning the distribution-invariant representation. 




 §.§ Conditional Diffusion Models



To achieve flexible and controllable generation, conditional diffusion methods combine the auxiliary information (e.g., text <cit.>, image <cit.>, etc.) to generate specific images. 
In particular, Stable Diffusion (SD) <cit.> successfully integrates the text CLIP <cit.> into latent diffusion. 

Given the efficiency of foundation models such as SD, most recent methods <cit.> resort to their powerful prior and introduce trainable prompts to encode different types of conditions as guidance information. 
For example, T2I-Adapter <cit.> enables rich controllability in the color and structure of the generated results by training lightweight adapters to align the internal knowledge with external control signals according to different conditions. 
Diff-Plugin <cit.> designs a lightweight task plugin with dual branches for a variety of low-level tasks, guiding the diffusion process for preserving image content while providing task-specific priors. 




§ PROPOSED METHOD




 §.§ Preliminary

Notation and Formulation. 







Formally, given the pre-trained teacher network N_T(·), knowledge distillation (KD) aims to learn a lightweight student network N_S(·) by minimizing the model discrepancy dis(N_T, N_S). With the original training data D = {(x_i, y_i)}_i=1^|D| (“|·|" is the data cardinality, x_i and y_i are the degraded image and clean image), traditional KD is usually achieved by minimizing the following loss: 

    L_kd(N_S) = 1/|D|∑_i=1^|D|[∥ N_T(x_i) - N_S(x_i)∥_2]
 

Problem Definition. 
In practice, the original training data D may be inaccessible due to transmission or privacy limitations, which hinders efficient model training. 
That means only the pre-trained teacher model is available.
Therefore, our D4IR aims to address two significant issues for data-free KD:
(1) how to capture the data for model optimization; (2) how to achieve effective knowledge transfer.

Technically, data-free KD methods simulate D with generated pseudo-data or web-collected data. 
To efficiently synthesize the domain-related images to the original degraded data for MWIR, we first analyze the mathematical and physical models <cit.> used in traditional IR method. 
The general formulation of the degraded image Y is assumed to be obtained by convolving a clean image X with a fuzzy kernel B and further adding noise n as follows:

    Y = X * B + n


where * denotes convolution operation. Inspired by disentangled learning <cit.>, we consider decoupling the low-quality images as degradation-aware (B, n) and content-related information (X) from web-collected degraded images D̅_X = {x̅_i }_i=1^|D̅_X| and unpaired clean images D̅_Y ={y̅_i}_i=1^|D̅_Y| to facilitate the pre-trained SD model to generate source domain-related degraded images.  





 §.§ Method Overview

As illustrated in Fig. <ref>, our method consists of three main components: degradation-aware prompt adapter (DPA), content-driven conditional diffusion (CCD), and pixel-wise knowledge distillation (PKD). These parts are collaboratively worked to generate data close to the source domain so as to achieve data-free distillation of MWIR.
 

First, DPA includes a lightweight learnable encoder Enc_DP, which is used to extract degradation-aware prompts Enc_DP(x̅) from the collected degraded images x̅. 
To learn task-specific and image-specific degradation representations across various images, Enc_DP is trained with contrastive learning <cit.>, i.e., the features of patches from the same image (q, k^+) are pulled closer to each other and pushed away from ones of other images (k_i^-). 








Then, CCD performs the diffusion process from the perturbed latent features z_T^' of the collected clean images y̅, which is designed to relieve the style shift between the original data and the images generated by frozen stable diffusion <cit.> starting from random noise.
Moreover, z_T^' is conditioned with the degradation-aware prompts Enc_DP(x̅) for synthesizing new domain-related images x̂. 



Finally, PKD is conducted with the generated images x̂. Without loss of generality, the student network is optimized with a pixel-wise loss L_kd between its output N_S(x̂) and the one of teacher network N_T(x̂). 

Note that L_kd is utilized to simultaneously optimize N_S(·) and Enc_DP. 
It aims to filter the degradation types domain-related to the original data from large-scale collected images for contributing to KD. 




 §.§ Degradation-aware Prompt Adapter 
 

As previously discussed, the degradation-aware prompt adapter (DPA) aims to extract the degradation representations that help the student network learn from the teacher network with web-collected low-quality images. 
To achieve this, the adapter needs to satisfy the following conditions. 

First, DPA expects to effectively learn diverse degradation representations across different images while focusing on the task-specific and image-specific degradation information that distinguishes it from other images for the input image. 
Therefore, we adopt contrastive learning <cit.> to optimize DPA to pull in the same degradation features and push away irrelevant features. 

Specifically, we randomly crop two patches x̅_q and x̅_k^+ from the collected degraded image x̅, which are considered to contain the same degradation information. 
Then, they are passed to a lightweight encoder Enc_DP with three residual blocks and a multi-layer perceptron layer to obtain the corresponding features q=Enc_DP(x̅_q) and k^+=Enc_DP(x̅_k^+). We treat q and k^+ as query and positive samples. 
On the contrary, the features k_i^-=Enc_DP(x̅_k_i^-) of the patches x̅_k_i^- cropped from other images are viewed as negative samples. All negative sample features are stored in a dynamically updated queue of feature vectors from adjacent training batches following MoCo <cit.>. Thus, the contrastive loss L_cl can be expressed as:

    L_cl(Enc_DP) = -logexp( q · k^+ / τ)/∑_i=1^K exp( q · k_i^- / τ)
 
where τ is a temperature hyper-parameter set as 0.07 <cit.> and K denotes the number of negative samples. 


Second, DPA needs to extract domain-related prompts to guide the diffusion model in synthesizing images that facilitate knowledge transfer. If we only use Eq. (<ref>) to optimize Enc_DP, the resulting prompts may overlook the degradation differences between the web-collected data and the original training data. 

This implies that DPA might only capture degradation features across different input images, leading to a distribution shift from the original data. 

To address this, we employ the distillation loss L_kd between the outputs of the student model and teacher model to simultaneously optimize the degradation prompt encoder and the student model. 


Replacing the text prompt encoder in the pre-trained SD model, we employ the DPA to align the internal knowledge prior with external encoded degradation-aware prompts by the cross-attention module <cit.> for generating images toward specific degradation-related images: 


    Attention(Q,K,V) = softmax(QK^T/√(d))· V

Q, K, and V projections are calculated as follows: 

    Q = W_Q^(i)·φ_i(z_t),     K = W_K^(i)· Enc_DP(x̅), 
    
        V = W_V^(i)   · Enc_DP(x̅)


where φ_i(z_t) denotes the intermediate representation of the UNet in SD. W_Q^(i), W_K^(i), and W_V^(i) are projection matrices frozen in SD. d is the scaling factor <cit.>. 










 §.§ Content-driven Conditional Diffusion


According to the degradation prompts, the diffusion models still cannot generate domain-related images. This is because they inevitably suffer from the content and style differences against the real images without specifying the content of the images. 
Therefore, it is necessary to address the content shift from the original degraded data while preserving the realism of the collected images.


Inspired by SDEdit <cit.>, we choose the noised latent features z_T^' encoded from the collected clean image y̅ instead of the random noise to synthesize domain-related images with realism. Specifically, we first encode the web-collected clean images y̅ into latent representations z_0 by the encoder Enc_SD frozen in SD via z_0 = Enc_SD(y̅).





Then, we replace the initial random Gaussian noise with the T^'-step noised features z_T^' of the latent features z_0 as the input to the diffusion model:  

    z_t = √(α̅_t)z_0 + √(1-α̅_t)ϵ_t,  t=T^'

where α̅_t is the pre-defined schedule variable <cit.>, ϵ_t ∼ N(0,1) is the random noise, T^' = λ * T, T is the total number of sampling steps in the diffusion model, and λ∈ [0,1] is a hyper-parameter indicating the degree of injected noise. 

With the learned conditional denoising autoencoder ϵ_θ, the pre-trained SD can gradually denoise z_T^' to z_0 conditioned with the degradation-aware prompts Enc_DP(x̅) via

    z_t-1 =    √(α̅_t-1)(z_t-√(1-α̅_t)ϵ_θ(z_t, t , Enc_DP(x̅))/√(α̅_t)) 
        
        + √(1-α̅_t)·ϵ_θ(z_t, t , Enc_DP(x̅))
 

Finally, the decoder Dec_SD reconstructs the image x̂ from the denoised latent feature z_0 as x̂ = Dec_SD(z_0).




As the noised input z_T^' to the diffusion model retains certain features of the real image y̅, the generated image x̂ closely aligns in style with the real image. More importantly, by starting from the partially noised features of the collected clean images, the pre-trained SD model can generate images D̂ = {(x̂_i)}_i=1^|D̂| that reflect the content and degradation characteristics of the original training data, when conditioned with degradation-aware prompts Enc_DP(x̅). 




 §.§ Pixel-wise Knowledge Distillation

Considering that image restoration focuses on pixel-level detail in an image, we calculate the distillation loss L_kd by the pixel-wise distance between the outputs of the student network and the teacher network as:



    L_kd(N_S, Enc_DP) = 1/|D̂|∑_i=1^|D̂|[∥ N_T(x̂_i) - N_S(x̂_i)∥_2]


where x̂_̂î denotes the synthesized images. For better generalization, we simply provide a simple way to conduct distillation, and other KD losses are also encouraged. 


Note that the distillation loss is used to optimize both the student network and the degradation prompt adapter. Therefore, the whole objective function is formulated as:

    L(N_S, Enc_DP) = L_kd(N_S, Enc_DP) + γ· L_cl(Enc_DP)

where γ is a regularization coefficient to balance the distillation loss and the contrastive loss.













§ EXPERIMENTS
 









 §.§ Experimental Settings

Datasets. Following the previous work in high-level tasks <cit.>, we introduce the web-collected data to synthesize data near the original distribution. Specifically, our datasets are as follows: 

1) Original Training Datasets: Here, we mainly consider the common weather following the representative AirNet <cit.>. The teacher networks are trained on Rain100L <cit.> for deraining, the Outdoor Training Set (OTS) <cit.> for dehazing, and Snow100K <cit.> for desnowing. 



2) Web-Collected Datasets: 
For image draining, we employ the training images from the large-scale deraining dataset Rain1400 <cit.> with 12,600 rainy-clean image pairs. 


For image dehazing, we adopt the training images from RESIDE <cit.> with 72,135 outdoor and 13,990 indoor hazy-clean image pairs. 
For image desnowing, we set the training images from the Comprehensive Snow Dataset (CSD) <cit.> with 8,000 snowy-clean image pairs. 
Note that the paired images are randomly shuffled during training to reach an unpaired configuration.

3) Test Datasets: 
Following the common test setting for different weather image restoration, we adopt Rain100L <cit.>, Synthetic Objective Testing Set (SOTS) <cit.>, and the test datasets of Snow100K for image deraining, dehazing and desnowing, respectively. 

Implementation Details. 
We employ the pre-trained AirNet as the teacher network and then halve the number of feature channels to obtain the student network. 
The initial learning rates of the student network N_S(·) and the degradation prompt encoder Enc_DP are set as 1×10^-3 and 1×10^-5, respectively, which are decayed by half every 15 epoch. 
Adam optimizer is used to train D4IR with β_1=0.9 and β_2=0.999. 
The specific sampling step of the latent diffusion <cit.> is 70. 
During training, the input RGB images are randomly cropped into 256×256 patches and the batch size is set following AirNet. 
To ensure the training stability, we first train N_S(·) and Enc_DP together as Eq. (<ref>) for 50 epochs, and then with the distillation loss as Eq. (<ref>) for 150 epochs. 
Besides, the hyperparameter λ in Eq. (<ref>) and the trade-off parameter γ in Eq. (<ref>) are set as 0.5 and 0.5, respectively (the analysis is shown in the supplementary material). 
All experiments are conducted in PyTorch on NVIDIA GeForce RTX 3090 GPUs. 


Evaluation Metrics. Peak signal-to-noise ratio (PSNR) <cit.> and structural similarity (SSIM) <cit.> are utilized to evalute the performance of our method. 

Besides, the parameters are used to evaluate model efficiency.









 §.§ Comparisons with the State-of-the-art 

To validate the effectiveness of our D4IR, we provide quantitative and qualitative comparisons for image deraining, dehazing, and desnowing. 
Here, we mainly compare our D4IR with four kinds of methods: 

1) directly train the student network with the original training data of the teacher network (Student). 
2) distill the student network with the original degraded data without the GT supervision (Data). 
3) distill the student network by DFSR <cit.> and DFMC <cit.>. 

Other data-free distillation methods are designed for high-level vision tasks, which cannot be applied to IR for comparison. 

4) the mainstream unsupervised methods that are trained on unpaired data. 

For Image Deraining. 
As shown in Tab. <ref>, it is observed that the performance of the student network obtained by our D4IR for image deraining improves by 0.91dB on PSNR and 0.023 on SSIM compared to “Data".


This benefits from the wider range of data synthesized by our D4IR, which is domain-related to the original degraded data so as to facilitate the student network to focus on the knowledge of the teacher network more comprehensively.
Besides, the performance of our D4IR also far exceeds that of the GAN-based DFSR and performs better than DFMC (0.44dB and 0.024 higher on PSNR and SSIM). 

Moreover, D4IR also performs better than most mainstream unsupervised image deraining methods and achieves comparable performance with Mask-DerainGAN with only the half parameters. 
The visual comparisons in Fig. <ref> show that D4IR achieves a significant rain removal effect and is better than DFMC, DFSR, and students distilled with original data for removing rain marks. 

For Image Dehazing. 
As shown in Tab. <ref>, our D4IR also outperforms the student distilled with the original degraded data (0.04dB higher on PSNR and 0.001 higher on SSIM) and performs much better than DFSR and DFMC, which lack specific degradation-related losses. 

Besides, compared to the popular unsupervised image dehazing methods, D4IR has a much smaller number of parameters in second place on PSNR and SSIM. 
The visual result is given in Fig. <ref>. It shows that our D4IR has a significant dehazing effect and is closer to the GT than DFMC, DFSR, and “Data".








In Fig. <ref>, we present visualized samples synthesized by DFMC, the pre-trained SD model, and our D4IR for image dehazing. The results indicate that GAN-based DFMC, which initiates from pure noise, struggles to produce images with semantic information. Additionally, generating images with rich texture and color details using simple textual prompts proves challenging for SD. In contrast, our D4IR method generates images with more detailed texture and semantic information compared to both DFMC and SD. 

The results for image desnowing are in the supplement.



 §.§ Ablation Studies 

Here, we mainly conduct the ablation experiments on the image deraining task as follows: 


Break-down Ablation. We analyze the effect of the degradation-aware prompt adapter (DPA) and content-driven conditional diffusion (CCD) by setting different input z_0 (noise and CCD) and prompts (none, textual features same as SD, content features encoded from clean images, and DPA) for frozen SD model in Tab. <ref>.
It is observed that the performance of M1 is slightly better than that of M2 since the “text-to-image" generative model is powerful in generating images with original textual prompts. 
Besides, the degradation-aware prompts can not work well without content-related information (M2) for the absence of content information compared with M3. 
Both textual degradation prompts (M5) and our proposed DPA (D4IR) effectively improve student models' performance compared with none prompts (M4). 
Our D4IR performs the best by jointly utilizing DPA and CCD to generate images close to the original degraded data. It improves 1.65dB on PSNR compared with the model relying solely on the pre-trained SD model (M1) and 1.34dB on PSNR compared with the model directly distilled with the web-collected data (M0). 




Real-world Dataset. 


For further general evaluation in practical use, we conducted experiments on the real-world rainy dataset SPA <cit.>. 
As shown in Tab. <ref>, our D4IR  also has comparable performance with the student distilled with original data in real-world scenarios (0.08dB higher on PSNR). 
More comparisons with other unsupervised methods are presented in the supplementary material. 




Different Backbones of Teacher Network. 
We also validate D4IR with a transformer-based teacher backbone Restormer <cit.> on Rain100L. Due to resource constraints, we use Restormer with halved feature channels (from 48 to 24) as our teacher network and a quarter of feature channels (from 48 to 12) as the student network. 
As shown in Tab. <ref>, the shrunk model capacity also leads to a large performance loss of the student network compared to the teacher network. 

Besides, it is observed that the performance of our D4IR is slightly lower than that of the student network distilled with the original degraded data. 
The reason lies in that the images generated by the diffusion model still differ from the real training data while the self-attention mechanism of the transformer pays more attention to the global contextual information of the images. 





§ CONCLUSION

This paper proposes a simple yet effective data-free distillation method with degradation-aware diffusion for MWIR. 
To achieve this, we mainly consider three concerns, including: 
1) investigate the application of the conditional diffusion model to solve the unstable training of the traditional GANs in data-free learning; 
2) introduce a contrast-based prompt adapter to extract degradation-aware prompts from collected degraded images; 
and 3) start diffusion generation from content-related features of collected unpaired clean images. 
Extensive experiments show that our D4IR obtains reliable student networks without original data by effectively handling the distribution shifts of degradation and content.
In future work, we will continue to study more effective prompt generation to enable efficient model learning.













