Title: The Best of Both Worlds: Amortized Variational Diffusion Posterior Sampling
Primary area: generative models
Keywords: diffusion models, inverse problems, amortized optimization

Abstract: Diffusion models pre-trained on large datasets are powerful priors for inverse problems such as super-resolution and inpainting. Zero-shot variational diffusion posterior sampling achieves state-of-the-art reconstructions without task-specific training but is slow due to costly test-time optimization. Supervised diffusion for inverse problems offers fast inference, yet demands large datasets and often fails under unseen degradations. We introduce a best-of-two-worlds strategy that jointly leverages upstream training and test-time likelihood guidance. An amortized inference model, trained on a small paired dataset, predicts a good initialization for a variational approximation problem involved in variational diffusion posterior sampling, while retaining the explicit use of the degradation operator to guide inference. This combination removes several gradient updates at inference, reducing inference time by up to 32 \% compared to zero-shot posterior sampling. Importantly, it remains robust against out-of-distribution degradation operators and training settings with limited data (e.g., 1 \% of the pre-training data), outperforming the supervised diffusion baselines in these scenarios. Our results show that coupling modest training with test-time operator knowledge can unlock fast, flexible, and high-quality diffusion reconstructions.

TL;DR: Amortized optimization for variational diffusion posterior sampling