Divide-and-Conquer Posterior Sampling for Denoising Diffusion Priors

Summary

The topic of using a pre-trained diffusion model as a prior for a Bayesian posterior sampling problem is considered. This formulation is compelling (particularly for images) as it combines the powerful generative capacities of large-scale diffusion models with the coherent inference under certainty provided by Bayes’ theorem. The paper refines the diffusion posterior sampler of Chung et al by replacing a point mass approximation with a variational Gaussian approximation to the conditional transition distributions. *Chung et al: Diffusion posterior sampling for general noisy inverse problems, ICLR, 2023*

Strengths

The paper tackles an interesting problem and is mathematically compelling as the point mass approximation is likely to be crude in practice. The experiments are practical and varied - investigating various likelihoods (linear+Gaussian, Poisson, JPEG decoding). And the results for the introduced sampler are promising.

Weaknesses

The paper is (perhaps understandably quite dense), in particular, I think some more clarity on the use of twisted potentials and how this differs (or to what extent) from the Chung et al approach would aid the reader. The resulting algorithm is somewhat complex and it is not clear how sensitive the algorithm is to the choice of twisting potential and hyperparameters, in particular the number of internal Langevin (M) and VI gradient steps (K). Additionally, as noted by the authors, the choice of twisting potentials considered is bespoke in that they have to be specified for each problem. A more general-purpose posterior sampling pipeline that fully separates likelihood and data specification from (automated) posterior sampling would be highly desirable. Although understandably as a future direction in the considered work. *Chung et al: Diffusion posterior sampling for general noisy inverse problems, ICLR, 2023*

Questions

Would MAP estimation (over posterior sampling) also be feasible/performative in this setting? And perhaps computationally easier? It's not clear to me how exactly the proposed DCPS algorithm is "divide-and-conquer". Over which axis does it divide? And which computational bottleneck does this speed-up? How does this compare to other computational bottlenecks? (I.e. is it dominating)

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors have addressed key limitations of the work. I see no potential for negative societal impact.

Summary

This work proposes a new posterior sampling scheme for unconditional diffusion models. The sampling algorithm, named divide-and-conquer posterior sampling (DCPS), combines Langevin Monte Carlo and Gaussian variational inference to operate the transitions between noise levels along the reverse diffusion process. The authors apply their sampling scheme to several inverse problems and compare their results with several previous works. DCPS seems to improve upon these works, both quantitatively and qualitatively.

Strengths

* The manuscript is well written for the most part. * The subject of posterior sampling algorithms in diffusion models is interesting and relevant for the ML and scientific communities. * The proposed algorithm seems sound, albeit too complex in my opinion (see weaknesses). * The experiments are well described and the results are clearly presented. * The code is provided, although it is not documented and instructions to reproduce experiments are minimal.

Weaknesses

* The proposed method, although presented quite differently, is very similar to the predictor-corrector approach proposed by Rozet et al. [1], where the predictor is a classic DDPM/DDIM step which approximates the transition from one noise level to the next, while the corrector is a number of Langevin Monte Carlo (LMC) iterations. The similarity is striking, especially when noticing that the DDPM step is a Gaussian approximation of the transition. The authors do not discuss nor mention Rozet et al. [1]. The authors should highlight the differences between their posterior sampling algorithm and the one of Rozet et al. [1] and compare their results. * Even if it were correct, the link between diffusion posterior sampling and Feynman-Kac (FK) models in Section 2 is not relevant and confusing. Unlike the authors state (line 110), the overwhelming majority of diffusion/score-based posterior sampling literature uses either a Markov chain or a stochastic differential equation (SDE) representation. Furthermore, this link with FK models is not even leveraged as part of the method, which only relies on Markov chain properties. * As mentioned by the authors, intermediate potentials should be designed for each problem, which could be difficult for some inverse problems. In addition the convolution of these potentials with the bridge kernel (Eq. (3.6)) should be (cheaply) tractable and differentiable. * There are two aspects to the proposed algorithm: the LMC iterations and the intermediate potentials/posteriors. It is hard to assess the contribution of each aspect in the results. My hypothesis is that most of the quality gains are due to the LMC iterations. My rationale is that the potentials $g^l_m(x_m) = \mathbb{E} [ g_{k_l}(x_{k_l}) \mid x_m ]$ are actually user-defined potentials, just like $g_m(x_m) = g_0(x_m)$ or $g_m(x_m) = \mathbb{E} [ g_0(x_0) \mid x_m ]$. As long as they converge to $g_0$ when $m \to 0$, the algorithm should work. ### Comments * I am not sure to understand the reference to "divide-and-conquer" from dynamic programming. * At line 80, $t$ is undefined. * At line 143, Finzi et al. [2] are the first to propose to use the Jacobian of the denoiser to estimate the covariance. Boys et al. [3] later proposed an approximation using the diagonal of the Jacobian. * In Algorithm 1, I think there is a typo in "for $j = \ell + 1 - 1$ to $\ell$ do". * Table 5 should be Figure 5. * At line 232, "resort to a optimizing" -> "resort to optimizing". * At line 258, $K$ is not defined in the main text. * At line 334, if I understand correctly DCPS is not limited to linear inverse problems. [1] Rozet et al. "Score-based Data Assimilation". In Advances in Neural Information Processing Systems. 2023. [2] Finzi et al. "User-defined Event Sampling and Uncertainty Quantification in Diffusion Models for Physical Dynamical Systems". In Proceedings of the 40th International Conference on Machine Learning. 2023. [3] Boys et al. "Tweedie Moment Projected Diffusions for Inverse Problems". 2024.

Questions

See weaknesses.

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

See weaknesses.

Summary

This paper aims to solve the inverse problem within the Bayesian framework. The authors propose DCPS, which approximately samples from the posterior. DCPS creates a sequence of intermediate posterior distributions to approach the target posterior step by step. The authors showcase the effectiveness of DCPS on inverse problems like image restoration and trajectory prediction.

Strengths

1. Good empirical results on Gaussian mixture problems. 2. DCPS is computationally efficient.

Weaknesses

1. No justification for the design of the intermediate potentials. They are approximations of $g^*_k$ defined in line 122. But why are they good approximations? The manuscript will benefit from a theoretical analysis of the approximation error. 2. In general, it is unclear how to design intermediate potential for nonlinear forward models. DCPS seems limited to linear problems. 4. Only LPIPS is reported in the image restoration tasks. The experiment section would benefit from adding evaluation metrics such as PSNR and SSIM. 5. DCPS generates pseudo observations in a similar way to the prior work FPS [1]. I believe this is an important baseline to discuss and compare. 6. Overall, the paper is hard to parse, with many long sentences, redundant phrases, and inconsistent notation usage. Just list some of the issues here. I believe a major revision is needed to make this paper more accessible. 1. Typos in Algorithm 1: "$X^l_{l+1}\leftarrow X_l+1$" and"for $j=l+1-1$ to $l$" 2. $l$ is used in subscript and superscript together, which is very confusing. For example, $X^l_{l+1}$ in Algorithm 1, $\pi^l_{k_l}$ in Eq. (3.4). 3. Line 66: remove the comma before "and illustrate" 4. Line 64: "describe theoretically" -> "theoretically describe" 5. Line 70: "generate efficiently" -> "efficiently generate" 6. Line 158: a comma is missing before "we find" 7. The authors use $N(x;\mu, \Sigma)$ to represent the Gaussian density. I suggest the conventional notation $\mathcal{N}$ from the literature for better consistency with the field. 8. Calling $g_k$ as the potential is confusing when it's actually a likelihood function. 9. The use of upper case $X$ and lower case $x$ is confusing throughout the paper. For example, in line 85, it's $X_k=x_k$. 10. Redundant phrases and long sentences should be revised for clarity. For example, line 62-71, line 96-99, line 151-153, line 160-163. 11. Line 102: "where $\mathcal{Z}=\int g_0(x_0)p_0(dx_0)$" -> "where $\mathcal{Z}=\int g_0(x_0)p_0(x_0)dx_0$" 12. Incorrect inline citations throughout the paper. [1]: Dou, Zehao, and Yang Song. "Diffusion posterior sampling for linear inverse problem solving: A filtering perspective." _The Twelfth International Conference on Learning Representations. 2024.

Questions

1. the unadjusted Langevin is a biased sampler for any finite time step. Why not include the Metroplis-Hasting correction? 2. Why is L much smaller than n? How does the value of L affect the final performance? 3. In the image restoration experiments, the authors claim to use only 5 Langevin steps. It is very unlikely that the algorithm is sampling from the right distribution within 5 Langevin steps. Langevin Monte Carlo is known for slow convergence for high-dimensional problems, which takes hundreds of steps even on simple low-dimensional problems. How can the proposed DCPS work with only 5 Langevin steps?

Rating

4

Confidence

5

Soundness

2

Presentation

1

Contribution

2

Limitations

1. Problem-specific intermediate potentials must be designed, limiting DCPS's application from a wider range of inverse problems. 2. While the goal of this paper is posterior sampling, the proposed DCPS is a biased sampler that does not converge to the true posterior.

Summary

This paper presents a sampling algorithm (DCPS) that leverages diffusion models as priors to solve linear inverse problems. DCPS defines and recursively solves a series of intermediate inverse problems that converge to the original inverse problem. In each iteration, a sample is drawn from the next posterior distribution following a combination of Langevin dynamics and an inhomogeneous Markov chain. DCPS shows nice results in synthetic inverse problems and also image restoration tasks.

Strengths

- DCPS is well-motivated by intuition to mitigate the approximation error of the conditional distribution, which is supported by solid theoretical evidence in Proposition 3.1. - The experimental results on synthetic data validate the claim of accurate posterior sampling.

Weaknesses

The experiments on images are not convincing enough. DCPS performs slightly better than DPS and DDRM, while recent works like DDNM[1], DiffPIR[2], and ReSample[3] have claimed much better posterior sampling quality on image restoration tasks but are not compared as baselines. [1] Yinhuai Wang, Jiwen Yu, and Jian Zhang. Zero-shot image restoration using denoising diffusion null-space model. ICLR, 2023. [2] Yuanzhi Zhu, Kai Zhang, Jingyun Liang, Jiezhang Cao, Bihan Wen, Radu Timofte, and Luc Van Gool. Denoising diffusion models for plug-and-play image restoration. CVPR, 2023. [3] Bowen Song, Soo Min Kwon, Zecheng Zhang, Xinyu Hu, Qing Qu, Liyue Shen. Solving Inverse Problems with Latent Diffusion Models via Hard Data Consistency. ICLR, 2024.

Questions

- It is noted for the Gaussian mixture experiments that DCPS does not beat MCGDiff in terms of accurate posterior sampling. However, it outperforms MCGDiff by a large margin in image restoration tasks. What are the pros and cons of DCPS compared to MCGDiff? And what makes synthetic inverse problems and image restoration tasks so different? - Would it be possible to apply DCPS to more realistic inverse problems in other domains? What are the possible limitations of this work when adapted to more complicated inverse problems?

Rating

7

Confidence

3

Soundness

3

Presentation

2

Contribution

4

Limitations

The authors address the limitations and provide possible future directions in their paper.