RestoreGrad: Signal Restoration Using Conditional Denoising Diffusion Models with Jointly Learned Prior
Authors: Ching-Hua Lee
Chouchang Yang
Jaejin Cho
Yashas Malur Saidutta
Rakshith Sharma Srinivasa
Yilin Shen
Hongxia Jin
Venue: ICML 2025 poster (2025)
Keywords: Denoising diffusion probabilistic model
prior distribution
posterior
speech enhancement
image restoration
Primary area: deep_learning->generative_models_and_autoencoders
TL;DR: This paper proposes an integration of conditional denoising diffusion probabilistic models (DDPMs) into the variational autoencoder (VAE) framework to jointly learn a more informative diffusion prior in signal restoration applications.

Abstract: Denoising diffusion probabilistic models (DDPMs) can be utilized to recover a clean signal from its degraded observation(s) by conditioning the model on the degraded signal. The degraded signals are themselves contaminated versions of the clean signals; due to this correlation, they may encompass certain useful information about the target clean data distribution. However, existing adoption of the standard Gaussian as the prior distribution in turn discards such information when shaping the prior, resulting in sub-optimal performance. In this paper, we propose to improve conditional DDPMs for signal restoration by leveraging a more informative prior that is jointly learned with the diffusion model. The proposed framework, called RestoreGrad, seamlessly integrates DDPMs into the variational autoencoder (VAE) framework, taking advantage of the correlation between the degraded and clean signals to encode a better diffusion prior. On speech and image restoration tasks, we show that RestoreGrad demonstrates faster convergence (5-10 times fewer training steps) to achieve better quality of restored signals over existing DDPM baselines and improved robustness to using fewer sampling steps in inference time (2-2.5 times fewer), advocating the advantages of leveraging jointly learned prior for efficiency improvements in the diffusion process.

RestoreGrad: Signal Restoration Using Conditional Denoising Diffusion Models with Jointly Learned Prior
Ching-Hua Lee
Chouchang Yang'
Jaejin Cho Yilin Shen
Yashas Malur Saidutta Rakshith Sharma Srinivasa Hongxia Jin
Abstract
able generative models consisting of i) the forward process, where the original data samples are gradually corrupted by adding Gaussian noise to eventually become a standard normal prior; ii) the reverse process, in which neural net- work is responsible for recovering the original data from the corrupted samples by learning to sequentially reverse the diffusion process. With their exceptional capabilities of gen- erating high-quality data, DDPMs can be applied to various signal restoration tasks recovering the missing compo- nents signal due to contamination (e.g., audio recorded with environmental noises (Lu et al., 2021; 2022; Tai et al., 2023b), images obstructed by various measurement noises (Özdenizei & Legenstein, 2023; Croitoru et al., 2023)), by conditioning the DDPM the degraded observations.
Denoising diffusion probabilistic models (DDPMs) can be utilized to recover a clean signal from its degraded observation(s) by conditioning the model on the degraded signal. The degraded signals themselves contaminated versions of the clean signals; due this correlation, they may encompass certain useful information about the target clean data distribution. However, existing adoption of the standard Gaussian as the prior distribution in turn discards such information when shaping the prior, resulting in sub-optimal performance. In this paper, we propose to improve conditional DDPMs for signal restoration by leveraging more informative prior that is jointly learned with the diffusion model. The proposed framework, called RestoreGrad, seamlessly integrates DDPMs into the variational autoencoder (VAE) framework, taking advantage of the correlation between the degraded and clean signals encode better diffusion prior. On speech and image restoration tasks, show that RestoreGrad demonstrates faster convergence (5-10 times fewer training steps) to achieve better quality of restored signals over existing DDPM baselines and improved robustness to using fewer sampling steps inference time (2-2.5 times fewer), advocating the advantages of leveraging jointly learned prior for efficiency improvements in the diffusion process.
However, for the diffusion model to adequately learn the reverse process, a large number of training iterations may be required, leading to potentially slow convergence. Such inefficiency was related to the discrepancy between the real data distribution and the accustomed choice of the standard Gaussian prior by Lee et al. (2022). They therefor proposed simple yet effective approach called PriorGrad, aiming to construct better prior by using rule- approaches to extract useful information from the conditioner data. How- ever, despite improving performance on some generative speech tasks, handcrafting a "better" prior requires certain knowledge about the data characteristics, and such guidance may always exist for a given task application.
In this paper, our main focus is investigate the question: Can we systematically obtain a better prior distribution that improves the efficiency of the diffusion generative process? In other words, we aim to develop a framework of learning- based diffusion priors for improved DDPM efficiency. high-level view of our approach depicted in Figure 1, where the conditional DDPM (parameterized by 0) samples the latent noise from learned prior distribution estimated by a prior encoder v. which takes the conditioner as input. The prior encoder is jointly trained with the DDPM to synthesize the data Xo, and a posterior encoder & that exploits information from both Xg and y, to align the prior and posterior distributions. The main idea here that, if there is certain correlation between the conditioner y and
1. Introduction
Denoising diffusion probabilistic models (DDPMs) (Ho et al., 2020; Sohl-Dickstein et al., 2015) are latent vari-
'Artificial Intelligence Center- Mountain View, Samsung Elec- tronics, Mountain View, CA, USA. Correspondence to: Ching- Hua Lee <chinghua. @ comm
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

RestoreGrad: Signal Restoration Using Conditional Denoising Diffusion Models with Jointly Learned Prior
Conditional DDPM
Speech enhancement (SE) comparison of conditional DDPMs
dala
forward process g(x|Xx. 1)
e
N I)
[slandard]
2.6
PP50 f
covi
VS.
30
X
X
XT
Pe (olely) [proposed]
25-
reverse process Px y).
jointly
1.
1.6
Baseline (CDIMUSE)
2.
Base ine (CDrfUuE)
postertor encoder
conditioner
prlor encoder
learned prior
distribution
+PriorGrad
RestoreGrad (ours]
15
+PrierGrad
RestoreGrad touns)
y
Traiiing epechs
Training spochs
qpleixa,y) y) = pely))
postarior distributions aligned
Image restoration (IR) comparison of conditional DDPMs
during model training
0:100
LPPS
HD 4
Figure 1. High-level view of the proposed method.
0.15
-
Basalina (RainDropo +RestoreGrad (ours)
Baseline (ialn
45
*RestoreGrad tours)
the target data Xo, e.g., in signal restoration problems where y is typically a degraded version of Xu. our framework can exploit such correlation construct a more informative prior in automatic and systematie manner.
0.155
150
0.145
0.140
20
60 Trainin spacts (hundreds)
80
To explore the idea, we introduce RestoreGrad, a new paradigm for improving conditional DDPM by learning the prior distribution in tandem with the diffusion model, fo- cusing on signal restoration applications. We apply Restore- Grad to speech enhancement (SE) and image restoration (IR) tasks to demonstrate its generality for signals of differ- ent nature. For SE, we compare with PriorGrad (Lee al., 2022) which provides guidance on handcrafting suitable pri- ors in the speech domain. For IR, we show that RestoreGrad serves as a promising solution for improving the baseline DDPM even in domain that lacks such a recipe for en- gineering the prior. As shown in Figure 2, models trained using RestoreGrad are more data and compute-efficien than the baseline DDPMs and PriorGrad; they converge faster to achieve higher quality of the restored signal. Further shown in Figure 3, the learned prior is more informative as it better correlates with the desired signal than an isotropic covariance, potentially simplifying the diffusion trajectory for improved model efficiency.
Orandhado)
Figure 2. Model learning performance. (Top) the speech domain, RestoreGrad outperforms PriorGrad (Lee et al., 2022), recently proposed improvement to the baseline DDPM (CDiffuSE (Lu et al., 2022)) which leverages handcrafted priors. (Bottom) In the image domain, RestoreGrad provides paradigm to improve the DDPM baseline (RainDropDiff (Özdenizei & Legenstein, 2023)) when there is no existing recipe for obtaining better priors.
Example leaned price SE
Example leaned grace tisk
Tarnet daan sphoch
Taspe clean image
Examod agror
Figure 3. Visualization of the learned prior distribution. Here, the distribution of the prior is modeled as: Po y) = N ((e 0, diag (f (y: ())))) where price is estimated by the prior encoder s with input y. It appears that peior follows the level vari- ation of the speech waveform (in SE) and preserves the structure of the original image (in IR). This indicates that an informative prior approximating the data distribution has been obtained, leading to improved efficiency of the diffusion process.
Our main contributions are summarized as follows:
We study the problem of learning the prior distribution jointly with the conditional DDPM for signal restora- tion applications. aiming at providing a more systematic, learning- based treatment to address the inefficiency in- curred by existing selections of the prior distribution.
Experiments demonstrate that the proposed paradigm is quite general and can benefit both training and sampling of DDPMs, achieving considerabl improvements with lightweight encoders in high quality signal restoration tasks of various modalities including images and audio.
We propose a new framework called RestoreGrad that learns the prior in conjuncture with the DDPM model through a prior encoder, by exploiting the correlation be- tween the targe signal and input degraded signal encoded by an auxiliary posterior encoder, for improved model efficiency. Our two-encoder learning framework is estab- lished based on novel evidence lower bound (ELBO) that seamlessly integrates the DDPM into the variational autoencoder (VAE) (Kingma& & Welling, 2014) to harness the advantages of both methodologies.
2. Backkround on DDPMs
2.1. Forward Process
DDPMs (Ho et al., 2020: Sohl-Dickstein et al., 2015) slowly corrupt the training data using Gaussian noise in the forward
2

---

RestoreGrad: Signal Restoration Using Conditional Denoising Diffusion Models with Jointly Learned Prior
process. Let Adata (dots(X)) be the data density of the original data Xo- The forward process is a fixed Markov Chain that sequentially corrupts the data Xo qdata (Xo) in T diffusion steps, by injecting Gaussian noise according to variance schedule (= E 1);
Based using a fixed covariance Ex(x,, () = ofl (e.g..
of 3,) as in Ho al. (2020), maximizing (3) corresponds to training network p4g (x,:x: f) that predicts ju, (X,, xo). Al-
ternatively, Ho et al. (2020) suggested the following repa- rameterization to rewrite the mean as a function of noise:
T g(x1T|(0)) = II (xxx
1 Vac X;
B. _x Ex(x
14g (x((xt,)) t)
(4)
VI
They train network Ex(x,, 1) to predict the real noise NO I) and use (4) to compute the mean. Practically it is carried out by minimizing a simplified training objective:
where (xx||x(-1)= = N(xx: v1 -, B2T) is the tran- sition probability at step L. It allows the direct sampling of X according to (x,, xo) = N (xxx vārxo, vI- v where 0. = II a; with a, == 1- B5; i.e., we can sample X2 = Vā + v1 ãye, where e .V(O I). A notable assumption is that with carefully designed variance sched- ule Br and large enough T, such that ãT is sufficiently small, q(x) Xb) converges to N M(x: O. I) so that the distribution of xT is well approximated by the standard Gaussian.
Eximlee (6) == Ex,,e,t [llc-co(x co t)11
(5)
which measures, for random time step L ~ u(II,.... T}), the distance between the actual noise and estimated noise.
2.4. Signal Restoration by Conditional DDPMs
2.2. Reverse Process
Signal restoration is concerned with recovering the origi- nal signals from their degraded observations, which are of paramount importance in reality but remaining challeng- ing, as noises are ubiquitous and may be strong enough to cause significant degradation of the signal quality. Recently, adoption of deep generative models (Kingma & Welling, 2014; Goodfellow et al., 2014; Ho al., 2020) for signal restoration tasks has considerably increased due to their re- markable capabilities of generating missing components in the data, with conditional DDPMs (Croitoru et al., 2023; Cao et al., 2024) demonstrating substantial promise. More formally, let y denote the degraded observation of the clean signal Xg. Recovering X,) given y by a model & can be cast as maximizing the conditional likelihood of data ps (xoly). The problem is in general intractable, but can be approxi- mated by using DDPM conditioned on y. The main idea is to learn diffusion model 0 with y provided as a conditioner in the reverse process (2):
One can generate new data samples from (Xa) by re- versing the predefined forward process utilizing the same functional form. That is, we can progressively transform a noise XT ~ p(xT) = (xT; 0, I) back into the data by ap- proximating the reverse of the forward transition probability. This process is defined by the joint distribution po (xo:r) of a Markov Chain with learned transitions:
T
(o(xur)) = p(xx) II t=1
po (x-- (1x))
(2)
where (x = (xxx, 1), Ex t)) is the
reverse transition probability parameterized by network 0.
2.3. DDPM Learning Framework
Ideally, we would train the model 0 with a maximum like- lihood objective such that the probability assigned by the model po(x)) to each training example is as large as pos- sible, which is unfortunately intractable (Croitoru et al., 2023). To circumvent such difficulty, DDPMs (Ho et al., 2020) instead maximize an ELBO of the data log-likelihood, by introducing a sequence of hidden variables X1:7 and the approximate variational distribution ((x1:T xx):
po(xorlly) = p(xT) II ( (6)
where Po(x;-1 /x..)))x))))))) =
y, y. t), I as-
suming a fixed covariance. In practice, noise estimator network eœ (x y, t) is adopted to predict the mean, follow- ing the practice in Ho et al. (2020).
po (Xxo:)
logp(xx) 2 E (x()))(())) log
(3)
3. Proposed Method: Integrating DDPM and VAE for Learnable Diffusion Prior
With the above parametric modeling of the forward and reverse processes, the ELBO (3) suggests training the net-
work 8 such that, each time step t. (x is as
We start with the conditional VAE (Sohn et al., 2015) for- mulation to maximize the conditional data log-likel log pog((x)))) ) = log ʃ p(x,, cly)de, where e is an introduced latent variable. To avoid intractable integral, in VAEs an ELBO is utilized as the surrogate objective by introducing
close as possible to the true forward process posterior con-
ditioned on Xo (Luo, 2022; Croitoru
al., 2023), i.c.,
q(x--1/X4 Xo) = (x PL, (x,, Xp), 3,I), where Be ==
1-3e-1 Be and (x,, xo) == val 19:
T-aa Xo + (ar(1-o) -o 1)
Xt-
T-o,
1-a;
3

---

RestoreGrad: Signal Restoration Using Conditional Denoising Diffusion Models with Jointly Learned Prior
an approximate posterior q(E) Xox y) (Harvey et al., 2022):
the posterior distribution with another network e (Posterior Ner), to incorporate richer information about the target sig- nal distribution into the learning of the prior. Together with (8), we introduce the new lower bound of the conditional data log-likelihood:
log (xxol)) 2 E B loog [ogg((xolx,e) c)]
reconstruction term
Dxi. (g(E) Y) {)lp((()))))))))
prior matching lerm
( log Po (X) ((X1 (xo)
logp(xuly) (e)p(r)))))))
The reconstruction and prior matching terms are typically realized by an incoder-decode architecture with c being the bottleneck representation sampled from the latent distri- bution. VAEs generally benefit from learnable latent spaces for good modeling efficiency. However, their generative per- formance often lags behind DDPMs that employ an iterative, more sophisticated decoding (reconstruction) process.
conditional DDPM
- -DiL. gas gs(elxo: y) || pop(ely)
Posterior Net
Prior Net
Based on the assumption Xy = 6, the conditional DDPM samples the latent noise X from the distribution of which is jointly estimated by the two encoders, ø and db. The two- encoder design is inspired by Kohl et al. (2018) for image segmentation with traditional U-Nets. Here, we adopt the idea in the context of DDPM for signal restoration, which is effective as it incorporates posterior information exploiting the correlation between clean and degraded signals. Based on (9), we introduce the training objective of RestoreGrad:
In this work, our aim is to embrace the best of both worlds, i.... remarkable generative power (DDPM) and modeling efficiency (VAE) to achieve improved output signal quality and training/sumpling efficiency simultaneously.
Proposition 3.1 (Incorporation of diffusion process into VAE). By introducing a sequence of hidden variables X1:T, under the setup of conditional diffusion models where the Markov Chain assumption is employed on the forward pro- cess q(x1r (xo)) = II (xxx(xt-1) and the reverse pro- cess Po(Xo:r (y] = p(xx) (T=iPo(x4- y) parame- terized by a DDPM 0, and assuming that XT 6. (i.e., the latent noise of DDPM samples from the VAE latent distri- bution), we have the lower bound on log p(xoly, e) in the reconstruction term of the VAE (7) as:
Proposition 3.2 (RestoreGrad). Assume the prior and posterior distributions are both zero-mean Gaussian, pa- rameterized as p.lely) = N(e; Sprior (y: +) and go (e(c(xo.)) = -N (xo Y; d)), respectively, where the covariances are estimated by the Prior Net v (taking y as input) and Posterior Net & (taking both Xo and y as input). Let us simply use Eprior and Eyost here- after to refer to Epriolly 1b) and p (Xo. y; 6) for concise notation. Then, with the direct sampling property in the for- ward path x, = Vaixo + vT - .4e at arbitrary timestep t where E 4o (E y), and assuming the reverse process has the same covariance as the true forward process poste- rior conditioned on xo, by utilizing the conditional DDPM eo (x,,y,t) f) as the noise estimator of the true noise e, we have the modified ELBO, [(( d, 1b), associated with (9):
log pixolxy 6) 2 E (xx)))))))) log po(Xo:T [y) q(x1:1 xo)
(8)
which is the ELBO of the conditional DDPM (i.e., the con-
ditional version of (3)).
The proof (based on Markov Chain property) is provided in Appendix A. Note that the assumption Xy = e follows the standard DDPM to sample xT from the distribution of the prior noise e, practically achieved by using large enough T and carefully designed variance schedule {8.)(= In our case, adopt the same assumption enable sampling Xr from the latent space of VAB. We interpret Proposition 3.1 as seamless integration of DDPM into the VAE framework to achieve improved generative (decoding) capabilities.
C(0, 6, v)
OT Exx X 2
log posst
Latent Regularization (1.R) terms
T
+
Yo
cx(X,, +) Spm
Denoising Matching (DM) terms
Having incorporated the DDPM as the decoder, we now discuss the encoding part, i.e., the prior matching term in (7). A straightforward design could be using a network vb (Prior Net) to parameterize the prior distribution as (cly)) while assuming the posterior to be fixed form of distribu- tion like the standard Gaussian. However, this may in turn discard any useful information inherent between Xn and y. To take advantage of the adequate correlation present in signal restoration settings, we propose to also parameterize
1
prior
+ 2
(log
+ ir(E prior Epon)) + C,
Epostl
Prior Matching (PM) terms
(10)
1>1 1
where 7t-
207a az) 2
are weighting factors,
1=1
| stant not depending on learnable parameters 0, d. th.
= x
x of
=
1-
Br and C is some con- I-d,
4

---

RestoreGrad: Signal Restoration Using Conditional Denoising Diffusion Models with Jointly Learned Prior
Training
S Sampling
Postestor Net
Conditional DDFM
estimele
noise
ilerate denoise
sample
clean target signal No
Postunior Net
Epors
Ra(E No.
Lam
semple
Conditional
SaN
confaminatol y
obseevation
Prior Nitt
Pylelyy
DDPM
Epnnr
RestoreGred losa tarms
restored signal
conraminnteed
Price Net
Pylely)
La latent l requlaization
conditioner
observation
confitioner
dencising meiching
Lou: prior matzhing
Figure 4, Proposed RestoreGrad, During training, the conditional DDPM 6 Prior Net v, and Posterior Net & are jointly optimized by (11). During inference, the DDPM samples the latent noise e from the jointly learned prior distribution synthesize the clean signal. (Summary of the algorithm details is presented in Appendix C.I.)
The derivation (see Appendix B) is based on combining VAE and the results in Lee et al. (2022). Notably, we join the conditional DDPM with the posterior/prior encoders and optimize all modules at once, by connecting the DDPM prior space with the latent space estimated by the encoders. (g(eeXxo33 y) is performed by
summation over t by sampling t ~ ~ T)) (exploit- ing the independency due to Markov assumption (Nichol Dhariwal, 2021)) in each training iteration. We have also introduced 7) > 0 for the LR terms and X for PM terms, to exert flexible control of the learned latent space.
To this end, the sampling the standard reparameterization trick as in VAEs, unlocking end-to-end training via gradient descent on the loss terms:
3.2. Sampling of RestoreGrad
In applications that RestoreGrad mainly concerned with, the target signal Xy is not available in inference time. As in Figure 4, the conditional DDPM then samples e ~ (ely) = (o, Eprir) from the Prior Net instead; the Posterior Net is no longer needed.
Latent Regularization (LR) terms: to help learn a reason- able latent space; e.g., minimizing log E avoids post from becoming arbitrary large due to the presence of inverse in the weighted norms.
Denoising Matching (DM) terms: responsible for train- ing the DDPM to predict the true noise.
3.3. The Role of Posterior Information
In the training stage of RestoreGrad, the latent code sam- ples from the posterior (x) which exploits both the ground truth signal X and conditioner y. is thus more advantageous than existing works on adaptive priors (e.g.. PriorGrad (Lee et al., 2022)) that only utilize the conditioner y. To observe the benefits brought by the posterior informa- tion, can make comparison with variant of RestoreGrad where the Posterior Net is excluded in training:
Prior Matching (PM) terms: to shape a desirable latent space by aligning the prior and posterior distributions. Note that we model the distributions as zero-mean, based that signals can be properly normalized.
3.1. Training of RestoreGrad
With the the conditional DDPM 0, Prior Net vb, and Poste- rior Net & defined in Proposition 3.2, optimization can be performed to learn the model parameters of 0, 1b, oi based on the modified ELBO. The RestoreGrad framework jointly trains the three neural network modules by minimizing (10) depicted in Figure 4. Following existing DDPM literature, approximate the objective by dropping the weighting constant 7 of the DM terms, leading to the simplified loss:
min n ||x 1 + log) Epicel + ||e co
(12)
which basically removes the Posterior Net & and only trains the Prior Net ei and DDPM A. Interestingly, our experimen- tal results show that RestoreGrad indeed performs better with the Posterior Net than without it in model training.
4. Experiments
6.0.0 min n( ār |xol s poi
+ log/ EpoaE + Eo I)
Gin
4.1. Application to Speech Enhancement (SE)
Σ Egriorl Eppce
+A [log
tr( 1--1 Epes)):
4.1.1. EXPERIMENTAL SETUP
Dataset: We validate performance on the benchmark SE dataset VoiceBank+ BeiceBank+DEMANDD (Valentini-Ba -otinnho et al., 2016), consisting of clean speech clips collected from the
Eyst
(11)
where we approximate the expectations by randomly sam- pling (Xo. y) ~ qdam (Xam and e ~ (O, Epma), ) and the
VoiceBank corpus (Veaux et al., 2013), mixed with ten types
5

---

RestoreGrad: Signal Restoration Using Conditional Denoising Diffusion Models with Jointly Learned Prior
of noise profiles from the DEMAND database (Thiemann al., 2013). Specifically, the training utterances from Voice- Bank are artificially contaminated with the noise samples from DEMAND at 0, 5, 10, and 15 dB signal-to-noise ratio (SNR) levels, amounting to 11 11,572 utterances. The testing utterances are mixed with different noise samples at 2.5, 7.5, 12.5, and 17.5 dB SNR levels, amounting to 824 utterances.
Infer wth stepp iifor with stego
2-2
Iifer with idep
oFfcr with Stops
2.1
2-5-
2.3
2.6
Bese OCON
#Retorobrdd
Bessine (COITuS))
(Prrrdrol
-Rotowored
We consider: PESQ: Perceptual Eval- uation of Speech Quality (ITU-T Rec. P.862.2, 2005). SI- SNR: Scale-Invariant SNR (Le Roux et al., 2019). SSNR: Segmental SNR (Hu & Loizou, 2007). CSIG, CBAK, COVL: Mean-opinion-score predictors of signal distortion, background-nowl noise intrusiveness, and overall signal quality, respectively (Hu & Loizou, 2007).
Evaluation Metrics:
Figure 5. Robustness the reduction in reverse sampling time steps for inference.
RestoreGrad performance as a function of
2.55
2.50
2.45
2:40
Models: The following models are compared:
2.30
Base inee (CDi(luSE)
Baseline DDPM: We adopt the CDiffuSE (Base) model from Lu et al. (2022), which is based on DiffWave (Kong al., 2021) with 4.28M learnable parameters.
2:25
+Prirttired
+RestoreGrad
30
10"
10th
10'
10'
Figure 6, Effect of latent regularization weight ny for LLR on SE.
PriorGrad: We implement the PriorGrad (Lee et al., 2022) on top of CDiffuSE by changing the prior distribu- tion from (0, I) to (0, Σy), where Σ, is the covari- ance of the data-dependent prior computed based on the conditioner y. using the rule-based estimation approach for the application to vocoder in Lee et al. (2022).
Robustness to Reduced Reverse Steps in Inference: Re- storeGrad can potentially reduce the inference complexity. In Figure 5, we show how the trained diffusion models with- stand the reduction in the number of inference steps. In each model, we trained the network for 96 epochs and then infer- enced with 3 reverse steps to compare with the originally adopted 6-step scheme in Lu et al. (2022). The noise sched- ule for the 3-step scheme was Simter [0.05, 0.2, 0.35], a subset of the 6-step schedule that resulted in best perfor- mance. We can see that the baseline DDPM is most sen- sitive to the step reduction, while PriorGrad shows certain resistance as leveraging closer-to-dat prior distribution. Finally, RestoreGrad barely degrades with reduced sam- pling steps, echoing that a better prior has been obtained as it recovers higher fidelity signal even in fewer reverse sleps.
We incorporate Prior Net and Posterior Net on top of CDiffuSE. Both modules adopt the ResNet-20 architect (He et al., 2016), suitably modified to 1-D con- volutions for waveform processing, each with only 93K learnable parameters (only 2% of the CDiffuSE model).
RestoreGrad:
Configurations: We adopted the basic configurations same as in Lu et al. (2022). The waveforms were processed at 16kHz sampling rate. The number of forward diffu- sion steps was T = 50. The variance schedule was 3. E [10-1,0.035] linearly spaced. The batch size was 16. The fast sampling scheme in Kong et al. (2021) was used in the reverse process with 6 steps to reduce infer- ence complexity, with the 6-step inference variance sched- ule ginter = [10-1, 10-3 0.01, 0.05, 0.2, 0.35]. Adam opti- mizer (Kingma & Ba, 2014) was utilized with a learning rate of 2 x 10-4, We set y = 0.1 and x = 0.5 for (11).
Effect of 77 An important factor in our prior learning scheme is the regularization weight r for LLR of the training loss. An appropriate value of 1) should be large enough to properly regularize the learned latent space for avoiding in- stability, while not adversely affecting signal reconstruction performance. It is thus interesting to see how the perfor- mance varies with the choice of 1). Empirically, we found the overall SE performance not to be very sensitive to the value of n across wide runge, as shown Figure 6: Roughly in the range of [10- 2, 2010] of the rj value we see that Restore- Grad gives better results over both PriorGrad and CDiffuSE.
4.1.2. RESULTS
Improved Model Convergence: As shown in Figure 2 (test set performance), RestoreGrad shows better conver- gence behavior over PriorGrad (using handcrafted prior) and CDiffuSE (using standard Gaussian prior). For exam- plc, PriorGrad reaches 2.4 in PESQ at 96 epochs, whereas RestoreGrad reaches it in (roughly) 10 epochs, indicating a 10x speed-up. The results suggest that jointly learning the prior distribution can be beneficial for conditional DDPMs.
We present in Table more detailed comparison of RestoreGrad with the baseline CDiffuSE. Here, the scores of CDiffuSE were di- rectly taken from the results reported in Lu et al. (2022) where the model has been fully trained for 445 epochs. For
Comparison to Fully-Trained CDiffuSE:
6

---

RestoreGrad: Signal Restoration Using Conditional Denoising Diffusion Models with Jointly Learned Prior
Table 1. Comparison with the fully-trained CDiffuSE model per-
Table 3. Performance of RestoreGrad models trained with and without using Posterior Net.
formance reported in Lu et al. (2022).
Methods
epochs
tman inter
stes
PESQ
1
CSIG
CBAK
COVL
SISKR
CDiffUSE (trained epochs) PriurGrad
SE models
PESQT
2.32
COVLt SSNRT
2,89
3.94
SI-SNR1
11.84
CDiffoSE
2.44
3,66
2,83
3.0
2.42
3.03
5,53
14.21
PrioeGrad
96
6 2.1233--3 3.6742--3
2.932
3.03-20-3
14.2142-33
RestoreGrad
3.14
5,92
14.74
RestoreGrnd
6 2.51+66-4
4804
3-00
S.14 44.744-3--4
RestoreGrad w/o Posterior Net
Iraining diverged
(ours)
3 2.50 3.75:22-4
2.99:77-4
3.11 14.65
RestoreGrad w/o Posterior Net (q) 0.01) 2.47
3.08
4,96
11.22
"Bold text for underlined for second values.
ResteroGrad w/o Posterior Net = 2.48
3.12
5.11
13.29
'Best values in bold.
4.2. Application to Image Restoration (IR)
Table 2. SE comparison of RestoreGrad models using encoder modules of different sizes and the corresponding latency and GPU
4.2.1. EXPERIMENTAL SETUP
memory usage (measured one NVIDIA Tesla V100 GPU) pre-
sented as the ratio of encoder to DDPM.
Following Özdenizei & Legenstein (2023), we consider the IR task of recovering clean images from their degraded versions contaminated by synthesized noises cor- responding to different weather conditions. Two datasets are considered, where one is a weather-specific dataset called RainDrop (Qian et al., 2018) and the other is multi- weather dataset named Al/Weather (Valanarasu et al., 2022). The RainDrop dataset consists of images captured with rain- drops on the camera sensor which obstruct the view. It has 861 training images and a test set of 58 images dedicated for quantitative evaluations. The Al/Weather dataset is a curated training dataset from Valanarasu et al. (2022), which has 18,069 samples composed of subsets of training images from Snow 100K (Liu al., 2018), Outdoor-Rad Rain (Li et al., 2019) and RainDrop (Qian et al., 2018), in order create balanced training set across three weather conditions.
Dataset:
Fncoder size
PESQU
COVL1
SSNR1
SI-SNR
Pmc.. Time
Memory
Tiny (24K)
2.48
3.11
5,10
13.74
1.9%
6.5%
Base (93K)
2.51
3,14
5.92
14.74
2.2%
10.3%
Large (370K)
2.54
3,16
6.15
15.01
2.6%
18.2%
PriorGrad and RestoreGrad we report the tsd com- puted based on results of 10 independent samplings. We can see that with RestoreGrad applied, the SE model can achieve better performance over the baseline CDiffuSE by only training for 96 epochs (4.6 times lesser) in all the met- rics. In addition, halving the number of reverse steps in inference time still maintains better performance than the fully-trained CDiffuSE and also the PriorGrad.
Signal Quality and Encoder Complexity Trade-Offs: We further present results using three different model sizes (24K, 93K, 370K) for the Prior and Posterior Nets (en- coders) in Table 2, along with latency and GPU memory usage (presented as the ratio of encoder to DDPM). The re- sults clearly show that the restored speech quality improves with increasing encoder size. This indicates that there is a trade-off between the restoration signal quality and encoder model complexity. Notably, the latency and memory usage of the encoder modules are relatively small compared to the DDPM decoding (<2.6% latency and <18.2% memory usage of the DDPM processing), suggesting that Restore- Grad is capable of achieving improved performance without incurring considerable increase in complexity compared to the adopted DDPM model.
Quantitative evaluations of restored images are performed via Peak Signal-to-Noise Ratio (PSNR) (Huynh-Thu & Ghanbari, 2008), Structural SIMilar- ity (SSIM) (Wang et al., 2004), Learned Perceptual Image Patch Similarity (LPIPS) (Zhang al., 2018), and Fréchet Inception Distance (FID) (Heusel et al., 2017).
Evaluation Metrics:
Models: The following IR models are compared:
We consider the RainDropDiffes and WeatherDiffas in Özdenizci & Legenstein (2023) trained on the RainDrop and AllWeather datasets, respectively, as baseline DDPMs. Our work is based on the imple- mentation provided by Özdenizei & Legenstein (2023).
Baseline DDPMs:
We incorporate the encoder modules, Prior Net and Posterior Net, on top of the baseline DDPMs Both encoder modules adopt the ResNet-20 architect (He et al., 2016) with only 0.27M learnable parameters, signifi- cantly smaller (< 0.3%) than the baseline DDPM models.
RestoreGrad:
Posterior Net Helps: Finally, we validate the benefits brought by employing Posterior Net in the training phase by comparing with the RestoreGrad models trained without Posterior Net as (12) for some 7j. For fairness, all models were trained with 96 epochs, inferred with 6 steps. In Table 3. we observe that RestoreGrud achieves better results with Posterior Net than without it, indicating the benefits of being informed of the target Xg by utilizing the Posterior Net. We also observe that without regularizing the latent space (i.c., with 7) = 0) it could lead to training divergence.
Configurations: We used Adam optimizer with learning rate of 2 x 10-5, An exponential moving average with a weight of 0.999 was applied. We used T = 1000 and linear noise schedule 3, E [10-4, 4, 0.02], same as Özdenizei & Legenstein (2023). batch size of was used.
7

---

RestoreGrad: Signal Restoration Using Conditional Denoising Diffusion Models with Jointly Learned Prior
Table Comparison with existing IR models. The multi-weather (MW) models were trained on the AllWeather training (Valanarasu et al., 2022) and tested on three different weather types: Snow100K L (Liu et al., 2018), Outdoor- Rain (Li et al., 2019), and Rain (Qian et al., 2018). Several spectir- (WS) models that were trained on individual weather types are also presented for reference.
Type Methods
Snow100K L
Methods
Outdoer Rain
Methods
RainDrop
PSNR 1 SSIM
PSNK 1 SSIM
PSNK SSIM
RESCAN (Li et al., 2018)
26.08
0.8108
HRGAN eLi 2019)
21.56
0,8550
AtcutivGAN (Qian 2018)
31.59
0.9170
DesnowNet 2018)
0.5963
PCNet Llinng et 2021)
26.19
RaindropAnu (Quas et 2019)
31.44
WS
DDMSNet (Zhang et al., 2021)
28.85
0.8772
MPRNet iZamir al., 2021)
28.03
0.9192
IDT (Xio el al., 2022)
31,87
0.9313
SnowDiff (Undenizci & Legenstein,, 2023)
30.43
0.9145
RainHazzeDiif (Ozdenizei Legonstcin, 2023)
28.38
0.9320
RainDroplhdf (Üzdenizzi & Logenstein, 2023)
32:29
0.5422
DTTM (Ye et 2024)
30.92
0,9174
DTPM (Yo 2024)
30,99
0.9340
DIPM al.. 2024)
32,72
0.9440
All One et 2020)
0.8820
All-in One et 2020)
24.71
Alli in (Li al., 2020)
31.12
0.9268
TramWWelher (Valammasu et al., 2022)
29.31
0.8879
Trans Weather (Nalaaarauu et al., 2022)
28.83
0.9000
Trans Weather (Valaarasu all, 2022)
30.17
0.9157
MW
WeatherDiif (Oadeniaci & Legenstein, 2023)
30.09
0
WeatherDiff (Ondenizei & Legenstein, 2023)
29.64
0.9312
WeatherDidff (Ondenizci & Legenstein, 2023)
30.71
0.9312
RettersGrad (ours)
30.82
0.9159
RestoreGrad (ours)
30.83
0.9111
RestureGrad (ours)
31.78
0,9394
ResteroGrad (ouns trained longer
0.9175
RestoreGrad (ouns) trained longer
30.70
RestureGrad (ours) trained longer
32:26
0.9414
*Bold text for best and underlined text for second best values.
Degraded Input Image [PSNR: 23,25 I SSIM:0.8943]
WeatherDiff
RestoreGrad
Ground Truth Image
[PSNR: 28.77 SSIM: 0.9619]
[PSNR: 30.17 SSIM: 0.9723]
TROUBLE
TROUBLE
TROUBLE
TROUBLE
Figure 7. Image restoration examples using a test image taken from the Snow 100K-L test set. We provide more examples, including other degradations (desnowing, deraining, raindrop removal and deblurring) in Appendix D.2.
4.2.2. RESULTS
epochs, which is 1.14x fewer than WeatherDiff), Restore- Grad achieves further improvements shown in the last row the table. On the other hand, our method, while trained on weather data, achieves a comparable performance to the recently proposed Diffusion Texture Prior Model (DTPM) (Ye et al., 2024) individually trained on the three weather-specific (WS) datasets. Notably, our method does not require a pretraining stage on a large dataset of high- quality images, unlike DTPM which is pretrained 55,000 images samples. This suggests the generality and effective-- ness of our method to improve baseline DDPM models.
Model Convergence: As presented in Figure 2 (test set performance), RestoreGrad demonstrates faster convergence and better restored image quality over the baseline DDPM (RainDropDifff) For example, RainDropDiff reaches 143 in LPIPS at 9.2K epochs, while RestoreGrad reaches it in 1.8K epochs only, indicating a 5x speed-up up due to the effectiveness of the joint prior learning scheme.
Comparison with Existing IR Models: We compare our method with existing models of the multi multi-weather (MW) type in Table 4, including All-in-One (Li et al., 2020) and TransWeather ( (Valanarasu et al., 2022) in addition to Weath- erDiff (Özdenizei & Legenstein, 2023). The models were all trained on the AllWeather dataset and the performance numbers were taken from Özdenizci & Legenstein (2023), where the WiftherDiff was trained for 1.775 epochs and inferenced with 25 steps. Our RestoreGrad was trained for only 887 epochs (2x fewer) and inferenced with 10 steps to already achieve the best performance in the MW cate- gory. In addition, when trained for more epochs (1,551
IR Example: Figure 7 presents examples of restored images by the models. It can be seen that RestoreGrad is able to better recover the original image, expecially in regions of the blue and red boxes where the baseline WeatherDiff fails to remove the snow obstructions. The higher PSNR and SSIM scores of RestoreGrad also reflect the improvements.
Other IR Tasks: Our method demonstrates the advantages of learnable priors in image deblurring and seper-vesolution tasks, suggesting its generality. See Appendix D.2.
8

---

RestoreGrad:
Signal Restoration Using Conditional Denoising Diffusion Models with Jointly Learned Prior
diffusion processes reduce the number of diffusion steps. Nichol & Dhariwal (2021) introduced a few simple modifi- cations to improve the log-likelihood. Pandey et al. (2022; 2021) used DDPMs to refine VAB-generated samples. Rom- bach et al. (2022) performed the diffusion process in the lower dimensional latent space of an autoencoder to achieve high-resolution image synthesis, and Liu et al. (2023b) stud- ied using such latent diffusion models for audio. Popov et al. (2021) explored using a text encoder extract better representations for continuous-time diffusion-based text- to-speech generation. More recently, Nielsen et al. (2024) explored using a time-dependent image encoder to param- eterize the mean of the diffusion process. Orthogonal to the above, PriorGrad (Lee et al., 2022) and follow-up work (Koizumi al., 2022) studied utilizing informative prior extracted from the conditioner data for improving learning efficiency. However, they become sub-optinal when the conditioner are degraded versions of the target data, posing challenges in applications like signal restoration tasks.
Table Evaluation on realistic image datasets of the IR models
trained on synthetic images of AllWeather training set.
RainDS. Real
Snow Real
Methods
Gen.
NIQE
NIQE
Trans Weather (Valanarasu et al.,, 2022)
N
4,005
3.161
WeatherDiff (Özdenizci & Legenstein, 2023)
Y
3,050
2.985
RestoreGrad (ours)
2,556
3.015
BBild text and underfined text for second values. column "Gen."
indicates the model geverative (Y) (N).
Table 6. Evaluation of SE models on CHiME-3 test set, where the models were trained on VoiceBank DEMAND training set.
Methods
Gen. PESOT
CSIG
CBAK
COVL
SIR
Unprocassed
1.27
2.61
1.93
1.88
7.51
Demucs (Defossez al., 2020)
1.38
2.50
2,08
1.88
WaveCRN (Hsieh et al., 2020)
1.43
2.53
2.03
1.91
DOSE (Tai et al., 2023a)
Y
1.52
2.71
2.15
2,06
CDiffuSE: (L. et ol., 2022)
Y
1.55
2.87
2.09
2.15
7.67
RestoroGrad (ours)
Y
1.54
2,88
2.14
2,16
8445
* 'Bold text lor best and undertined text for second best values. The column "Gen."
indicates if the model is generative (Y) (N).
Diffusion-Basses Signal Restoration: Built on top of the diffusion models for audio generation, e.g., Kong et al. (2021); Chen et al. (2020); Leng et al. (2022), many SE models have been proposed. The pioneering work of Lu et al. (2022) introduced conditional DDPMs to the SE task and demonstrated the potential. Other works (Serrà et al., 2022; Welker et al., 2022; Richter et al., 2023; Yen et al., 2023; Lemercier et al., 2023; Tai et al., 2023a) have also attempted to improve SE by exploiting diffusion models, In the vision domain, diffusion models have demonstrated impressive performance for IR tasks (Li et al., 2023; Zhu et al., 2023; Huang et al., 2024; Luo al., 2023; Xia et al., 2023; Fei et al., 2023; Hurault et al., 2022; Liu et al., 2023a; Chung et al., 2023b:a; Zhou et al., 2024; Xiao et al., 2024; Zheng et al., 2025; Liu et al., 2024; Ye et al., 2024), and a comprehensive view of recent advancements is provided by He et al. (2025). A notable IR work is by Özdenizei & Legenstein (2023) that achieved impressive performance on several benchmark datasets for restoring vision in adverse weather conditions. Despite showing promising results, ex- isting works have not fully exploited prior information about the data as they mostly settle on standard Gaussian priors.
4.3. Generalization to Out-of-Distrihution (OOD) and Realistic Data
We have so far evaluated the models on in-domain scenarios with synthetic noisy data where RestoreGrad has shown substantial improvements. A natural question is whether the demonstrated improvements have actually come at the expense of the model's generalizability to unseen or realistic data. To address the concern, we evaluate the IR models on two additional RainsDS (2021) and Snow/00K Wetw//0KKReeea -Real from Liu et al. (2018) consist- ing of real-world images, using the reference-free Natural Image Quality Evaluator (NIQE) metric (Mittal et al., 2012) (a lower score indicates better quality). In Table 5 we see that RestoreGrad is able to perform on par with or better than WiftherDiff and the non-generative model of Tran- s sWeather. For OOD testing, we evaluate the SE models on the CHiME-3 dataset (Barker et al., 2017) unseen during model training. Table 6 compares RestoreGrad with CDif- fuSE also trained for 96 epochs, DOSE (Tai et al., 2023a), and two discriminative SE models, Demues (Defossez et al., 2020) and WaveCRN (Hsieh et al., 2020). We can see that RestoreGrrd is able to perform equally well as the CDif- fuSE while outperforming DOSE and the non-generative SE models. The results show that RestoreGrad is capable of im- proving in-domain performance while maintaining desirable generalizability of generative models.
from Quan et al.
6. Conclusion
We investigated the potential of learning the prior distribu- tion in tandem with the conditional DDPM for improved efficiency. We demonstrated the advantages of RestoreGrad that leverages lcarning-based priors, providing a more sys- tematic way of estimating the prior than existing selections. A limitation of the current work is that it focuses on signal restoration applications, where we suitably assume a zero- mean Gaussian prior and only learn its covariance. In the future, it can be interesting explore using a more generic prior form and extend the idea to other applications.
5. Related Work
Diffusion Efficiency Improvements: Das et al. (2023) uti- lized the shortest path between two Gaussians and Song et al. (2020) generalized DDPMs via a class of non-Mann non-Markovian

---

RestoreGrad: Signal Restoration Using Conditional Denoising Diffusion Models with Jointly Learned Prior
Impact Statement
Das, A., Fotiadis, S., Batra, A., Nabiei, F., Liao, F., Vakili, S., Shiu, D.-s., and Bernacchia, A. Image generation with shortest path diffusion. In International Conference Machine Learning (ICML), pp. 7009-7024, 2023.
This paper presents work whose goal to advance the field of Machine Learning. There are many potential societal consequences of our work, none of which feel must be specifically highlighted here.
Defosse7, A., Synnaeve, G., and Adi, Y. Real time speech enhancement in the waveform domain. In Annual Con- ference of the International Speech Communication Asso- ciation (Interspeech), pp. 3291 3291-3295, 2020.
References
Abd El-Fattah, M., Dessouky, M. I., Diab, S., and Abd El-Samie, F. Speech enhancement using an adaptive Wiener filtering approach. Progress in Electromagneties Research M, 4:167-184, 2008.
Fei, B., Lyu, Z., Pan, L.., Zhang. J., Yang, W., Luo, T. Zhang, B., and Dai, B. Generative diffusion prior for unified image restoration and enhancement. In /EEE/CVF Conference on Computer Vision and Pattern Recognition 2023.
Agustsson, E. and Timofte, R. NTIRE 2017 Challenge on Single Image Super-Resolutio: Dataset and Study. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR) Workshops, 2017.
(CVPR), pp.
Freirich, D., Michaeli, T., and Meir, R. A theory of the Eestortion-pereepee tradeoff in wasserstein space. In Advances in Neural Information Processing Systems (NeurIPS), pp. 25661-25677 2021.
Barker, J.. Marxer, R., Vincent, E., and Watanabe, S. The third 'CHiME' speech separation and recognition chal- lenge: Analysis and outcomes. Computer Speech Lan- guage, 46:605-623 626.6666622222262226 2017.
Goodfellow, I., Pouget-Abadie J., Mirza, M., Xu, B., Warde-Farley, D., Ozair, S., Courville, A., and Bengio, Y. Generative adversarial nets. In Advances in Neural Information Processing Systems (NIPS), 2014.
Benita, R., Elad, M., and Keshet, J. DiffAR: Denoising diffusion autoregressive model for raw speech waveform generation. In International Conference on Learning Representations (ICLR), 2024.
Gulati, A., Qin, J., Chiu, C.-C., Parmar, N., Zhang, Y., Yu, J., Han, W., Wang, S., Zhang, Z., Wu, Y., and Pang, R. Conformer: Convolation-augmente transformer for speech recognition. In Annual Conference of the Interna- tional Speech Communication Association (Interspeech), pp. 5036-50000000 5040, 2020.
Blau, Y. and Michaeli,
T. The perception-distorine trade- off. In IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 6228-6222 0088-637,77 2018.
Cao, H., Tan, C., Gao, Z., Xu, Y., Chen, G., Heng, P.-A., and Li, S. Z. A survey generative diffusion models. IEEE Transactions on Knowledge and Data Engineering, 2024.
Harvey, W., Naderiparizi, S., and Wood, F. Conditional im- age generation by conditioning variational auto-encoders. International Conference on Learning Representations (ICLR), 2022.
Chen, N., Zhang, Y., Zen, H., Weiss, R. J., Norouzi, M., and Chan, W. WaveGrad: Estimating gradients for waveform generation. In International Conference on Learning Representations (ICLR), 2020.
He, C., Shen, Y., Fang, C., Xiao, F., Tang, L.,, Zhang, Y., Zuo, W., Guo, Z., and Li, X. Diffusion models in low- level vision: A survey, IEEE Transactions on Pattern Analysis and Machine Intelligence, 2025.
Chung, H., Kim, J., Mccann, M. T., Klasky, M. L., and Ye, J. C. Diffusion posterior sampling for general noisy in- verse problems. In International Conference on Learning Representations (ICLR), 2023a.
He, K., Zhang, X., Ren, S., and Sun, Deep residual learning for image recognition. In IEEE Conference Computer Vision and Pattern Recognition (CVPR), pp. 770-778, 2016.
Chung, H., Kim, J., and Ye, J. C. Direct diffusion bridge using data consistency for inverse problems. In Advances in Neural Information Processing Systems (NeurIPS), 2023b.
Heusel, M., Ramsauer, H., Unterthiner, T.,, Nessler, B., and Hochreiter, S. GANs trained by two time-scale update rule converge to a local Nash equilibrium. In Advances in Neural Information Processing Systems (NIPS), 2017.
Croitoru, F-A., Hondru, V., Ionescu, R. T., and Shah, M. Diffuusion models in vision: A survey. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2023.
Ho, J., Jain, A., and Abbeel, P. Denoising diffusson prob- abilistic models. In Advances in Neural Information Processing Systems (NeurIPS), pp. 6840 6851, 2020.
10

---

RestoreGrad: Signal Restoration Using Conditional Denoising Diffusion Models with Jointly Learned Prior
Hsieh, T.-A., Wang, H.-M., Lu, X., and Tsao, Y. Wavecrn: An efficient convolutional recurrent neural network for end-to-end speech enhancement. IEEE Signal Processing Letters, 27:2149 2153, 2020.
Kupyn, O., Martyniuk, T., Wu, J., and Wang. Z. DeblurGAN-v2: Deblurring (orders-of-magnitude) faster and better. In IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp.
2019.
Hu, Y. and Loizou, P. C. Evaluation of objective quality measures for speech enhancement. IEEE Transactions Audio, Speech, and Language Processing, 16(1):229-235) 2007.
Le Roux, Wisdom, S., Erdogan, H., and Hershey, J. SDR- HoRRhhaff-baked well done? In IEEE International Conference Acoustics, Speech and Signal Processing (ICASSP), pp. 626-630, 2019.
Huang, Y., Huang, J., Liu, J., Yan, M., Dong. Y., Lyu, J., Chen, C., and Chen, S. WaveDM: Wavelet-based diffu- sion models for image restoration. IEEE Transactions Multimedia, 2024.
Lee, S. Kim, H.,, Shin, C., Tan, X., Liu, C., Meng, Q., Qin, T., Chen, W., Yoon, S., and Liu, T.-Y. PriorGrad: Improving conditional denoising diffuusion models with data-dependent adaptive prior. In International Confer- ence on Learning Representations (ICLR), 2022.
Hurault, S., Leclaire, A., and Papadakis, N. Gradient step denoiser for convergent plug-and-play. In International Conference Learning Representations (ICLR), 2022.
Lemercier, J.-M., Richter, Welker, S., and Gerkmann, StoRM: A diffusion- based stochastic regeneration model for speech enhancement and dereverberation. IEEE/ACM Transactions on Audio, Speech, and Language Process- ing, 2023.
Huynh-Thu, Q. and Ghanbari, M. Scope of validity psnr in image/video quality assessment. Electronics letters, (13):800-801, 2008.
ITU-T Rec. P.862.2. Wideband extension to recommen- dation P.862 for the assessment of wideband telephone networks and speech codecs. International Telecommuni- cation Union, 2005.
Leng, Y., Chen, Z., Guo, J., Liu, H., Chen, J., Tan, X., Mandic, D., He, L.,, Li, X., Qin, T., zhao, and Liu, T.-Y. BinauralGrad: A two-stage conditional diffusion proba- bilistic model for binaural audio synthesis. In Advances in Neural Information Processing Systems (NeurlPS), pp. 23689-23700, 2022.
Jiang, K., Wang. Z., Yi, P., Chen, C., Wang, Z., Wang, X., Jiang, J., and Lin, C.-W. Rain-free and residue hand-in- hand: A progressive coupled network for real-time image deraining. JEEE Transactions on Image Processing. 30: 7404-7418, 2021.
Li, R., Cheong, L.-F., and Tan, R. T. Heavy rain image restoration: Integrating physics model and conditional ad- versarial learning. In IEFE/CVF Conference on Computer Vision and Puttern Recognition (CVPR), pp. 1633-1642, 2019.
Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. In International Conference on Learning Representations (ICLR), 2014.
Li, R., Tan, R. T., and Cheong, L.-F. All in one bad weather removal using architectural search. In /EFF/CVF Con- ference on Computer Vision and Pattern Recognition (CVPR), pp. 3175-3185, 2020.
Kingma, D. P. and Welling, M. Auto-encoding variational bayes. In International Conference on Learning Repre- sentations (ICLR), 2014.
Kohl, S., Romera-Parcee B., Meyer, C., De Fauw, J., Led- sam, J. R., Maier-Hein, K., Eslami, S., Jimenez Rezende, D., and Ronneberger, O. A probabilistic U-Net for seg- mentation of ambiguous images. In Advances in Neural Information Processing Systems (NIPS), 2018.
Li, X., Wu, J., Lin, 7., Liu, H., and Zha, H. Recurrent squeeze-and-excitatio context aggregation net for sin- gle image deraining. In Proceedings of the European Conference on Computer Vision (ECCV), pp. 259-269, 2018.
Koizumi, Y., Zen, H., Yatabe, K., Chen, N., and Bacchiani, M. SpecGrad: Diffusion probabilistic model based neural vocoder with adaptive noise spectral shaping. In Annual Conference of the International Speech Communication Association (Interspeech), pp. 803-807, 2022.
Li, X., Ren, Y., Jin, X., Lan, C., Wang, X., Zeng, W., Wang, X., and Chen, Z. Diffusion models for image restora- tion and enhancement-A comprehensive survey, arXiv preprint arXiv:2308.0 09388, 2023.
Kong, Z., Ping, W., Huang. J., Zhao, K., and Catanzaro, B. DiffWave: A versatile diffusion model for audio synthesis. In International Conference on Learning Representations (ICLR), 2021.
Liu, G.-H., Vahdat, A., Huang, D.-A., Theodorou, E., Nie, W., and Anandkumar, A. I'SB: Image-to-image Schrödinger bridge. In International Conference on Ma- chine Learning (ICML), pp. 22042-2022 -22062, 2023a.

---

RestoreGrad: Signal Restoration Using Conditional Denoising Diffusion Models with Jointly Learned Prior
Liu, H., Chen, Z., Yuan, Y., Mei, X., Liu, X., Mandie, D., Wang, W., and Plumbley, M. D. AudioLDMM: Text-to- audio generation with latent diffusion models. In Interna- tional Conference on Machine Learning (ICML), 2023b.
In International Conference on Learning Reprexentations (ICLR), 2024.
Özdenizci, O. and Legenstein. R. Restoring vision in ad- verse weather conditions with patch-based denoising dif- fusion models. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2023.
Liu, J., Wang, Q. Fan, H., Wang, Y., Tang. Y. and Qu, L.. Residual denoising diffusion models. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition (CVPR), pp. 2773 2024.
Pandey, K., Mukherjee, A., Rai, P., and Kumar, A. VAEs meet diffusion models: Efficient and high-fidelity gener- ation. In NeurlPS 2021 Workshop on Deep Generative Models and Downstream Applications, 2021.
Liu, X., Suganuma, M., Sun, Z., and Okatani, T. Dual residual networks leveraging the potential of paired oper- ations for image restoration. In IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 7007-7016, 2019.
Pandey, K., Mukherjee, A., Rai, P., and Kumar, A. Diffuse- VAE: Efficient, controllable and high-fidelity generation from now-dimensional latents. Transactions on Machine Learning Research, 2022.
Liu, Y.-F., Jaw, D.-W., Huang, S.-C., and Hwang, J.-N. DesnowNet: Context-aware deep network for snow re- moval. IEEE Transactions on Image Processing, 27(6): 2018.
Pascual, S., Bonafonte, A., and Serra, J. SEGAN: Speech
enhancement generative adversarial network. In Annual Conference of the International Speech Communication Association (Interspeech), pp. 3642-3646, 2017.
Lu, Y.-J., Tsao, Y., and Watanabe, S. A study on speech enhancement based diffusion probabilistic model. In Asia-Pacific Signal and Information Processing Associa- tion Annual Summit and Conference (APSIPA ASC), pp. 659-666, 2021.
Phan, H., McLoughlin, I. V., Pham, I.., Chén, O. Y., Koch, P., De Vos, M., and Mertins, A. Improving gans for speech enhancement. IEEE Signal Processing Letters, 27:1700-1704, 2020.
Lu, Y.-.,, Wang, Z.-Q., Watanabe, S., Richard, A., Yu, C.,, and Tsao, Y. Conditional diffusion probabilistic model for speech enhancement. In IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 7402-7406, 2022.
Popov, V., Vovk, I., Gogoryan, V., Sadekova, T., and Kudi- nov, M. Grad-TTS: A diffusion probabilistic model for text-to-speech, In International Conference on Machine Learning (ICML), pp. 8599-88098 8608,8688,80 2021.
Qian, R. Tan, R. T., Yang, W., Su, J., and Liu, J. Attentive generative adversarial network for raindrop removal from a single image. In IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 2491, 2018
Luo, C. Understanding diffusion models: A unified perspec- tive. arXiv preprini arXiv:2208.1 1970, 2022.
Luo, Z. Gustafsson, K., Zhao, Z., Sjölund, J., and Schön, T. B. Refusion: Enabling large-size realistic im- age restoration with latent-space diffusion models. In IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 1680-1691, 2023.
Quan, R., Yu, X., Liang, Y., and Yang, Y. Removing rain- drops and rain streaks in one go. IEEE/CVF Con- ference on Computer Vision and Pattern Recognition (CVPR), pp. 9147-9156. 2021.
Majumdar, S., Balam, J., Hrinchuk, O., Lavrukhin, V., Noroozi, V., and Ginsburg, B. Citrinet: Closing the gap between non-autoregressive and autoregressive end- to-end models for automatic speech recognition. arXiv preprint arXiv:2 2021.
Quan, Y., Deng, S., Chen, Y., and Ji, H. Deep learning for seeing through window with raindrops. In IEEE/CVF International Conference on Computer Vision (ICCV), pp- 2463 2203-2271,222 2019.
Mittal, A., Soundararajan, R., and Bovik, A. C. Making "completely blind" image quality analyzer. IEEE Signal Processing Letters, 20(3):209 212, 2012.
Richter, J., Welker, S., Lemercier, J.-M., Lay, B., and Gerk- mann, T. Speech enhancement and dereverberation with diffusion- based generative models. IEEF/ACM Transac- tions on Audio, Speech, and Language Processing. 2023.
Nichol, A. Q. and Dhariwal, P. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning (ICML), pp. 8162-8171, 2021.
Rim, J., Lee, H., Won, J., and Cho, S. Real-world blur dataset for learning and benchmarking deblurring algo- rithms. In European Conference on Computer Vision (ECCV), 2020.
Nielsen, B. M. G., Christensen, A., Dittadi, A., and Winther, O. DiffEnc: Variational diffusion with a learned encoder.
12

---

RestoreGrad: Signal Restoration Using Conditional Denoising Diffusion Models with Jointly Learned Prior
Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B. High-resolution image synthesis with latent diffusion models. In IEEE/CVF Conference Com- puter Vision and Pattern Recognition (CVPR), pp. 10684- 10695, 2022.
Valanarasu, J. M. Yasarla, R., and Patel, V. M. Tran- sWeather: Transformer-basee restoration of images de- graded by adverse weather conditions. In IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 2353-2363, 2022.
Serrà, J., Pascual, S., Pons, J., Araz, R. O., and Scaini, D. Universal speech enhancement with score-based diffu- sion. arXiv preprint arXiv:2206.03005 2022.
Valentini-Botinhar C., Wang, X., Takaki, S., and Yamag- ishi, J. Investigating RNN-based speech enhancement methods for noise-robust text-to-speech. In ISCA Work- shop on Speech Synthesis Workshop (SSW), pp. 146-152, 2016.
Sohl-Dickstein, J., Weiss, E., Maheswaranathan, N., and Ganguli, S. Deep unsupervised learning using nonequi- librium thermodynamics. In International Conference on Machine Learning (ICML), pp. 2256-2265. 2015.
Veaux, C., Yamagishi, J., and King, S. The voice bank corpus: Design, collection and data analysis of large regional accent speech database. In International Confer- ence Oriental COCOSDA held jointly with 2013 Confer- ence on Asian Spoken Language Research and Evaluation (O-COCOSDA/CASLRE), 2013.
Sohn, K., Lee, H., and Yan, Learning structured output representation using deep conditional generative models. In Advances in Neural Information Processing Systems (NIPS), 2015.
Song. J., Meng, C., and Ermon, S. Denoising diffusion implicit models. In International Conference on Learning Representations (ICLR), 2020.
Wang, Z., Bovik, A. C., Sheikh, H. R., and Simoncelli, E. P. Image quality assessment: from error visibility to struc- tural similarity. JEEE ransactions on Image Processing. 13(4):600-6) 612, 2004.
Stoica, P., Moses, R. L., et al. Spectral analysis of signals, volume 452. Pearson Prentice Hall Upper Saddle River, NJ, 2005.
Welker, S., Richter, J., and Gerkmann, T. Speech enhance- ment with score-based generative models in the complex STFT domain. In Anmual Conference of the International Speech Communication Association (Interspeech), pp.
Strauss, M. and Edler, B. A low-based neural network for time domain speech enhancement. In IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 5754-5758, 2021.
2022.
Xia, B., Zhang, Y., Wang, S., Wang, Y., Wu, X., Tian, Y., Yang, W., and Van Gool, L. DifIRR: Efficient dif- fusion model for image restoration. In IEEE/CVF In- lernational Conference on Computer Vision (ICCV), pp. 13195-13055, 2023.
Tai, W., Lei, Y., Zhou, F., Trajcevski, G., and Zhong, T. DOSE: Diffusion dropout with adaptive prior for speech enhancement. In Advances in Neural Information Pro- cessing Systems (NeurIPS), 2023a.
Tai, W. Zhou, F., Trajcevski. G., and Zhong. T. Revisiting denoising diffusion probabilistic models for speech en- hancement: Condition collapse, efficiency and refinement. In AAAI Conference on Artificial Intelligence (AAAI), vol- ume 37, pp. 13627-13635, 2023b.
Xiao, J., Fu, X., Liu, A., Wu, F., and Zha, Z.-J. Image de-raining transformer. IEEE Transactions on Pallern Analysis and Machine Intelligence, 2022.
Xiao, J., Feng, R., Zhang, H., Liu, Z., Yang, Z., Zhu, Y., Fu, X., Zhu, K., Liu, Y., and Zha, Z.-J. DreamClean: Restoring clean image using deep diffusion prior. In International Conference on Learning Representations (ICLR), 2024.
Tao, X., Gao, H., Shen, X., Wang, J., and Jia, J. Scale- recurrent network for deep image deblurring. In IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 8174-8188 2018.
Thiemann, J., Ito, N., and Vincent, E. The diverse environ- ments multi-channel acoustic noise database (DEMAND): A database of multichannel environmental noise record- ings. In Proceedings of Meetings on Acoustics, 2013.
Ye, T., Chen, S., Chai, W., Xing, Z., Qin, J., Lin, G., and Zhu, L. Learning diffusion texture priors for image restora- tion. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 2524-2534, 2024.
Timofte, R., Agustsson, E., Van Gool, I.., Yang, M.-H., Zhang, L., Lim, B., et al. NTIRE 2017 Challenge on Single Image Super-Resolution In IEEE Conference on Computer Vision and Pattern Recognition (CVPR) Workshops, 2017.
Yen, H., Germain, F. G., Wichern, G., and Le Roux, J. Cold diffusion for speech enhancement. In IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), 2023.
Methods and Results.
13

---

RestoreGrad: Signal Restoration Using Conditional Denoising Diffusion Models with Jointly Learned Prior
Zamir, S. W., Arora, A., Khan, S., Hayat, M., Khan, F. S., Yang, M.-H., and Shao, L. Multi-stage progressive image restoration. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 14821-14831, 2021.
Zhang, K., Li, R., Yu, Y., Luo, W., and Li, C. Deep dense multi-scale network for snow removal using semantic and depth priors. IEEE Transactions on Image Processing, 30:7419-7431, 2021.
Zhang, R., Isola, P., Efros, A. A., Shechtman, E., and Wang, O. The unreasonable effectiveness of deep features perceptual metric. In /EEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 586-595, 2018.
Zheng, K.,, He, G., Chen, J., Bao, F., and Zhu, J. Diffusion bridge implicit models. In International Conference Learning Representations (ICLR), 2025.
Zhou, L., Lou, A., Khanna, S., and Ermon, S. Denoising diffusion bridge models. In International Conference Learning Representations (ICLR), 2024.
Zhu, Y., Zhang, K., Liang, J., Cao, J.. Wen, B., Timofte, R., and Van Gool, L. Denoising diffusion models for plug- and-play image restoration. In IEEE/CVF Conference Computer Vision and Pullern Recognition (CVPR), pp. 1219-1229, 2023.
14

---

RestoreGrad: Signal Restoration Using Conditional Denoising Diffusion Models with Jointly Learned Prior
A. Proof of Proposition 3.1
Proposition 3.1 (Incorporation of diffusion process into VAE). By introducing a sequence of hidden variables X1.Ts under the setup of conditional diffusion models where the Markov Chain assumption employed on the forward process q(XI:T Xo) II q(xx{x{-1) and the reverse process pe (Xo:Tly ly) p(xT) II pe D0(Xt-1)as y) parameterized by DDPM 0, and assuming that XT (i.e., the latent noise of DDPM samples from the VAE latent distribution), we have the lower bound on log p(xoly, 6) in the reconstruction term of the VAE (7) as:
log po(xo:rly) ((xI.T) Xo
logo(xoly, E) 2 Be(xi bxo)
which is the ELBO of the conditional DDPM (i.e., the conditional version of (3)).
Proof:
We lower bound the conditional data log-likelihood log p(xoly, e) by utilizing sequence of hidden latent representations X1.7 and the approximate variational distribution g
/ px(xully, e)dx1:1
log p(xuxx, 6) = log
/
po(xo.Tly, ())(() (x0)
= log
dxier
(x1,r|/00) Xo)
(13)
po (xorrly
-log (x(xur((x)) (x0)
(x1::1 (Xu)
Po (xo.rly, E)
2 xo)
log
(x Xo)
where the inequality obtained by applying Jensen's inequality. The reverse diffusion process incorporating e given as:
T
(x(:Tolyy <) =p(xx)
po (x N,,Y, c)
t=1
=p(xx) (x))
Po [x,, xX)
(14)
t=1
T
=p(xT)
II
po (x:-1 Xt, =: (oxo:ll))))
by using the assumption that the diffusion prior samples from the VAE latent space, i.e., Xy = e, and utilizing the Markov Chain property on the reverse transition probability remove the Xy as condition. This indicates that:
Po Po (xo:rlyy
Po (o(xoolly) log (X1:r Xn)
(x
log
Eg(xur (xo) (x))
(15)
(xx:r)(x))
Together, (13) and (15) lead to the result Proposition 3.1.
B. Derivation of Proposition 3.2
Proposition 3.2 (RestoreGrad), Assume the prior and posterior distriuttions are both zero-mean Gaussian, parameterized respectively, where the covariances are estimated by the Prior Net ip (taking y as input) and Posterior Net & (taking both Xn and y as input). Let us simply use Eprior and Epost hereafter to refer to Eprior (y: tb) and Epos (Xy,X8 So.yy o) for concise notation. Then, with the direct sampling (0 property the forward path x, = 00 + v1 646 arbitrary timestep where e ( Xo. Xo,Y), and assuming the
as (el)) = (N(e;; 0, Eprior (y: 1b) and qo (e(E)xo.3 y) = NN(e; 0, (Xo.
15

---

RestoreGrad: Signal Restoration Using Conditional Denoising Diffusion Models with Jointly Learned Prior
reverse process has the same covariance the true forward process posterior conditioned Xg, by utilizing the conditional DDPM co Eo as the noise estimator of the true noise have the modified ELBO, - -[00, Φ, 1b), associated with (9):
ÜT Exx
[(0, 0, ()
+
log/ Epost
Y (e 6o (x 1)
2
Latent Regularization (IR) terms
Denoising Matching (DM)
prior
loog
+ tr
D prior Epost) + C,
2
Σ possi
Prior Matching (PM)
t> 1
where 20
are weighting factors, |xx| =x X X. o, =² I-as 3r and C is some constant
=
1=1
2cn not depending on learnable parameters 0, Φ, th.some constant not depending on the learnable parameters 9, o and t/.
Derivation:
Recall our proposed lower bound in (9) to incorporate the conditional DDPM into the VAE framework is given as:
(Xo:r ly)
Bygle (x),)) Byx(xer (xo)
DRL. (1a(e)xo.v)llnnoeee))
(16)
((x1:11 xx)
-L(9,d)
Note that as assumed in standard DDPMs, the forward diffusion process gradually corrupts the data distribution into the prior distribution, which can be achieved by carefully designing the variance schedule for the forward pass, i.e., {B:)²= such that Xy (as result of oT 0). More specifically, the q(xx (x,) of the forward diffusion process converges in distribution to the approximate posterior qo (c(xo. y) from the posterior encoder şi Then, the term L(e, Φ) in (16) suggests training a conditional diffusion model 0 to reverse the diffusion trajectory from the estimated distribution of given by the posterior encoder á back to the target data distribution of Xg-
The ELBO C(0, Φ) can be shown to expanded as (Luo, 2022; Sohl-Dickstein al., 2015; Ho et al., 2020):
log (X (x1:7 (x))
C(0,0) 6) = Eg(xi (xo)
p(xT) ITT Pe (x 1 ()..)
Eg(x1r120) (x00)
log
I1 9 (xx xx-1)
p(x p(xT)po(Xol))) X1. I Po 27((xx (x- Xt,
= =- B (xo)
log
a(xi xn)
I1= a (xee(x(-1)
(17)
log p(XT)
log Pa (x,x 1 (x,,))
--g(xrr/x)) (xog x)] Bx((1.)(x))
Ex( (xo)
q(2T (x0)
9 (xx-1 X2, xx)
B xx) dog Ex(x)(x0) log p(xT)
po (x |x,, y) g(x+-1 X1, xu)
Σ
E Ex No)
log
((x)(x0))
-Lo + Lr + Σ Lr-1,
1=2
where
Lo == Ex(x, (x)) [log po [xo]xx,))]
(18)
h-1 = Ex(x//x0) [DKL (g(x1_1/x0) xo) (x v))l
(19)
CT DKL. (g(x7/x0)llp(x)))
(20)
16

---

RestoreGrad: Signal Restoration Using Conditional Denoising Diffusion Models with Jointly Learned Prior
According to Lee al. (2022), the terms of the loss function for training the noise estimator network 0 of the conditional DDPM for arbitrary (0, Σ) can be explicitly written
1 Ex en - (e:ooxilll-oxxxx,,,1111 Co = log (12T
2
+ 2011
1
L--1
E
20
(4-
(21)
3
Ex Ex cx(x, y, t)1ls-1,
a - a+)
αT Ex |x0 |3
d
CT
(ār log -r)),
2
with 0 II 0; and 0ç 1 33 for t 1, the parameter freedom and of
T where {B+ is the noise variance schedule hyperparameter, d
T- 1-ae-1 Br-
In our case, we have assumed modeling of the posterior distribution where the e sampled from as the zero-mean Gaussian (O0, poss where the covariance Epest == Epoa(Xu) y; Ø) is estimated by the Posterior Net o, taking both the ground
truth data Xo and the conditioner y input. By directly plugging in E = Epest for each term in (21), we obtain:
T
[(0, 6)
aT Exo |xo| ب + log ppst
Σ
c0(va,x0 + v dre y. Spa
+ C, (22)
2
1=1
X;
where
B3 20fae
t 1
7/
āz)
t=1
201
and C' is some constant not depending on the learnable parameters.
For the prior matching term (16), we can utilize the analytic form of the KL divergence between two Gaussians which leads to:
Σ Priiol
DKL.
log
tr(
Ar(E-1
peior Epoct)
(23)
Σ puot
where the covariances Eprior Eprior (y; vb) and Epast Dpusa (xo, Y; 1).
Combining (22) and (23), we have obtained the ELBO L(0, d, eb) of Proposition 3.2.
C. Implementation Details
C.1. Algorithms
Algorithm Training of RestoreGrad
Algorithm Sampling of RestoreGrad
for i= 0, 1,2..., Nier do
Sample (XD, y) Geana (Xu, Exiir Prior Net(y; (b) Expo Posterior Net(xo. y: y) Sample e ~ (0,, Epen) and (((1...., T))
Exaiar + Prior Net(y; v
Sample Xr ~ (0, E
for f = T.T 1,... do
ift > then
Sample NO00 Epaxe)
VouXo II are
else
CIR âr ||x0|| 1 log Epoal
E =
Lim ||e cx y, t) t)]3?_ D
end
X vor (xe
L xx( ose VI-a,
LyM log Exual
Equil tr( E p Epen)
Update 0, th, Φ with Vo.u.o nLLR Lom ALPM
end for return
end for
17

---

RestoreGrad: Signal Restoration Using Conditional Denoising Diffusion Models with Jointly Learned Prior
C.2. Experiments on Speech Enhancement (SE)
C.2.1. DATASET
We used the VoiceBank+DEMAND dataset (Valentini-ootinhao et al., 2016) with the same experimental setup as in previous work (Paseual et al., 2017; Phan et al., 2020; Strauss & Edler, 2021; Lu al., 2022) to perform direct comparison. The clean speech and noise recordings were provided from the VoiceBank corpus (Veaux et al., 2013) and the Diverse Environments Multichannel Acoustic Noise Database (DEMAND) (Thiemann et al., 2013), respectively, each recorded with sampling rate of 48kHz. Noisy speech inputs used for training were composed by mixing the two datasets with four signal-to-noiss ratio (SNR) settings from {0, 5, 10, 15} dB, using 10 types of noise (2 artificially generated + 8 real recorded from the DEMAND dataset) and 28 speakers from the Voice Bank corpus. The test set inputs were made with four SNR settings different from the training set, i.e., (2.5, 7.5, 12.5, 17.5) dB, using the remaining 5 noise types from DEMAND and 2 speakers from the VoiceBank corpus. There are totally 11,527 utterances for training and 824 for testing. Note that the speaker and noise classes were uniquely selected for the training and test sets. The dataset is publicly available at: https: / da aaaasare.... ed. ac. ww//handee/112222222 2200022232222222 2826. 6. In our experiments, the audio steams were resampled to 16kHz sampling rate.
C.2.2. MODEL ARCHITECTURE
Baseline DDPM-based SE Model: The baseline SE model considered in this work, i.e., CDiffuSE (Lu et al., 2022), performs enhancement in the time domain. We utilized the CDiffuSE base model, which has approximately 4.28M learnable parameters, from the implementation at: https: / / g t.hbb. - com/nei 1 Ll23/CDD fuSE. The model is implemented based on Wave (Kong et al., 2021), versatile diffusion probabilistic model for conditional and unconditional waveform generation. The basic model structure of CDiffuSE is similar to that of DiffWave. However, since the target task is SE, CDiffuSE uses the noisy spectral features as the conditioner, rather than the clean Mel-spectral features used in Wave utilized for vocoders. After the reverse process is completed, the enhanced waveform further combine the observed noisy signal y with the ratio 0.2 to recover the high frequency speech in the final enhanced waveform, as suggested Abd El-Fattah et al. (2008); Defossez et al. (2020).
PriorGrad: We implemented the PriorGrad (Lee al., 2022) top of the CDiffuSE model by using a data-dependent prior (0, Ey). where Σ, is the covariance of the prior distribution computed based on using the mel-spectrogram of the noisy input y. Following the application to vocoder in Lee et al. (2022), we leveraged normalized frame-level energy of the mel-spectrogram for acquiring data-dependent prior, exploiting the fact that the spectral energy contains an exact correlation to the waveform variance (by Parseval's 's theorem (Stoica et al., 2005)). More specifically, we computed the frame-level energy by taking the square root of the sum of eYp(Y) over the frequency axis for each time frame, where Y the mel-spectrogram of the noisy input y from the training data. We then normalized the frame-level energy to a range of (0, 1] to acquire the data-dependent diagonal variance Σy. Then we upsampled Ey in the frame level Ey in the waveform-level using the given hop length of computing the mel-spectrogram, We imposed the minimum standard deviation of the prior to 0.1 through clipping to ensure numerical stability during training, as suggested in Lee et al. (2022).
Prior Net and Posterior Net for RestoreGrad: The additional encoder modules for the RestoreGrad adopt the ResNet- 20 architect (He et al., 2016) using the implementation from: https: / /tt///itttti git com/ akamast /ttsttteepppttttooc resnet_cifan O. We suitably modified the original 2-D convolutions in 2e to I-D convolutions for waveform processing. The modified ResNet-20 model has only 93K learnable parameters (only 2% of the size of CDiffuSE model). The Prior Net takes the noisy speech waveform y as input, while the Posterior Net takes both the clean and noisy waveforms, Xo and y. as input, which are concatenated along the channel dimension. We employed the exponential nonlinearity at the network output for estimating the variances of the prior and posterior distributions.
C.2.3. OPTIMIZATION AND INFERENCE
We used the same configurations of CDiffuSE (Base) (Lu et al., 2022) for optimizing all the models, where the batch size was 16, the Adam optimizer was used with learning rate of 2 x 10-1, 1, and the diffusion steps T = 50 with linearly spaced 3, [10-4, 0.035). For RestoreGrad, we imposed the minimum standard deviation amin = 0.1 by adding it to the output of the Prior Net and Posterior Net to ensure stability during training. The fast sampling scheme in Kong et al. (2021) was used
in the reverse processes with 6-step inference schedule ginfer [10- 4, 10-3 0.01, 0.05, 0.2, 0.35). The models were trained on one NVIDIA Tesla V100 GPU of 32 GB CUDA memory and finished training for 96 epochs in 1 day.
18

---

RestoreGrad: Signal Restoration Using Conditional Denoising Diffusion Models with Jointly Learned Prior
C.2.4. EVALUATION METRICS
PESQ: speech quality measure using the wide-band version recommended in ITU-T P.862.2 (ITU-T Rec. P.862.2, 2005). It basically models the mean opinion scores (MOS) that cover scale from 1 (bad) to 5 (excellent). We used the Python-based PESQ implementation from: https: : / github. com/ Budllwww wwt.llwws/pppt/ee http//ytttohpppppeeeeeee peeq..
SI-SNR: a variant of the conventional SNR measure taking into account the scale-invariance of audio signals. The SI-SDR is a more robust and meaningful metric than the traditional SNR for measuring speech quality. A higher SI-SNR score indicates better perceptual speech quality. We adopted the SI-SNR implementation from: hitps: / Bocs/torchmetricss staalel
watii//scale_ivvee aerpiinnttniggll_aeeee www.innnnnnnneeeaaaaaeeee cww.sseerrari____. html.
SSNR: SNR measure, instead of working the whole signal, that calculates the average of the SNR values of short segments (segment length = 30 msec, 75% overlap, SNR -10 dB, SNR 35 dB). We use the Python-based SSNR implementation from: https: /itthb.. com/ htthhmmmmpppp/p/pppprpmmm
CSIG: The mean opinion score (MOS) prediction of the signal distortion (from to 5, the higher the better) (Hu & Loizou, 2007). We used the implementation from: https: / /gggtttthh./...///g/ com/ cetoom//chhphh// ppysepm.
CBAK: MOS prediction of the intrusiveness background noises (from 1 to 5, the higher the better) (Hu & Loizou, 2007). We used the implementation from: https:: //g thub. com/ schmi 2h ppyspp..
COVL: MOS prediction of the overall effect (from to 5, the higher the better) (Hu & Loizou, 2007). We used the implementation from: https: /ttthub.ttttt//tt// http//hb.coooc/occccoccocccocece htmm schmiph2 2am22/pysppp..
C.3. Experiments on Image Restoration (IR)
C.3.1. DATASETS
We used three standard benchmark image restoration datasets considering adverse weather conditions of snow, heavy rain with haze, and raindrops on the camera sensor, following Özdenizei & Legenstein (2023).
Snow100K (Liu et al., 2018): a dataset for evaluation of image desnowing models. The images are split into approximately equal sizes of three Snow synthetic snow strength imposed via snowflake sizes (light/mid/heavy). We used the Snow100K-L sub-test set for evaluation. The dataset can be downloaded from: https: /sites....s/.../////////////ttt/ google. com/v wtw.eeewwwnnuuuneeeeeeeeeeeeeeee
sub-test sets (with 16,61 1,6,,,8/8888/ samples, respectively), indicating the
Outdoor-Rai Rain (Li et al., 2019): dataset of simultancous rain and fog which exploits physics-based generative model to simulate not only dense synthetic rain streaks, but also incorporating more realistic scene views, constructing an inverse problem of simultaneous image deraining and dehazing. We used the test set, denoted in Li et al. (2019) as Test1, which is of size 750 for quantitative evaluations. The dataset can be accessed at: https: / /g thub com/ 1 ruoteng/ HleayRRaiRemoval
RainDrop (Qian al., 2018): a dataset of images with raindrops introducing artifacts on the camera sensor and obstructing the view. It consists of 861 training images with synthetic raindrops, and test set of 58 images dedi- cated for quantitative evaluations, denoted in Qian et al. (2018) RainDrop-A. The dataset is provided at: https: www.:ririi9991//. Deto///eaainnoroorrrrr
/ //githubs
In addition, we also used the composite dataset for weatheweather IR model training:
AllWeather (Valanarasu et al., 2022): a dataset of 18,069 samples composed of subsets of training images from the training sets of the three datasets above, in order to create balanced training set across three weather conditions with a similar approach to Li et al. (2020). The dataset is publicly available at: https: //tpthub..ttttt//// com/ jeya maria josee/ TransWeather.
C.3.2. MODEL ARCHITECTURE
Baseline DDasd IR Models: The baseline IR models considered in this work, i.e., the RainDropDiff and WeatherDiff from Özdenizci& & Legenstein (2023), perform patch-based diffusive restoration the images. The models perform diffusion process at the patch level, where overlapping p x p patches are taken as input. When sampling, all p x p patches extracted from the image with hop size are processed by the DDPM model, utilizing the mean estimated noise based sampling updates for the overlapping pixels to synthesize the clean image. In this work, considered p = 64 and 16, which
19

---

RestoreGrad: Signal Restoration Using Conditional Denoising Diffusion Models with Jointly Learned Prior
correspond to the RainDrooDiffszand WeatherDiffe: models (with 110M and 82 M learnable parameters, respectively) provided by the authors at:
f fus
Prior Net and Posterior Net for RestoreGrad:
The additional encoder modules for the RestoreGrad adopt the ResNet- com/ akamaster
enwieee
20 architect (He et al.. 2016) using the implementation from: https: resnet_ci farll, O. The ResNet-20 model has 0.27M learnable parameters, which is less than 0.3% of the size of RainDropDiff and WeatherDiff. The Prior Net takes the noisy image y as input, while the Posterior Net takes both the clean and noisy images, Xy and y, as input, which are concatenated along the channel dimension. We employed the exponential nonlinearity at the network output for estimating the variances of the prior and posterior distributions.
C.3.3. OPTIMIZATION AND INFERENCE
We used the same configurations of Özdenizci & Legenstein (2023) for optimizing all the models, except the batch size was 4 instead of 16 due to GPU memory limitation. The Adam optimizer with a fixed learning rate of x 10-5 was used for training models without weight decay, and an exponential moving average with a weight of 0.999 was applied during parameter updates. The number of diffusion steps was T = 1000 and the noise schedule was B, E [10- 0.02], linearly spaced. For inference, we used the sampling scheme with 10 timesteps for each model that we trained on our own. We did not use the deterministic implicit sampling scheme in Özdenizci & Legenstein (2023) for our RestoreGrrad-base found using the normal stochastic sampling scheme actually works better. The models were trained on 2 NVIDIA Tesla V100 GPU of 32 GB CUDA memory and finished training for 9,261 epochs on the RainDrop dataset in 12 days and 887 epochs on the Alleeather dataset in 21 days.
DDPM models
C.3.4. EVALUATION METRICS
PSNR: a non-linear full-reference metric that compares the pixel values of the original reference image to the values of the degraded image based on the mean squared error (Huynh-Thu & Ghanbari, 2008). A higher PSNR indicates better reconstruction quality of images in terms of distortion. PSNR can be calculated for the different color spaces. We followed Özdenizei& & Legenstein (2023) to compute PSNR based on the luminance channel Y of the YCbCr color space. We used the implementation form https: : / / thub. com/ Ji mnyyuniiagg / Swi i n IR for PSNR calculation.
SSIM: a non-linear full-reference metric compares the luminance, contrast and structure of the original and degraded image (Wang et al., 2004). It provides a value from 0 to 1, the closer the score is to 1, the more similar the degraded image is to the reference image. We followed Özdenizei & Legenstein (2023) to compute SSIM based on the luminance channel Y of the YCbCr color space. We used the implementation form https: / //ggittthhs com JingyunLiang/ SwinIR for SSIM.
LPIPS (Zhang et al., 2018) and FID (Heusel al., 2017): to provide better quantification of perceptual quality over the traditional distortion measures of PSNR and SSIM (Blau & Michaeli, 2018; Freirich et al., 2021). For the LPIPS we used the implementation from https: /tap::::///tie implementation from https: / tig.t...................... com wtt.oo/egg//ggaaayorrrroo In both metrics, a lower score indicates better perceptual quality of the restored image.
com/ ceom/ricchin http://chhhann/eeeeetttaaaaaaaaa sammmmimmmmmmmmmm larity, and for FID we used the
C.4. Experiments on Generalization to OOD and Realistic Data
C.4.1. DATASETS
The additional datasets considered for experiments on realistic data for the IR task are:
ReinDS-Real (Qian et al., 2018): is the raindrop removal test subset of the RainDS dataset presented in Qian et al. (2018). consists of 98 real-world captured raindrop obstructed images. The dataset is publicly available at: https: //github. cww www.:oooooooooo.oe Sw/g//onngoorrrr wwttooorrrrrriannaoeaeae
(Liu et al., 2018): is the subset of the Snow 100K dataset (Liu al., 2018) that consists of 1.329 realistic snowy images for testing real-world restoration cases. The dataset can be accessed at: https: /tt:://ssiiiee googl//sttss...llle.eeeeeooo.gee com/
The additional dataset considered for experiments on OOD data of the SE task is:
CHiME-3 (Barker al., 2017): is 6-channel microphone recording of talkers speaking in a noisy environment, sampled at 16 kHz. It consists of 7,138 and 1,320 simulated utterances for training and testing, respectively, which are generated
20

---

RestoreGrad: Signal Restoration Using Conditional Denoising Diffusion Models with Jointly Learned Prior
by artificially mixing clean speech data with noisy backgrounds of four types, i.e. cafe, bus, street, and pedestrian area. In this paper, we only take the 5-th channel recordings for the experiments. The dataset information be found at: ht https: /tww//wwww////// chwmechhaaacccccccccmcmcccccc.ce org.eecaaaaeeggeeoocaaaaaaaaaaaa
C.4.2. EVALUATION METRICS
The additional evaluation metric used in the corresponding section is:
NIQE: a reference-free quality assessment of real-world restoration performance introduced by Mittal et al. (2012) which measures the naturalness of given image without using any reference image. lower NIQE score indicates better perceptual image quality. We used the NIQE implementation from: https: peettts://gthhieee /github....../ com/ wtt. Saan//caaffegg//gg// IQA PoreeeennreeAAAAPAooorrr
C.5. Applications to Other IR tasks
C.5.1. DATASETS
The datasets considered for experiments on image deblurring and super-resolution tasks are:
RealBlur (Rim et al., 2020): large-scale dataset of real-world blurred images and ground truth sharp images for learning and benchmarking single image deblurring methods. The images were captured both in the camera raw and JPEG formats, leading to two datasets: RealBlur-R R from the raw images and RealBlur-J from the JPEG images. Each training set consists of 3,758 image pairs and each test set consists of 980 image pairs. The dataset can be downloaded from: httpps://// /cg. cg. wwwpppsssss...o.o..o.o.o..ppp... ac. kr
DIV2K (Agustsson & Timofte, 2017; Timofte al., 2017): dataset of 2K resolution high quality images collected from the Internet as part of the NTIRE 2017 super-resolution challenge. There are 800, 100, and 100 images for training, validation, and testing, respectively. The dataset provides x2, x3, and x4 downscaled images with bicubic and unknown downgrading operations. The dataset can be downloaded from: https: // //////daat /aaaaa vision. ww..eoeeeeeeeeeee.eeee....ee.... ethz. ch/ cvvl /DIV2K/.
C.5.2. MODEL ARCHITECTURE
The baseline conditional DDPM (cDDPM) implements the same architecture as the patch-hased denoising diffusion model of WeatherDiff (Özdenizci & Legenstein, 2023). The Prior Net and Posterior Net of RestoreGrad also adopt the same ResNet models as in the IR experiments under adverse weather conditions. For more details please refer to Appendix C.3.2.
C.5.3. OPTIMIZATION AND INFERENCE
The models were optimized and inferenced using the same configurations and settings given in Appendix C.3.3 for the IR experiments under adverse weather conditions. The models were trained on 2 NVIDIA Tesla V100 GPU of 32 GB CUDA memory and finished training for 853 epochs on the RealBlur-(R.) dataset each in 5 days and 2000 epochs on the DIV2K-(x2x4] dataset each in 3 days.
D. Additional Experimental Results
D.I. Additional Results on SE
Model Learning Performance in Terms of Other Metrics: In addition to the results evaluated by PESQ and COVL in Figure 2, we provide the learning curves in terms of the CSIG, CBAK, and SI-SNR metrics in Figure 8, to further support the advantages of RestoreGrad over the baseline DDPM and PriorGrad for improved training behavior and efficiency.
Performance with Using Different Numbers of Inference Steps: In Figure 9, show how the trained diffusion models perform with respect to using different numbers of reverse steps for inference. Specifically, in each case of CDiffuSE, PriorGrad, and RestoreGrad, we trained the model for 96 epochs and then inferenced with S E {3, 4, 5) reverse steps to compare with the originally adopted S = 6 steps in Lu et al. (2022). We used ginfer = [10 4 10 0.05, 0.2, 0.35) for = 5, ginfor = [10-4, 0.05, 0.2, 0.35] for S = 4, and Sinner = [0.05, 0.2, 0.35] for S = 3. These choices were selected from the subsets of the original noise schedule for S = 6, i.e., gimter [10-1 10-3, 0.010 0.05, 0.2, 0.35), that resulted in best performance of the models. For the figure we can see that as becomes smaller, the baseline CDiffuSE degrades
21

---

RestoreGrad: Signal Restoration Using Conditional Denoising Diffusion Models with Jointly Learned Prior
CSIG
CBAK
SI-SNR
3.0
14-
2.8
12.
3.0
2.4
2.5
2.2
Baseline (CDiffuSE)
2.0
Base ine (CDIffuSE)
Baseline (CDiffuSE)
PriorGrad
1.8
Prroorrrad
PriorGrad
RestoreGrad
RRetooreGrad
RestoreGrad
40
100
60
40
60
100
Training epochs
Taining spachs
Training epochs
Figure Model learning performance terms of CSIG, CBAK, and SI-SNR metries. Improved training behavior of RestoreGrad over CDiffuSE and PriorGrad is ohserved among all metries.
Baseline (CDIfTuSE)
PrforGrdd
RestoraGrad
infer
steps infer with steps
2.5
infer with 6 stegs
2,5
2.5
0.51
001
2.4
-2.08
2.3
-5 DE
2.3
2:3
0.20
2.1
2.1
21-
Infer with
infer with 6 steps
2.0
5-3
5-4
5-5
5-3
5
Figure 9. Effect of using reduced numbers of sampling steps in inference on the SE performance, in terms of PESQ. RestoreGrad demonstrates the strongest resistance the reduction in reverse sampling steps for inference.
Table 7. Performance comparison of RestoreGrad with the baseline DDPM (CDiffuSB) and PriorGrad for using various numbers of sampling steps during inference.
SI-SNR
CSIG
CBAK
COVL 1
Methods
8=6
S=5
S=4
8=3
S=6
8=5
S=4
S=3
S=6
S=5
S=4
S=3
S=6
S=5
S=4
S=3
CDiffuSE (Lu et al., 2022)
11.84
11.46
11.32
11.28
3.52
3.44
3.15
3.13
2.76
2.72
2.64
2.63
2.89
2.82
2.60
2.58
+ PriorGrad (Lee et al., 2022)
14.21
13.98
13.93
13.93
3.67
3.61
3,56
3.54
2.93
2.90
2.88
2.88
3.02
2.97
2.93
2.92
+ RestoreGrad (ours)
14.74
14.66
14.64
14.65
3.80
3.77
3.75
3.75
3.00
2.99
2.99
2.99
3.14
3.12
3.11
3.11
"Best values are indicated with bold text.
considerably, while PriorGrad shows certain resistance, and RestoreGrad manages to maintain the high performance. We present more comparison in Table 7 in terms of SI-SNR, CSIG, CBAK, and COVL metrics. The results further support that RestoreGrad is much more robust to the reduction in sampling steps, achieving the best quality scores in all the metrics over the baseline DDPM and PriorGrad across all sampling steps considered.
how the latent noise prior that has been learned by RestoreGrad looks like and how it compares with that of the PriorGrad. In Figure 10 we present an example of randomly chosen noisy speech waveform and the corresponding latent noise E, = diag (a } of PriorGrad and that of RestoreGrad (with (n, A) = (0.1, 0.5) for (1))). It can be seen that the variances of the pre-defined (PriorGrad) and learned (RestoreGrad) latent noise distributions are actually quite different, though both show the trend of following the variation the conditioner signal level. This trend indicates that both latent distributions aim to better approximate the true signal distribution in more informative manner for improved efficiency, as against the standard Gaussian prior used in the original DDPM. Note that in the RestoreGrad training, we have chosen a proper KL weight À so that the Prior Net distribution matches the Posterior Net distribution reasonably well without harming the reconstruction performance of the DDPM model. On the other hand, using a too large À might lead to a collapsed latent space as the optimization could have put too much emphasis matching the prior and posterior distribution, discarding the contribution of the reconstruction loss term. In contrast, using a too small might result in large discrepancy between the learned prior and posterior distributions, as also illustrated in Figure 10. Empirically, we found a naive choice of works reasonably well and also for similar values, c.g., 0.5, 10, etc., as similarly observed in the VAE-type model of Kohl et al. (2018).
Visualizing the Learned Prior: It would be interesting
22

---

RestoreGrad: Signal Restoration Using Conditional Denoising Diffusion Models with Jointly Learned Prior
PriorGrad
RestoreGrad
Too large
0.4
0.4
Condition dgnal eac
Amplituee
02
Amelituaee
0.7
Learned a
0.0
0.0
A4 Mg
-12
-12
Condition signal
Too small A.
Condition
Learned Meno
Pre-computed
Sum
00 0.2
0.8 3.0
1.
0.0
aigu
Time Bec)
Time Gac)
Learnel Wage
Figure 10. An example of learned latent distribution variances, Sprie = diag(o (r and Ejos = diag (oou) by RestoreGrad, and the effect of the KI. weight λ of the prior matching loss Lrs on the resulting latent distribution variances. The pre-computed variance of the handcrafted prior using PriorGrad is also presented for reference purposes.
Comparison to Existing Waveform-Domall Generative SE Models: In Table we benchmark RestoreGrad with several generative SE approaches. Note that although RestoreGrad performs slightly inferior to DOSE, a recent SE model also based on DiffWave (Kong et al., 2021), was actually achieved with 4.6x fewer training epochs.
Table 8. Comparison with existing time-domain generative SE models.
Methods
PESQ
CSIG
CBAK
COVL1
Unprocessed
1.97
3.35
2.44
2.63
SEGAN (Pascual al., 2017)
2.16
3.48
2.94
2.80
DSEGAN (Phan al., 2020)
2.39
3.46
3.11
2.90
SD-Flow (Struuss & Edler, 2021)
2.28
3.70
3.03
2.97
DOSE (Tai et al., 2023a)
2.56
3.83
3.27
3.19
CDiftuSE (Lu el al., 2022)
2.44
3.66
2.83
3.03
RestoreGrad (ours)
2.51
3.80
3.00
3.14
"Best values beld and second best values underlined.
Evaluation Using Automatic Speech Recognition (ASR): Following Benita et al. (2024) who perform evaluation of diffusion-based speech generation using ASR, we evaluate the SE model as a front-end denoiser for ASR under noisy environments. To this end, we pre-process the noisy VoiceBandDDEMAND test data samples through the well-trained SE model and feed the denoised audio separately to two pre-trained ASR engines taken from the NVIDIA NeMo toolkit Conformer-traassuucer--
(Gulati et al., 2020) and Citrinet-1024 (Majumdar et al., 2021). We report the word error rate (WER) and character eтror rate (CER) for each ASR engine outcome, where the lower WER / CER indicate better SE performance. The results are presented in Table with all the SE models trained after 96 epochs. inferred using 6 steps. It interesting to see that CDiffuSE and PriorGrad actually lead to worse performance than the unprocessed speech case for Citrinet ASR. Our RestoreGrad is able to achieve the lowest WER and CER for both ASR models, demonstrating its efficacy for enhancing machine learning capabilities under noisy environments.
Enhanced Speech Examples: We present several audio examples in Figure 11 to facilitate the comparison of the baseline DDPM and our RestoreGrad. It can be seen the RestoreGrad is able to recover a better speech signal closer to the target clean speech, which is also reflected by the higher PESQ scores obtained.
D.2. Additional Results on IR
Comparison to Existing IR Models on RainDrop Dataset: We compare our method with existing IR models including AttentiveGAN (Qian et al., 2018), DuRN (Liu et al., 2019), RaindropAttn (Quan al., 2019), IDT (Xiao et al., 2022) and RDDM (Liu et al., 2024) in Table 10 on the RainDrop dataset (Qian al., 2018), where the models were all trained and
23

---

RestoreGrad: Signal Restoration Using Conditional Denoising Diffusion Models with Jointly Learned Prior
Table Following Benita al. (2024) who perform evaluation of diffusion- based speech generation using ASR, evaluate SE models two ASR engines (Conformer, Citrinet) for the VoieeBandDDEMAND test set. The results further confirm the superiority of RestoreGrad over the baseline DDPM (CDiffuSE) and PriorGrad.
ASR: WER & (%) CER 4 (%)
SE model
Conformer (Gulati et al., 2020) Citrinet (Majumdar et al., 2021)
Unprocessed
6.62/ 6.15
8.69 6.86
CDiffuSE (Lu al., 2022)
6.55 6.01
9.77 / 7.41
+ PriorGrad (Lee et al., 2022)
6.13 5.70
9.15 / 7.00
+ RestoreGrad (Ours)
5.07 5.27
8.15 / 6.51
Best values are indicated with bold text.
CDiffuSE
RestoreGrad
PESQ: 2.48
clean enhanced
PESQ:
clean enhanced
Ampltude
0.25
Amplitude
0.25
0.00
0.00
-0.25
-0.25
-0.50
-0.50
0.0
0.5
1.0
1.5
2.0
2.5
3.0
0.0
0.5
1.0
1.5
2.0
2.5
3.0
Time (sac)
Time (soc)
CDiffuSE
RestoreGrad
0.50
PESQ: 2.85
0.50
PESQ: 3.16
enhanced
entanced
Amglitue
0.25
Anglitude
0.25
0.00
0.00
-0.25
-0.25
-0.50
-0.50
0.4
0.8
1.2
1.4
0.4
0.8 (sec)
1.0
1.2
1.4
(sec)
CDiffuSE
RestoreGrad
0.50
PESQ: 2.60
dcan
PESQ: 3.05
clean enhanced
enhanced
Amplitude
0.25
Anglitude
0.25
0.00
-0.25
-0.25
-0 50
-0.50
0.0
1.0
1.5
2.0
2.5
3.0
3.5
4.0
0.0
0.5
1.0
1.5
2.0
2.5
3.5
4.0
Time (sec)
Time (sec)
CDiffuSE
RestoreGrad
PESQ: 1.97
dean anhanced
PESQ: 2.17
clean
enhanced
Ampltude
0.25
Amplitude
0.25
0.00
-0.25
-0.25
-0.50
0.0
0.2
0.6
0.8
1.0
12
1.4
0.0
0.2
0.4
0.6
0.8
1.0
1.2
1.4
Time (5ec)
Time (sac)
Figure 11. Enhanced speech examples of the haseline DDPM (CDiffuSE) and the proposed RestoreGrad for several noisy samples taken from the VoiceBank+DEMMNN test
tested on the same training and test samples. The results the compared models were taken from Özdenizci Legenstein (2023) and Liu et al. (2024), where the baseline RainDropDiff was trained for 37,042 epochs. Our RestoreGrad was only trained for 9 epochs (4x fewer than RainDropDiff), and has achieved higher PSNR and SSIM scores than RainDropDiff
24

---

RestoreGrad: Signal Restoration Using Conditional Denoising Diffusion Models with Jointly Learned Prior
(here we report + std of RestoreGrad based results of 10 independent samplings). Moreover, our performance is comparable the recent state-of-the-an approach of RDDM, further suggesting the potential of our method to effectively improve baseline DDPM approaches (e.g., RainDropDiff).
Table 10. Weather-specific (RainDrop dataset) model comparison.
RainDrop
Methods
PSNR t
SSIM
AttentiveAAN (Qian et al., 2018)
31.59
0.9170
DuRN (Liu et al., 2019)
31.24
0.9259
RaindropAttn (Quan et al., 2019)
31.44
0.9263
IDT (Xiao al., 2022)
31.87
0,9313
RDDM (Liu ct al., 2024)
32.51
0.9563
DrffDro (Ózdenizci & Legenstein, 2023)
32.29
0,9422
RestoreGrad (ours)
32.69=0..03
0.9
Bestt values in bold and second best values underlined.
Visualizing the Learned Prior: We visualize the learned prior distribution variances for chosen image input with various 7) values in Figure 12 since we are interested in the effect of this newly introduced hyperparameter. We plot the results for the first channel of the image. The original contaminated image (i.e., the conditioner to the DDPM model) is also presented for reference purposes. As expected for the latent space regularization effect, a large results smaller variances as enforcing stronger regularization, while a small 17 leads to larger variances, as observed in the plots. Moreover, the learned prior appears to preserve structure the image, indicating that tends to learn prior distribution that approximates the data distribution.
Degraded image (conditioner)
Learned o prior (1 0.005)
3.0
25
100
200
300
400
200
300
400
500
600
700
05
Learned prior (n = 0.001)
Learned prior (n 0.1)
006
300
004
400
500
600
700
Figure 12. Visualization of learned prior distribution variances with various for a sample image taken from the RainDrop test set (Qian al., 2018). Mind the magnitude color bar of each figure. We that larger v) results in smaller variance of the prior distribution. while a smaller leads to larger variance.
25

---

RestoreGrad: Signal Restoration Using Conditional Denoising Diffusion Models with Jointly Learned Prior
Restoration Performance vs. 1/ and A: We also study the IR performance of the RestoreGrad models trained across various combinations of and À in Table 11, where the models were trained and tested the RrinDrop dataset. The results show that RestoreGrad works effectively for a wide range of 1 and À values to outperform the baseline DDPM model, RainDropDiff from Özdenizei & Legenstein (2023), which utilizes the standard Gaussian prior for the diffusion process.
Table 11. RestoreGrad performance for various and A, where the models were trained on the Dron training set (Qian et al., 2018) for 9,261 epochs and evaluated on the test set. The bascline RainDropDiff model scores reported in the original paper of Özdenizei & Legenstein (2023) (which was trained for 37,042 epochs, 4 times more than our RestoreGrad models) are also presented here for comparison purposes.
Model
7
À
PSNR
SSIM
0.05
32.55
0.9440
0.01
32.73
0,9448
RestoreGrad (ours)
0.005
0.1
32.69
0.9441
0.001
32.63
0.9404
0.0005
32.50
0.9405
10
32.74
0.9442
I
32.72
0.9441
RestoreGrad (ours)
0.005
0.1
32.69
0.9441
0.01
32.41
0.9417
DropDif (Özdenizei & Legenstein, 2023)
32.29
0,9422
*Values in bold text indicate better scores than the bascline ReainDropDiff model.
Images Generated from Different Prior Noise Samples: In Figure 13, we present example images generated for the same conditioner y using different random seed values (1, 10, and 20) when sampling the latent noise c, for both the baseline DDPM (RainDropDiff) and our RestoreGrad. In the figure, although it is challenging perceive the difference between the results of different seed values simply by inspecting the images visually, the quality (PSNR and SSIM scores) of the images produced by Grst is consistently better than the baseline DDPM among the different seed values used, further demonstrating the superiority of our method in restoring higher fidelity signals.
Signal Quality and Encoder Complexity Trade-Offs: Similar to the computational complexity analysis provided Table 2 for the SE task, we also present the results for the IR task on the RainDrop dataset in Table 12. From the table, we again observe that the performance of our approach slightly improves with the of a larger encoder, while also with increasing complexity. However, the computational overhead of our approach, i.c., the additional complexity due to the encoder module(s), is again relatively smaller compared to the adopted DDPM module (i.e., RainDropDiff (Özdenizei & Legenstein, 2023)), in terms of processing latency (<1.3% of DDPM) and memory usage (<30% of DDPM).
Table 12. IR comparison of RestoreGrad models (on RainDrop datasct) using encoder modules of different sizes and the corresponding latency and GPU memory usage (measured on one NVIDIA Tesla V100 GPU) presented the ratio of encoder to DDPM.
Encoder size
PSNRT
SSIM
Proc. Time
Memory
Base (0.27M)
32.65
0.9414
0.9%
16%
Large (1.9M)
32.77
0.9444
1.3%
30%
Visualization of Diffusion Processes: We present a qualitative analysis of the diffusion processes by showing a sequence of images from to / T for the baseline conditional diffusion model (WeatherDiff) and with the proposed method (RestoreGrad) in Figure 14. It can be seen that, at a given timestep t, the generated image by using RestoreGrad is cleaner than that of WeatherDiff, indicating a better diffusion trajectory developed by using our approach.
26

---

RestoreGrad: Signal Restoration Using Conditional Denoising Diffusion Models with Jointly Learned Prior
Seed:
Seed: 10
Seed: 20
[PSNR: 32.92 SSIMM:09469]
[PSNR: 33.04| SSIM:0. 9471]
[PSNR: 33.04 SSIM:0.9475]
RainDropDiff
[PSNR: 34.51| SSIM:0.9520]
[PSNR: 34.62 SSIM:0.9523]
[PSNR: 34.69| SSIM:0.9523)
RainDropDiff + RestoreGrad
Seed:
Seed: 10
Seed: 20
[PSNR: 33.73| SSIM:009448]
[PSNR: 33.64| SSIM:0.9448]
[PSNR: 33.42 | SSIM:0.9442]
RainDropDiff
[PSNR: 34.03| SS.M:
[PSNR: 34.12 0S.M:0
[PSNR: 33.97 SSIM:0.9477]
DainDropDiff
+ RestoreGrad
Figure 13. Example restored images from different prior noise samples (by using different random seed values) for the baseline DDPM (RainDropDifff) and our approach (RestoreGrad).
27

---

RestoreGrad: Signal Restoration Using Conditional Denoising Diffusion Models with Jointly Learned Prior
forward process
t=T
WeatherDiff
WeatherDiff RestoreGrad
t=0
reverse process
Figure 14. Visualization of diffusion processes for the baseline DDPM (WeatherDifh) and our method (RestoreGrad).
Experiments on Image Deblurring: We apply RestoreGrad to the baseline conditional DDPM (cDDPM) which implements the same architecture as the patch-based DDPM of Özdenizei & Legenstein (2023) used for weather degradations for deblurring. We trained the baseline cDDPM and RestoreGrad models and validated their performance on the RealBlur dataset (Rim et al., 2020), a large-scale dataset of real-world blurred images captured both in the camera raw and JPEG formats, leading to two sub-datasets: RealBlur-R from the raw images and RealBlur-J from the JPEG images. Each training set consists of 3,758 image pairs and each set consists of 980 image pairs. In Table 13, we present results the bascline cDDPM and RestoreGrad models trained after 853 epochs. We also include scores of two existing models, SRN-DeblurNet (Tao et al., 2018) and DeblurGAN-v2 (Kupyn et al., 2019), which performed similarly to the baseline cDDPM (taken from results by Rim et al. (2020)), as references for comparison. We can see that, except for LPIPS and FID on RealBlur-J, RestoreGrad is able to achieve improved scores than the baseline cDDPM, and outperform the compared methods.
Table 13. Image deblurring of realistic blurred images.
RealBlur-J
RealBlur-R
Methods
PSNRT
SSIM
LPIPS}
FID
PSNRT
SSIM
LPIPS)
FIDI
SRN-Deblurett
31.38
0.9091
38.65
0.9652
DeblurGAN-v2
29.69
0.8703
36.44
0.9347
Baseline cDDPM
30.69
0.9043
0.220
15.17
37.71
0.9777
0.626
14.46
RestoreGrad (ours)
31.51
0.9095
0.224
15.53
38.78
0.9796
0.122
13.61
*Bold text for best and underlined text for second best values.
Experiments on Image Super- Resolutio:: We further study the benefits of RestoreGrad the baseline conditional DDPM (cDDPM) model on image super-resolution tasks with the DIV2K dataset (Agustsson & Timofte, 2017; Timofte et al., 2017). We compare RestoreGrad with the baseline cDDPM model (the same architecture of the patch-based DDPM of WeatherDiff (Özdenizci & Legenstein, 2023)) for x2 and x4 downscale factor subsets (with bicubic downgrading operators). There are 800 images for training and 100 images for validation in each subset. For both subsets, we trained a baseline cDDPM and the RestoreGrad models for 2000 epochs on the training set and evaluated their performance on the corresponding validation set. The results are presented in Table 14, where we can see that except for the LPIPS metric, RestoreGrad is more beneficial then the baseline cDDPM in terms of achieving better scores in the other three metrics.
28

---

RestoreGrad:
Signal Restoration Using Conditional Denoising Diffusion Models with Jointly Learned Prior
Table 14. Comparison of baseline conditional DDPM (cDDPM) and the RestoreGrad image super-resolution tasks.
DIV2K x2
DIV2K x4
Methods
PSNR
SSIM 1
LPIPS 4
FID 4
PSNR 1
SSIM f
LPIPS
FID 4
Baseline cDDPM
27.40
0.9291
0.127
7.577
25.18
0.8064
0.269
7.849
+ RestoreGral (ours)
27.56
0,9341
0.136
7,547
25.56
0.8228
0.290
7.839
Better values are indicated with bold text.
More Image Restoration Examples: We provide more examples Figures 15, 16, 17, 18, 19, for comparing our RestoreGrad with the baseline DDPM approaches (i.e., RainDropDiff, WiatherDiff) of Özdenizei & Legenstein (2023). The restored image of RainDropDiff in Figure 15 was obtained by using the model weights trained by ourselves. The restored images of WeatherDiff in Figures 16, 17, 18, 19 were obtained by using the trained model weights provided by Özdenizei & Legenstein (2023) at https:: : / / g thub. Cwo./IIIITGGaa. CaG.TTGGraz/Maattaai /eeat Efusion, We also provide examples of image deblurring in Figures 20 and 21.
Degraded Input Image [PSNR: 24.201 SSIM. 88455
RainDropDiiff
RestoreGrad [PSNR: 31.85
Ground Truth Image
[PSNR: 30.60 SSIM: 0.9167J
SSIM: 0.9246]
Figure 15. Restored images by RainDropDiff (Özdenizei & Legenstein, 2023) and RestoreGrad (ours) for sample from the RainDrop test set.
Degraded Input Image
WiatherDiff
RestoreGrad
Ground Truth Image
[PSNR: 17.00 SSIM:O 5015)
[PSNR: 27.11 SSIM: 0.8344)
[PSNR: 27.86 SSIM: 0.8528)
FRE
FIE
2009
ZOT
Figure 16. Restored images by WeatherDiff (Özdenizci & Legenstein, 2023) and RestoreGrad (ours) for a test sample from the Outdoor- Rain test
29

---

RestoreGrad: Signal Restoration Using Conditional Denoising Diffusion Models with Jointly Learned Prior
Degraded Input Image [PSNR: 19.69 SSIM:0.6346]
WeatherDiff
RestoreGrad
Ground Truth Image
[PSNR: 30.38 SSIM: 0.9230]
[PSNR: 31,44 SSIM: 0.9356]
Figure Restored images by WiftherDiff (Özdenizci Legenstein, 2023) and RestoreGrad (ours) test sample from the Snow 100K-
I test set.
Degraded Input Image
WiatherDiff
RestoreGrad
Ground Truth Image
[PSNR: 23.81 SSIM:O
[PSNR: 30.88 SSIM: 0.9178)
[PSNR: 32.74 SSIM: 0.9309)
Figure Restored images by WeatherDiff (Özdenizci Legenstein, 2023) and RestoreGrad (ours) for a test sample from the RainDrop
test set.
Degruded Input Image [PSNR: 22.75 SSIM:0.8653
WeatherDiff
RestoreGrad
Ground Truth Image
[PSNR: 31 .28 SSTM: 0.9497]
[PSNR: 32.97 SSIM: 0.9578)
Figure 19. Restored images by WeatherDiff (Özdenizei Legenstein, 2023) and RestoreGrad (ours) for a test sample from the RainDrop test set.
30

---

RestoreGrad: Signal Restoration Using Conditional Denoising Diffusion Models with Jointly Learned Prior
Degruded Input Image IPSNR: 25.42 SSIM:0.7025]
Baseline cDDPM [PSNR: 28.61 SSIM: 0.8526]
RestoreGrad
Ground Truth Image
[PSNR: 29.90 SSIM: 0.8726]
MENU
APRICI
ARRIENN
6606
CLID DCO 619-7354
253-5534 CITD
D53-619 103-019-7361 -T3S
DCD 253 923--9335
C5D 918-9918-144
HIOH
MENU
MENU
MENU
Figure 20. Image deblurring examples using test image taken from the RealBlur test set.
Degraded Input Image [PSNR: 26.38 SSIM:0.7594)
Baseline cDDPM
RestoreGrad
Ground Truth Image
[PSNR: 29.68 SSIM: 0.8891]
[PSNR: 30.77 SSIM: 0.9012J
IN
IN
Figure 21. Image deblurring examples using a test image taken from the RealBlur test set.
31


Reviews:
[{"review_id": "u4r8N1yVWy", "reviewer": "Reviewer_2bPR", "summary": "This paper proposes RestoreGrad to restore signals with diffusion models. In order to achieve this, the paper proposes to train a prior net and a posterior net jointly with the diffusion model so that the noise can be sampled from the noise distribution output by the prior net. During training, three losses are applied to optimize the three networks: the latent regularization loss, the prior matching loss and denoising matching loss which is the original diffusion loss. During inference, the posterior net is removed and the prior net is used to output the distribution that is conditioned on the input so that the noise could be sampled directly from this distribution. The experiments are performed mainly on speech enhancement and image restoration.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We greatly appreciate that the reviewer found our method interesting and the quantitative results competitive. In this rebuttal, we provide our responses to address the main concerns and questions.\n\n---\n\n## Relation To Broader Scientific Literature\n\n> **The idea ...**\n\nWe would like to take this opportunity to highlight that, *the authors of PriorGrad were not able to demonstrate benefits of using learning-based priors, and they eventually just settled on using a handcrafted prior*. In other words, *the idea of learned priors was not yet well studied in the work of PriorGrad*. In contrast, our framework successfully demonstrates the **feasibility of leveraging learnable priors** for improving the efficiency of conditional DDPMs, including studies on practical nuances that make our proposed idea work in practice, which PriorGrad failed to demonstrate.\n\n---\n\n## Other Strengths And Weaknesses\n\n### Weaknesses\n\n> **The proposed ...**\n\nWe sincerely thank the reviewer for mentioning the comparison of our method with PriorGrad. We further clarify their differences.\n\n1. PriorGrad only utilizes information from the condition signal $\\mathbf{y}$ to compute the data-dependent prior. **As such, PriorGrad is not optimal for the applications where the condition signal contains certain degradations, e.g., in speech enhancement where $\\mathbf{y}$ is a noisy version of the target $\\mathbf{x}\\_0$**. In contrast, *our RestoreGrad exploits information from the target signal $\\mathbf{x}\\_0$ in addition to $\\mathbf{y}$ when learning to predict the prior.*\n\n3. We have introduced *a new loss function for jointly learning priors by employing the prior and posterior encoders, $\\psi$ and $\\phi$*. Let us compare ours with PriorGrad: \n\n- PriorGrad:\n$$\n\\min\\_{\\theta} \\\\,\\\\,\\\\, ||\\boldsymbol{\\epsilon}-\\boldsymbol{\\epsilon}\\_\\theta(\\mathbf{x}\\_t,\\mathbf{y},t)||^2\\_{\\boldsymbol{\\Sigma}^{-1}\\_{y}}, \\\\,\\\\,\\\\, \\text{where} \\\\,\\\\,\\\\, \\boldsymbol{\\Sigma}\\_{y}=f(\\mathbf{y}) \\\\, \\text{and} \\\\, f(\\cdot) \\\\, \\text{is a pre-defined function that maps} \\\\, \\mathbf{y} \\\\, \\text{into the prior distribution.}\n$$\n\n- RestoreGrad (ours):\n$$\n\\min_{\\theta,\\phi,\\psi} \\\\,\\\\,\\\\, \\eta\\bigr(\\bar{\\alpha}\\_T||\\mathbf{x}\\_0||^2\\_{\\boldsymbol{\\Sigma}^{-1}\\_{\\text{post}}}+\\log|\\boldsymbol{\\Sigma}\\_{\\text{post}}|\\bigr)+||\\boldsymbol{\\epsilon}-\\boldsymbol{\\epsilon}\\_\\theta(\\mathbf{x}\\_t,\\mathbf{y},t)||^2\\_{\\boldsymbol{\\Sigma}^{-1}\\_{\\text{post}}}+\\lambda \\bigr(\\log\\frac{|\\boldsymbol{\\Sigma}\\_{\\text{prior}}|}{|\\boldsymbol{\\Sigma}\\_{\\text{post}}|}+\\text{tr}(\\boldsymbol{\\Sigma}\\_{\\text{prior}}^{-1}\\boldsymbol{\\Sigma}\\_{\\text{post}})\\bigr).\n$$\n\nWe can see that our framework takes care of the prior estimation by utilizing encoders, by-passing the requirement to search for a suitable mapping function $f(\\cdot)$ requiring the domain knowledge of a specific task (e.g., spectral analysis in speech tasks). Our framework is thus **more general and applicable to more modalities** as we demonstrated.\n\n> **The model ...**\n\nIn this work, we have shown that our framework achieves significant improvements by using *standard*, *lightweight* architectures (ResNets) for the Prior and Posterior Nets. Specifically, the encoder modules we used were only 2\\% and 0.3\\% of the DDPM models in size for the SE and IR tasks, respectively. Such small encoders have already led to considerable performance boost for the DDPM backbone, and the computational overhead for the encoders is relatively small compared to the DDPM itself (e.g., see Table 2 of the submitted manuscript and new results for IR at [Link](https://openreview.net/forum?id=NbjrGgxLPi&noteId=Dc1KIYzvCF)). \n\n---\n\n## Questions For Authors\n\n> **I'm wondering ...**\n\nAs the idea of exploiting jointly learnable priors by combining DDPMs and VAEs is quite general, we think it has the potential to be applied to other advanced models. However, in the current work we focus on signal restoration using standard conditional DDPMs. If we want to apply our idea to other variants, such as SDXL (Podell et al., 2024) that utilizes latent diffusion models (LDMs) (Rombach et al., 2022), more effort may be needed. To elaborate on it, note that our approach depends on the decent correlation between the target and the conditioning data. However, exploiting correlation in a more implicit domain, such as in the hidden space of an LDM, could still be possible while may also require more effort on studying and analyzing the characteristics of the signals in the hidden representation space. It might also call for proper design of specialized architectures or even loss functions for different encoders/decoders used in LDMs. Our current work has not addressed such aspects, but we think it will be very interesting to explore relevant research in future work.\n\n**Refs:**\n- Podell et al. SDXL: Improving latent diffusion models for high-resolution mage synthesis. In *ICLR*, 2024.\n- Rombach et al. High-resolution image synthesis with latent diffusion models. In *CVPR*, 2022.", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "I checked the proof of Propositions 3.1 and 3.2 in the appendix roughly and they appear good to me.", "experimental_designs_or_analyses": "The soundness/validity of the experimental designs looks good.", "relation_to_broader_scientific_literature": "The idea of learned prior is well studied in previous works like PriorGrad. The key idea here is to match the prior distribution with the posterior distribution so that during inference, the learned condition prior distribution could be used to sample the noise.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Strengths\n1. The paper is clearly written and the proposed method is interesting by utilizing the posterior net and the prior net.\n2. The quantitative results look competitive.\n\nWeaknesses\n1. The proposed method seems incremental compared with PriorGrad. The key idea is similar, with the main difference being the jointly trained prior and posterior encoders. \n2. The model requires the joint training of prior encoder and posterior encoder, which might be costly if the model or the input signal is large.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "I'm wondering if the method could be utilized for more state-of-the-art and more foundational image diffusion models e.g., SDXL? In this way.  a much powerful image distribution could be utilized.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "TdyiM0fw6U", "reviewer": "Reviewer_wfM7", "summary": "The paper proposes a novel framework to enhance the efficiency of conditional Denoising Diffusion Probabilistic Models (DDPMs) for signal restoration tasks, such as speech enhancement (SE) and image restoration (IR). By jointly learning a data-informed prior distribution using two encoder networks (Prior Net and Posterior Net) alongside the DDPM, the method aims to address the inefficiency of traditional DDPMs caused by the mismatch between the standard Gaussian prior and real data distributions. The authors integrate DDPM into a variational autoencoder (VAE) framework, deriving a new evidence lower bound (ELBO) to optimize the model. Extensive experiments on benchmark datasets demonstrate improved convergence speed, signal quality, and robustness compared to baseline DDPMs and PriorGrad.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for acknowledging that our method is a significant contribution to address key limitation in existing works, and is theoretically sound and creative. In the following, we provide our responses to address the main questions and concerns.\n\n---\n\n## Experimental Designs Or Analyses\n\nOur framework aims to provide a general technique to improve conditional DDPM based Signal Restoration approaches. We have provided evidence of improved performance across multiple signal restoration tasks. We believe the capabilities and generality of our approach have the potential to extend to other applications like image denoising, and even to other signal modalities. \n\nWe would like to point out that the baseline image restoration model we used, i.e., WeatherDiff (Özdenizci and Legenstein, 2023), was published in 2023 and considered a reasonably strong baseline which achieved SOTA performance in several restoration tasks. Also, our main goal is to show the feasibility of learnable priors in a generic setting, rather than pursuing the SOTA performance in a specific task. However, as the reviewer suggested, it is still a good idea to compare with the latest methods, e.g., Diffusion Ye et al., (2024) and Liu et al., (2024). We have further provided such comparison in Tables A and B through the anonymous [Link](https://anonymous.4open.science/r/restoregrad-8A09/comparison.pdf).\n\n**Refs:**\n\n- Özdenizci and Legenstein. Restoring vision in adverse weather conditions with patch-based denoising diffusion models. *IEEE TPAMI*, 2023.\n- Ye et al. Learning diffusion texture priors for image restoration. In *CVPR*, 2024.\n- Liu et al. Residual denoising diffusion models. In *CVPR*, 2024.\n\n---\n\n## Essential References Not Discussed\n\nAs the reviewer mentioned, a significant amount of research has been conducted in the field of IR, and we might have missed some of the important ones. We will include more discussion on recent diffusion-based approaches and refer the reader to latest survey papers, e.g., He et al. (2025). As this is a rapidly developing field with many emerging approaches, we may not be able to discuss them all. If the reviewer has any specific works in mind that are important but we have missed, please kindly let us know.\n\n**Ref:**\n\n- He et al. Diffusion models in low-level vision: A survey. *IEEE TPAMI*, 2025.\n\n---\n\n## Other Strengths And Weaknesses\n\n### Weaknesses\n\n> **No. 5**\n\nIn fact, the assumption is used in the standard DDPM by assuming a large enough $T$ and a carefully designed variance schedule $\\\\{\\\\beta\\_t\\\\}\\_{t=1}^T$. Thus, it is not new here and has been adopted in DDPM literature to sample $\\mathbf{x}_T$ from the latent variable space of $\\boldsymbol{\\epsilon}$. What is new in our approach is that we explore the feasibility of having the distribution of $\\boldsymbol{\\epsilon}$ learnable. We will explicitly mention the conditions to improve the rigor of Proposition 3.1.\n\n> **No. 6**\n\nWe have followed the common practice in the literature to model the reverse process with a fixed variance schedule. In the future, it is also possible to combine our method with alternative approaches, such as Improved DDPM (Nichol \\& Dhariwal, 2021) which learns the variances instead of using a fixed noise schedule, to further improve the modeling efficiency.\n\n**Ref:**\n\n- Nichol and Dhariwal. Improved denoising diffusion probabilistic models. In *ICML*, 2021.\n\n> **No. 7 and No. 8** \n\nAlready responded in Experimental Designs Or Analyses.\n\n> **No. 9**\n\nFrom Table 2 of our submitted paper, we observed that the computational overhead due to the encoder module is relatively small compared to the DDPM module, in terms of processing latency (2\\% of DDPM) and memory usage (10\\% of DDPM). We have further conducted similar experiment on the IR task (on the RainDrop dataset) to measure the latency and GPU memory usage (presented as the ratio of encoder to DDPM) shown in the table below:\n\n| Encoder size | PSNR$\\uparrow$ | SSIM$\\uparrow$ | Proc. Time | Memory |\n|--------------|----------------|----------------|------------|--------|\n| Base (0.27M) | 32.65          | 0.9414         | 0.9\\%      | 16\\%   |\n| Large (1.9M) | 32.77          | 0.9444         | 1.3\\%      | 30\\%   |\n\n> **No. 10**\n\nOur method is quite general in the sense that it can potentially be adopted in existing conditional DDPM models while introducing minimal additional complexity. However,  there are definitely limitations. First, we have focused on signal restoration where we suitably assume a zero-mean Gaussian prior distribution. Second, our approach improves the conditional DDPM by exploiting the decent correlation between the conditioner signal $\\mathbf{y}$ and the target signal $\\mathbf{x}_0$. However, if such correlation is rather weak or very unstructured, our current encoding scheme may not be able to directly extract useful information from there. In this case, a more sophisticated encoding strategy may be required.", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "Almost correct.", "theoretical_claims": "Almost correct.", "experimental_designs_or_analyses": "1. The paper claims image restoration. It also compares many different tasks, but I am curious about the results of image denoising. Can the author make further comparisons? In addition, the baseline selected for comparison seems to be insufficient and a bit old.\n\n2. The methods compared in the article seem to be very old algorithms, such as WeatherDiffusion, which seems to be a method from 23 years ago. The latest CVPR'24 and ECCV'24 also have newer algorithms. The author must compare them to prove the superiority of the methods and their similarities and differences.", "relation_to_broader_scientific_literature": "It advances the field of generative modeling by enhancing the efficiency of DDPMs. It introduces a systematic, learning-based prior, aligning with trends in data-driven generative models and hybrid frameworks.", "essential_references_not_discussed": "The authors should provide a comparison section to discuss the work of image restoration in detail, because there are already a lot of works in this category, otherwise, the content of this section is insufficient.", "other_strengths_and_weaknesses": "Strengths:\n\n1. The proposal of systematically learning the prior distribution for DDPMs, rather than relying on handcrafted priors, is a significant contribution that addresses a key limitation in prior work (e.g., PriorGrad). This fills a gap in improving DDPM efficiency in a domain-agnostic manner.\n\n2. Embedding DDPM into a VAE framework with a novel ELBO formulation is theoretically sound and creative, effectively combining the generative power of DDPMs with the modeling efficiency of VAEs.\n\n3. The method’s ability to reduce training time (e.g., 10× faster convergence in SE) and maintain quality with fewer inference steps enhances its applicability in real-world scenarios.\n\n4. The article compares multiple data sets and can almost prove the effectiveness and robustness of the method.\nWeaknesses:\n\n1. The assumption that the final diffusion step $\\mathbf{x}_T$ equals the latent variable $\\epsilon$ (Proposition 3.1) is critical but lacks rigorous justification. Its validity may depend on the number of diffusion steps $T$ and the noise schedule $\\beta_t$, which are not thoroughly analyzed.\n\n2. The fixed covariance assumption ($\\sigma_t^2 = \\hat{\\beta}_t$) in the reverse process simplifies training but may limit the model’s flexibility in capturing complex data distributions. Alternative approaches are not explored.\n\n3. The paper claims image restoration. It also compares many different tasks, but I am curious about the results of image denoising. Can the author make further comparisons? In addition, the baseline selected for comparison seems to be insufficient and a bit old.\n\n4. The methods compared in the article seem to be very old algorithms, such as WeatherDiffusion, which seems to be a method from 23 years ago. The latest CVPR'24 and ECCV'24 also have SOTA diffusion-based algorithms. The author must compare them to prove the superiority of the methods and their similarities and differences.\n\n5. More comparisons on computational overhead are needed. Please provide more details on this part.\n\n6. The authors should further analyze the limitations of this method and the error cases, and explain why.", "other_comments_or_suggestions": "Please refer to the weaknesses.", "questions_for_authors": "Please refer to the weaknesses.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "tVNGLWXJB7", "reviewer": "Reviewer_JrVk", "summary": "The authors present a new type of conditional diffusion model with an application image restoration.\nAn image $x_0$ is predicted conditioned on a degraded version $y$.\nThe presented method is based on interpreting the diffusion model as a variational autoencoder (VAE). Based on this idea, they allow the prior over the noise $\\epsilon$ at $t=T$ to be conditioned on $y$ (implemented by a CNN that outputs per-pixel standard deviations).\nAdditionally, they model the posterior $q(\\epsilon|x_0, y)$ as conditioned on $y$.\nThe authors apply their method to image and audio data and show that it produces superior results for especially with a reduced number of iterations during inference and allows for faster convergence during training.\n\n## update after rebuttal\n\nI appreciate the detailed response by the authors and will stick with my original 'accept' rating.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We greatly appreciate that the reviewer found our method elegant and sensible, and that our experimental validation makes sense and is convincing. Below, we address the main concerns and questions.\n\n---\n\n## Methods And Evaluation Criteria\n\n> **The learned ...**\n\nThank you for carefully checking our Figures 3 and 12 and observing the details (grid patterns). Please note that the figures are not the restored images, **but a visualization of the learned priors**. The purpose of visualizing the priors is to demonstrate that they are informative and structured. In terms of the grid patters, we point out that the priors are actually estimated in a **patch-based manner**, following the practice of our conditional DDPM backbone, WeatherDiff (Özdenizci \\& Legenstein, 2023). As a result, the estimated priors might present the grid patterns between patches due to potential discontinuity. However, this is alleviated in the later stage when merging the restored image patches through the *mean estimated noise based sampling updates for the overlapping pixels* used in WeatherDiff. Therefore, the final restored images (our Figures 7, 13-19) are free from grid structures.\n\n> **I am missing ...**\n\nPlease refer to the figure via the anonymous [Link](https://anonymous.4open.science/r/restoregrad-8A09/diffusion_process.png). \n\n> **I wonder ...**\n\nThe applications considered in this work, i.e., restoring a signal from degraded observations, prioritize the output signal quality over diversity. Specifically, the models are trained to generate signals close to the original clean signal $\\mathbf{x}_0$ given the conditioner $\\mathbf{y}$, instead of producing diverse signals.\n\n---\n\n## Theoretical Claims\n\n> **I checked ...**\n\nFor the claim: ``However, existing adoption ...'', we meant to say that **in the stage of estimating the prior distribution** that the information from $\\mathbf{y}$ is discarded. We will make sure to rectify the claim for accuracy.\n\n> **Looking at ...**\n\nWe thank the reviewer for providing the interesting connection between our approach and the conditional VAE shared in the blog post. Our reason for why the dependent scheme may be preferred over the independent scheme is that, a data-dependent prior can potentially improve the diffusion trajectory of DDPMs in signal restoration. Specifically, as we can exploit the information from $\\mathbf{y}$ which adequately correlates with $\\mathbf{x}_0$, it is possible to form a prior distribution that is closer to the target distribution, improving the trajectory of the diffusion process for more efficient training and inference. \n\n----\n\n## Experimental Designs Or Analyses\n\n> **I agree ...**\n\nAs shown in the table below, the RestoreGrad model is still able to improve over the baseline DDPM (CDiffuSE) considerably when Posterior Net is conditioned only on $\\mathbf{x}_0$. This makes sense, as the information of $\\mathbf{x}_0$ is still available to the model during training which is important for estimating a desirable prior. \n\n| SE model                                              | PESQ $\\uparrow$ | COVL $\\uparrow$ | SSNR $\\uparrow$ | SI-SNR $\\uparrow$ |\n|-------------------------------------------------------|-----------------|-----------------|-----------------|-------------------|\n| CDiffuSE (96 epochs)                                  | 2.32            | 2.89            | 3.94            | 11.84             |\n| + RestoreGrad                                         | 2.51            | 3.14            | 5.92            | 14.74             |\n| + RestoreGrad (only $\\mathbf{x}_0$ for Posterior Net) | 2.52            | 3.13            | 5.58            | 14.48             |\n\nOur framework concentrates on signal restoration to generate a signal as close to the original signal as possible. In terms of the unconditional setting, if the task emphasizes more on diversity of generated data, then the data-dependent prior may not benefit that much. Unfortunately, our current work have not explored the unconditional setting and we do not have a concrete answer, but we sincerely hope that our work can be a good starting point for the research going forward.\n\n--- \n\n## Questions For Authors\n\n> **Parametrization ...**\n\nIn theory we can also parameterize the prior and posterior encoders to predict the mean in addition to the (co)variance, which may in turn require more capacity of the encoders to estimate both entities adequately. In this work, we have chosen to demonstrate the idea of learnable priors by estimating only the variance for simplicity. Our reason for that is, as the mean term does not contribute to the randomness of the samples (which only shifts the distribution), we reserve the capacity of the encoder modules for estimating the variance term that actually controls the shape (spread) of the distribution for sampling. In signal restoration applications, we also exploit normalization of the signals to zero mean and perform sampling in the normalized domain for efficiency.", "claims_and_evidence": "**\"We study the problem of learning the prior distribution jointly with the conditional DDPM for signal restoration applications, aiming at providing a more systematic, learning-based treatment to address the inefficiency incurred by existing selections of the prior distribution.\"**\nThis is correct, they paper looks at this problem and finds a sensible solution.\n\n**\"We propose a new framework called RestoreGrad that learns the prior in conjuncture with the DDPM model through a prior encoder\"**\nThis is also correct. The method is clearly described.\n\n**Performance claims:**\n\n**\"faster convergence (5-10 times fewer training steps)\"**\nThis claim is experimentally validated.\n\n**\"improved robustness to using fewer sampling steps in inference time (2-2.5 times fewer)\"**\nThis claim is experimentally validated", "methods_and_evaluation_criteria": "I believe this is an elegant and sensible approach. The methodology follows directly form viewing the diffusion model as a VAE.\nI think the evaluation datasets and baselines make a lot of sense and are convincing.\n\nHowever, I think the qualitative evaluation and discussion could be improved and makes me a bit worried:\n\n* The learned prior in Figure 3 shows grid patterns, especially visible in the top right corner, even though there are no grid structures in the image. These are also visible in supplementary Figure 12. The authors do not comment on this grid pattern. It would be interesting to see examples of learned priors for other images and to hear a discussion on why these artefacts occur.\n\n* I am missing a qualitative analysis of the generation process. It would be great to see a sequence of images showing the generation form t=T to t=0 for the proposed method and a vanilla conditional diffusion model.\n\n* I wonder of the learned prior effects the diversity of generated images. This could be qualitatively checked by showing multiple images generated for the same $y$ for the proposed method and a vanilla baseline.", "theoretical_claims": "I checked the equations in the main paper and they seem valid as far as I can tell.\nI am not completely sure, I a agree with the reasons the authors give for why their method gives improved results.\nFor example they say:\n*\"However, existing adoption of the standard Gaussian as the prior distribution in turn discards such information* [form $y$] *, resulting in sub-optimal performance.\"*\nAt each step of the generation with a vanilla conditional diffusion model, y is provided as input, so all information is available (albeit not in the noisy image but in the ) and no information is disregarded.\n\nLooking at this from the perspective of conditional VAEs, I think the proposed method corresponds to Figure 2 in [1] versus the vanilla conditional diffusion model that corresponds to Figure 1 in [1].\nThese are two different types of models and it's not obvious for me form a theoretical perspective why one is to be preferred over the other.\nCan the authors comment on this?\n\n\n \n\n\n[1]: https://beckham.nz/2023/04/27/conditional-vaes.html#sec_zy_independent", "experimental_designs_or_analyses": "I agree with the experimental design and especially appreciate the authors checking the effect of $\\nu$ and the importance of training the posterior net.\nI wonder how well the method would preform if the posterior net would only be conditioned on $x_0$ and not on $y$.\nThis could even be an interesting experiment in a non conditional setting. could a non conditional diffusion model benefit from using a posterior net conditioned on  $x_0$?", "relation_to_broader_scientific_literature": "The authors apply a VAE perspective to conditional diffusion models, which enables them to apply naturally concepts from literature on conditional VAEs to diffusion models, yielding a way to train a prior and posterior together with the normal diffusion unet.", "essential_references_not_discussed": "I did not notice any missing references to related literature.", "other_strengths_and_weaknesses": "I don't have additional points beyond the ones made above.", "other_comments_or_suggestions": "**Parametrization of prior and posterior net:**\nIs there a reason why the prior and posterior net predict only the standard deviations for each pixel and not the mean?\nWould this not make the method potentially more powerful?\n\n**Typos:**\n\n\"To avoid intractable integral\" -> \"To avoid the intractable integral\"", "questions_for_authors": "I don't have additional questions beyond the ones made above.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper seamlessly integrates DDPMs into the variational autoencoder framework and exploits the correlation between the degraded and clean signals to encode a better diffusion prior. The authors validate their method of speech enhancement and image restoration. This paper is reviewed by three experts in the field. The recommendations are: two weak accept and one accept. As a result, I think this paper can be accepted in the current form.
Decision: Accept (poster)