Divide-and-Conquer Posterior Sampling for Denoising Diffusion priors

Dear Reviewer GK7d,

We thank you for your time, your feedback and suggestions. Please find our response to your questions below. In the PDF attached to the main rebuttal you can also find an update of Table 2 in the main paper with 4 more competitors suggested by the other Reviewers. These new results further demonstrate the strength of our method in comparison with the state-of-the-art.

More precision on the difference with Chung et al.: The potentials, and the subsequent approximations we introduce in this paper, differ in many aspects from the approximation developed in Chung et al. While [1] derives a new Gaussian approximation of $p(y|x_t)$, we take a different approach by introducing likelihood functions at various timesteps $(k \_\ell) \_{\ell=0}^L$ in the diffusion process. This transforms the initial sampling problem into sampling from intermediate posteriors, $\pi \_{k_\ell}$, defined by these new likelihoods. DCPS then iteratively samples from $\pi \_{k_\ell}$ using approximate samples from the previous posterior $\pi_{k_{\ell+1}}$. Essentially, we reduce the initial problem to solving $L$ simpler posterior sampling problems , which only require approximating $p_{k_\ell|t}(\cdot|x_t)$ for $t \in [k_{\ell}, k_{\ell+1}]$ (hence the divide-and-conquer terminology), rather than the difficult-to-estimate $p_{0|t}(\cdot|x_t)$ on which DPS but also many recent works in the literature have focused on deriving accurate estimates. When unrolling the backward transitions involved in our algorithm, it can be seen that these are fundamentally different from those of DPS. We will add a side by side comparison in the revised version of the manuscript.

Choice of twisting potentials: In the paper, we proposed choices of potentials for three important problems, all following the same principle: each potential is derived by annealing the initial potential and applying it to a rescaled input, i.e., $g_k(x_k) = g_0 \left( \frac{x_k}{\beta_k} \right)^{\gamma_k}$, where $\beta_k$ and $\gamma_k$ are tunable parameters. Specifically, we used $\beta_k = \sqrt{\alpha_k}$ and $\gamma_k = \alpha_k$ for linear inverse problems and JPEG dequantization. For the Poisson problem, we used $\gamma_k = \sqrt{\alpha_k}$ and the normal approximation of the Poisson distribution for $g_0$. We believe that this principle can serve as a general recipe for constructing potentials, leaving the tuning of the parameters to the users. Future works include deriving informed choices of hyperparameters in the general case. We have added a discussion of this matter in the revision of the manuscript.

Sensitivity to hyperparameters: Please note that in the Gaussian mixture experiment we examine the impact of the number of Langevin steps, which can be seen to significantly improve the performance when using a large number of steps. Similarly, we have found experimentally that increasing the number of gradient steps improves the performance.

MAP estimation: We believe that the principles we presented in this paper can be extended to MAP estimation; one could also design a sequence of intermediate posteriors well-suited for MAP estimation such that the terminal distribution is the posterior of interest, and sequentially solve these problems. The MAP at the previous iteration should be a good initialization for the optimization in the next step. An initial guess would be to start off with the sequence of distributions used in our paper.

Dear Reviewer af3x,

We would like to thank you for the time you took to review our manuscript and for your feedback and suggestions. Please find below our response to what we believe is the most crucial point in your review. Please consider reading the main rebuttal containing comparisons with new competitors and comments on the choice of potentials.

Similarity with [1]: We humbly disagree with the claim that our algorithm is similar to SDA [1]. To clarify this misunderstanding, let us first give a brief description of SDA. The authors propose an approximation of $p(y|x_t)$ ($g_t ^{*} (x_t)$ with our notation). The resulting density approximation is $\hat{p}(y|x_t)=N(y;A \hat{x}\_{0|t}^\theta(x_t), \sigma^2 \_y + \frac{\gamma (1 - \alpha_t)}{\alpha_t} AA^T )$. It should be noted that it is similar to the one used in PGDM [3], but with a slightly different choice of variance. After sampling $X_{t-1}$ given $X_t$ by leveraging the approximate conditional score, the authors perform few corrections steps using Langevin Diffusion targeting the marginal distribution $\propto \hat{p}(y|x_{t-1}) p_{t-1}(x_{t-1})$. As the authors acknowledge, this PC step is not new and already appears in the seminal paper of Song et al. 2021

We now proceed to a step-by-step comparison with [1].

• Intermediate posteriors (Please see also the second point in the main rebuttal): While [1] focuses on deriving a new approximation of $p(y|x_t)$ using a Gaussian approximation of $p_{0|t}(\cdot x_t)$, we take a different approach and introduce likelihood functions $g_{k_\ell}$ at different timesteps $(k_\ell) \_{\ell = 0} ^L$ of the diffusion process. By introducing these likelihood functions, we transform the initial problem of sampling from the posterior of interest to that of sampling from the intermediate posteriors, denoted by $\pi_{k_\ell}$ in our paper. DCPS is then an iterative procedure where at each step we aim to sample from $\pi_{k \_\ell}$ using approximate samples from the previous posterior $\pi \_{k \_{\ell+1}}$. In short, we reduce the initial problem to that of solving $L$ simpler posterior sampling problems. These problems are simpler since, in contrast with the initial problem which required deriving an approximation of the difficult to estimate distribution $p_{0|t}(\cdot | x_t)$, now we only require approximating the simpler distributions $p_{k_\ell | t}(\cdot | x_t)$ for $t \in [k_{\ell}: k_{\ell+1}]$.

• Backward transition approximation: The conditional score in [1] is used within DDIM whereas we posit a Gaussian approximation with learnable mean and diagonal covariance matrix and perform KL minimization on the conditional target distribution that immediately precedes Eq. 3.5. This is step 2 (see line 185). This step does not target the same distribution as [1] and relies on a different approximation.

• Langevin steps: In [1], $C$-Langevin steps are performed directly after each DDIM step. On the other hand, we resort to Langevin steps only to move from one block to the next, i.e., to obtain an approximate sample of $\pi^\ell _{k _{\ell+1}}$ that initializes the next block; see Eq. 3.4. This is step 1, (see line 183). We thus use significantly less LMC steps than SDA and the performance gains are not simply due to them, as opposed to [1].

Our algorithm DCPS and the algorithm in [1] are thus entirely different. We have implemented SDA and compared it with ours on linear inverse problem imaging experiments using the authors' implementation. We used two Langevin correction steps, $\tau = 0.1$ for the Langevin step size, and a diagonal approximation of $\gamma = 0.1$ after tuning. While SDA performs well on FFHQ, we couldn't find optimal parameters for ImageNet, resulting in degraded performance. DCPS outperforms it in both average and median LPIPS scores. Please see the PDF in the main rebuttal for detailed results and comparisons with other prominent algorithms.

Choice of potentials: Please see the second point of the main rebuttal. The reviewer is right in that the choice of potentials is arbitrary and any sequence should perform well as long as it converges to the inverse problem likelihood function. However, in practice, a mischosen potential will require a very large number of Langevin steps or may even result in significant instabilities. We illustrate the latter point. Take as example the choice of potential $g_m(x_m) = g_0(x_m)$ that you propose. Assume that $x \in R^2$ and $g_0(x) = N(y; x_1, \sigma^2 \_y)$. Assume that $x \in R^2$ and $g_0(x) = N(y; x_1, \sigma^2 \_y)$. We also assume that the prior distribution is for example a mixture with well separated modes far from 0 and that we observe $y >> 0$, the first coordinate of either one of the modes. During the early steps of the diffusion, when $\sigma_y \approx 0$, samples from the intermediate posteriors have first coordinates centered around $y$, while the second coordinate remains nearly Gaussian. This leads to instabilities in subsequent steps. When moving to the next step, i.e., sampling $X_{k-1}$ given $X_k$, the denoiser network is evaluated at $X_k$. However, at the $k$-th noise step, the denoising network has been trained to process samples that are almost white noise, which is not the case for our current sample $X_k$, which has a large first component. As a result, the network returns an incoherent value, and since this happens at every step of the diffusion, the errors accumulate, leading to significant instabilities.

As for the second example $g_m(x_m) = \mathbb{E}[g_0(X_0) | X_m = x_m]$, this is a valid argument and we believe that it can work well in practice. However, it should be noted that using this method requires a significant amount of memory and further increases the runtime. Namely, differentiating the loss in Eq. A.8 with respect to both parameters requires differentiating the composition of two Denoisers, which is very computationally intensive.

Dear Reviewer xni9,

We would like to thank you for the time you took to review our manuscript and for your feedback and suggestions. Please find below our response to what we believe are the most crucial points.

On the approximation of $g_k ^{*}$ and design of the intermediate potentials: We would like to emphasize that in our methodology the functions $g_k$ are completely user-defined. This is exactly the main point of our paper; we propose a principled method that deals with the intractability of the true potentials and allows the use of problem-informed, user-defined potentials. Therefore, it is not particularly relevant to question whether the functions $g_k$ are precise approximations of the functions $g_k^*$. As already emphasized in the paper, our method does not require any knowledge of $g^* _k$.

In the paper we have proposed choices of potentials for three important problems and their design follows the same principle; please see the second point in the main rebuttal for more details. We do not claim that these potentials are good approximations of $g^*_k$ as this is not a requirement of our method; nevertheless, we demonstrate that although they may not be accurate approximations, our method still demonstrates strong performance; as evidenced by our experiments. As an illustration, you can find in the PDF of the main rebuttal a plot of the gradient field of our log potentials compared with those of the true potentials $g^* _k$ (which are computable in a closed form in the Gaussian mixture case). Notably, these are very different and yet DCPS achieves a strong performance, as you stated in the strengths.

Finally, many recent published papers in top conferences [1, 2, 3, 4, 5, 6, 7] (including the FPS paper you suggest we should compare with) has focused only on linear inverse problems, which continues to be a highly relevant and challenging problem class, and as far as we are concerned there is still no gold-standard algorithm that is widely accepted for its robust performance. Besides linear problems, our study also covers two other significant inverse problems for which we also design potentials $g_k$ and empirically assess and illustrate their benefits. Therefore, while it would undoubtedly be ideal to have a parameter-free general solution, up to our knowledge only no existing works have been able to tackle this degree of generality while showing a performance comparable to more “specialized” schemes.

Comparison to FPS: We agree with the reviewer that both DCPS and FPS rely on “pseudo-observations,” as acknowledged in the paper (see the discussion in Section A.4 of the appendix where we already describe FPS). For the sake of completeness, we have added a comparison with FPS on imaging experiments but also with three other competitors suggested by the other reviewers, see the main rebuttal. Clearly, DCPS outperforms FPS by a significant margin. While SMC-based algorithms like FPS and MCGDiff have strong theoretical guarantees, their performance scales poorly with increasing dimensions, requiring exponentially more memory. In contrast, DCPS performs well across both low and high dimensions without incurring such high memory costs.

LPIPS only: We want to emphasize a key point. Our goal is to develop a sampling algorithm that efficiently generates diverse samples over the entire posterior distribution. For the multimodal examples in our paper, high-quality samples should be varied and differ significantly from the true image. A pixel-by-pixel comparison with the true image is not meaningful for posterior sampling, as it doesn't reflect sampler performance. Since the true image is near the maximum a posteriori (MAP) estimator, good pixel-by-pixel performance might indicate the algorithm generates non-diversified draws from dominant modes, failing to explore the full posterior support.

Why no Metropolis-Hastings correction? The author is right in thinking about the use of a MH correction. However, note that we do not actually have access to (unnormalized versions of) $\pi^\ell \_{k_{\ell+1}}$ since we cannot evaluate $p_k$. Therefore we cannot use MH type algorithms. LMC is of course biased when run for a finite number of steps, but as we show with the Gaussian mixture example, it can still perform decently, very close to that of the asymptotically exact algorithm MCGDiff.

Typos: We would like to thank you for drawing our attention to the notation problems present in Algorithm 1 and we apologize for these. Please find a corrected version of our algorithm in the PDF attached to the main rebuttal. Thank you also for pointing out other typos and suggesting corrections. Note that throughout the paper we use upper case for random variables and lower case for their realizations. With this in mind, saying that $p_{0|k}(\cdot | x_k)$ is the conditional density of $X_0$ given $X_k = x_k$ is not confusing.

bias of our algorithm: To our knowledge, all methods for sampling from posterior distributions with probabilistic diffusion priors are biased with finite iterations. This includes algorithms based on importance sampling, like SMC, which are also biased with limited particles. Regarding LMC, similarly to SMC using an infinite number of particles, if we were using decreasing step sizes converging to $0$, it is also asymptotically unbiased by [8, 9]. DCPS can be made asymptotically unbiased by running LMC targeting the posterior with the final sample from Algorithm 1. The key issue is whether our method’s samples are adequate initializations. While we have shown this empirically, a quantitative analysis of our algorithm’s accuracy remains for future work. Therefore, based on the current existing literature, pointing out that our method is biased (without any further indications) as a limitation seems unjustified in our opinion. We humbly ask the reviewer to elaborate more on their comment.

[1] Wang, Y., Yu, J. and Zhang, J. Zero-shot image restoration using denoising diffusion null-space model. International Conference on Learning Representations (ICLR) 2023.

[2] Kawar, B., Elad, M., Ermon, S. and Song, J., 2022. Denoising diffusion restoration models. Advances in Neural Information Processing Systems.

[3] Zhang, G., Ji, J., Zhang, Y., Yu, M., Jaakkola, T. and Chang, S., 2023, July. Towards Coherent Image Inpainting Using Denoising Diffusion Implicit Models. In International Conference on Machine Learning (ICML) 2023

[4] Song, J., Vahdat, A., Mardani, M. and Kautz, J., 2023, May. Pseudoinverse-guided diffusion models for inverse problems. In International Conference on Learning Representations.

[5] Dou, Z. and Song, Y., 2024. Diffusion posterior sampling for linear inverse problem solving: A filtering perspective. In The Twelfth International Conference on Learning Representations.

[6] Cardoso, G., Le Corff, S. and Moulines, E., 2023. Monte Carlo guided Denoising Diffusion models for Bayesian linear inverse problems. In The Twelfth International Conference on Learning Representations.

[7] Trippe, B.L., Yim, J., Tischer, D., Baker, D., Broderick, T., Barzilay, R. and Jaakkola, T.S., Diffusion Probabilistic Modeling of Protein Backbones in 3D for the motif-scaffolding problem. In The Eleventh International Conference on Learning Representations.

[8] Durmus, Alain, and Eric Moulines. "Nonasymptotic convergence analysis for the unadjusted Langevin algorithm." (2017): Annals of Applied Probability, 1551-1587.

[9] Lamberton, Damien, and Gilles Pages. “Recursive computation of the invariant distribution of a diffusion.” (2002): Bernoulli, 367-405.

Dear Reviewer vN6C,

We would like to thank you for your time, your feedback and suggestions. We answer your questions below.

Further experiments: We have followed your suggestion and proceeded to compare our algorithm with DDNM [1] and DiffPIR [2] on the imaging experiment. We did not compare with ReSample [3] since it is specialized for latent diffusion, which is not the focus of the present paper and is left for future work. We have also added a discussion of these papers to the manuscript. Following the suggestions of Reviewer xni9 and af3x we have also added FPS [4] and SDA [3] to our benchmark. The results can be found in the PDF in the main rebuttal. Both algorithms you suggested show strong performance, with DDNM coming second in our benchmark, and DiffPIR scoring third in terms of median performance and fourth in terms of average. Our algorithm, DCPS, still ranks first. Finally, we would like to emphasize that it ranks first or second consistently on 14 of the 16 tasks we have considered (across low and large observation standard deviation), which is not the case for any other algorithm under consideration.

Comparison with MCGDiff: Please note that SMC-based methods like MCGDiff heavily suffer from the curse of dimensionality. More specifically, these methods iteratively evolve a set of $N$ samples approximating the final posterior. The number of samples required to obtain a good approximation to the posterior distribution should be an increasing function of the dimension of the problem and may even grow exponentially with this quantity (see [5]), requiring a prohibitive amount of memory. Typically, for a number of samples $N$ which has been chosen too small relative to the dimension, all samples typically collapse into a single one, so that the final posterior distribution is approximated by basically a Dirac mass.

In our high-dimensional experiments, we used 32 particles for MCGDiff, which is far from sufficient, but this was the maximal number we were able to tackle due to memory constraints. The returned approximation is therefore relatively poor. For the Gaussian mixture experiment, the dimension is low, and this is exactly the setting where MCGDiff performs best and is thus better than DCPS. Nonetheless, note that if the computational budget of DCPS is increased, it can still perform better than MCGDiff, as we show in the $10$-dimensional example. To summarize, the comparison with MCGDiff provides a convincing argument in favor of DCPS; in low dimensions it comes very close to the performance of SMC methods, and, crucially, it performs much better for high-dimensional problems.

Other applications: As an example of other possible applications, we are currently applying DCPS to zero-shot conditional sampling for traffic simulation. In this setting, one has access to a diffusion model for the joint distribution of $N$ interacting agents and the aim is to be able to sample joint agent trajectories that satisfy user-specified constraints. Two notable examples are the simulation of joint trajectories that do not collide or trajectories that satisfy certain speed limits, see Table I in [6]. Both cases define specific potentials $g_0: \mathbb{R}^{N \times T \times d} \to \mathbb{R}$ where $N$ is the number of agents, $T$ is the length of the trajectory and $d$ is the number of trajectory features considered. DCPS applies successfully by considering potentials built following the same recipe as in the paper; please see the second point in the main rebuttal for more details. We have found that this general principle works well for various inverse problems and as a future work we aim to understand how to select these hyperparameters in a principled manner.

[3] Rozet et al. "Score-based Data Assimilation". In Advances in Neural Information Processing Systems. 2023.

[4] Dou, Zehao, and Yang Song. "Diffusion posterior sampling for linear inverse problem solving: A filtering perspective." _The Twelfth International Conference on Learning Representations. 2024.

[5] Bickel, Peter, Bo Li, and Thomas Bengtsson. "Sharp failure rates for the bootstrap particle filter in high dimensions." In Pushing the limits of contemporary statistics: Contributions in honor of Jayanta K. Ghosh, vol. 3, pp. 318-330. Institute of Mathematical Statistics, 2008.

[6] Zhong, Ziyuan, Davis Rempe, Danfei Xu, Yuxiao Chen, Sushant Veer, Tong Che, Baishakhi Ray, and Marco Pavone. "Guided conditional diffusion for controllable traffic simulation." In 2023 IEEE International Conference on Robotics and Automation (ICRA), pp. 3560-3566. IEEE, 2023.