RestoreGrad: Signal Restoration Using Conditional Denoising Diffusion Models with Jointly Learned Prior

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.

---

Methods And Evaluation Criteria

The learned ...

Thank 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.

I am missing ...

Please refer to the figure via the anonymous Link.

I wonder ...

The 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.

---

Theoretical Claims

I checked ...

For 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.

Looking at ...

We 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.

---

Experimental Designs Or Analyses

I agree ...

As 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.

Our 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.

---

Questions For Authors

Parametrization ...

In 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.

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.

---

Experimental Designs Or Analyses

Our 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.

We 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.

Refs:

• Özdenizci and Legenstein. Restoring vision in adverse weather conditions with patch-based denoising diffusion models. IEEE TPAMI, 2023.
• Ye et al. Learning diffusion texture priors for image restoration. In CVPR, 2024.
• Liu et al. Residual denoising diffusion models. In CVPR, 2024.

---

Essential References Not Discussed

As 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.

Ref:

• He et al. Diffusion models in low-level vision: A survey. IEEE TPAMI, 2025.

---

Other Strengths And Weaknesses

Weaknesses

No. 5

In 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.

No. 6

We 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.

Ref:

• Nichol and Dhariwal. Improved denoising diffusion probabilistic models. In ICML, 2021.

No. 7 and No. 8

Already responded in Experimental Designs Or Analyses.

No. 9

From 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:

No. 10

Our 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.

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.

---

Relation To Broader Scientific Literature

The idea ...

We 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.

---

Other Strengths And Weaknesses

Weaknesses

The proposed ...

We sincerely thank the reviewer for mentioning the comparison of our method with PriorGrad. We further clarify their differences.

1. 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.

2. 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:

• PriorGrad:

$$\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.}$$

• RestoreGrad (ours):

$$\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).$$

We 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.

The model ...

In 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).

---

Questions For Authors

I'm wondering ...

As 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.

Refs:

• Podell et al. SDXL: Improving latent diffusion models for high-resolution mage synthesis. In ICLR, 2024.
• Rombach et al. High-resolution image synthesis with latent diffusion models. In CVPR, 2022.