Amortized Posterior Sampling with Diffusion Prior Distillation
Abbas Mammadov, Hyungjin Chung, Jong Chul Ye
TL;DR
This paper introduces Amortized Posterior Sampling (APS), an unsupervised variational inference method using diffusion models that efficiently generates diverse posterior samples for inverse problems across various domains.
Contribution
The paper presents a novel, unsupervised, diffusion-based amortized sampler that requires no paired data and generalizes across Euclidean and non-Euclidean domains.
Findings
APS outperforms existing methods in computational efficiency.
APS achieves competitive reconstruction quality.
APS enables real-time solutions to inverse problems.
Abstract
We propose Amortized Posterior Sampling (APS), a novel variational inference approach for efficient posterior sampling in inverse problems. Our method trains a conditional flow model to minimize the divergence between the variational distribution and the posterior distribution implicitly defined by the diffusion model. This results in a powerful, amortized sampler capable of generating diverse posterior samples with a single neural function evaluation, generalizing across various measurements. Unlike existing methods, our approach is unsupervised, requires no paired training data, and is applicable to both Euclidean and non-Euclidean domains. We demonstrate its effectiveness on a range of tasks, including image restoration, manifold signal reconstruction, and climate data imputation. APS significantly outperforms existing approaches in computational efficiency while maintaining…
Peer Reviews
Decision·ICLR 2025 Conference Withdrawn Submission
- The paper is very well written and mathematically sound. - The results cover diverse settings and show clear favorable runtime for APS. - Related work is addressed well and the paper’s contribution is put in proper context.
- The method requires training at test time (e.g. on the test set). This stands in contradiction to the speed claim of the method being a fast sampler. The question then becomes: How large should the test set be to compensate for this prolonged training? - The method is claimed to produce high-quality and diverse posterior samples. Nonetheless, the resulting samples look rather noisy (e.g. CelebA denoising in Fig. 1, row 1), and they seem to have little to no diversity. - The datasets used
- Clarity: The paper is clearly written and well-structured, making the methodology and key contributions easy to understand. - Significance: Previous approaches to solving inverse problems often require multiple neural network evaluations, leading to high computational costs. This paper’s proposed method addresses this limitation by using an amortized variational inference framework, enabling efficient, single-step sampling. This improvement in computational efficiency represents a meaningful
- Limited novelty: The primary contribution of the paper is the use of amortization and diffusion model distillation to enable efficient, single-step sampling. However, this concept has already been explored in recent work [1] with a very similar approach, also leveraging **unsupervised diffusion prior distillation** with **amortized variational inference** for inverse problems. This challenges the paper's claim of novelty (i.e., "the first diffusion prior distillation" in line 101 or "the first
The paper is well-written and easy to follow. It employs a conditional normalizing flow model to achieve one-step posterior inference. Experiments on both Euclidean and non-Euclidean datasets across various tasks—including inpainting, denoising, Gaussian deblurring, and super-resolution—demonstrate the proposed distillation method's fast speed and competitive performance.
1. Though the proposed method enables one-step inference, training such a model using equation (11) in the main text is time-consuming. Additionally, it requires training on test images, whereas previous methods like DPS perform zero-shot inference on test images and achieve strong performance. Given the strong performance in the field of diffusion-based inverse problems, I find APS's performance underwhelming. Furthermore, including more challenging tasks on higher-resolution images would be be
The method allows one-step inference for posterior sampling, offering substantial time savings over traditional DIS methods APS is tested on both Euclidean and non-Euclidean data, showcasing its potential in various fields. The integration of conditional NFs with diffusion priors, particularly in unsupervised settings, is promising.
Limited to specific models, such as MCG, DPS, and Noise2Score, but does not benchmark against other efficient variational methods recently introduced for inverse problems. Also benchmarking against similar methods would help demonstrate the behavior of this approach to others: DDNM, Pi-GDM, FPS, FPS-SMC, ect. For example, it’s unclear how the model can fail and what the failure cases may look like.
Videos
No videos yet. Explain this paper in a talk, walkthrough, or lecture? Add one.
Taxonomy
TopicsAdvanced Statistical Process Monitoring
MethodsDiffusion · Variational Inference