Research map / Image generation / Fast diffusion sampling
RestoreGrad: Signal Restoration Using Conditional Denoising Diffusion Models with Jointly Learned Prior
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes RestoreGrad, which jointly learns a data-dependent Gaussian prior for conditional DDPMs via a VAE-style ELBO with prior and posterior encoders, speeding convergence and enabling fewer sampling steps for speech enhancement and image restoration.
Problem
Conditional DDPMs for signal restoration converge slowly and need many sampling steps because the standard Gaussian prior discards information in the degraded conditioner about the clean signal.
Approach
Integrates the conditional DDPM into a VAE framework and jointly learns a zero-mean Gaussian prior (from the degraded signal via a Prior Net) and a posterior (from clean plus degraded signal via a Posterior Net), so the DDPM samples its initial noise from the learned prior instead of N(0, I).
Key results
- On VoiceBank+DEMAND, RestoreGrad reaches PESQ 2.51 in 96 epochs versus 2.44 for fully trained CDiffuSE (445 epochs), and keeps 2.50 PESQ with 3 sampling steps.
- On image restoration, RainDropDiff reaches LPIPS 0.143 at 9.2K epochs while RestoreGrad reaches it at 1.8K epochs, a 5x speed-up.
- On the multi-weather AllWeather benchmark, RestoreGrad (WeatherDiff backbone) reaches 30.82 PSNR on Snow100K-L and 31.78 on RainDrop, the best multi-weather results in the comparison table.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Jointly learn a prior and posterior with the conditional DDPM | CDiffuSE on VoiceBank+DEMAND · VoiceBank+DEMAND test | PESQ: 2.44 -> 2.51 | Helped |
| Use RestoreGrad instead of PriorGrad | CDiffuSE on VoiceBank+DEMAND · VoiceBank+DEMAND test | PESQ: 2.42$\pm$3e-3 -> \textbf{2.51}$\pm$6e-4 | Helped |
| Use RestoreGrad for fewer inference steps | CDiffuSE on VoiceBank+DEMAND · VoiceBank+DEMAND test | PESQ: 2.51 -> 2.50 | Helped |
| Add jointly learned prior and posterior encoders | WeatherDiff on Snow100K-L · Snow100K-L | PSNR: 30.09 -> 30.82 | Helped |
| Add jointly learned prior and posterior encoders | WeatherDiff on Outdoor-Rain · Outdoor-Rain | PSNR: 29.64 -> 30.83 | Helped |
| Add jointly learned prior and posterior encoders | WeatherDiff on RainDrop · RainDrop | PSNR: 30.71 -> 31.78 | Helped |
| Train RestoreGrad longer | WeatherDiff on Outdoor-Rain · Outdoor-Rain | PSNR: 30.83 -> 30.70 | Hurt |
| Add jointly learned prior and posterior encoders | WeatherDiff on real-world image data · RainDS-Real | NIQE: 3.050 -> 2.556 | Helped |
| Add jointly learned prior and posterior encoders | WeatherDiff on real-world image data · Snow100K-Real | NIQE: 2.985 -> 3.015 | Hurt |
| Add jointly learned prior and posterior encoders | CDiffuSE on CHiME-3 · CHiME-3 test | PESQ: 1.55 -> 1.54 | No clear effect |
| Add jointly learned prior and posterior encoders | CDiffuSE on CHiME-3 · CHiME-3 test | SI-SNR: 7.67 -> 8.45 | Helped |
| Increase encoder size from Tiny (24K) to Base (93K) (ablation) | RestoreGrad on VoiceBank+DEMAND · VoiceBank+DEMAND test | PESQ: 2.48 -> 2.51 | Helped |
Surprising findings
- Training with the Posterior Net gives better results than the prior-only variant, and removing the latent regularization (eta=0) led to training divergence.
- Performance is reported to be not very sensitive to the latent regularization weight eta across roughly [0.01, 10].
Limitations
- Restricted to signal restoration with a zero-mean Gaussian prior where only the covariance is learned.
- Image restoration results are compared mostly against baselines trained on the same synthetic data; PSNR/SSIM gap to DTPM on individual weather types is not closed.
- Speech experiments use a single VoiceBank+DEMAND setup and 1-day training on one V100 for 96 epochs.
Experimental setup
Datasets: VoiceBank+DEMAND RainDrop AllWeather Snow100K Outdoor-Rain CHiME-3 RainDS-Real Snow100K-Real
Benchmarks: VoiceBank+DEMAND CHiME-3 Snow100K-L Outdoor-Rain RainDrop RainDS-Real Snow100K-Real
Metrics: PESQ SI-SNR SSNR CSIG CBAK COVL PSNR SSIM LPIPS FID
Baselines: CDiffuSE PriorGrad DOSE Demucs WaveCRN WeatherDiff RainDropDiff DTPM TransWeather All-in-One
Builds on: DDPM PriorGrad CDiffuSE WeatherDiff VAE
Closest papers by meaning
- Denoising as Adaptation: Noise-Space Domain Adaptation for Image Restoration (ICLR 2025 · similarity 0.89)
- SeNM-VAE: Semi-Supervised Noise Modeling with Hierarchical Variational Autoencoder (CVPR 2024 · similarity 0.89)
- Adapt and Diffuse: Sample-adaptive Reconstruction via Latent Diffusion Models (ICML 2024 · similarity 0.88)
- Diffusion Priors for Variational Likelihood Estimation and Image Denoising (NeurIPS 2024 · similarity 0.88)
- Residual Denoising Diffusion Models (CVPR 2024 · similarity 0.88)
- Stochastic Deep Restoration Priors for Imaging Inverse Problems (ICML 2025 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.