Practical and Asymptotically Exact Conditional Sampling in Diffusion Models

Summary

This paper focuses on solving inverse problems using diffusion based probabilistic models. More precisely, it is only assumed that one has access to a diffusion model for the prior distribution and a likelihood, so that no additional training is needed. The aim of the present paper is to provide an asymptotically exact method. The method builds up on the idea of reconstruction guidance which aims at approximating the score involved in the diffused posterior, $\nabla \log p_t(x_t | y)$ which is equal to $\nabla \log p(y|x_t) + \nabla \log p_t(x_t)$ where the second term is given by the prior diffusion model. The first term is however intractable but can be approximated by noting that $p(y | x_t) = \int p(y|x_0) p(x_0 | x_t) dx_0$. This integral is then approximated by simply taking the mean a posteriori i.e. $p(y | x_t) \approx p(y | E(x_0 | x_t))$ which itself can be approximated using Tweedie's formula and the prior score. The authors use this idea to obtain particle approximations of the marginal of the backward diffusion of the posterior. They do so using twisting functions. This allows them to define a principled SMC samplers for the target distribution. By standard SMC results, the obtained particle approximation of the posterior converges to the posterior of the diffusion model. The extension to Riemannian manifolds is also provided.

Strengths

- The contribution of this paper is original. Unlike other methods trying to solve inverse problems, this method is theoretically grounded and is guaranteed to be reliable in relatively complex problems. This is for example not the case of DPS [1] which fails to sample the posterior even in this simplest Bayesian settings. - I also acknowledge that the paper is well written, I enjoyed reading it. - The numerical experiments are sound and the application to protein design is interesting.

Weaknesses

- While I find the approximations of the optimal twisting functions reasonable, I am not sure if this is the best idea if we take into account the computational cost. Indeed, sampling from the transition requires computing the gradient of the score network with respect to the input. In large scale applications (like images with 3x256x256 dimensions) this severly limits the number of particles that can be used, since the gradient will be computed for each particle. At the end of the day, the number of particles used is very important even if the transition kernels and weights are near optimal; if the posterior is highly multimodal then we inevitably need many particles in order to populate each mode. SMCDiff on the other hand should work better on these problems since one can use a larger number of particles given a fixed computational budget. - The comparison with SMCDiff is in my opinion unfair. SMCDiff performs particle filtering only on the unobserved part of the state and does not require computing the gradient of the score network. Therefore the computational time and memory cost of TDS are much larger than SMCDiff. This is not taken into account in the numerical experiments. Furthermore, it seems that SMCDiff outperforms TDS in terms of ESS. I understand that SMCDiff relies on a stringent assumption that does not hold in practice, but in my opinion it is as reasonable as saying that the instrumental kernel and weights proposed for TDS are a good approximation. In short, i am not really convinced that TDS is better than SMCDiff for inpainting tasks. For noisy inverse problems of course the comparison is not relevant since SMCDiff is not designed for such problems. - As far as I can tell, the impact of the dimension of the observation is not discussed. As is widely known in the SMC community, this makes the weights degenerate. I believe that a discussion on this matter should be added so that readers are not misled into thinking that in high dimensional problems this method will provide $K$ diversified samples from the posterior. It should be emphasized that this method works best on low dimensional problems.

Questions

It is claimed in the appendix that the assumptions required to prove the convergence are mild. However, assuming that the ratio of the twisting function (assumption (b), line 588-589) is bounded is not a mild assumption since it doesn't hold in the simplest cases. It is claimed that this holds if the prediction $\hat{x}_0(x_t)$ is compactly supported but I am not sure how this is possible even if the data distribution is compactly supported. Indeed, the marginals of the forward process $q_t$ are the result of the convolution of the data distribution with a Gaussian kernel so that it cannot be compactly supported. Since $\hat{x}_0(x_t)$ contains the score, it cannot be compactly supported too. Next, assuming that the score is bounded is also not reasonable. Finally, assuming that $p(y; x^0)$ is bounded away from zero works but this does not hold for Gaussian linear inverse problems, which is quite unfortunate. In short, both assumptions (b) and (c) are strong assumptions and it should be stated explicitely. There is a typo in assumption (c); its not the gradient that should be continuous with bounded gradients.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

See weaknesses.

Summary

The paper proposes an SMC algorithm to draw conditional samples form a diffusion model. Specifically, they wish to draw samples p(x0|y) given a diffusion model p(x0) and likelihood p(y|x0). Existing techniques to do so rely on expensive training of conditional diffusion models or heuristics which do not sample from the correct conditional distribution.

Strengths

- Theoretically well-founded method. The paper does a good job of describing the flaws of commonly used methods for conditional generation and proposes a technique to target the true conditional distribution without needing task-specific training. - State-of-the-art results for protein motif scaffolding, which is a very reasonable application of the proposed exact targeting of conditional distributions

Weaknesses

- For some of the experiments (such as inpainting), a reasonable baseline to compare against would be that proposed in Section 3.1 of VDM (https://arxiv.org/abs/2204.03458). This performs conditional generation with an unconditional diffusion model using a heuristic that they show improves on naive guidance (although doesn't have theoretical guarantees like the proposed SMC method). A comparison against this would make it clearer when Twisted SMC is indeed the best option. - Section 4.3 is hard to read and lacking in plots or concrete results. Adding something like Figure J from the appendix would make the results more convincing for readers who do not venture to the appendix to search for plots. - How accurate is the inferred conditional distribution when less sampling steps are used? With progressive distillation or more modern ODE/SDE solvers it is now common to sample from diffusion models with tens of integration steps. The performance of SMC presumably degrades when less steps are used, since each individual step must be bigger. Can the authors comment on this or quantify this effect? E.g. how would Twisted SMC perform relative to the baselines if the number of diffusion steps used for motif-scaffolding was halved?

Questions

See weaknesses.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Yes

Summary

This paper proposes a practical approach for achieving asymptotically exact inference from diffusion models through exact conditional sampling in terms of Sequential Monte Carlo (SMC) . This discovers the connection between SMC and diffusion models, and one of the key feature is to approximate the optimal twisting functions by som tractable twisting alternatives. The proposed method is extended to tackle inpainting, inpainting with degrees of freedom problems, and make diffusion model compatible on Riemannian manifolds. The experiments involve synthetic diffusion models, class-conditional sample generation on MNIST dataset, and motif-scaffolding problem, which is relevant to protein design.

Strengths

1. The inspiration of the paper is good, which tries to achieve exact conditional sampling from diffusion models by making connection to SMC, and extend to Riemannian diffusion models. 3. (More like an) Ablation study in the experiment portion to examine whether the proposed method works better.

Weaknesses

Based on understandings, 1. The novelty of this paper is not clear to me. The strong and close connection between diffusion models and the SMC technique is not crystal clear to me. Even though the reverse process with T diffusion steps from the diffusion models and sampling across T steps in terms of SMC could be related somehow, the significance of this connection is not apparent. The authors also point out that Classifier guidance[1] requires additional classifier training, but no computational cost and sample quality comparison is provided; Classifier-free guidance [2] requires classifier information as an additional input but again, none of the aforementioned result is presented and class information is still needed by the proposed TDS method (from my understanding but could be wrong). Equation 10, which is referred to as key insight in the paper, not been explained well, and actually has been proposed in [5] with even further clear explanation. The previous paper is not cited either. The title is also not informative enough and the emphasis on SMC is not evident. 2. Lack of real benchmark for the first two experiments, particularly the examination on the MNIST dataset. I am also concerned that the "Guidance" method performs significantly worse than the other four methods without any justification, and no pre-trained diffusion model is served as a benchmark here. Three experiments are all competing to themselves. It would be more convincing for the image synthesis experiments if they were conducted on more diverse and challenging datasets such as CIFAR-10, instead of solely relying on the MNIST dataset, which consists of grayscale images with only 10 classes. 3. Not sure why but the ``Related work`` is added in the appendix. MNIST class-conditional generation experiment part is in the paper, the inpainting is put in the appendix as well (which is mentioned in the ``abstract`` and ``conclusion``). I do not think it is pretty common but I would suggest to move them into the main paper from appendix. 4. There is a big room to improve for writing. Too many typos in both grammar and math expressions, not being cited properly, i.e., Tweedie's formula in the ``background``; unclear descriptions, i.e., the diffusion models described in the ``background`` and ``methodology`` are not referred to VE or VP diffusion models until the appendix, and mathematical notations appear without proper explanation, i.e. $\nu_t(x^{t:T})$ the first occurrence on line 94, where its meaning is not clarified. The lack of coherence between sentences and paragraphs further hampers the clarity of the paper. [1]: Prafulla Dhariwal and Alexander Nichol. Diffusion models beat GANs on image synthesis. Advances in Neural Information Processing Systems, 2021. [2]: Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12598, 2022. [3]: Hyungjin Chung, Jeongsol Kim, Michael T Mccann, Marc L Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems. In International Conference on Learning Representations, 2023.

Questions

1. ``Choosing weighting function $w_T(x^T) = w_t(x^t, x^{t+1}) = 1$ for $t = 1, \dots, T$'' constantly appearing in the paper to leave the proposal functions only in the equation. Are you assuming no resampling here and no particle is being discarded? If so, please make it more clear. 2. For weighting function on line 95, $w_k^t$ is a function related to $x_k^t, x_k^{t+1}$ and no $y$ is involved, so what is the point of not including $y$ here? When twisted weighting function is defined on equation 15, it is related to $y$ by $\tilde{p}_\theta(y | x^t)$ though. 3. On line 141, ``recall that $\hat{x_0}(x^t) = \mathbb{E}_{p_\theta}[x^0 | x^t]$ if $p_\theta$ is optimized to exactly match the true distribution $q$''. What if the assumption does not held perfectly? Is there any statistical analysis to justify this? 4. On equation 10, ``the approximated optimal twisting functions is $\tilde{p}_\theta(y | x^t) := p(y; \hat{x}_0(x^t, t))$''. Is there any analysis (if this idea is from [5], please ignore this question about analysis) or experiment result for not denoising $x^t$ back to $x^0$ for comparison? 5. Is the first experiment two-dimensional? What if it is high-dimensional? 6. What's the computational cost for various number of particles? Is there any variance reported for different number of particles for each experiment?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

I would expect the authors address the weakness and questions aforementioned.

Summary

This paper addresses the challenge conditioning in unconditionally-trained diffusion models. The most successful approaches often require explicitly training on conditional data. This paper frames sampling from such conditionals as an SMC procedure and proposes Twisted Diffusion Sampler (TDS), a method derived from the twisting technique in SMC literature. TDS uses a classifier trained on the clean, non-noisy data to construct a sub-optimal proposal distribution for sampling from the diffusion process. This proposal is then used in the twisted SMC framework to get weighted particles that more closely approximate the conditional distribution of interest. The authors prove that their proposed SMC method targets the correct distribution and is asymptotically exact. Furthermore, they show how to extend TDS to work for various inpainting problems and Reimannian diffusion models. Finally, the authors empirically verify TDS's correctness on a simple synthetic diffusion model with tractable score functions. Then they show its effectiveness on more realistic tasks of class-conditional MNIST image generation and inpainting. They lastly show TDS achieves state-of-the-art results on motif-scaffolding problem in protein design.

Strengths

The central contribution of the paper is introducing the twisted SMC framework for sampling from diffusion models. While the twisted SMC is not novel, applying it to diffusion models and demonstrating use cases is significant. Moreover, the extension of TDS to some settings other than class-conditional framework was very interesting to me. In terms of clarity, this paper is mostly clearly written. There are some typos and clarifications required that I have listed in the questions section.

Weaknesses

I think the main weakness of the paper lies in the experimental section. First, the experiments are rather small scale. It would be nice to see how their method perform on larger scales standard tasks such as class-conditional CIFAR or ImageNet generation. Second, I expected to see more baselines in the experiments. The paper is motivated by pointing out that many existing methods require additional training. I believe it such methods should be included as baseline. For example, classifier-guidance with a proper classifier trained on noisy data or classifier-free guidance would be good baselines to compare against. Admittedly, they might outperform TDS, particularly with smaller K. However, it would be a good upper bound to have. Moreover, it gives insights on the tradeoff between a large K or training a separate model.

Questions

1. I am not sure if what the authors refer to "reconstruction guidance" is correct. Quoting from the related work section "the use of denoising estimate $\hat{x}_0(x^t, t)$ to form an approximation $\tilde{p}_\theta(x^t | x^{t+1}, y)$ to $p_\theta(x^t | x^{t+1}, y)$ is called _reconstruction guidance_." I am not sure if this is correct. As far as I know, reconstruction guidance is proposed in [1] and improves the replacement method [2] which is for inpainting applications. Moreover, the other two citations are [3,4] (citation numbers 5, 25 in the paper) which are unrelated to reconstruction guidance (this might be related to the question 3 below though) 2. In the related work section (App. B), on line 523, there is a "[33]" citation which does not exist. The reference 29 in the paper ([5] here) seems to be what it meant to refer to. This paper, however, does not do reconstruction guidance as far as I can tell. 3. Looking at the rest of the related work section, seems like most citation numbers are incorrect. For example citations 26 and 18 on lines 524 and 528, respectively do not seem to refer to the correct paper. I suggest double-checking the citations. 4. In section 3.3, it is not clear to me how Eq. (18) is derived and whether it is a convenient likelihood that often works in practice or it is a general result. I did not see an explanation in the appendices either. I think it would be helpful to clarify this. 5. In the experiments section 1. "Guidance" is defined as "TDS" with one sample. If I understand it correctly, it would be equivalent to "TDS-IS" with one sample. Is this correct? If so, I expect it to be outperformed by "TDS-IS" with more samples while the middle and right panel of Figure 1 shows otherwise. Additionally, in Figure 2(a), I do not understand why would "TDS-IS" show less diversity than "guidance". Why would more particles destroy the diversity? 2. I am wondering if the authors have insights on why in Figure 2(b) ESS jumps back up very close to t=0. 6. Minor issues - Section 3 and subsection 3.2 have the same name. - The background section on diffusion models explains diffusion processes without scaling (e.g. the variance-preserving process).It is worth mentioning in the background section that not all diffusion models follow this exact framework. - I think it's nicer if Eq. 11 is written to define $\tilde{p}_\theta(x^{t-1} | x^{t}, y)$ instead of $\tilde{p}_\theta(x^t | x^{t+1}, y)$ to be similar to Eq. 3. 7. Typos - Line 9: the -> that - Line 99: provide -> provides - Line 148: eq. -> Eq. - Line 153: I might be wrong, but I guess $s_\theta(x^t, y) = \nabla_{x^t} \log q(x^t, y)$ should actually be $s_\theta(x^t) = \nabla_{x^t} \log q(x^t)$. Otherwise, the assumption of $\tilde{p}_\theta(y | x^t) = p_\theta(y | x^t)$ becomes unnecessary for Eqs. (13,14). - Line 157: $w_T(x^T) := \tilde{p}_\theta(x^T)$ -> $w_T(x^T) := \tilde{p}_\theta(y | x^T)$ - Linr 218: sytematic -> systematic - Lien 570: systematica -> systematic - In the second line of equations under line 578: $p(y | x^{t'})$ -> $p(y | x^{t})$ [1] Jonathan Ho, Tim Salimans, Alexey A Gritsenko, William Chan, Mohammad Norouzi, David J Fleet. "Video diffusion models". NeurIPS 2022. [2] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, Ben Poole. "Score-based generative modeling through stochastic differential equations". ICLR 2020. [3] Hyungjin Chung, Jeongsol Kim, Michael T Mccann, Marc L Klasky, Jong Chul Ye. "Diffusion posterior sampling for general noisy inverse problems". ICML 2023. [4] Jue Wang, Sidney Lisanza, David Juergens, Doug Tischer, Joseph L Watson, Karla M Castro, Robert Ragotte, Amijai Saragovi, Lukas F Milles, Minkyung Baek, et al. "Scaffolding protein functional sites using deep learning". Science 2022. [5] Guanhua Zhang, Jiabao Ji, Yang Zhang, Mo Yu, Tommi Jaakkola, and Shiyu Chang. "Towards coherent image inpainting using denoising diffusion implicit models". arXiv preprint, 2023.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

Limitations

The main limitation of the proposed method is the additional sampling cost due to multiple SMC particles which is mentioned in the discussion section of the paper. Otherwise, I do not see other limitations.

Response to Authors

I would like to thank the authors for their answers and the additional information. Most of my questions and concerns were touched and addressed. For the additional experiments on ImageNet, I can see a big improvement on classification accuracy when the classifier guidance can already achieve 99% accuracy; however, it is noteworthy that the quality and diversity of the generated samples, as indicated by the FID and Inception scores, have notably decreased.. These induced to two questions: 1) Is it justifiable to sacrifice sample quality and diversity to improve classification accuracy by applying TDS? This is without considering the increased time consumption resulting from the higher number of particles. 2) I may still not fully understand it, but can it be considered a valid assumption to employ TDS performed upon a pre-trained classifier guidance, particularly when the training a classifier is time-intensive? I would be happy to update my score once my concern is addressed.

Response to Authors

I appreciate the authors' effort to conduct further experiments and their clarifications. Most of my questions are answered and I will raise my score. However, the new results do not demonstrate a clear advantage from the techniques introduced. In particular, the scores of TDS with one particle (which is equivalent to "naive" guidance) very close to TDS with 16 particles. The exception is inception score where TDS (P=16) shows a higher score suggesting more sample diversity, but the diversity of reported images is very limited. Nonetheless, I still believe the proposed method and the details of implementing it in different application is elegant and fruitful to the community.