<figure><img src="image_1.png" /></figure>

HAL
open science

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

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

► To cite this version:

Léon Zheng, Thomas Hirtz, Yazid Janati, Eric Moulines. Fast and Robust Likelihood-Guided Diffusion Posterior Sampling with Amortized Variational Inference. 2026. hal-05498000

HAL Id: hal-05498000

https://hal.science/hal-05498000v1

Preprint submitted on 6 Feb 2026

HAL is a multi-disciplinary open access
archive for the deposit and dissemination of sci-
entific research documents, whether they are pub-
lished or not. The documents may come from
teaching and research institutions in France or
abroad, or from public or private research centers.

L'archive ouverte pluridisciplinaire HAL, est
destinée au dépôt et à la diffusion de documents
scientifiques de niveau recherche, publiés ou non,
émanant des établissements d'enseignement et de
recherche français ou étrangers, des laboratoires
publics ou privés.

<figure><img src="image_2.png" /></figure>

Distributed under a Creative Commons CC BY 4.0 - Attribution - International License
# Fast and Robust Likelihood-Guided Diffusion Posterior Sampling with
## Amortized Variational Inference

Léon Zheng<fnref n="1" /> Thomas Hirtz<fnref n="1" /> Yazid Janati² Eric Moulines³

## 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.

## 1. Introduction

An inverse problem aims to recover an unknown signal from incomplete, noisy observations. Such problems arise in a wide range of applications, including imaging tasks such as super-resolution, inpainting, and deblurring. These problems are typically difficult because they are ill-posed: a single perturbed observation may correspond to multiple plausible reconstructions. To overcome this difficulty, recent work has employed diffusion models (Ho et al., 2020; Nichol & Dhariwal, 2021) pre-trained on large-scale image datasets. These models serve as powerful priors that guide the solution of inverse problems towards reconstructions that are both realistic and consistent with the observed data.

A prominent line of work is zero-shot diffusion posterior sampling (Kadkhodaie & Simoncelli, 2021; Song et al.,

2021; Kawar et al., 2022; Lugmayr et al., 2022; Wang et al., 2023; Song et al., 2023; Chung et al., 2023), which encompasses a family of methods that require no additional training beyond a pre-trained diffusion model. At test time, these approaches aim to sample from the posterior distribution of reconstructions conditioned on the degraded observation, assuming that the underlying degradation process is known and can be expressed through a likelihood probability density function. By design, such methods are flexible and can accommodate arbitrary degradations at inference. This stands in contrast to supervised diffusion methods (Saharia et al., 2022a;b), which achieve strong performance on reconstruction tasks involving degradation operators seen during training on clean-degraded pairs, but typically struggle to generalize to unseen degradations during training.

The flexibility of current zero-shot diffusion posterior sampling methods comes, however, at the cost of increased inference-time computation. This overhead stems from the need to repeatedly evaluate an approximation of the so-called likelihood guidance (Chung et al., 2023; Song et al., 2023; Rozet et al., 2024) at each step of the reverse diffusion process in order to steer samples toward the target posterior distribution. Different zero-shot approaches mainly differ in how this likelihood guidance term is heuristically approximated, leading to varying trade-offs between inference computational cost and reconstruction quality.

To address the high inference cost of zero-shot posterior sampling, a natural idea is to leverage amortized optimization (Amos, 2023), which shifts computation from inference to an upstream training phase. However, existing amortized diffusion approaches (Lee et al., 2024; Feng et al., 2023; Feng & Bouman, 2024) typically replace likelihood-based posterior sampling with an implicit inference model. While this enables fast inference for degradations seen during training, it inherently ties the method to a fixed set of operators and severely limits robustness to unseen degradations.

This exposes a fundamental tension between efficiency and flexibility in diffusion-based inverse problems, and motivates the following question: Can zero-shot diffusion posterior sampling be accelerated through amortization, while preserving its ability to handle arbitrary (unseen) degradation operators at test time?

<footnote n="1">¹Huawei Lagrange Mathematics and Computing Research Center, Paris, France ²Institute of Foundation Models, MBZUAI ³Mohamed Bin Zayed University of AI, Chair of the ML Department, EPITA, Laboratoire Recherche de l'EPITA. Correspondence to: Léon Zheng <leon.zheng@polytechnique.org>.</footnote>
<figure><img src="image_3.png" /><figcaption>**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. Section C.1). (Middle) Reconstruction results for ×4 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. Section 5.1), while preserving flexibility to handle out-of-distribution operators (cf. Section 5.2).</figcaption></figure>

**Contributions.** We introduce a novel amortization strategy for diffusion posterior sampling that reconciles inference-time efficiency with the flexibility of zero-shot methods (Figure 1). Rather than replacing likelihood-based posterior sampling with a fully implicit inference model, we amortize the inner optimization problems arising at each timestep of so-called variational diffusion posterior sampling methods (Janati et al., 2025a; Moufad et al., 2025) (cf. Section 3 for more details), while preserving explicit likelihood guidance at test time. This design enables significant acceleration for degradation operators seen during training, without sacrificing the ability to handle arbitrary (and previously unseen) degradations at inference. As a result, our method, named LAVPS—standing for Likelihood-guided Amortized Variational Posterior Sampling—strictly improves over zero-shot approaches under limited inference budgets for in-distribution operators, while retaining their robustness to out-of-distribution degradations, where supervised and fully amortized diffusion methods typically fail. To the best of our knowledge, this is the first work to demonstrate that amortization and likelihood-guided diffusion can be combined to jointly achieve efficiency and robustness in diffusion-based inverse problems.

## 2. Background and Motivation

**Notations.** The Gaussian distribution of mean $\mu$ and covariance $\Sigma$ is denoted by $\mathcal{N}(\mu, \Sigma)$. Its density is denoted as $x \mapsto N(x; \mu, \Sigma)$. The identity matrix of size $d$ is denoted as $\mathbf{I}_d$. $\mathbf{X} \sim p$ indicates that $\mathbf{X}$ is a random variable whose probability distribution has density $p$. We use $\propto$ to define a probability density up to a normalization constant.

## 2.1. Diffusion Models for Inverse Problems

**Diffusion Models.** Denoising Diffusion Models (DDMs) (Sohl-Dickstein et al., 2015; Song & Ermon, 2019; Ho et al., 2020) define a generative process for a data distribution $p_0 := p_{\text{data}}$, initialized from a non-informative base distribution $p_T := \mathcal{N}(0, \mathbf{I}_d)$. The model then sequentially samples from increasingly structured intermediate distributions $p_t$ in reverse order until $p_0$ is recovered. The probability path $(p_t)_{t=0}^T$ is determined by the forward diffusion process: $p_t(\mathbf{x}_t) := \int q_{t|0}(\mathbf{x}_t|\mathbf{x}_0)p_0(\mathbf{x}_0) \, d\mathbf{x}_0$ where we define for all $s < t$,

$$q_{t|s}(\mathbf{x}_t|\mathbf{x}_s) := \mathcal{N}(\mathbf{x}_t; \alpha_{t|s}\mathbf{x}_s, \sigma_{t|s}^2 \mathbf{I}_d), \quad (1)$$

with $\alpha_{t|s} := \alpha_t/\alpha_s$, $\sigma_{t|s}^2 := \sigma_t^2 - \alpha_{t|s}^2 \sigma_s^2$ where $(\alpha_s)_{s=0}^T$, $\alpha_0 = 1$ and $(\sigma_t)_{t=0}^T$, $\sigma_0 = 0$ are positive sequences chosen so that the signal-to-noise ratio $(\alpha_t^2/\sigma_t^2)_{t>1}$ decreases monotonically. A popular choice of schedule is $\alpha_t = 1 - t/T$ and $\sigma_t = t/T$ (Lipman et al., 2023; Esser et al., 2024). For $s < t$, define the Gaussian transition $q_{s|0,t}(\mathbf{x}_s|\mathbf{x}_0, \mathbf{x}_t) \propto q_{s|0}(\mathbf{x}_s|\mathbf{x}_0)q_{t|s}(\mathbf{x}_t|\mathbf{x}_s)$ and the reverse transitions $(0 < s < t)$

$$p_{s|t}(\mathbf{x}_s|\mathbf{x}_t) := \int q_{s|0,t}(\mathbf{x}_s|\mathbf{x}_0, \mathbf{x}_t) p_{0|t}(\mathbf{x}_0|\mathbf{x}_t) \, d\mathbf{x}_0, \quad (2)$$

which can be shown to satisfy the recursion $p_s(\mathbf{x}_s) = \int p_{s|t}(\mathbf{x}_s|\mathbf{x}_t)p_t(\mathbf{x}_t) \, d\mathbf{x}_t$. Given a sample $\mathbf{X}_t \sim p_t$, drawing $\mathbf{X}_s \sim p_{s|t}(\cdot|\mathbf{X}_t)$ yields a sample from the more structured distribution $p_s$. However, these backward transition kernels are intractable due to the integration over the posterior $p_{0|t}(\cdot|\mathbf{x}_t)$. Following (Ho et al., 2020), the latter is approximated with a Dirac mass at $D_{0|t}^\theta(\mathbf{x}_t)$, the neural
network approximation of its expectation, i.e., the denoiser $\mathbb{E}[X_0|X_t = x_t]$ where $X_0 \sim p_0$ and $X_t \sim q_{t|0}(\cdot|X_0)$. This yields the transition approximation

$$p_{s|t}^{\theta}(x_s|x_t) := q_{s|0,t}(x_s|D_{0|t}^{\theta}(x_t), x_t). \quad (3)$$

The parameters $\theta$ of the denoiser are typically learned by minimizing a weighted $\ell_2$ denoising loss over several time steps:

$$\min_{\theta} \sum_{t=1}^{T} w_t \mathbb{E}[\|D_{0|t}^{\theta}(\mathbf{X}_t) - \mathbf{X}_0\|_2^2], \quad (4)$$

where $(w_t)_{t=1}^T$ denotes a sequence of non-negative weights. At inference time, the generative process starts with $\hat{\mathbf{X}}_T \sim \mathcal{N}(0, I_d)$ and iteratively samples $\hat{\mathbf{X}}_{t-1} \sim p_{t-1|t}(\cdot|\hat{\mathbf{X}}_t)$ for $t > 1$. Finally, $\hat{\mathbf{X}}_0 := D_{0|1}^{\theta}(\hat{\mathbf{X}}_1)$ is used as an approximate sample from $p_0$.

**Bayesian Inverse Problems.** In the Bayesian formulation of inverse problems, the following observation model is considered:

$$\mathbf{Y} = \mathcal{A}(\mathbf{X}_0) + \sigma_y \mathbf{W}, \quad \mathbf{X}_0 \sim p_0, \quad \mathbf{W} \sim \mathcal{N}(0, I_{d'}), \quad (5)$$

where $\mathcal{A}$ is a known degradation operator (also called forward model) acting on the space of the signal to recover (e.g., the pixel space for image inverse problems), $\sigma_y$ is the noise level and $\mathbf{X}_0$ and $\mathbf{W}$ are independent. The probability of observing $\mathbf{Y} = \mathbf{y}$ given $\mathbf{X}_0 = \mathbf{x}_0$ and a degradation operator $\mathcal{A}$ is defined as

$$\ell_0(\mathbf{y}|\mathbf{x}_0, \mathcal{A}) := \mathbb{N}(\mathbf{y}; \mathcal{A}(\mathbf{x}_0), \sigma_y^2 I_{d'}). \quad (6)$$

In the Bayesian paradigm, the reconstruction of $\mathbf{X}_0$ from an observation $\mathbf{Y} = \mathbf{y}$ and the known degradation operator $\mathcal{A}$ amounts to sampling from the posterior distribution associated with (5):

$$\pi_0(\mathbf{x}_0|\mathbf{y}, \mathcal{A}) \propto \ell_0(\mathbf{y}|\mathbf{x}_0, \mathcal{A}) p_0(\mathbf{x}_0). \quad (7)$$

## 2.2. Diffusion Posterior Sampling

The core idea of zero-shot diffusion posterior sampling algorithms (Daras et al., 2024; Janati et al., 2025b; Chung et al., 2025) is to approximate the dynamics of a Denoising Diffusion Model (DDM) targeting the posterior distribution (7) by suitably adapting, at inference time, a DDM trained on the prior $p_0$. A key property of these methods is that they require no additional training on data generated from the observation model (5), beyond the training (4) of the DDM model targeting the prior $p_0$. Constructing a DDM for the posterior distribution $\pi_0(\cdot|\mathbf{y}, \mathcal{A})$ for a given known $\mathbf{y}$ and $\mathcal{A}$ amounts to approximately sampling along the distributional path $(\pi_t(\cdot|\mathbf{y}, \mathcal{A}))_{t=0}^T$, defined analogously to the unconditional probability path $(p_t)_{t=0}^T$:

$$\pi_t(\mathbf{x}_t|\mathbf{y}, \mathcal{A}) := \int q_{t|0}(\mathbf{x}_t|\mathbf{x}_0)\pi_0(\mathbf{x}_0|\mathbf{y}, \mathcal{A}) \mathrm{d}\mathbf{x}_0, \quad (8)$$

This is achieved by having a parametric approximation of the posterior denoiser $\mathbb{E}[X_0 | X_t = x_t, Y = \mathbf{y}]$, where $X_0 \sim p_0$, $Y \sim \ell_0(\cdot|X_0, \mathcal{A})$, and $X_t \sim q_{t|0}(\cdot|X_0)$. It satisfies the identity (Daras et al., 2024, Eq. 2.15 and 2.17)

$$\mathbb{E}[X_0 | X_t = x_t, Y = \mathbf{y}] = \mathbb{E}[X_0 | X_t = x_t] + \frac{\sigma_t^2}{\alpha_t} \nabla_{x_t} \log \ell_t(\mathbf{y}|x_t, \mathcal{A}), \quad (9)$$

with

$$\ell_t(\mathbf{y}|\mathbf{x}_t, \mathcal{A}) := \int \ell_0(\mathbf{y}|\mathbf{x}_0, \mathcal{A}) p_{0|t}(\mathbf{x}_0|\mathbf{x}_t) \mathrm{d}\mathbf{x}_0. \quad (10)$$

Hence, assuming that we already have pre-trained a neural network $D_{0|t}^{\theta}(\mathbf{x}_t)$ approximating the unconditional denoiser $\mathbb{E}[X_0 | X_t = x_t]$, estimating the posterior denoiser $\mathbb{E}[X_0 | X_t = x_t, Y = \mathbf{y}]$ reduces to approximating the intractable log-gradient on the right-hand side of (9). A widely used approximation (Ho et al., 2022; Chung et al., 2023) replaces $p_{0|t}(\cdot|\mathbf{x}_t)$ in the integral by a Dirac mass centered at the prior denoiser $D_{0|t}^{\theta}(\mathbf{x}_t)$, leading to

$$\nabla_{x_t} \log \ell_t(\mathbf{y}|\mathbf{x}_t, \mathcal{A}) \approx \nabla_{x_t} \log \hat{\ell}_t(\mathbf{y}|\mathbf{x}_t, \mathcal{A}) \quad (11)$$

with

$$\hat{\ell}_t(\mathbf{y}|\mathbf{x}_t, \mathcal{A}) := \ell_0(\mathbf{y}|D_{0|t}^{\theta}(\mathbf{x}_t), \mathcal{A}). \quad (12)$$

This likelihood-guided correction is the key source of flexibility of zero-shot methods, but also their dominant computational bottleneck, and will be the focus of the amortization strategy introduced in the next section.

## 2.3. Supervised Diffusion and Fully Amortized Methods

To reduce inference-time computation, an alternative to zero-shot posterior sampling is to rely on an upstream training phase, as done in supervised diffusion and fully amortized inference methods. While these approaches enable fast inference, they fundamentally trade flexibility for efficiency, and typically fail to generalize to degradation operators not seen during training.

**Supervised Diffusion.** Supervised diffusion methods (Saharia et al., 2022b;a) directly learn a neural approximation of the posterior denoiser $\mathbb{E}[X_0 | X_t = x_t, Y = \mathbf{y}]$ by training a conditional denoiser on paired clean-degraded data. Inference then proceeds via ancestral sampling using learned, observation-conditioned backward transitions. Recent extensions condition the denoiser additionally on the degradation operator $\mathcal{A}$ to handle multiple operators within a single model (Elata et al., 2025). However, since the likelihood correction is entirely internalized by the network during training, these methods remain tightly coupled to the distribution of operators seen during training and exhibit poor robustness to unseen degradations (see Section C.1 for an experimental illustration).
**Fully Amortized Methods.** Fully amortized approaches (Lee et al., 2024; Feng et al., 2023; Feng & Bouman, 2024) rely on variational inference to learn an implicit distribution $r_\phi(\cdot|y, \mathcal{A})$ that approximates the target posterior $\pi_0(\cdot|y, \mathcal{A})$. By parameterizing this distribution as the push-forward of a simple base distribution through a neural network, they enable one-step inference at test time. However, similarly to supervised diffusion, the absence of explicit likelihood guidance at inference makes these methods highly dependent on the training distribution of degradation operators, resulting in limited out-of-distribution generalization.

In contrast to both supervised diffusion and fully amortized inference, our approach retains explicit likelihood-guided posterior sampling at test time, thereby preserving robustness to arbitrary, previously unseen degradation operators.

## 3. Method

We now present our approach to accelerate diffusion posterior sampling while preserving out-of-distribution robustness. Our key insight is to introduce amortization not at the level of the posterior itself, but within the inner variational inference problems solved at each timestep of so-called variational diffusion posterior sampling methods. This design leverages upstream training to reduce inference-time computation while retaining explicit likelihood guidance at test time, yielding an amortization strategy that, to our knowledge, has not been explored in prior work.

### 3.1. Variational Diffusion Posterior Sampling

Variational diffusion posterior sampling (Janati et al., 2025a; Moufad et al., 2025) denotes a class of zero-shot posterior sampling methods that rely on variational inference to approximate intractable intermediate distributions—such as selected conditional or marginal distributions arising along the diffusion trajectory—thereby enabling approximate backward diffusion toward the target posterior while retaining explicit likelihood guidance. To illustrate our approach, we focus on a representative instance of this class, MGDM (Janati et al., 2025a), but emphasize that the amortization strategy proposed in this work applies more broadly to variational diffusion posterior sampling methods.

**Variational Inference in MGDM.** In Janati et al. (2025a), the key idea is to approximate the marginal distribution $\pi_t(\cdot|y, \mathcal{A})$ in (8) by a mixture $\hat{\pi}_t(\cdot|y, \mathcal{A})$ of midpoint-based approximations $\{\hat{\pi}_t^s(\cdot|y, \mathcal{A})\}_{s=1}^{t-1}$, i.e., $\pi_t(x_t|y, \mathcal{A}) \approx \hat{\pi}_t(x_t|y, \mathcal{A}) := \sum_{s=1}^{t-1} \omega_t^s \hat{\pi}_t^s(x_t|y, \mathcal{A})$, where $(\omega_t^s)_{s=1}^{t-1}$ are non-negative weights summing to 1, and

$$ \hat{\pi}_t^s(x_t|y, \mathcal{A}) \propto \left( \int \hat{\ell}_s(y|x_s, \mathcal{A}) p_{s|t}(x_s|x_t) dx_s \right) p_t(x_t), \quad (13) $$

is an approximation of $\pi_t(x_t|y, \mathcal{A})$ based on a midpoint timestep $s < t$. In Janati et al. (2025a), approximate sampling from the target posterior distribution $\pi_0(\cdot|y, \mathcal{A})$ is performed sequentially through the intermediate distributions $\hat{\pi}_T(\cdot|y, \mathcal{A})$, $\hat{\pi}_{T-1}(\cdot|y, \mathcal{A})$, ..., $\hat{\pi}_0(\cdot|y, \mathcal{A})$. At each step $t$, sampling from the mixture $\hat{\pi}_t(\cdot|y, \mathcal{A})$ proceeds by first drawing a midpoint index $s < t$ according to the categorical distribution with weights $(\omega_t^s)_{s=1}^{t-1}$, and then sampling from the corresponding component $\hat{\pi}_t^s(\cdot|y, \mathcal{A})$ in (13) using a Gibbs sampler with the following data augmentation scheme:

$$ \bar{\pi}_{0,s,t}(x_0, x_s, x_t | y, \mathcal{A}) \propto p_{0|s}(x_0|x_s) \hat{\ell}_s(y|x_s, \mathcal{A}) p_{s|t}(x_s|x_t) p_t(x_t). \quad (14) $$

By construction, the $x_t$-marginal of this joint distribution is exactly $\hat{\pi}_t^s(\cdot|y, \mathcal{A})$. The Gibbs sampler (Gelfand & Smith, 1990) repeats for $R$ steps the following updates: given $(X_0^r, X_s^r, X_t^r)$,

1. Sample $X_s^{r+1}$ from

$$ \bar{\pi}_{s|0,t}(x_s|X_0^r, X_t^r, y, \mathcal{A}) \propto \hat{\ell}_s(y|x_s, \mathcal{A}) q_{s|0,t}(x_s|X_0^r, X_t^r). \quad (15) $$

In practice, this step is implemented using a Gaussian variational distribution, as we describe below. This approximation is precisely what places the method of Janati et al. (2025a) within the class of variational posterior sampling approaches.

2. Sample $X_0^{r+1}$ from $p_{0|s}(\cdot|X_s^{r+1})$ using the approximate reverse process of the pre-trained DDM that targets the prior $p_0$, based on the kernels (3).

3. Sample $X_t^{r+1}$ from $q_{t|s}(\cdot|X_s^{r+1})$, which is a forward Gaussian noising process (1) in DDMs.

**Variational Approximation.** Given the context variable $c := (x_0, x_t, s, t, y, \mathcal{A})$, the variational approximation consists in minimizing

$$ \begin{aligned} & \min_{\mu, \rho} \mathcal{L}(\mu, \rho; c) \\ & := D_{\text{KL}}(\text{N}(\cdot; \mu, \text{diag}(\rho)) \| \bar{\pi}_{s|0,t}(\cdot|x_0, x_t, y, \mathcal{A})), \end{aligned} \quad (16) $$

where $D_{\text{KL}}$ denotes the Kullback-Leibler divergence, $\bar{\pi}_{s|0,t}(x_s|x_0, x_t, y, \mathcal{A})$ is defined in (15), and $\text{diag}(\rho)$ is the diagonal matrix with entries given by the vector $\rho$. This KL divergence is in practice estimated using the reparameterization trick (Kingma & Welling, 2013) with a single Monte Carlo sample. Achieving a sufficiently small value of this objective is critical for obtaining high-quality reconstructions in variational diffusion posterior sampling algorithms (Janati et al., 2025a). However, in practice this minimization
is performed at inference time using first-order optimization, which typically requires multiple gradient updates and costly backpropagations through the pre-trained unconditional denoiser $D_{0|t}^\theta$. The repeated numerical resolution of this variational optimization problem at inference time constitutes the main computational bottleneck of MGDM, and is the target of the amortization strategy introduced next.

## 3.2. Amortized Variational Inference

We propose an amortized optimization approach (Amos, 2023) that transfers the burden of minimizing the variational objective (16) from inference to a dedicated upstream training phase.

**Inference Model Training.** The inference model is defined as $c \mapsto \lambda_\varphi(c) := (\mu_\varphi(c), \rho_\varphi(c))$, with parameters $\varphi$, to predict a good initialization to address the variational approximation problem (16), given a new context $c$. The model is implemented using a single forward pass through a neural network and trained by minimizing the objective-based loss (Amos, 2023):

$$ \min_{\varphi} \mathbb{E}_c [\mathcal{L}(\mu_{\varphi}(c), \rho_{\varphi}(c); c)], \quad (17) $$

where $\mathcal{L}$ is the variational objective from (16). The training loss (17) is minimized by stochastic gradient descent. The context $c := (x_0, x_t, s, t, y, A)$ in the expectation of (17) is a random variable drawn from a distribution designed to match those encountered at inference, i.e., during posterior sampling with MGDM. We find the following sampling procedure for $c$ to be effective in practice:

$$ (x, y, A) \sim p_{\text{data}}(x) \ell_0(y|x, A) p_{\text{op}}(A), \quad (18) $$

$$ t \sim \text{Unif}(\mathcal{T}), \quad s | t \sim \text{Categorical}((\omega_t^{s'})_{s'=1}^{t-1}), $$

$$ x_t | x \sim q_{t|0}(\cdot|x), \quad x_0 := D_{0|t}^\theta(x_t). \quad (19) $$

where $p_{\text{data}}(x)$ denotes the distribution of clean images, $p_{\text{op}}(A)$ denotes the distribution of degradation operators modeled during training, and $\mathcal{T} \subseteq \{1, \dots, T\}$ denotes the subset of timesteps on which the inference model is applied during test-time inference.

**Inference Model Parameterization.** The inference model aims to predict a near minimizer of problem (16), so it is natural to start from the initialization used in MGDM (Janati et al., 2025a), where the optimization is initialized with the statistics of the Gaussian bridge transition $q_{s|0,t}(x_s|x_0, x_t) \propto q_{s|0}(x_s|x_0)q_{t|s}(x_t|x_s)$ and then refined by gradient descent. To mimic this procedure in a single forward pass, we parameterize the network so that it learns the *residual* between these prior statistics and the posterior parameters that approximately minimize (16).

**Inference Model Architecture.** Concerning the architecture of this neural network, note that minimizing the variational objective (16) amounts to approximately sampling from the conditional law $\bar{\pi}_{s|0,t}(x_s|x_0, x_t, y, A) \propto \hat{\ell}_s(y|x_s, A) q_{s|0,t}(x_s|x_0, x_t)$ defined in (15), given a context $c = (x_0, x_t, s, t, y, A)$. In other words, the task is to denoise from $x_t$ down to an intermediate state $x_s$ with $s < t$, while conditioning on $x_0$ and, crucially, on the observation $y$ and the operator $A$. This conditioning naturally suggests drawing inspiration from the conditional denoisers used in supervised diffusion models. To efficiently account for multiple degradations during training, we recommend adopting the InvFussion (Elata et al., 2025) architecture, which is expressly designed to handle external inputs such as $y$ and $A$. Therefore, the neural-network architecture of our inference model can be derived from the backbone of any existing conditional denoisers, with the following minor modifications: (i) concatenate $x_0, x_t$ along the channel dimension; (ii) early fusion of the embedding of the timesteps $t$ and $s$; (iii) double the number of output channels so the network jointly predicts variational parameters that approximately minimize (16).<fnref n="1" /> Further details are deferred to Section A.

## 3.3. Sampling with Amortized Initialization

We now summarize the new sampling algorithm, called LAVPS (Likelihood-guided Amortized Variational Posterior Sampling), as it uniquely combines amortization and explicit likelihood-guidance.

**Amortized Warm-start with Likelihood Guidance.** At each diffusion timestep, we use the trained inference model to warm-start the optimization of (16). Given a context $c := (x_0, x_t, s, t, y, A)$, the model predicts $(\mu_\varphi(c), \rho_\varphi(c))$ in a single forward pass, which is used to initialize a first-order optimizer. This initialization is tailored to the inverse problem at hand by explicitly accounting for the observation $y$ and degradation operator $A$, and allows the variational objective (16) to be minimized using only a small number of likelihood-guided gradient steps, without compromising final reconstruction quality.

**Safeguarded Initialization for Robustness.** To preserve robustness when the degradation operator at test time is out-of-distribution, we introduce a simple safeguard at initialization. Specifically, we compare the variational objective evaluated at the warm-start initialization, denoted by $\mathcal{L}_{\text{ws}}$, with its value under the standard zero-shot initialization used in MGDM (Janati et al., 2025a), i.e., an initialization derived from the unconditional prior bridge and independent of the observation $y$ and degradation operator $A$, denoted by $\mathcal{L}_{\text{zs}}$.

<footnote n="1">¹These minor changes are compatible with all existing standard conditional denoiser architectures, e.g., UNet (Ho et al., 2020), DiT (Peebles & Xie, 2023), HDiT (Crowson et al., 2024).</footnote>
<figure><img src="image_4.png" /><figcaption>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.</figcaption></figure>

If $\mathcal{L}_{ws} > \mathcal{L}_{zs}$, indicating that the inference model may be unreliable, we discard the warm-start and revert to the zero-shot initialization. In practice, our experiments show that this fallback mechanism is sufficient to ensure robustness to arbitrary degradation operators, including those unseen during training, and guarantees performance no worse than that of MGDM in terms of reconstruction quality.

# 4. Related Work

Two recent works (Elata et al., 2025; Mbakam et al., 2025) share with us the idea of training a component that can be reused at inference to remain flexible to different likelihoods (equivalently, different degradation operators), but they differ in how the likelihood is ultimately exploited: they predict a likelihood-conditioned denoised sample directly through a network that takes the likelihood parameters as input. Elata et al. (2025) adapt deep unrolling (Gregor & LeCun, 2010) by training a conditional diffusion model with an operator-aware block that applies the degradation operator $\mathcal{A}$ and its pseudo-inverse directly to the activations, enabling the network to predict a likelihood-conditioned denoised sample. Mbakam et al. (2025) unfold the LATINO Langevin sampler (Spagnoletti et al., 2025) into a sequence of learnable steps and train it with a consistency trajectory loss so that the network likewise predicts a likelihood-conditioned denoised sample. Because these methods rely on the network to internalize the likelihood correction, they may require large paired datasets and generalize poorly to degradation operators not covered during training.

In contrast, our approach keeps likelihood-guided optimization explicit: the inference model provides an approximate initialization, while the final correction is obtained by directly minimizing the variational objective (16) under the true test-time likelihood. This combination of amortized initialization and explicit likelihood guidance enables both fast inference and improved robustness to previously unseen operators, as demonstrated in the next section.

# 5. Experiments

Our protocol follows Elata et al. (2025). The unconditional denoiser $D_{0It}^{\theta}$ is an HDiT (Crowson et al., 2024) denoiser. We validate our method both on facial images using the FFHQ dataset (Karras et al., 2019) and natural images using the ImageNet (Deng et al., 2009), at the resolution of 256 x 256. Methods are evaluated on diverse image restoration tasks: super-resolution, inpainting, and motion deblurring. The noise level in (5) is fixed at $\sigma_y = 0.05$. The evaluation is performed on 300 test images, unseen during training or hyperparameter tuning. The reconstruction quality is measured with per-image metrics comparing the similarity between the reconstructed image and a reference one, such as LPIPS (Zhang et al., 2018), SSIM (Wang et al., 2004) and PSNR. We report the average of these metrics with their standard deviations over all test images. To measure the photorealism of the reconstructed images, we also report the CMMD<fnref n="2" /> (Jayasumana et al., 2024), which is particularly appropriate for inpainting tasks, where multiple valid reconstructions exist and per-instance fidelity metrics are inherently ambiguous. All details about the experimental protocol are deferred to Section B.

## 5.1. Acceleration for In-distribution Operators

We first show that LAVPS improves the trade-off between reconstruction quality and inference speed compared to zero-shot posterior sampling methods, thanks to the proposed amortization strategy. To isolate this effect, we first consider a simplified setting in which amortization and evaluation are performed on the *same* degradation operator. The more general scenario—where amortization is carried out over a *family* of operators and robustness to out-of-distribution degradations is assessed—is deferred to Section 5.2.

<footnote n="2">²With only 300 images, we found FID (Heusel et al., 2017) to be unstable and poorly aligned with perceptual quality. We therefore rely on CMMD, which is more sample-efficient and better suited to our experiments.</footnote>
<table><caption>**Table 1.** Inference speedup of LAVPS vs. MGDM for matched reconstruction quality levels on FFHQ and ImageNet (256 × 256). We report the sampling time required to reach a specific quality target (LPIPS, or CMMD denoted by †) measured on a single NVIDIA V100.</caption><thead><tr><th>Task</th><th>Dataset</th><th>Target Quality ↓</th><th>Inference time of MGDM</th><th>Speedup with LAVPS</th></tr></thead><tbody><tr><td rowspan="3">SR × 4</td><td rowspan="2">FFHQ</td><td>0.102</td><td>29.3s</td><td>1.73x</td></tr><tr><td>0.096</td><td>44.5s</td><td>1.41x</td></tr><tr><td>ImageNet</td><td>0.275<br>0.261</td><td>66.1s<br>86.1s</td><td>2.01x<br>2.15x</td></tr><tr><td rowspan="3">Deblurring</td><td rowspan="2">FFHQ</td><td>0.106</td><td>19.2s</td><td>1.27x</td></tr><tr><td>0.062</td><td>48.3s</td><td>1.25x</td></tr><tr><td>ImageNet</td><td>0.257<br>0.237</td><td>39.6s<br>43.8s</td><td>1.20x<br>1.34x</td></tr><tr><td rowspan="3">Inpainting</td><td rowspan="3">ImageNet</td><td>0.174†</td><td>46.2s</td><td>1.43x</td></tr><tr><td>0.133†</td><td>65.9s</td><td>2.03x</td></tr><tr><td>0.116†</td><td>86.0s</td><td>2.21x</td></tr></tbody></table>

As shown by Janati et al. (2025a), MGDM is a strong baseline that achieves a good trade-off between quality and speed among zero-shot posterior sampling methods (cf. Section C.2 for details). Therefore, our goal here is to show that LAVPS achieves a better trade-off than MGDM.

**Pareto Frontier Analysis.** Figure 2 shows that LAVPS consistently shifts the Pareto frontier of MGDM downward for the ImageNet dataset, on super-resolution ×4, motion deblurring with a kernel size 21 × 21, and the inpainting task where the mask is a center box of size 128 × 128. Section C.4 shows the same conclusion for the FFHQ dataset. The gains of LAVPS are most pronounced in the low-to-medium computational budget regime<fnref n="3" />, where amortized initialization enables high-fidelity reconstructions at substantially reduced sampling cost. This is illustrated in Figure 1, where LAVPS recovers crisp high-frequency details—such as the bird’s feathers and the shark’s teeth—whereas the zero-shot baseline produces significantly blurred or noisy reconstructions that require considerably more computational budget to reach a similar level of fidelity. See Section C.5 for an ablation study analyzing the role of amortized initialization in reducing the number of inner-loop gradient steps required to achieve high-quality reconstructions.

**Quantifying Speedups.** In Table 1, we quantify the acceleration of LAVPS at inference compared to MGDM, when considering different reconstruction quality target along the Pareto frontier of MGDM. Specifically, we identify the fastest LAVPS configuration that matches the quality of representative points from the baseline, validating performance

<table><caption>**Table 2.** Motion deblurring on ImageNet: Improvement on in-distribution while keeping out-of-distribution robustness.</caption><thead><tr><th rowspan="2">Method</th><th colspan="2">In-distribution operators<br>(kernel size 21 × 21)</th><th colspan="2">Out-of-distribution operators<br>(kernel size 35 × 35)</th></tr><tr><th>Time (s)</th><th>LPIPS ↓</th><th>Time (s)</th><th>LPIPS ↓</th></tr></thead><tbody><tr><td colspan="5"><strong>Inference time ≤ 35.1 s</strong></td></tr><tr><td>MGDM</td><td>35.1</td><td>0.250</td><td>35.1</td><td><strong>0.278</strong></td></tr><tr><td>LAVPS</td><td><strong>33.2</strong></td><td><strong>0.236</strong></td><td><strong>34.1</strong></td><td>0.279</td></tr><tr><td colspan="5"><strong>Inference time ≤ 47.3 s</strong></td></tr><tr><td>MGDM</td><td>45.9</td><td>0.210</td><td>39.5</td><td>0.256</td></tr><tr><td>LAVPS</td><td><strong>40.3</strong></td><td><strong>0.188</strong></td><td><strong>37.5</strong></td><td><strong>0.238</strong></td></tr></tbody></table>

<table><caption>**Table 3.** Inpainting on ImageNet: Improvement on in-distribution while keeping out-of-distribution robustness.</caption><thead><tr><th rowspan="2">Method</th><th colspan="2">In-distribution operators<br>(missing rectangles)</th><th colspan="2">Out-of-distribution operators<br>(missing pixels)</th></tr><tr><th>Time (s)</th><th>CMMD ↓</th><th>Time (s)</th><th>CMMD ↓</th></tr></thead><tbody><tr><td colspan="5"><strong>Inference time ≤ 34.9 s</strong></td></tr><tr><td>MGDM</td><td>34.8</td><td>0.199</td><td>34.9</td><td>0.336</td></tr><tr><td>LAVPS</td><td><strong>32.3</strong></td><td><strong>0.124</strong></td><td><strong>32.7</strong></td><td><strong>0.335</strong></td></tr><tr><td colspan="5"><strong>Inference time ≤ 102.6 s</strong></td></tr><tr><td>MGDM</td><td>66.3</td><td>0.113</td><td>102.6</td><td><strong>0.130</strong></td></tr><tr><td>LAVPS</td><td><strong>43.8</strong></td><td><strong>0.105</strong></td><td><strong>93.2</strong></td><td>0.131</td></tr></tbody></table>

parity through a statistical non-inferiority test. This procedure ensures that the reported speedups represent true efficiency gains without sacrificing reconstruction quality. Our results demonstrate systematic speedups across all tasks:

* On FFHQ, LAVPS achieves targeted qualities 1.25× to 1.73× faster for motion deblurring and super-resolution.

* On ImageNet, evaluations show higher speedups, up to 2.03× for inpainting and up to 2.15× for super-resolution.

These inference-time gains come at the cost of an additional amortization training phase, whose computational overhead is quantified in Section C.3. In practice, this cost is modest—only a few hours on a single V100 GPU—making the resulting inference speedups highly favorable.

## 5.2. Robustness to Out-of-distribution Operators

We now show that the improvement of the trade-off between quality and inference speed with LAVPS does not come at the expense of losing the robustness to out-of-distribution degradation operators.

Our protocol consists in training an inference model by amortizing over a family of degradation operators, extending the setting of Section 5.1, where amortization was performed for a single fixed operator. Formally, during training, operators are sampled from a fixed distribution $p_{op}(\mathcal{A})$, which we refer to as the in-distribution setting. Degradation operators encountered at test time that fall outside this distribution are referred to as out-of-distribution. At inference time, we are given a reconstruction task for which the degradation operator $\mathcal{A}$ is known, consistent with standard

<footnote n="3">³As the inference budget increases, the performance of both methods converges, which is expected: given sufficient compute, the iterative optimization in MGDM can eventually compensate for its coarse initialization and reach the same variational optimum as LAVPS.</footnote>
<figure><img src="image_5.png" /><figcaption>Figure 4. Inpainting on ImageNet. The inference time is constrained to be lower than 34.9s, cf. Figure 10 for more images.</figcaption></figure>

assumptions in zero-shot diffusion posterior sampling. However, we do not assume access to information about whether $\mathcal{A}$ is in-distribution or not, as is standard when evaluating robustness to distributional shift. Evaluation of LAVPS and MGDM are then performed on diverse degradation operators (in-distribution or out-of-distribution).

**Tasks.** We consider two reconstruction tasks on ImageNet. The first task is motion deblurring, where degradation operators are convolutions with kernels generated from random motion trajectories whose length and curvature are controlled by an intensity parameter.<fnref n="4" /> In-distribution kernels have size 21 × 21 and intensity 0.9, while out-of-distribution kernels have size 35 × 35 and intensity 0.3. The second task is inpainting, where in-distribution masks consist of rectangular regions at random locations, with width and height uniformly sampled between 0.4 and 0.6 times the image resolution, and out-of-distribution masks are generated by independently removing individual pixels according to a Bernoulli distribution with parameter $p \in [0.8, 0.85]$.

**Results.** When amortizing over a family of degradation operators, LAVPS continues to achieve a better trade-off between reconstruction quality and inference time than MGDM for in-distribution operators. This is shown by the Pareto fronts reported in Figure 8 (Section C.6) and is consistent with the fixed-operator setting of Section 5.1,

showing that the efficiency gains of amortization are preserved when training over multiple operators rather than a single one.

Tables 2 and 3 provide quantitative comparisons under fixed inference-time budgets for both in-distribution and out-of-distribution operators, with an illustration in Figures 3 and 4. The conclusion is that LAVPS consistently improves reconstruction quality over MGDM on in-distribution operators for comparable or lower inference times, and the performance on out-of-distribution kernels remains comparable to MGDM across budgets, with no systematic degradation. Overall, these results indicate that amortizing the inner optimization in variational diffusion posterior sampling improves efficiency on in-distribution operators without compromising the robustness to previously unseen degradation operators that characterizes zero-shot posterior sampling.

## 6. Conclusion

We introduced a novel amortization strategy that accelerates diffusion posterior sampling methods while preserving their ability to handle arbitrary degradation operators at test time. Our approach amortizes the inner variational inference problems arising in variational diffusion posterior sampling, enabling a principled combination of amortization with explicit likelihood guidance. Overall, this work demonstrates that amortization and likelihood-guided sampling can be effectively reconciled, paving the way for diffusion-based inverse problem solvers that are both fast at inference and flexible across a wide range of degradation operators.

<footnote n="4">⁴Kernels are rendered as piecewise-smooth paths, smoothed, downsampled, and normalized to form valid point spread functions.</footnote>
**Impact Statement**

This paper presents work whose goal is to advance the field of machine learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.

**References**

Amos, B. Tutorial on amortized optimization. *Foundations and Trends in Machine Learning*, 16(5):592–732, 2023.

Chung, H., Kim, J., Mccann, M. T., Klasky, M. L., and Ye, J. C. Diffusion posterior sampling for general noisy inverse problems. In *ICLR*, 2023.

Chung, H., Kim, J., and Ye, J. C. Diffusion models for inverse problems. *arXiv preprint arXiv:2508.01975*, 2025.

Crowson, K., Baumann, S. A., Birch, A., Abraham, T. M., Kaplan, D. Z., and Shippole, E. Scalable high-resolution pixel-space image synthesis with hourglass diffusion transformers. In *ICML*, 2024.

Daras, G., Chung, H., Lai, C.-H., Mitsufuji, Y., Ye, J. C., Milanfar, P., Dimakis, A. G., and Delbracio, M. A survey on diffusion models for inverse problems. *arXiv preprint arXiv:2410.00083*, 2024.

Deng, J., Dong, W., Socher, R., Li, L.-J., Li, K., and Fei-Fei, L. Imagenet: A large-scale hierarchical image database. In *CVPR*, 2009.

Elata, N., Chung, H., Ye, J. C., Michaeli, T., and Elad, M. InvFussion: Bridging supervised and zero-shot diffusion for inverse problems. *arXiv preprint arXiv:2504.01689*, 2025.

Esser, P., Kulal, S., Blattmann, A., Entezari, R., Müller, J., Saini, H., Levi, Y., Lorenz, D., Sauer, A., Boesel, F., Podell, D., Dockhorn, T., English, Z., Lacey, K., Goodwin, A., Marek, Y., and Rombach, R. Scaling rectified flow transformers for high-resolution image synthesis. In *ICML*, 2024.

Feng, B. and Bouman, K. Variational bayesian imaging with an efficient surrogate score-based prior. *Transactions on Machine Learning Research*, 2024.

Feng, B. T., Smith, J., Rubinstein, M., Chang, H., Bouman, K. L., and Freeman, W. T. Score-based diffusion models as principled priors for inverse imaging. In *ICCV*, 2023.

Gelfand, A. E. and Smith, A. F. M. Sampling-based approaches to calculating marginal densities. *Journal of the American Statistical Association*, 85(410):398–409, 1990.

Gregor, K. and LeCun, Y. Learning fast approximations of sparse coding. In *ICML*, 2010.

Heusel, M., Ramsauer, H., Unterthiner, T., Nessler, B., and Hochreiter, S. Gans trained by a two time-scale update rule converge to a local nash equilibrium. In *NeurIPS*, 2017.

Ho, J., Jain, A., and Abbeel, P. Denoising diffusion probabilistic models. *NeurIPS*, 2020.

Ho, J., Salimans, T., Gritsenko, A., Chan, W., Norouzi, M., and Fleet, D. J. Video diffusion models. *NeurIPS*, 2022.

Janati, Y., Moufad, B., Abou El Qassime, M., Durmus, A. O., Moulines, E., and Olsson, J. A mixture-based framework for guiding diffusion models. In *ICML*, 2025a.

Janati, Y., Moulines, E., Olsson, J., and Oliviero-Durmus, A. Bridging diffusion posterior sampling and Monte Carlo methods: A survey. *Philosophical Transactions A*, 383 (2299):20240331, 2025b. doi: 10.1098/rsta.2024.0331.

Jayasumana, S., Ramalingam, S., Veit, A., Glasner, D., Chakrabarti, A., and Kumar, S. Rethinking FID: Towards a better evaluation metric for image generation. In *CVPR*, 2024.

Kadkhodaie, Z. and Simoncelli, E. P. Stochastic solutions for linear inverse problems using the prior implicit in a denoiser. In *NeurIPS*, 2021.

Karras, T., Laine, S., and Aila, T. A style-based generator architecture for generative adversarial networks. In *CVPR*, 2019.

Kawar, B., Elad, M., Ermon, S., and Song, J. Denoising diffusion restoration models. In *NeurIPS*, 2022.

Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. In *ICLR*, 2015.

Kingma, D. P. and Welling, M. Auto-encoding variational bayes. *arXiv preprint arXiv:1312.6114*, 2013.

Lee, S., Park, D., Kong, I., and Kim, H. J. Diffusion prior-based amortized variational inference for noisy inverse problems. In *ECCV*, 2024.

Lipman, Y., Chen, R. T. Q., Ben-Hamu, H., Nickel, M., and Le, M. Flow matching for generative modeling. In *ICLR*, 2023.

Loshchilov, I. and Hutter, F. Decoupled weight decay regularization. In *ICLR*, 2019.

Lugmayr, A., Danelljan, M., Romero, A., Yu, F., Timofte, R., and Van Gool, L. Repaint: Inpainting using denoising diffusion probabilistic models. In *CVPR*, 2022.
Mbakam, C. K., Spence, J., and Pereyra, M. Learning few-step posterior samplers by unfolding and distillation of diffusion models. *arXiv preprint arXiv:2507.02686*, 2025.

Moufad, B., Janati, Y., Bedin, L., Durmus, A. O., Douc, R., Moulines, E., and Olsson, J. Variational diffusion posterior sampling with midpoint guidance. In *ICLR*, 2025.

Nichol, A. Q. and Dhariwal, P. Improved denoising diffusion probabilistic models. In *ICML*, 2021.

Peebles, W. and Xie, S. Scalable diffusion models with transformers. In *CVPR*, 2023.

Rozet, F., Andry, G., Lanusse, F., and Louppe, G. Learning diffusion priors from observations by expectation maximization. In *NeurIPS*, 2024.

Saharia, C., Chan, W., Chang, H., Lee, C., Ho, J., Salimans, T., Fleet, D., and Norouzi, M. Palette: Image-to-image diffusion models. In *SIGGRAPH*, 2022a.

Saharia, C., Ho, J., Chan, W., Salimans, T., Fleet, D. J., and Norouzi, M. Image super-resolution via iterative refinement. *IEEE Transactions on Pattern Analysis and Machine Intelligence*, 45(4):4713-4726, 2022b.

Sohl-Dickstein, J., Weiss, E., Maheswaranathan, N., and Ganguli, S. Deep unsupervised learning using nonequilibrium thermodynamics. In *ICML*, 2015.

Song, J., Vahdat, A., Mardani, M., and Kautz, J. Pseudoinverse-guided diffusion models for inverse problems. In *ICLR*, 2023.

Song, Y. and Ermon, S. Generative modeling by estimating gradients of the data distribution. In *NeurIPS*, 2019.

Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. Score-based generative modeling through stochastic differential equations. In *ICLR*, 2021.

Spagnoletti, A., Prost, J., Almansa, A., Papadakis, N., and Pereyra, M. Latino-PRO: Latent consistency inverse solver with prompt optimization. *arXiv preprint arXiv:2503.12615*, 2025.

Wang, Y., Yu, J., and Zhang, J. Zero-shot image restoration using denoising diffusion null-space model. In *ICLR*, 2023.

Wang, Z., Bovik, A. C., Sheikh, H. R., and Simoncelli, E. P. Image quality assessment: from error visibility to structural similarity. *IEEE Transactions on Image Processing*, 13(4):600-612, 2004.

Zhang, R. Y., Isola, P., Efros, A. A., Shechtman, E., and Wang, O. The unreasonable effectiveness of deep features as a perceptual metric. In *CVPR*, 2018.
## A. Additional Details on the Method (Section 3)

We recall that the context variable $c := (x_0, x_t, s, t, y, A)$ in (16) is a tuple containing images $x_0$ and $x_t$ at time 0 and $t$, the two timesteps $t$ and $s$, the observation $y$ and the degradation operator $A$.

**Parameterization of the Inference Model.** We parameterize the inference model using a neural network $c \mapsto f_{\varphi}(c) := (f_{\varphi}^{\text{mean}}(c), f_{\varphi}^{\text{var}}(c))$ that predicts the residual between the mean, covariance statistics $(\mu_{s|0,t}(x_0, x_t), \sigma_{s|0,t}^2 I_d)$ of the prior Gaussian transition $q_{s|0,t}(x_s|x_0, x_t) \propto q_{s|0}(x_s|x_0)q_{t|s}(x_t|x_s)$, and the target posterior statistics $(\mu_{\varphi}(c), \rho_{\varphi}(c))$ that warm-start the variational approximation problem (16):

$$ 
\begin{aligned}
\mu_{\varphi}(c) &:= \mu_{s|0,t}(x_0, x_t) + f_{\varphi}^{\text{mean}}(c), \\
\rho_{\varphi}(c) &:= \sigma_{s|0,t}^2 \mathbf{1}_d + f_{\varphi}^{\text{var}}(c),
\end{aligned}
\qquad (20) $$

where $\mathbf{1}_d$ denotes the vector full of ones,

$$ \mu_{s|0,t}(x_0, x_t) := \gamma_{t|s} \alpha_{s|0} x_0 + (1 - \gamma_{t|s}) \alpha_{t|s}^{-1} x_t \qquad (21) $$

with $\gamma_{t|s} := \sigma_{t|s}^2 / \sigma_{t|0}^2$, and

$$ \sigma_{s|0,t}^2 := \sigma_{t|s}^2 \sigma_{s|0}^2 / \sigma_{t|0}^2, \qquad (22) $$

because by definition, $q_{s|0,t}(x_s|x_0, x_t) \propto q_{s|0}(x_s|x_0)q_{t|s}(x_t|x_s)$, and $q_{s|0}(x_s|x_0) := N(x_s; \alpha_{s|0}x_0, \sigma_{s|0}^2 I_d)$ and $q_{t|s}(x_t|x_s) := N(x_t; \alpha_{t|s}x_s, \sigma_{t|s}^2 I_d)$, by definition (1).

The choice of this parameterization is motivated by the fact the MGDM (Janati et al., 2025a) proposes to initialize the variational approximation problem (16) at inference time using the prior statistics $\mu_{s|0,t}(x_0, x_t)$ and $\sigma_{s|0,t}^2$, which are updated using gradient descent updates rules.

**Architecture of the Inference Model.** The architecture of the neural network model $c \mapsto f_{\varphi}(c) := (f_{\varphi}^{\text{mean}}(c), f_{\varphi}^{\text{var}}(c))$ is inspired by the standard design of conditional denoisers used in supervised diffusion. Such denoisers take as inputs the noised image $x_t$, the timestep $t$, the observation $y$ and the degradation operator $A$, and output a prediction of the denoised image (Saharia et al., 2022b;a; Elata et al., 2025). In our setting, the neural network model involved in the parameterization of the inference model requires two additional inputs: the midpoint timestep $s < t$ and the condition on the reconstructed image $x_0$ at timestep 0. Its outputs are the residual $f_{\varphi}^{\text{mean}}(c)$ of the approximate Gaussian distribution and the residual diagonal entries $f_{\varphi}^{\text{var}}(c)$ of its covariance, both having the same dimensionality as an image. These minor changes are compatible with all existing standard conditional denoiser architectures, e.g., UNet, DiT, HDiT, etc.

There are many possible modifications of the standard conditional denoiser to accommodate these changes in inputs and outputs. We found the following design to be effective in practice: (i) concatenate $x_0$ and $x_t$ channel-wise, thereby doubling the input channels; (ii) embed the midpoint timestep $s$ analogously to the embedding of $t$, but independently (e.g., using a separate linear layer), and add both timestep embeddings at the beginning of the architecture (early fusion); (iii) double the number of output channels to jointly predict the Gaussian mean and diagonal covariance entries.

In all our experiments, the backbone of the neural network for the inference model is an HDiT (Crowson et al., 2024). Specifically, in Section 5.1, we adopt the design of Palette (Saharia et al., 2022a) using the HDiT architecture. In Section 5.2, where specified, we follow the InvFussion approach (Elata et al., 2025) by incorporating so-called feature degradation layers into the HDiT backbone. These layers allow the architecture to adapt to different degradation operators during training and inference.

**Training Algorithm and Design of the Context Distribution.** The training algorithm of the inference model is summarized in Algorithm 1. We now comment the choice of the context distribution involved in the training objective function (17) for amortization. The sampling of $(x, y, A)$ creates a tuple of a clear image, its degraded observation along with the corresponding degradation operator, following the model (5). The sampling of the timesteps $(t, s)$ with $s < t$ aims to match those encountered during posterior sampling, cf. Section 2. Finally, the choice of the distribution for $(x_t, x_0)$ aims to capture the fact that, during posterior sampling with MGDM, $x_t$ is a noised sample, and $x_0$ is a denoised version of $x_t$: for this reason, we choose to set $x_0 := D_{0|t}^0(x_t)$ rather $x_0 := x$, where $x$ is sampled directly from $p_{\text{data}}$, the distribution of clean images.
**Algorithm 1** Training fully-amortized inference model

**Require:** Dataset of clean images $\mathcal{D}$, distribution of operators $p_{op}$, subset $\mathcal{T} \subseteq \{1, \dots, T\}$ of timesteps, learning rate $\alpha > 0$

**Ensure:** Trained inference model $\lambda_{\varphi}$

1: $\varphi \leftarrow$ initialization from scratch

2: **for** each iteration **do**

3: Sample clean images $\mathbf{x} \sim \mathcal{D}$

4: Sample degradation operator $\mathcal{A} \sim p_{op}$

5: Sample observation $\mathbf{y} \sim \ell_0(\mathbf{y}|\mathbf{x}, \mathcal{A})$

6: Sample timestep $t \sim \text{Unif}(\mathcal{T})$

7: Sample midpoint timestep $s \sim \text{Categorical}((\omega_t^{s'})_{s'=1}^{t-1})$,

8: Sample $\mathbf{x}_t \sim q_{t|0}(\cdot|\mathbf{x})$ \quad \{Forward noising\}

9: $\mathbf{x}_0 \leftarrow D_{0|t}^{\theta}(\mathbf{x}_t)$ \quad \{Pre-trained unconditional denoiser\}

10: $\mathbf{c} \leftarrow (\mathbf{x}_0, \mathbf{x}_t, s, t, \mathbf{y}, \mathcal{A})$

11: Compute gradient $g \leftarrow \nabla_{\varphi} \mathcal{L}(\mu_{\varphi}(\mathbf{c}), \rho_{\varphi}(\mathbf{c}); \mathbf{c})$

12: Update parameters $\varphi \leftarrow \varphi - \alpha g$

13: **end for**

14: **return** Inference model $\lambda_{\varphi} := (\mu_{\varphi}(\mathbf{c}), \rho_{\varphi}(\mathbf{c}))$

**Inference Algorithm.** Algorithm 2 summarizes how the trained inference model is used during our method called LAVPS. The three conditional distributions involved in the Gibbs sampling of Section 2, as described below (14), are implemented in Algorithm 2.

## B. Details of the Experimental Protocol (Section 5)

This section regroups all the important details to reproduce the results in Section 5.

**Training, Validation and Test Datasets.** All experiments are conducted at a resolution of 256 x 256, with images resized and center-cropped accordingly.

For FFHQ, the images used for training the unconditional denoiser correspond to indices ranging from 0 to 59 999. The validation set consists of images with indices from 61 000 to 61 031, and the test set (used for the final evaluation of all methods) consists of images with indices from 69 000 to 69 299.

For ImageNet, we use a deterministic subset corresponding to 1% of the official training split to train the inference model, ensuring that these images are fixed across runs and do not overlap with the validation or test sets. We use only images from the official validation split to construct the evaluation sets. We construct two disjoint subsets by selecting at most one image per class: the validation set is formed by taking the first image of each class, while the test set is formed by taking the last image of each class. We retain 32 images for validation and 300 images for testing.

**Normalization of Pixel Values.** To give a meaning to the relative noise level $\sigma_y$ in (5), we assume that the pixel values are normalized between -1 and 1.

**Inference Model Training.** The inference model used in LAVPS is trained by minimizing the population loss (17) using AdamW (Loshchilov & Hutter, 2019). We consider two training setups depending on the target dataset.

* **FFHQ.** The inference model is trained on the full FFHQ training set (100%, 60 000 images) for 5 epochs, with a learning rate of $1 \times 10^{-4}$ and a weight decay of $1 \times 10^{-5}$.

* **ImageNet.** The inference model is trained on 1% of the ImageNet training set (12 800 images) for 20 epochs, with a learning rate of $5 \times 10^{-5}$ and no weight decay.

In both cases, training is performed with a batch size of 16, using images at a resolution of 256 x 256, using a single NVIDIA V100 or H100 GPU for FFHQ and ImageNet, respectively. The subset of timesteps $\mathcal{T} \subseteq \{1, \dots, T\}$ that is chosen in practice is of the form $\mathcal{T} := \{1, \dots, \lceil(1 - r_{\text{switch}})T\rceil\}$, for a certain parameter $r_{\text{switch}} \in [0, 1]$ that we tune.
<caption>**Algorithm 2** Likelihood-guided Amortized Variational Posterior Sampling (LAVPS)</caption>

**Require:** Likelihood $\ell_0(y|\mathcal{A})$ defined from (6) with observation $y$ and operator $\mathcal{A}$, unconditional denoiser $(\boldsymbol{x}, t) \mapsto D_{0|t}^{\theta}(\boldsymbol{x})$ pre-trained via (4), inference model $c \mapsto \lambda_{\varphi}(c)$ trained via (17), number of timesteps $K$, timesteps $(t_k)_{k=0}^K$ with $t_K = T$ and $t_0 = 0$, Gibbs repetition $R$, number of denoising steps $M$, learning rates $\eta$, gradient steps $(G_k)_{k=1}^K$, probabilities $(\omega_k^\ell)_{k \in \{2, \dots, K-1\}, \ell \in \{1, \dots, t_{k-1}\}}$, subset of timesteps $\mathcal{T} \subseteq \{1, \dots, T\}$ for which we apply the inference model

**Ensure:** Approximate sample of the posterior distribution $\pi_0(\cdot|y, \mathcal{A})$ defined in (7)

1: $\boldsymbol{x}_t \sim \mathcal{N}(0, I_d)$
2: $\boldsymbol{x}_0 \leftarrow D_{0|T}^{\theta}(\boldsymbol{x}_t)$
3: **for** $k = K-1 \rightarrow 2$ **do**
4: &nbsp;&nbsp;&nbsp;&nbsp; $s \sim \text{Categorical}((\omega_k^\ell)_{\ell=1, \dots, t_{k-1}})$
5: &nbsp;&nbsp;&nbsp;&nbsp; $\boldsymbol{x}_t \sim q_{t_k|0, t_{k+1}}(\cdot|\boldsymbol{x}_0, \boldsymbol{x}_t)$
6: &nbsp;&nbsp;&nbsp;&nbsp; **for** $r = 1 \rightarrow R$ **do**
7: &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; $c \leftarrow (\boldsymbol{x}_0, \boldsymbol{x}_t, s, t_k, y, \mathcal{A})$
8: &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; $(\boldsymbol{\mu}_{zs}, \boldsymbol{\rho}_{zs}) \leftarrow (\boldsymbol{\mu}_{s|0, t_k}(\boldsymbol{x}_0, \boldsymbol{x}_t), \sigma_{s|0, t_k}^2 \mathbf{1}_d)$
9: &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; **if** $t_k \in \mathcal{T}$ **then**
10: &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; $(\boldsymbol{\mu}, \boldsymbol{\rho}) \leftarrow \lambda_{\varphi}(c)$
11: &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; **if** $\mathcal{L}(\boldsymbol{\mu}, \boldsymbol{\rho}; c) > \mathcal{L}(\boldsymbol{\mu}_{zs}, \boldsymbol{\rho}_{zs}; c)$ **then**
12: &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; $(\boldsymbol{\mu}, \boldsymbol{\rho}) \leftarrow (\boldsymbol{\mu}_{zs}, \boldsymbol{\rho}_{zs})$
13: &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; **end if**
14: &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; **else**
15: &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; $(\boldsymbol{\mu}, \boldsymbol{\rho}) \leftarrow (\boldsymbol{\mu}_{zs}, \boldsymbol{\rho}_{zs})$, cf. (21) and (22)
16: &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; **end if**
17: &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; **for** $g = 1 \rightarrow G_k$ **do**
18: &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; $(\boldsymbol{\mu}, \boldsymbol{\rho}) \leftarrow (\boldsymbol{\mu}, \boldsymbol{\rho}) - \eta \nabla_{(\boldsymbol{\mu}, \boldsymbol{\rho})} \mathcal{L}(\boldsymbol{\mu}, \boldsymbol{\rho}; c)$
19: &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; **end for**
20: &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; $\boldsymbol{x}_s \sim \mathcal{N}(\boldsymbol{\mu}, \text{diag}(\boldsymbol{\rho}))$
21: &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; $\boldsymbol{x}_0 \leftarrow \text{DDIM}(\boldsymbol{x}_s, s, M)$
22: &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; $\boldsymbol{x}_t \sim q_{t_k|s}(\cdot|\boldsymbol{x}_s)$
23: &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; **end for**
24: &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; **end for**
25: &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; **return** $\boldsymbol{x}_0$

{Zero-shot initialization}

{Warm-start candidate}

{Safeguard fallback}

{Likelihood-guidance}

{Sampling $\boldsymbol{x}_s$ knowing $\boldsymbol{x}_0$ and $\boldsymbol{x}_t$}

{Sampling $\boldsymbol{x}_0$ knowing $\boldsymbol{x}_s$}

{Sampling $\boldsymbol{x}_t$ knowing $\boldsymbol{x}_s$}

**Hyperparameters for LAVPS and MGDM.** We describe the hyperparameters of Algorithm 2 (LAVPS) and MGDM that we consider in our experiments. We set the weights $(\omega_t^s)_{s=1}^{t-1}$ of the approximation mixture $\pi_t(\boldsymbol{x}_t|y, \mathcal{A}) \approx \hat{\pi}_t(\boldsymbol{x}_t|y, \mathcal{A}) := \sum_{s=1}^{t-1} \omega_t^s \hat{\pi}_t^s(\boldsymbol{x}_t|y, \mathcal{A})$ (see Section 2) to $\omega_t^s = 0$ for $s < t-1$ and $\omega_t^{t-1} = 1$, as this configuration proved most effective in practice. We always choose a number of $K=100$ timesteps, and the number of Gibbs repetition is $R=1$. The timesteps $(t_k)_{k=0}^K$ is a uniform grid between $t_0=0$ and $t_K=T$. The choice of the subset of timesteps $\mathcal{T} \subseteq \{1, \dots, T\}$ for which we apply the inference model is the same as the one considered during training: $\mathcal{T} := \{1, \dots, \lceil(1-r_{\text{switch}})T\rceil\}$, for a certain parameter $r_{\text{switch}} \in [0, 1]$ that we tune. The number of gradient steps $(G_k)_{k=1}^K$ are chosen as follows:

$$ G_k := \begin{cases} G_{\text{start}} & \text{if } t_k \notin \mathcal{T}, \\ G_{\text{end}} & \text{otherwise,} \end{cases} \qquad (23) $$

where $G_{\text{start}}$, $G_{\text{end}}$ are hyperparameters to tune. The number of denoising steps $M$ is tuned as well. We use the Adam optimizer (Kingma & Ba, 2015) without weight decay to minimize the variational objective function (16), with a learning rate $\eta$ to tune.

**Time Measurements.** All runtimes are measured on a single NVIDIA V100 GPU with a batch size of 1. For each image, we measure the wall-clock time of the full sampling loop of the selected method using Python's high-resolution timer (*time.perf_counter*). Timing starts immediately before the sampling call and stops right after it returns; consequently, preprocessing steps (e.g., degradation setup and noise generation), metric computation, and file I/O are excluded. All experiments are conducted using Python 3.12, PyTorch 2.6, and CUDA 12.4.
**Image Reconstruction Tasks.** We describe the degradation operators corresponding to the different considered reconstruction tasks: super-resolution, inpainting and motion deblurring. We rely on the implementations of degradation operators provided in the codebase of Janati et al. (2025a); Elata et al. (2025).

* **Super-resolution:** The degradation operator is a uniform downsampling operation, by averaging over non-overlapping patches of size $r \times r$. We set $r = 4$, corresponding to the task of $\times 4$ super-resolution.

* **Inpainting:** The degradation operator is a masking operator, where only pixels corresponding to unmasked indices are retained, and the inverse problem is to reconstruct the full image. We consider the following masks. Center inpainting (for Section 5.1 and Section 5.2): we remove a central square region covering half of the image ($I/2 \times I/2$ pixels for an image of size $I \times I$). Missing rectangle (for Section 5.2): we choose to mask a rectangle of random shape and random position, as implemented by Elata et al. (2025). Missing pixels (for Section 5.2): we choose to mask some randomly selected pixels, following the implementation of MissingPatch in Elata et al. (2025).

* **Motion deblurring:** The degradation operator is a convolution with a fixed motion blurring kernel. We use a kernel of size $21 \times 21$ with intensity 0.9 in Section 5.1, following the implementation of MotionBlur in Janati et al. (2025a). We also use kernel of size $35 \times 35$ with intensity 0.3 in Section 5.2.

## B.1. Protocol for the Inference Time Acceleration with Amortization (Section 5.1)

For the MGDM baseline, we perform an extensive hyperparameter search on 32 validation images, evaluating more than 90 combinations to determine the Pareto front of reconstruction quality (measured with LPIPS, a perceptual similarity metric well aligned with human visual judgments) versus inference time. This Pareto front represents the set of configurations that yield the best trade-off between reconstruction fidelity and inference speed for MGDM. Using the same hyperparameter grid, we tune LAVPS to identify configurations surpassing this Pareto front.

**Hyperparameter Grid.** To characterize the trade-off between reconstruction quality and inference speed for both MGDM and LAVPS, we evaluate a broad grid of hyperparameter configurations. We adopt the naming conventions introduced in Algorithm 2. The grid comprises:

* $G_{start} \in \{1, 3\}$

* $G_{end} \in \{0, 1, 3, 10\}$

* Learning rate $\eta \in \{0.01, 0.03\}$

* Number of denoising steps $M \in \{1, 5\}$

* Switch ratio $r_{switch} \in \{70\%, 80\%, 90\%\}$

**Pareto Front Construction.** We trace the MGDM Pareto front by sweeping the hyperparameters presented previously on a 32-image validation set, thereby characterizing the trade-off between reconstruction quality and inference time. We then take the best zero-shot configurations (highest quality at a given speed) and evaluate it on 300 images from the test set.

**Statistical Test.** We compare LAVPS to the zero-shot baseline using a paired non-inferiority t-test on per-image LPIPS scores (denoted by LPIPS<sub>ZS</sub> for the zero-shot baseline, and LPIPS<sub>AM</sub> for LAVPS, lower is better). For each test image i, we form the paired difference

$$ \delta_i := \text{LPIPS}_{\text{ZS},i} - \text{LPIPS}_{\text{AM},i}, \qquad (24) $$

so that larger $\delta_i$ favors LAVPS.

With non-inferiority margin $\Delta > 0$, the hypotheses are

$$ H_0: \delta \le -\Delta \qquad (\text{LAVPS worse than MGDM by at least } \Delta), \qquad (25) $$

$$ H_1: \delta > -\Delta \qquad (\text{LAVPS is not inferior to MGDM}). \qquad (26) $$
Let $\bar{\delta}$ and $\bar{s}$ be the empirical mean and standard deviation of the samples $\{\delta_i\}_{i=1}^n$. The statistical test is

$$T = \frac{\bar{\delta} + \Delta}{\bar{\delta}/\sqrt{n}}, \qquad (27)$$

We reject $H_0$ at level $\alpha$ if

$$T > t_{1-\alpha, n-1}, \qquad (28)$$

where $t_{1-\alpha, n-1}$ denotes the $(1 - \alpha)$-quantile of the Student $t$ distribution with $n - 1$ degrees of freedom. Equivalently, we conclude non-inferiority if the one-sided $(1 - \alpha)$ lower confidence bound for $\delta$ exceeds $-\Delta$.

We use $\alpha = 0.05$ and $\Delta = 0.003$. Configurations of LAVPS that reject $H_0$ are considered to be configurations that does not degrade the reconstruction quality compared to the considered MGDM sampling. They are therefore retained for speedup reporting when aiming for a certain target quality, as it is done in Table 1.

## B.2. Protocol for Out-of-distribution Robustness (Section 5.2)

The architecture of the inference model for the inpainting task in Section 5.2 is based on Palette (Saharia et al., 2022a), and the one for the motion deblurring task is based on InvFussion (Elata et al., 2025).

Given a reconstruction task and a fixed inference computational budget, each method (LAVPS and MGDM) is tuned over the same hyperparameter grid described in Section B.1. The optimal hyperparameter configuration is selected on a validation set of 32 images by maximizing a chosen quality metric (e.g., LPIPS or CMMD). Final performance is then reported on a held-out test set of 300 images.

## C. Additional Experiments

This section contains additional experiments to reinforce the motivation of this paper and to further support our claims.

### C.1. Lack of Out-of-distribution Robustness for Supervised Diffusion

Our method relies on explicit likelihood guidance, as in posterior sampling, which we show to provide an advantage in robustness to out-of-distribution operators compared to InvFussion (Elata et al., 2025)—a supervised diffusion baseline that embeds the degradation operator into the conditional denoiser. While this design is effective when degradations match those seen during training, we hypothesize that it struggles to generalize when the inference-time operator deviates from this distribution. In contrast, likelihood guided methods enforces data consistency with the actual operator at inference time, thereby avoiding this generalization bottleneck.

To test our hypothesis, we reproduced the supervised baseline InvFussion (Elata et al., 2025) following its exact setup for image inpainting, considering no more a fixed mask, but different families of masks during training and inference. During training, the learned conditional denoiser only observes masks that remove a single randomly placed rectangle (in-distribution degradations). At evaluation, we introduce a different family of masks in which individual pixels are randomly removed according to an i.i.d. Bernoulli sampling with parameter $p \in [0.9, 0.92]$, creating out-of-distribution degradations. In order to provide a reference, we evaluate MGDM on the same tasks.

**Training and Inference of InvFussion.** The training of the supervised baseline InvFussion (Elata et al., 2025) uses the full training set of FFHQ at resolution 64 × 64 (60000 images), following the exact protocol of Elata et al. (2025). In their protocol, they use an Adam optimizer without weight decay, with a reference learning rate of 5 × 10⁻⁵ that follows a warm-up-and-decaying schedule. The model is trained on 32768 kims with a batch size of 512. At inference, we use the DDIM sampler with 500 steps to generate reconstructions with the supervised diffusion baseline InvFussion (Elata et al., 2025). For comparison of the reconstruction quality, we compare with the zero-shot posterior sampling baseline MGDM (Janati et al., 2025a). Its hyperparameters are the following: we use $K = 100$ diffusion steps, the number of gradient steps are $G_{start} = 3$ and $G_{end} = G$ with $G \in \{0, 30\}$ gradient updates applied to the variational approximation during the final 10 steps, i.e, for timesteps $t_k$ with $k \in \{1, \dots, 10\}$. The learning rate $\eta$ for the variational approximation problem is set to 0.01. The number of Gibbs repetition is $R = 1$, and the number of denoising steps is $M = 1$.
**Results.** Table 4 shows that the supervised baseline InvFussion (Elata et al., 2025) performs well on in-distribution masks but collapses under out-of-distribution masks, despite explicitly encoding the degradation operator. See Figure 1 for an illustration: the reconstruction of InvFussion on out-of-distribution masks has artefacts. This illustrates its lack of robustness when the inference-time degradation departs from what was modeled during training. In contrast, MGDM maintains strong reconstruction quality across both mask types. This is because MGDM is a zero-shot method whose performance is driven by the test-time guidance term, which ensures that the generated samples remain consistent with the actual observation regardless of the specific mask geometry.

These findings highlight the central role of test-time likelihood guidance in achieving robustness to out-of-distribution degradations. They also illustrate the limitations of supervised baselines, which fail when the degradation operator is unseen during training.

<table><caption>Table 4. Supervised diffusion InvFussion (Elata et al., 2025) vs. MGDM (Janati et al., 2025a) on inpainting with out-of-distribution masks (missing pixels). Evaluation on 300 images of FFHQ64. Models were trained on in-distribution masks (single missing rectangle). The best metric is in bold.</caption><thead><tr><th></th><th>LPIPS ↓</th><th>SSIM ↑</th><th>PSNR ↑</th><th>CMMD ↓</th></tr></thead><tbody><tr><td colspan="5">In-distribution masks (missing rectangle)</td></tr><tr><td>InvFussion (Elata et al., 2025)</td><td>0.015 ± 0.027</td><td>0.947 ± 0.075</td><td>33.2 ± 8.1</td><td><strong>0.010</strong></td></tr><tr><td>MGDM (Janati et al., 2025a)</td><td><strong>0.014 ± 0.024</strong></td><td><strong>0.952 ± 0.066</strong></td><td><strong>33.3 ± 7.1</strong></td><td>0.011</td></tr><tr><td colspan="5">Out-of-distribution masks (missing pixels)</td></tr><tr><td>InvFussion (Elata et al., 2025)</td><td>0.403 ± 0.118</td><td>0.427 ± 0.077</td><td>16.8 ± 1.4</td><td>1.029</td></tr><tr><td>MGDM (Janati et al., 2025a)</td><td><strong>0.040 ± 0.022</strong></td><td><strong>0.846 ± 0.051</strong></td><td><strong>23.6 ± 2.5</strong></td><td><strong>0.077</strong></td></tr></tbody></table>

## C.2. Comparison with Other Zero-shot Posterior Sampling Methods

We put into perspective the trade-off between reconstruction quality and inference time of MGDM studied in Section 5.1 compared to other existing zero-shot posterior sampling algorithm, such as DPS (Chung et al., 2023).

**Protocol.** The zero-shot DPS baseline is tuned on the validation set (32 images). We vary the number of sampling steps ($T \in 10, 30, 100, 300, 1000$) and the step sizes ($\zeta \in 0.1, 0.3, 1, 3$), and retain for each $T$ the best-performing configuration. The final evaluation is conducted on the test set (300 images). This procedure yields the reconstruction-time trade-off for DPS, which we compare against MGDM trade-off studied in Section 5.1.

<figure><img src="image_6.png" /><figcaption>Figure 5. Trade-off between reconstruction quality (LPIPS) and inference time of MGDM (Janati et al., 2025a) vs. DPS (Chung et al., 2023), for motion deblurring in FFHQ. Configurations of MGDM and DPS are Pareto-optimal.</figcaption></figure>

**Results.** Figure 5 shows that MGDM consistently achieves a better trade-off between reconstruction quality and inference time than DPS. Across the full range of tested compute budgets, MGDM obtains lower LPIPS at equal inference time.
For example, for inference times below 30 seconds, DPS reaches an average LPIPS of 0.205, whereas MGDM attains a substantially lower LPIPS of 0.044. Even when allowing longer runtimes, DPS only improves to an average LPIPS of 0.153 at 98 seconds of inference. These results are consistent with prior observations that MGDM provides higher reconstruction fidelity than other zero-shot sampling methods (Janati et al., 2025a).

### C.3. Training Cost vs. Inference Efficiency Analysis in LAVPS

We analyze the computational investment required to train the inference model in LAVPS compared to the MGDM baseline. We demonstrate that the training overhead for the amortization module is modest and is rapidly compensated by the significant acceleration achieved during inference.

**Protocol.** We compare MGDM and LAVPS under configurations selected to achieve comparable high-fidelity reconstruction quality (LPIPS ≈ 0.045).

* **MGDM:** We use inference hyperparameters $G_{end} = 10$, $G_{start} = 1$, $M = 1$, learning rate $\eta = 0.03$, and a switch ratio $r_{switch} = 80\%$.

* **LAVPS:** We utilize the same inference settings but increase the switch ratio to $r_{switch} = 90\%$ to leverage the amortized initialization. The inference model is trained on 10% of the FFHQ dataset.

<table><caption>Table 5. Training time, inference time (average) and LPIPS (average and standard error) for MGDM and LAVPS on motion deblurring in FFHQ. Metrics are computed over 300 test images.</caption><thead><tr><td>Method</td><td>Training time (GPUh)</td><td>Inference time / image (s)</td><td>LPIPS ↓</td></tr></thead><tbody><tr><td>MGDM</td><td>0</td><td>28.7</td><td>0.044 ± 0.018</td></tr><tr><td>LAVPS</td><td>6.6</td><td>20.5</td><td>0.045 ± 0.017</td></tr></tbody></table>

**Results.** As reported in Table 5, training the inference model for LAVPS requires approximately 6.6 GPU-hours. Although MGDM has no initial training cost, it is slower at inference (28.7 s/image) compared to LAVPS (20.5 s/image). By defining the total cost as the sum of training time and cumulative inference time, we find that the initial training overhead is fully amortized after approximately $N = 2,900$ images. Consequently, for any workload exceeding this relatively small number of reconstructions, LAVPS is strictly more computationally efficient than the zero-shot baseline.

### C.4. Extended Pareto Frontier Analysis on FFHQ

Under the same protocol as in Section 5.1, Figure 6 expand our Pareto frontier analysis in Section 5.1 to include reconstruction tasks on the FFHQ dataset, such as motion deblurring and super-resolution. This further validates the improved trade-off between reconstruction quality and inference speed for LAVPS compared to MGDM.

<figure><img src="image_7.png" /><figcaption>Figure 6. Trade-off between reconstruction quality and inference time of LAVPS vs. MGDM on the FFHQ dataset. The configurations are Pareto-optimal.</figcaption></figure>
### C.5. Ablation Study: Impact of Gradient Steps on Convergence

We analyze how the number of inference gradient steps, denoted $G_{end}$, affects both reconstruction quality and inference speed for LAVPS and MGDM. This comparison highlights the efficiency of the likelihood-aware warm-start.

**Protocol.** MGDM hyperparameters are fixed as follows: switch ratio $r_{switch} = 80\%$; the number of gradient steps for timesteps outside $\mathcal{T} := \{1, \dots, \lfloor(1 - r_{switch})T\rfloor\}$ is $G_{start} = 1$; the learning rate is $\eta = 0.03$; and the number of denoising steps is $M = 1$. We then vary the number of gradient steps $G_{end} \in \{0, 1, 3, 10\}$ timesteps in $\mathcal{T}$, to study its role in both LVAPS and MGDM. We considered here 32 images from the validation set.

<figure><img src="image_8.png" /></figure>

<table><caption>Figure 7. Impact of the number of gradient steps $G_{end}$ on reconstruction performance and inference time for motion deblurring on FFHQ. Left: LPIPS vs. inference time for LAVPS and MGDM. Right: corresponding LPIPS values (mean ± std) for the different values of number of gradient steps $G_{end}$.</caption><thead><tr><th>G<sub>end</sub></th><th>MGDM</th><th>LAVPS</th></tr></thead><tbody><tr><td>0</td><td>0.212 ± 0.080</td><td>0.072 ± 0.027</td></tr><tr><td>1</td><td>0.156 ± 0.057</td><td>0.063 ± 0.023</td></tr><tr><td>3</td><td>0.085 ± 0.034</td><td>0.050 ± 0.019</td></tr><tr><td>10</td><td>0.044 ± 0.018</td><td>0.039 ± 0.015</td></tr></tbody></table>

**Results.** As shown in Figure 7, for the same choice of hyperparameter, LAVPS is slightly slower than MGDM due to the additional forward pass through the inference model used to warm-start the variational problem. However, this small overhead is outweighed by a substantial improvement in reconstruction quality, confirmed quantitatively in Figure 7. Overall, LAVPS maintains high reconstruction quality even with very few gradient steps, whereas MGDM performs poorly in the regime with a low number of gradient steps. This highlights that the learned inference model enables LAVPS to operate with significantly fewer gradient steps (hence, a lower inference time) while preserving performance, leading to a better quality-time trade-off.

### C.6. Pareto Frontier Analysis When Amortizing on a Family of Operators

Figure 8 shows that, for the reconstruction tasks considered in Section 5.2, the improved trade-off between reconstruction quality and inference speed achieved by LAVPS over MGDM persists even when the inference model is trained on a family of operators. The protocol used to construct the Pareto front is identical to that of Section 5.1, except that the inference model is trained on a family of operators rather than a single operator.

As claimed in Section 5.2, this improvement on in-distribution opetors does not sacrifice the flexibility on arbitrary opetors of posterior sampling methods: LAVPS remains robust to out-of-distribution operators. We complement Figures 3 and 4 with additional reconstruction samples in Figures 9 and 10. Under a certain target computational budget at inference, LAVPS achieves better reconstructions than MGDM, e.g., reconstructions are sharper and hallucinate less, while being able to recover the reconstruction quality of MGDM on out-of-distribution operators.
<figure><img src="image_9.png" /></figure>

<figure><img src="image_10.png" /><figcaption>Figure 8. Efficiency-quality trade-off for LAVPS vs. MGDM, on in-distribution degradation operators. The inference model in LAVPS is trained on a family of degradation operators. Markers represent Pareto-optimal configurations for each method across various restoration tasks on the Imagenet dataset.</figcaption></figure>

<figure><img src="image_11.png" /></figure>

<figure><img src="image_12.png" /><figcaption>Figure 9. Motion deblurring on ImageNet. The inference time is constrained to be lower than 47.3s.</figcaption></figure>
<figure><img src="image_13.png" /><figcaption>**Figure 10.** Inpainting on ImageNet. The inference time is constrained to be lower than 34.9s.</figcaption></figure>