Table of Contents
Fetching ...

Fast and Robust Likelihood-Guided Diffusion Posterior Sampling with Amortized Variational Inference

Léon Zheng, Thomas Hirtz, Yazid Janati, Eric Moulines

TL;DR

An amortization strategy for diffusion posterior sampling is introduced that preserves explicit likelihood guidance by amortizing the inner optimization problems arising in variational diffusion posterior sampling, which accelerates inference for in-distribution degradations while maintaining robustness to previously unseen operators, thereby improving the trade-off between efficiency and flexibility in diffusion-based inverse problems.

Abstract

Zero-shot diffusion posterior sampling offers a flexible framework for inverse problems by accommodating arbitrary degradation operators at test time, but incurs high computational cost due to repeated likelihood-guided updates. In contrast, previous amortized diffusion approaches enable fast inference by replacing likelihood-based sampling with implicit inference models, but at the expense of robustness to unseen degradations. We introduce an amortization strategy for diffusion posterior sampling that preserves explicit likelihood guidance by amortizing the inner optimization problems arising in variational diffusion posterior sampling. This accelerates inference for in-distribution degradations while maintaining robustness to previously unseen operators, thereby improving the trade-off between efficiency and flexibility in diffusion-based inverse problems.

Fast and Robust Likelihood-Guided Diffusion Posterior Sampling with Amortized Variational Inference

TL;DR

An amortization strategy for diffusion posterior sampling is introduced that preserves explicit likelihood guidance by amortizing the inner optimization problems arising in variational diffusion posterior sampling, which accelerates inference for in-distribution degradations while maintaining robustness to previously unseen operators, thereby improving the trade-off between efficiency and flexibility in diffusion-based inverse problems.

Abstract

Zero-shot diffusion posterior sampling offers a flexible framework for inverse problems by accommodating arbitrary degradation operators at test time, but incurs high computational cost due to repeated likelihood-guided updates. In contrast, previous amortized diffusion approaches enable fast inference by replacing likelihood-based sampling with implicit inference models, but at the expense of robustness to unseen degradations. We introduce an amortization strategy for diffusion posterior sampling that preserves explicit likelihood guidance by amortizing the inner optimization problems arising in variational diffusion posterior sampling. This accelerates inference for in-distribution degradations while maintaining robustness to previously unseen operators, thereby improving the trade-off between efficiency and flexibility in diffusion-based inverse problems.
Paper Structure (63 sections, 25 equations, 10 figures, 5 tables, 2 algorithms)

This paper contains 63 sections, 25 equations, 10 figures, 5 tables, 2 algorithms.

Figures (10)

  • Figure 1: Summary of the novel amortization strategy. (Left) Illustration of the lack of robustness of supervised diffusion models to out-of-distribution operators, shown on the inpainting task using the FFHQ dataset (cf. \ref{['app:lack_robustness_ood_sft']}). (Middle) Reconstruction results for $\times4$ super-resolution and motion deblurring obtained via zero-shot posterior sampling on ImageNet, at different sampling times, highlighting the high inference cost of these methods. (Right) Our amortization strategy significantly accelerates inference compared to zero-shot posterior sampling (cf. \ref{['subsec:acceleration_id']}), while preserving flexibility to handle out-of-distribution operators (cf. \ref{['subsec:robustesse_ood']}).
  • Figure 2: Efficiency-quality trade-off for LAVPS vs. MGDM, on in-distribution degradation operators. Markers represent Pareto-optimal configurations for each method across various restoration tasks on the Imagenet dataset.
  • Figure 3: Motion deblurring on ImageNet. The inference time is constrained to be lower than $47.3$s, cf. \ref{['fig:mb_qualitative_47p3_appendix']} for more images.
  • Figure 4: Inpainting on ImageNet. The inference time is constrained to be lower than $34.9$s, cf. \ref{['fig:inpainting_qualitative_34p9_appendix']} for more images.
  • Figure 5: Trade-off between reconstruction quality (LPIPS) and inference time of MGDM janati2025a vs. DPS chung2023diffusion, for motion deblurring in FFHQ. Configurations of MGDM and DPS are Pareto-optimal.
  • ...and 5 more figures