Annealed Langevin Posterior Sampling (ALPS): A Rapid Algorithm for Image Restoration with Multiscale Energy Models

TL;DR

The paper introduces Annealed Langevin Posterior Sampling (ALPS), a rapid algorithm for image restoration that distills diffusion-model knowledge into multi-scale Energy-Based Models (EBMs). By formulating a family of static posteriors p_t(x|y) ∝ p(y|x) p_t(x) and applying preconditioned Langevin dynamics across annealed scales, ALPS achieves efficient MAP, MMSE, and uncertainty estimation for inverse problems such as inpainting and MRI reconstruction. Distillation preserves the probabilistic interpretability and compositionality of EBMs while retaining high-quality sample generation and posterior consistency, outperforming or matching diffusion-based baselines with fewer steps. The framework supports out-of-domain detection, model-mismatch diagnostics, and practical deployment considerations via principled posterior energies and tailored preconditioners. Overall, ALPS offers a scalable, principled approach to Bayesian imaging with competitive accuracy and reduced inference cost.

Abstract

Solving inverse problems in imaging requires models that support efficient inference, uncertainty quantification, and principled probabilistic reasoning. Energy-Based Models (EBMs), with their interpretable energy landscapes and compositional structure, are well-suited for this task but have historically suffered from high computational costs and training instability. To overcome the historical shortcomings of EBMs, we introduce a fast distillation strategy to transfer the strengths of pre-trained diffusion models into multi-scale EBMs. These distilled EBMs enable efficient sampling and preserve the interpretability and compositionality inherent to potential-based frameworks. Leveraging EBM compositionality, we propose Annealed Langevin Posterior Sampling (ALPS) algorithm for Maximum-A-Posteriori (MAP), Minimum Mean Square Error (MMSE), and uncertainty estimates for inverse problems in imaging. Unlike diffusion models that use complex guidance strategies for latent variables, we perform annealing on static posterior distributions that are well-defined and composable. Experiments on image inpainting and MRI reconstruction demonstrate that our method matches or surpasses diffusion-based baselines in both accuracy and efficiency, while also supporting MAP recovery. Overall, our framework offers a scalable and principled solution for inverse problems in imaging, with potential for practical deployment in scientific and clinical settings. ALPS code is available at the GitHub repository \href{https://github.com/JyoChand/ALPS}{ALPS}.

Figures (18)

• Figure 1: Overview of the proposed Annealed Langevin Posterior Sampling (ALPS) algorithm for inverse problems using multi-scale EBM regularizers. Please see Algorithm 1 and text for details. Current diffusion-based inverse problem solvers traverse through samples from annealed prior distributions $p_t({\boldsymbol{x}})$, which correspond to noise corrupted images. These schemes carefully add data consistency terms so that the trajectory does not deviate from the learned paths. The compositional property of EBMs enable us to define a family of time-dependent posterior distributions, each of which are sampled using preconditioned Langevin dynamics. The samples in these distributions are NOT noise perturbed versions from the time-dependent priors, but from well defined posterior distributions (bottom row marked by ${\boldsymbol{x}}_t \sim p_t({\boldsymbol{x}}|{\boldsymbol{y}}) :\propto p({\boldsymbol{y}}|{\boldsymbol{x}})p_t({\boldsymbol{x}})$). Breaking from the conventional diffusion sampling setting enables simpler algorithms that alternate between denoising using the score model at a specific scale to obtain ${{\boldsymbol{d}}_t}$ from ${\boldsymbol{x}}_t$, data consistency enforcement using quadratic optimization to obtain $\tilde{{\boldsymbol{x}}_t}$ from ${\boldsymbol{d}}_t$, and forward-model dependent noise addition (more noise in the null-space than range space of ${\mathbf{A}}$) to yield the posterior samples ${\boldsymbol{x}}_t$. The posterior samples ${\boldsymbol{x}}_{t}$ derived after $K$ iterations at one scale are used as initialization for Langevin dynamics at the next scale $p_{t+\delta}$, denoted by the curved yellow arrows. This approach results in a smoother trajectory from a good initialization ${\boldsymbol{x}}_0$ (top left corner) related to the least square solution to the final one, unlike most diffusion inverse solvers that start with pure noise. The EBM in study was distilled from a diffusion model, trained on the AFHQ ($64 \times 64$) dataset.
• Figure 2: Unconditional generation of multi-scale EBM using noise predictor on CIFAR-10 dataset at $32 \times 32$ resolution.
• Figure 3: Image inpainting using multi-scale EBMs: Illustration of recovery of images in the context of image inpainting problem for two different masks in (a) and (b). Top row in each figure shows the original image, the corresponding measurements, MAP, MMSE, and uncertainty estimates - whose values are higher in the masked regions. Second and the third rows presents several samples from the corresponding posterior distribution. We also show the negative log-posterior and negative log-prior values on top of each of the samples. Samples in Fig. (a) shows diversity in eyes and lips areas while the samples in Fig. (b) shows diversity in the eyes region.
• Figure 4: Comparison of performance of multi-scale energy models with diffusion models on the T2-weighted fastmri brain data set using (a) 4x Cartesian and (b) 8x Poisson undersampling masks. In each figure, first row shows the MAP, MMSE, and uncertainty estimates of the multi-scale EBM. Second and third rows shows the MMSE and uncertainty estimates obtained using DPS and DAPS algorithm, respectively. First image in the second row shows the mask employed for undersampling.
• Figure 5: Learned EBM using score-based learning from the moons dataset. In the top row we show the data samples (blue points) and the score-learned energy at different time-scales. At $t=\sigma_{max}$, the energy is roughly quadratic corresponding to a Gaussian prior. As time approaches to $0$, the energy becomes more representative of the true distribution. We also assume the measurements to be defined by the red straight line. The posterior energies at different time scales are shown in the bottom row; the maximum of the negative log-posterior is clipped for improved visualization. As time $t\rightarrow 0$, the posterior evolves from the likelihood to the true multi-modal posterior. The evolution of the samples offered by the ALPS algorithm is shown in Fig. \ref{['fig:alps_vs_daps']}