Variational Diffusion Posterior Sampling with Midpoint Guidance

## Abstract

Abstract Diffusion models have recently shown considerable potential in solving Bayesian inverse problems when used as priors. However, sampling from the resulting
denoising posterior distributions remains a challenge as it involves intractable terms. To tackle this issue, state-of-the-art approaches formulate the problem as that of sampling from a surrogate diffusion model targeting the posterior and decompose its scores into two terms: the prior score and an intractable guidance term. While the former is replaced by the pre-trained score of the considered diffusion model, the guidance term has to be estimated. In this paper, we propose a novel approach that utilises a decomposition of the transitions which, in contrast to previous methods, allows a trade-off between the complexity of the intractable guidance term and that of the prior transitions. We validate the proposed approach through extensive experiments on linear and nonlinear inverse problems, including challenging cases with latent diffusion models as priors, and demonstrate its effectiveness in reconstructing electrocardiogram (ECG) from partial measurements for accurate cardiac diagnosis.

## 1 Introduction

Inverse problems aim to reconstruct signals from incomplete and noisy observations and are prevalent across various fields. In signal and image processing, common examples include signal deconvolution, image restoration, and tomographic image reconstruction Stuart (); Idier (). Other applications extend to protein backbone motif scaffolding Watson et al. () and urban mobility modeling Jiang et al. (). Due to the presence of noise and the inherent complexity of the measurement process, these problems are typically ill-posed, meaning they have an infinite number of possible solutions. While many of these solutions may fit the observed data, only a few align with the true underlying signal.
Bayesian inverse problems provide a principled framework for addressing the challenges of signal reconstruction by incorporating prior knowledge, enabling more plausible and meaningful solutions. The *a priori* knowledge about the signal $x$ to be reconstructed is captured in the prior distribution $q(x)$, while the information about the observation $y$ is encoded in the likelihood function $p(y|x)$. Given these components, solving the inverse problem boils down to sampling from the posterior distribution $\pi(x)\propto p(y|x)q(x)$, which integrates both prior knowledge and observational data.

The choice of the prior distribution is crucial for achieving accurate reconstructions. If the goal is to reconstruct high-resolution data, the prior must be able to model data with similar fine detail and complexity. Recently, denoising diffusion models (DDM) Sohl-Dickstein et al. (); Song & Ermon (); Ho et al. (); Song et al. () have emerged as a powerful approach in this context. These models can generate highly realistic and detailed reconstructions and are becoming increasingly popular as priors.
Using a DDM, a sample $X_{0}$ being approximately distributed according to a given data distribution $q$ of interest can be generated by iteratively denoising an initial sample $X_{n}$ from a standard Gaussian distribution.
The denoising process $(X_{k})_{k=n}^{0}$ consists of the iterative refinement of $X_{k}$ for $k\in\llbracket 1,n\rrbracket$ based on a parametric approximation of the score $\nabla\log q_{k}$, where $q_{k}(x_{k})=\int q(x_{0})q_{k|0}(x_{k}|x_{0}),\mathrm{d}x_{0}$, with $q_{k|0}(\cdot|x_{0})$ being a Gaussian noising transition.

Denoising Diffusion Models (DDMs) form a highly powerful class of prior distributions, but their use introduces significant challenges in posterior sampling. Specifically, since a DDM prior does not admit an explicit and tractable density, conventional Markov chain Monte Carlo (MCMC) methods, such as the Metropolis–Hastings algorithm and its variants Besag (); Neal (), cannot be applied in general. Furthermore, gradient-based MCMC methods often prove inefficient, as they tend to get trapped in local modes of the posterior.

The problem of sampling from the posterior in Bayesian inverse problems with DDM priors has recently been addressed in several papers Kadkhodaie & Simoncelli (); Song et al. (); Kawar et al. (; ); Lugmayr et al. (); Ho et al. (); Chung et al. (); Song et al. (). These methods typically modify the denoising process to account for the observation $y$. One principled approach for accurate posterior sampling involves skewing the denoising process at each stage $k$ using the score $\nabla\log\pi_{k}$, where $\pi_{k}$ is defined analogously to $q_{k}$, with $q$ replaced by the posterior $\pi$. This guides the denoising process in a manner that is consistent with the observation. Notably, the posterior score $\nabla\log\pi_{k}$ decomposes into two components: the prior score and an additional intractable term commonly referred to as *guidance*. Previous works leverage *pre-trained* DDMs for the prior score, while providing various *training-free* approximations for the guidance term Ho et al. (); Chung et al. (); Song et al. (); Finzi et al. (). This framework enables solving a wide range of inverse problems for a given prior. However, despite the notable success of these methods, the approximations often introduce significant errors, leading to inaccurate reconstructions when the posterior distribution is highly multimodal, the measurement process is strongly nonlinear, or the data is heavily contaminated by noise.

##### Our contribution.

We begin with the observation that the posterior denoising step, which transforms a sample $X_{k+1}$ into a sample $X_{k}$, does not necessarily require the conditional scores $\nabla\log\pi_{k+1}$ or $\nabla\log\pi_{k}$.
Instead, we demonstrate that this step can be decomposed into two intermediate phases: first, denoising $X_{k+1}$ into an intermediate *midpoint* state $X_{\ell_{k}}$, where ${\ell_{k}}<k$, and then noising back to obtain $X_{k}$, unconditionally on the observation $y$.
This decomposition introduces an additional degree of freedom, as it only requires the estimation of the guidance term at the intermediate step ${\ell_{k}}$ rather than step $k+1$.
Building on this insight, we introduce Midpoint Guidance Posterior Sampling  (MGPS), a novel diffusion posterior sampling scheme that explicitly leverages this approach.
Our algorithm develops a principled approximation of the denoising transition by utilizing a Gaussian variational approximation, combined with the guidance approximation proposed by Chung et al. () at the intermediate steps ${\ell_{k}}$.
The strong empirical performance of our method is demonstrated through an extensive set of experiments.
In particular, we validate our approach on a Gaussian mixture toy example, as well as various linear and nonlinear image reconstruction tasks–inpainting, super-resolution, phase retrieval, deblurring, JPEG dequantization, high-dynamic range–using both pixel-space and latent diffusion models (LDM).
We also apply our posterior sampling method to imputation tasks for 12-lead electrocardiograms (ECGs).
Our results demonstrate that *training-free* posterior sampling performs on par with diffusion models specifically trained for ECG imputation Alcaraz & Strodthoff ().
Moreover, our method demonstrates significantly improved performance compared to existing approaches.

## 2 Posterior sampling with DDM prior

##### Problem setup.

In this paper, we focus on the approximate sampling from a density of the form

$$ $\pi(x)\vcentcolon=g(x)q(x)/\mathcal{Z}\,,$ (2.1) $$

where $g:\mathbb{R}^{{d}}\to\mathbb{R}$ is a *non-negative* and *differentiable* potential that *can be evaluated pointwise*, $q$ is a prior distribution, and $\smash{\mathcal{Z}\vcentcolon=\int g(x)q(x)\,\mathrm{d}x}$ is the normalizing constant.
We are interested in the case where a DDM $p^{\theta}_{0}$ for the prior has been pre-trained for the prior based on i.i.d. observations from $q$.
As a by-product, we also have parametric approximations $(s^{\theta}_{k})_{k=1}^{n}$ of the *scores* $(\nabla\log q_{k})_{k=1}^{n}$, where the marginals $(q_{k})_{k=1}^{n}$ are defined as $q_{k}(x_{k})\vcentcolon=\int q(x_{0})\,q_{k|0}(x_{k}|x_{0})\,\mathrm{d}x_{0}$ with $q_{k|0}(x_{k}|x_{0})\vcentcolon=\mathcal{N}(x_{k};\sqrt{\alpha_{k}}x_{0},v_{k}
\mathrm{I}_{d})$ and $v_{k}\vcentcolon=1-\alpha_{k}$.
Typically, the sequence $(\alpha_{k})_{k=0}^{n}$ is a decreasing sequence with $\alpha_{0}=1$ and $\alpha_{n}$ approximately equals zero for $n$ large enough.

These score approximations enable the generation of new samples from $q$ according to one of the DDM sampling schemes Song & Ermon (); Ho et al. (); Song et al. (; ); Karras et al. ().
All these approaches boil down to simulating a Markov chain $(X_{k})_{k=n}^{0}$ backwards in time starting from a standard Gaussian $p_{n}\vcentcolon=\mathcal{N}(\mathrm{0},\mathrm{I}_{{d}})$ and following the Gaussian transitions
$p^{\theta}_{k|k+1}(x_{k}|x_{k+1})=\mathcal{N}(x_{k};\smash{m^{\theta}_{k|k+1}}
(x_{k+1}),v_{k|k+1}\mathrm{I}_{{d}})$,
where the mean functions $\smash{m^{\theta}_{k|k+1}}$ can be derived from the learned DDM, and $v_{k|k+1}>0$ are fixed and pre-defined variances.
The backward transitions are designed in such a way that the marginal law of $X_{k}$ approximates $q_{k}$; see Section A.1 for more details on the form of the backward transitions.

In our problem setup, we assume we only have access to the approximated scores of the prior and no observation from neither posterior $\pi$ nor the prior $q$.
This setup encompasses of solving an inverse problem without training a conditional model from scratch based on a paired signal/observation dataset; a requirement typically imposed by conditional DDM frameworks Song et al. (); Batzolis et al. (); Tashiro et al. (); Saharia et al. ().
On the other hand, we require access to $g$.
This setup includes many applications in Bayesian inverse problems, *e.g.*, image restoration,
motif scaffolding in protein design Trippe et al. (); Wu et al. (), and trajectory control in traffic-simulation frameworks Jiang et al. ().

##### Conditional score.

Since we have access to a DDM model for $q$, a natural approach to sampling from $\pi$ is to define a DDM approximation based on the pre-trained score approximations for $q_{k}$.
Following the DDM approach, the basic idea here is to sample sequentially from the marginals $\pi_{k}(x_{k})\vcentcolon=\int\pi(x_{0})q_{k|0}(x_{k}|x_{0})\,\mathrm{d}x_{0}$ by relying on approximations of the conditional scores $(\nabla\log\pi_{k})_{k=1}^{n}$.
The latter can be expressed in terms of the unconditional scores by using

$$ $\pi_{k}(x_{k})\propto\int g(x_{0})q(x_{0})q_{k|0}(x_{k}|x_{0})\,\mathrm{d}x_{0 }\,,$ (2.2) $$

where after defining the backward transition kernel $q_{0|k}(x_{0}|x_{k})\vcentcolon=q(x_{0})q_{k|0}(x_{k}|x_{0})\big{/}q_{k}(x_{k})$ yields

$$ $\pi_{k}(x_{k})\propto g_{k}(x_{k})q_{k}(x_{k})\,,\quad\text{where}\quad g_{k}( x_{k})\vcentcolon=\int g(x_{0})q_{0|k}(x_{0}|x_{k})\,\mathrm{d}x_{0}\,.$ (2.3) $$

It then follows that

$$ $\nabla\log\pi_{k}(x_{k})=\nabla\log q_{k}(x_{k})+\nabla\log g_{k}(x_{k})\,.$ $$

While we have a parametric approximation of the first term on the r.h.s., the bottleneck is the second term, which is intractable due to the integration of $g$ against the conditional distribution $\smash{q_{0|k}(\cdot|x_{k})}$.

##### Existing approaches.

To circumvent this computational bottleneck, previous works have proposed approximations that involve replacing $q_{0|k}(\cdot|x_{k})$ in (2.3) with either a Dirac delta mass Chung et al. () or a Gaussian approximation Ho et al. (); Song et al. ()

$$ $\smash{q_{0|k}(x_{0}|x_{k})\approx\mathcal{N}(x_{0};\smash{m^{\theta}_{0|k}}(x _{k}),v_{0|k}\mathrm{I}_{d})}\,,$ (2.4) $$

where $\smash{m^{\theta}_{0|k}}$ is a parametric approximation of $\smash{\smash{m_{0|k}}(x_{k})\vcentcolon=\int x_{0}\,q_{0|k}(x_{0}|x_{k})\,
\mathrm{d}x_{0}}$ and $v_{0|k}$ is a tuning parameter. Using Tweedie’s formula Robbins (), it holds that $\smash{m_{0|k}}(x_{k})=(-x_{k}+\sqrt{\alpha_{k}}\,\nabla\log q_{k}(x_{k}))/v_{k}$, which suggests the approximation $\smash{m^{\theta}_{0|k}}(x_{k})\vcentcolon=(-x_{k}+\sqrt{\alpha_{k}}s^{\theta}
_{k}(x_{k}))/v_{k}\,.$
As for the variance parameter $v_{0|k}$, Ho et al. () suggest setting $v_{0|k}=v_{k}/\alpha_{k}$, while Song et al. () use $v_{0|k}=v_{k}$.
When $g$ is the likelihood of a linear model, *i.e.*, $g(x)\vcentcolon=\mathcal{N}(y;\mathrm{A}x,\sigma_{y}^{2}\mathrm{I}_{d_{y}})$, for $\mathrm{A}\in\mathbb{R}^{{d_{y}}\times{d}}$ and $\sigma_{y}>0$, it can be exactly integrated against the Gaussian approximation (2.4), yielding $g_{k}(x_{k})\approx\mathcal{N}(y;\mathrm{A}\smash{m^{\theta}_{0|k}}(x_{k}),
\sigma_{y}^{2}\mathrm{I}_{d_{y}}+v_{0|k}\mathrm{A}\mathrm{A}^{\intercal})$.
Chung et al. (), on the other hand, use the pointwise approximation, yielding $g_{k}(x_{k})\approx g(\smash{m^{\theta}_{0|k}}(x_{k}))$.
We denote by $\tilde{g}_{k}$ the resulting approximation stemming from any of these methods. Approximate samples from the posterior distribution are then drawn by simulating a backward Markov chain $(X_{k})_{k=n}^{0}$, where $X_{n}\sim p_{n}$ and then, given $X_{k+1}$, $X_{k}$ is obtained via the update

$$ $X_{k}\vcentcolon=\tilde{X}_{k}+w_{k+1}(X_{k+1})\nabla\log\tilde{g}_{k+1}(X_{k+ 1})\,,\quad\text{where}\quad\tilde{X}_{k}\sim p^{\theta}_{k|k+1}(\cdot|X_{k+1} )\,,$ (2.5) $$

and $w_{k+1}$ is a method-dependent weighting function. For instance, Chung et al. () use $w_{k+1}(X_{k+1})=\zeta/\|y-\mathrm{A}\smash{m^{\theta}_{0|k+1}}(X_{k+1})\|$, with $\zeta\in[0,1]$.
We note that the update (2.5) in general involves a vector-Jacobian product, *e.g.*, considering the approximation suggested by Chung et al. (), $\nabla_{x_{k}}\log\tilde{g}_{k}(x_{k})=\nabla_{x_{k}}\smash{m^{\theta}_{0|k}}(
x_{k})^{\intercal}\nabla_{x_{0}}\left.\log g(x_{0})\right|_{x_{0}=\smash{m^{
\theta}_{0|k}}(x_{k})}$, which incurs an additional computational cost compared to an unconditional diffusion step.

## 3 The MGPS algorithm

In this section we propose a novel scheme for the approximate inference of $\pi$. We start by presenting a midpoint decomposition of the backward transition that allows us to trade adherence to the prior backward dynamics for improved guidance approximation.

##### Midpoint decomposition.

First, consider the DDPM forward process
$q_{k|j}(x_{k}|x_{j})\vcentcolon=\mathcal{N}(x_{k};\sqrt{\alpha_{k}/\alpha_{j}}
x_{j},(1-\alpha_{k}/\alpha_{j})\mathrm{I}_{d})$
defined for all $(j,k)\in\llbracket 0,n\rrbracket^{2}$ such that $j<k$. Define also the joint distribution obtained by initializing the same process with the posterior of interest (2.1)

$$ $\pi_{0:n}(x_{0:n})\vcentcolon=\pi(x_{0})\prod_{k=0}^{n-1}q_{k+1|k}(x_{k+1}|x_{ k})\,.$ (3.1) $$

Note that the identity $\smash{q_{k|0}(x_{k}|x_{0})=\int\prod_{j=0}^{k-1}q_{j+1|j}(x_{j+1}|x_{j})\,
\mathrm{d}x_{1:k-1}}$ implies that the $k$-th marginal of this distribution is $\pi_{k}$ as defined in (3.1).
The time-reversed decomposition of (3.1) writes

$$ $\pi_{0:n}(x_{0:n})\vcentcolon=\pi_{n}(x_{n})\prod_{k=0}^{n-1}\smash{\pi}_{k|k+ 1}(x_{k}|x_{k+1})\,,$ (3.2) $$

where, more generally, for all $(i,j)\in\llbracket 0,n\rrbracket^{2}$ such that $i<j$,

$$ $\smash{\pi}_{i|j}(x_{i}|x_{j})=\pi_{i}(x_{i})q_{j|i}(x_{j}|x_{i})\big{/}\pi_{j }(x_{j})\,.$ (3.3) $$

Note that following (2.2), we have that $\smash{\pi}_{i|j}(x_{i}|x_{j})\propto g_{i}(x_{i})q_{i|j}(x_{i}|x_{j})$.
From this decomposition an exact sample from (2.1) is obtained by first sampling $X_{n}$ from $\pi_{n}$ and then recursively sampling $\smash{X_{k}\sim\smash{\pi}_{k|k+1}(\cdot|X_{k+1})}$ for $k\in\llbracket 0,n-1\rrbracket$. However, the backward-transition kernels (3.3) are difficult to approximate, as they also require knowledge of the intractable potentials $g_{k}$ according to (2.2).
To mitigate the effect of the introduced approximation error associated with these potentials, we consider a new approach to the approximate sampling from $\smash{\pi}_{k|k+1}$ based on the decomposition (3.4) given in Lemma 3.1. First, define, for $\ell\in\llbracket 1,k\rrbracket$,
the bridge kernel $q_{k|\ell,k+1}(x_{k}|x_{\ell},x_{k+1})\propto q_{k|\ell}(x_{k}|x_{\ell})q_{k+1
|k}(x_{k+1}|x_{k})$, which is the density of a Gaussian distribution; see Section A.1.
By convention we assume that $q_{k|k,k+1}(\cdot|x_{k},x_{k+1})=\delta_{x_{k}}$.

###### Lemma 3.1 .

For all $k\in\llbracket 1,n-1\rrbracket$ and ${\ell_{k}}\in\llbracket 0,k\rrbracket$ it holds that

$$ $\smash{\pi}_{k|k+1}(x_{k}|x_{k+1})=\int q_{k|{\ell_{k}},k+1}(x_{k}|x_{\ell_{k} },x_{k+1})\smash{\pi}_{{\ell_{k}}|k+1}(x_{\ell_{k}}|x_{k+1})\,\mathrm{d}x_{ \ell_{k}}\,.$ (3.4) $$

The proof is provided in Section A.3.
From an approximate inference standpoint, this alternative decomposition is appealing for sampling from $\smash{\pi}_{k|k+1}$.
First, notice that sampling from the bridge kernel $q_{k|{\ell_{k}},k+1}$ is straightforward as it involve sampling from a Gaussian distribution, hence we only need to focus our efforts on approximate sampling from $\smash{\smash{\pi}_{{\ell_{k}}|k+1}}$.
This observation is the basis of the introduction of the following surrogate model approximating $\pi_{0:n}$ that we propose. For all $k\in\llbracket 1,n-1\rrbracket$ and ${\ell_{k}}\in\llbracket 0,k\rrbracket$, define the approximation

$$ $\smash{\hat{\pi}^{\ell}_{k|k+1}(x_{k}|x_{k+1})}\vcentcolon=\int q_{k|{\ell_{k} },k+1}(x_{k}|x_{\ell_{k}},x_{k+1})\smash{\hat{\pi}^{\theta}_{{\ell_{k}}|k+1}(x _{\ell_{k}}|x_{k+1})}\,\mathrm{d}x_{{\ell_{k}}}$ (3.5) $$

of (3.4), where

$$ $\smash{\hat{\pi}^{\theta}_{{\ell_{k}}|k+1}(x_{\ell_{k}}|x_{k+1})}\propto\hat{g }^{\theta}_{{\ell_{k}}}(x_{\ell_{k}})p^{\theta}_{{\ell_{k}}|k+1}(x_{\ell_{k}}| x_{k+1})\,,\quad\hat{g}^{\theta}_{{\ell_{k}}}(x_{\ell_{k}})\vcentcolon=g( \smash{m^{\theta}_{0|{\ell_{k}}}}(x_{\ell_{k}}))\,,$ (3.6) $$

and $\smash{p^{\theta}_{{\ell_{k}}|k+1}(\cdot|x_{k+1})\vcentcolon=q_{{\ell_{k}}|0,k
+1}(\cdot|\smash{m^{\theta}_{0|k+1}}(x_{k+1}),x_{k+1})}$ is the DDPM Gaussian approximation of the prior backward transition $q_{{\ell_{k}}|k+1}(\cdot|x_{k+1})$ Ho et al. (). In short, we combine the DPS approximation of $g_{{\ell_{k}}}$ and the DDPM approximation of $q_{{\ell_{k}}|k+1}$. Finally, for any sequence $({\ell_{k}})_{k=1}^{n}$ satisfying ${\ell_{k}}\leq k$, we define our surrogate model for $\pi_{0:n}$ and the posterior (2.1) as

$$ $\smash{\hat{\pi}^{\ell}_{0:n}}(x_{0:n})\vcentcolon=\smash{\hat{\pi}_{n}}(x_{n} )\prod_{k=0}^{n-1}\smash{\hat{\pi}^{\ell}_{k|k+1}(x_{k}|x_{k+1})}\,,\quad \smash{\hat{\pi}^{\ell}_{0}}(x_{0})\vcentcolon=\int\smash{\hat{\pi}^{\ell}_{0: n}}(x_{0:n})\,\mathrm{d}x_{1:n}\,,$ (3.7) $$

where $\smash{\hat{\pi}_{n}}(x_{n})=\mathcal{N}(x_{n};\mathrm{0}_{d},\mathrm{I}_{d})$ and $\smash{\hat{\pi}^{\ell}_{0|1}(\cdot|x_{1})}\vcentcolon=\delta_{\smash{m^{
\theta}_{0|1}}(x_{1})}$ (similarly to Ho et al., ; Song et al., ).
Before providing details on the method we propose to sample from $\smash{\hat{\pi}^{\ell}_{0:n}}$, and consequently $\smash{\hat{\pi}^{\ell}_{0}}$, we give an insight into the choice of the sequence $({\ell_{k}})_{k=1}^{n-1}$.

###### Lemma 3.1 .

##### Choice of the sequence ( ℓ k ) k = 1 n − 1 superscript subscript subscript ℓ 𝑘 𝑘 1 𝑛 1 ({\ell_{k}})_{k=1}^{n-1} ( roman_ℓ start_POSTSUBSCRIPT italic_k end_POSTSUBSCRIPT ) start_POSTSUBSCRIPT italic_k = 1 end_POSTSUBSCRIPT start_POSTSUPERSCRIPT italic_n - 1 end_POSTSUPERSCRIPT .

The accuracy of the surrogate model (3.7) ultimately depends on the design of the sequence $({\ell_{k}})_{k=1}^{n-1}$.
Indeed, for a given $k$, note that decreasing ${\ell_{k}}$ ensures that the approximation $g\circ\smash{m^{\theta}_{0|{\ell_{k}}}}$ of $g_{{\ell_{k}}}$ becomes more accurate. For instance, setting ${\ell_{k}}=0$

Figure: Figure 1: Right: average $W_{2}$ with $10\%$–$90\%$ quantile range. Left: distribution of the minimizing $\eta^{*}$.
Refer to caption: x1.png

eliminates any error of the approximate potential (since $g_{0}=g$); however, decreasing ${\ell_{k}}$ can cause the backward transition $\smash{\smash{\pi}_{0|k+1}(\cdot|x_{k+1})}$ to become highly multimodal, especially in the early stages of the diffusion process, making the DDPM Gaussian approximation less accurate.
Conversely, setting ${\ell_{k}}$ closer to $k$, similar to the DPS approach of Chung et al. (), improves the accuracy of the transition approximation $\smash{q_{{\ell_{k}}|k+1}}$ but worsens the estimate of $\smash{g_{{\ell_{k}}}}$ by $\hat{g}^{\theta}_{{\ell_{k}}}$ as in (3.6).
The choice of ${\ell_{k}}$ is thus subject to a particular trade-off that we illustrate using the following Gaussian toy example. When considering the parameterization $\ell^{\eta}_{k}=\lfloor\eta k\rfloor$ with $\eta\in[0,1]$, we show that the Wasserstein-2 distance between $\pi$ and $\smash{\hat{\pi}^{\ell(\eta)}_{0}}$ reaches its minimum w.r.t. $\eta$ for $\eta$ around $0.5$,
confirming that the additional flexibility provided by the introduction of a midpoint can provide a surrogate model (3.7) with reduced approximation error.

###### Example 3.2 .

Consider a linear inverse problem $Y=\mathrm{A}X+\sigma_{y}Z$ with Gaussian noise and prior $q=\mathcal{N}(m,\Sigma)$, where $\mu\in\mathbb{R}^{d}$ and $\Sigma\in\mathcal{S}^{++}_{d}$.
In this setting, we have access to the exact denoiser $\smash{m_{0|k}}{}$ in a closed form, which allows us to take $\smash{m^{\theta}_{0|k}}=\smash{m_{0|k}}$.
Moreover, since the transition (3.6) admits an explicit expression in this case, we can quantify the approximation error of (3.7) w.r.t. $\pi$ for any sequence $({\ell_{k}})_{k=1}^{n-1}$.
Note that in this case, although the true backward transitions $q_{{\ell_{k}}|k+1}$ are Gaussian and have the same mean as the Gaussian DDPM transitions, they differ in their covariance.
All necessary derivations and computational details can be found in Appendix B.
We consider sequences $(\ell^{\eta}_{k})_{k=1}^{n}$ with $\ell^{\eta}_{k}=\lfloor\eta k\rfloor$, $\eta\in[0,1]$ and compute the Wasserstein-2 distance $W_{2}(\pi,\smash{\hat{\pi}^{\ell(\eta)}_{0}})$ as a function of $\eta$ on the basis of $500$ samples $(\mathrm{A},\mu,\Sigma)$ for ${d}=100$.
The left-hand side plot in Figure 1 displays the average error $W_{2}(\pi,\smash{\hat{\pi}^{\ell(\eta)}_{0}})$ as a function of $\eta$ whereas the right-hand side plot displays the associated distribution of $\eta^{*}:=\operatorname*{argmin}_{\eta\in[0,1]}W_{2}(\pi,\smash{\hat{\pi}^{
\ell(\eta)}_{0}})$.
Figure 1 illustrates that the smallest approximation error for the surrogate model (3.7) is reached at intermediate values close to $\eta=0.5$ rather than $\eta=1$, which corresponds to a DPS-like approximation.

###### Example 3.2 .

##### Variational approximation.

To sample from $\smash{\hat{\pi}^{\ell}_{0:n}}$ in (3.7),
we focus on generating approximate samples from the sequence $(\smash{\hat{\pi}^{\ell}_{k}})_{k=0}^{n}$ of marginal distributions of $\smash{\hat{\pi}^{\ell}_{0:n}}$ in (3.5), by backward recursion over $k$. Ideally, for $k=n$, we simply sample from $\smash{\hat{\pi}^{\theta}_{n}}$; then, recursively, assuming that we have access to an approximate sample $X_{k+1}$ from $\smash{\hat{\pi}^{\ell}_{k+1}}$, the next state $X_{k}$ is drawn from $\smash{\hat{\pi}^{\ell}_{k|k+1}(\cdot|X_{k+1})}$.
However, as $\smash{\hat{\pi}^{\ell}_{k|k+1}(\cdot|X_{k+1})}$ is intractable due to its definition (3.5), we propose to use the Gaussian approximation that we specify next.

In the following, let $k\in\llbracket 1,n\rrbracket$ and ${\ell_{k}}\in\llbracket 0,k\rrbracket$ be fixed. For $\smash{\varphi=(\hat{\mu}_{{\ell_{k}}},\hat{\rho}_{\ell_{k}})\in(\mathbb{R}^{d
})^{2}}$, we consider the Gaussian variational distribution

$$ $\displaystyle\lambda^{\varphi}_{k|k+1}(x_{k}|x_{k+1})$ $\displaystyle\vcentcolon=\int q_{k|{\ell_{k}},k+1}(x_{k}|x_{\ell_{k}},x_{k+1}) \lambda^{\varphi}_{{\ell_{k}}|k+1}(x_{\ell_{k}})\,\mathrm{d}x_{\ell_{k}}\,,$ (3.8) $\displaystyle\lambda^{\varphi}_{{\ell_{k}}|k+1}(x_{\ell_{k}})$ $\displaystyle\vcentcolon=\mathcal{N}(x_{\ell_{k}};\hat{\mu}_{\ell_{k}},\mathrm {diag}(\mathrm{e}^{2\hat{\rho}_{\ell_{k}}}))\,.$ (3.9) $$

In definition (3.9) the exponential is applied element-wise to the vector $\smash{\hat{\rho}_{\ell_{k}}}$ and $\smash{\mathrm{diag}(\mathrm{e}^{\hat{\rho}_{\ell_{k}}})}$ is the diagonal matrix with diagonal entries $\smash{\mathrm{e}^{\hat{\rho}_{\ell_{k}}}}$. Based on the family $\smash{\{\lambda^{\varphi}_{k|k+1}\,:\,\varphi\in(\mathbb{R}^{d})^{2}\}}$ and an approximate sample $X_{k+1}$ from $\smash{\hat{\pi}^{\ell}_{k+1}}$, we then seek the best fitting parameter $\varphi$ that minimizes the upper bound on the backward KL divergence
obtained using the data-processing inequality

$$ $\mathsf{KL}(\lambda^{\varphi}_{k|k+1}(\cdot|X_{k+1})\parallel\smash{\hat{\pi}^ {\ell}_{k|k+1}(\cdot|X_{k+1})})\leq\mathsf{KL}(\lambda^{\varphi}_{{\ell_{k}}|k +1}\parallel\smash{\hat{\pi}^{\theta}_{{\ell_{k}}|k+1}(\cdot|X_{k+1})})= \vcentcolon\mathcal{L}_{k}(\varphi;X_{k+1})\,.$ $$

Here the gradient of the upper bound
is given by

$$ $\nabla_{\varphi}\mathcal{L}_{k}(\varphi;X_{k+1})=-\mathbb{E}\big{[}\nabla_{ \varphi}\log\hat{g}^{\theta}_{{\ell_{k}}}(\hat{\mu}_{{\ell_{k}}}+\mathrm{diag} (\mathrm{e}^{\hat{\rho}_{\ell_{k}}})Z)\big{]}+\nabla_{\varphi}\mathsf{KL}( \lambda^{\varphi}_{{\ell_{k}}|k+1}\parallel p^{\theta}_{{\ell_{k}}|k+1}(\cdot| X_{k+1}))\,,$ $$

where $Z\sim\mathrm{N}(\mathrm{0}_{d},\mathrm{I}_{d})$ is independent of $X_{k+1}$.
The first term is dealt with using the reparameterization trick Kingma & Welling () and can be approximated using a Monte Carlo estimator based on a single sample $Z_{k+1}$; we denote this estimate by $\smash{\nabla_{\varphi}\widetilde{\mathcal{L}}_{k}(\varphi,Z_{k+1};X_{k+1})}$.
The second term is the gradient of the KL divergence between two Gaussian distributions and can thus be computed in a closed form.
Similarly to the previous approaches of Ho et al. (); Chung et al. (); Song et al. (), our gradient estimator involves a vector-Jacobian product of the denoising network.

Figure: Algorithm 1 Midpoint Guidance Posterior Sampling

We now provide a summary of the MGPS algorithm, whose pseudocode is given in Algorithm 1. Given $({\ell_{k}})_{k=1}^{n-1}$, MGPS proceeds by simulating a Markov chain $(X_{k})_{k=n}^{0}$ starting from $X_{n}\sim\mathcal{N}(\mathrm{0}_{d},\mathrm{I}_{d})$.
Recursively, given the state $X_{k+1}$, the state $X_{k}$ is obtained by

- 1.
minimizing $\mathcal{L}_{k}(\cdot;X_{k+1})$ by performing $M$ stochastic gradient steps using the gradient estimator $\nabla_{\varphi}\widetilde{\mathcal{L}}_{k}(\cdot;X_{k+1})$, yielding a parameter $\varphi^{*}_{k}(X_{k+1})$,
- 2.
sampling $\hat{X}_{\ell_{k}}\sim\lambda^{\varphi^{*}_{k}}_{{\ell_{k}}|k+1}$, where we drop the dependence on $X_{k+1}$ in $\varphi^{*}_{k}$, and then $X_{k}\sim q_{k|{\ell_{k}},k+1}(\cdot|\hat{X}_{\ell_{k}},X_{k+1})$.

###### Remark 3.3 .

While conditionally on $\varphi^{*}_{k}$ we have that $\smash{\hat{X}_{\ell_{k}}\sim\lambda^{\varphi^{*}_{k}}_{{\ell_{k}}|k+1}}$, the actual law of $\hat{X}_{\ell_{k}}$ given $X_{k+1}$ is not a Gaussian distribution as we need to marginalize over that of $\varphi^{*}_{k}$ due to the randomness in the stochastic gradients.

A well-chosen initialization for the Gaussian variational approximation parameters is crucial for achieving accurate fitting with few optimization steps, ensuring the overall runtime of MGPS remains competitive with existing methods. Given $\hat{X}_{\ell_{k+1}}$ sampled from $\smash{\lambda^{\varphi^{*}_{k+1}}_{{\ell_{k+1}}|k+2}}$ at the previous step, we choose the initial parameter $\varphi_{k}$ so that

$$ $\lambda^{\varphi_{k}}_{{\ell_{k}}|k+1}(x_{k})=\smash{q_{{\ell_{k}}|0,k+1}(x_{k }|\smash{m^{\theta}_{0|{\ell_{k+1}}}}(\hat{X}_{\ell_{k+1}}),X_{k+1})}\,.$ $$

This initialization is motivated by noting that $q_{{\ell_{k}}|0,k+1}(\cdot|m^{\pi}_{0|k+1}(X_{k+1}),X_{k+1})$, the DDPM approximation of $\smash{\pi}_{{\ell_{k}}|k+1}(\cdot|X_{k+1})$, where $m^{\pi}_{0|k+1}$ is a supposed denoiser for the posterior, is a reasonable candidate for the initialization. As it is intractable, we replace it with $\smash{m^{\theta}_{0|{\ell_{k+1}}}}(\hat{X}_{\ell_{k+1}})$, our best current guess.

Finally, in Section C.2 we devise a warm-start approach that improves the initialization during the early steps of the algorithm. This has proved advantageous in the most challenging problems.

###### Remark 3.3 .

##### Related works.

We now discuss existing works that have similarities with our algorithm, MGPS.
In essence, our method tries to reduce the approximation error incurred by DPS-like approximations.
This has also been the focus of many works in the literature, which we review below.

When $g$ is a likelihood associated to a linear inverse problem with Gaussian noise, Finzi et al. (); Stevens et al. (); Boys et al. () leverage the fact that the Gaussian projection of $\smash{q_{0|k}(\cdot|x_{k})}$ minimizing the forward KL divergence
can be approximated using the pre-trained denoisers; see Meng et al. (, Theorem 1).
In the considered Gaussian setting, the potential $g$ can be exactly integrated against the estimated Gaussian projection.
However, this involves a matrix inversion that may be prohibitively expensive.
Boys et al. () circumvent the latter by using a diagonal approximation which involves the computation of ${d}$ vector-Jacobian products.
For general likelihoods $g$, Song et al. () use the Gaussian approximations of Ho et al. (); Song et al. () to estimate $g_{k}$ for non-linear potentials $g$ using a vanilla Monte Carlo estimate.
Zhang et al. () also consider the surrogate transitions (3.5) with ${\ell_{k}}=k$, and, given an approximate sample $X_{k+1}$ from $\hat{\pi}^{\ell}_{k+1}$, estimate the next sample $X_{k}$ maximizing $x_{k}\mapsto\smash{\hat{\pi}^{\theta}_{{k}|k+1}(x_{k}|X_{k+1})}$ using gradient ascent.
They also consider improved estimates of the potential $g_{k}$ running a few diffusion steps.
However, it is shown (Zhang et al., , last subtable in Table 2) that this brings minor improvements at the expense of a sharp increase in computational cost, which is mainly due to backpropagation over the denoisers.

Finally, a second line of work considers the distribution path $(\hat{\pi}_{k})_{k=n}^{0}$, where $\hat{\pi}_{k}(x_{k})\propto g(\smash{m_{0|k}}(x_{k}))q_{k}(x_{k})$ for $k\in\llbracket 0,n-1\rrbracket$ and $\hat{\pi}_{n}=\mathcal{N}(\mathrm{0}_{d},\mathrm{I}_{d})$.
This path bridges the Gaussian distribution and the posterior of interest.
Furthermore, if one is able to accurately sample from $\hat{\pi}_{k+1}$, then these samples can be used to initialize a sampler targeting the next distribution $\hat{\pi}_{k}$. As these distributions are expected to be close, the sampling from $\hat{\pi}_{k}$ can also be expected to be accurate. Repeating this process yields approximate samples from $\pi$ *regardless* of the approximation error in the potentials.
This approach is pursued by Rozet & Louppe (), who combines the update (2.5) with a Langevin dynamics targeting $\hat{\pi}_{k}$.
Wu et al. () use instead sequential Monte Carlo Gordon et al. () to recursively build empirical approximations of each $\hat{\pi}_{k}$ by evolving a sample of $N$ particles.

## 4 Experiments

Figure: Figure 2: MGPS sample images for half mask (left), expand task, Gaussian blur and motion blur (right) on the ImageNet dataset.
Refer to caption: extracted/5901560/figures/mgps_ddm.jpeg

We now evaluate our algorithm on three different problems and compare it with several competitors.
We begin in benchmarking our method on toy Gaussian-mixture targets and image experiments with both pixel space and latent diffusion models. We review the latter in Section A.2.
Finally, we perform inpainting experiments on ECG data.
For experiments based on pixel-space diffusion models, MGPS is benchmarked against state-of-the art methods in the literature: DPS Chung et al. (), PGDM Song et al. (), DDNM Wang et al. (), DiffPir Zhu et al. (), and RedDiff Mardani et al. ().
Regarding experiments using latent diffusion models, we compare against PSLD Rout et al. () and ReSample Song et al. (). Full details on the hyperparameters can be found in Section D.1.

##### Gaussian mixture.

We first evaluate the accuracy of our method and the impact of the hyperparameters on a toy linear inverse problem with Gaussian mixture (GM) as $q$ and a Gaussian likelihood $g(x)\vcentcolon=\mathcal{N}(y;\mathrm{A}x,\sigma_{y}^{2}\mathrm{I}_{d})$, where $\mathrm{A}\in\mathbb{R}^{{d_{y}}\times{d}}$ and $\sigma_{y}=0.05$. We repeat the experiment of Cardoso et al. (, App B.3),
where the prior is a GM with 25 well-separated strata.
For this specific example, both the posterior and the denoiser $\smash{m_{0|k}}$ are available in a closed form. In particular, the posterior can be shown to be a GM. The full details are provided in Section D.3. We consider the settings $(20,1)$ and $(200,1)$ for the dimensions $({d},{d_{y}})$. For each method, we reconstruct $1000$ samples and then compute the sliced Wasserstein (SW) distance to exact samples from the posterior distribution. We repeat this procedure $100$ times with randomly drawn matrices $\mathrm{A}$, and consider the resulting averaged SW distance.
In every replication, the observation is generated as $Y=\mathrm{A}X+\sigma_{y}Z$ where $Z\sim\mathcal{N}(\mathrm{0}_{d_{y}},\mathrm{I}_{d_{y}})$ and $X\sim q$. The results are reported in Section 4.
It can be observed that MGPS with $\eta=0.75$ outperforms all baselines in both settings. Although we have tuned the parameters of DPS  it still exhibits considerable instability and often diverges. To account for these instabilities we set the SW to $10$ when it diverges.

Figure: Figure 3: Left: SW as a function of $\eta$ with ${\ell_{k}}=\lfloor\eta k\rfloor$. Right: SW as a function of the number of gradient steps, for a specific choice of $({\ell_{k}})_{k}$.
Refer to caption: x2.png

*Ablations.*  Next, we perform an ablation study on the total number of gradient steps and the choice of the sequence $({\ell_{k}})_{k=1}^{n-1}$. The right-hand plot in Figure 3 shows the average SW distance as a function of the number of gradient steps per denoising step when using ${\ell_{k}}=\lfloor 3k/4\rfloor$.
For ${d}=20$, we observe a monotonic decrease of the SW distance. For ${d}=200$, the SW distance decreases and then stabilizes after 10 gradient steps, indicating that in this case, MGPS performs well without requiring many additional gradient steps.
We also report the average SW as a function of $\eta\in[0,1]$ in the left-hand plot of Figure 3, following the same sequence choices as in Example 3.2, *i.e.*, $\ell^{\eta}_{k}=\lfloor\eta k\rfloor$. In both settings, the best SW is achieved at $\eta=0.75$. In dimension $d=200$, the SW at $\eta=0.75$ is nearly twice as good as at $\eta=1$, which corresponds to a DPS-like approximation. This demonstrates that the introduced trade-off leads to non-negligeable performance gains.

##### Images.

We evaluate our algorithm on a wide range of linear and nonlinear inverse problems with noise level $\sigma_{y}=0.05$.
As linear problems we consider: image inpainting with a box mask of shape $150\times 150$ covering the center of the image and half mask covering its right-hand side; image Super Resolution (SR) with factors $\times 4$ and $\times 16$; image deblurring with Gaussian blur; motion blur with $61\times 61$ kernel size.
We use the same experimental setup as Chung et al. (, Section 4) for the last two tasks.
Regarding the nonlinear inverse problems on which we benchmark our algorithm, we consider: JPEG dequantization with quality factor $2\%$; phase retrieval with oversampling factor $2$; non-uniform deblurring emulated via the forward model of Tran et al. (); high dynamic range (HDR) following Mardani et al. (, Section 5.2).
Since the JPEG operator is not differentiable, we instead use the differentiable JPEG framework from Shin & Song (). We note that only DPS and RedDiff apply to nonlinear problems.

*Datasets and models.* We test our algorithm on the $256\times 256$ versions of the FFHQ Karras et al. () and ImageNet Deng et al. () datasets.
We chose randomly $50$ images from each to be used in the evaluation of the algorithms.
We leverage pre-trained DDMs that are publicly available.
For FFHQ, we use the pixel-space DDM of Choi et al. () and the LDM of Rombach et al. () with VQ4 autoencoder.
We use the model by Dhariwal & Nichol () for ImageNet.

*Evaluation.* We use LPIPS Zhang et al. () as the metric to evaluate the quality of the reconstructions. Given that the posterior distribution is expected to be highly multimodal for most of the tasks, pixel-wise comparisons with the ground truth, such as PSNR and SSIM, provide limited insight into reconstruction quality.
For all the considered algorithms, we use the hyperparameters proposed in their official implementations if available, otherwise we manually tune them to achieve the best reconstructions. We defer implementation details to Section D.1.
For each inverse problem, we compute the LPIPS between the ground-truth image and a reconstruction and then average over the dataset. For phase retrieval, however, due to the task’s inherent instability, we follow the approach of Chung et al. () by selecting the best reconstruction out of four; further discussion on this can be found in Chung et al. (, Appendix C.6).

*Results.*  We report the results in Table 2 for FFHQ and ImageNet with DDM prior, and Section 4 for FFHQ with LDM prior.
These reveal that MGPS consistently outperforms the competing algorithms across both linear and nonlinear problems. Notably, on nonlinear tasks, MGPS achieves up to a twofold reduction in LPIPS compared to the other methods with DDM prior. It effectively manages the additional nonlinearities introduced by using LDMs and surpasses the state-of-the-art, as demonstrated in Section 4. Reconstructions obtained with MGPS are displayed in Figure 2 and Figure 4. Further examples and comparisons with other competitors are provided in Section D.6. It is seen that MGPS provides high quality reconstructions even for the most challenging examples. For the DDM prior, the results are obtained using the warm-start approach, see Algorithm 3, and setting ${\ell_{k}}=\lfloor k/2\rfloor$ on all the tasks based on the FFHQ  dataset. As for the ImageNet dataset, we use the same configuration on all the tasks, except for Gaussian deblur and motion deblur. For these tasks we found that using ${\ell_{k}}=\lfloor k/2\rfloor\mathbbm{1}_{k\geq\lfloor n/2\rfloor}+k\mathbbm{1
}_{k<\lfloor n/2\rfloor}$ improves performance. We use a similar strategy for the LDM prior. In Section D.2.1, we measure the runtime and GPU memory requirements for each algorithm. The memory requirement of our algorithm is the same as that of DPS and PGDM, but the runtime with $n=300$ is larger. With fewer diffusion steps, the runtime is halved compared to DPS and PGDM and comparable to that of DiffPir, DDNM and RedDiff. Still, MGPS remains competitive and consistently outperforms the baselines, particularly on nonlinear tasks. See Table 7 and 8 for detailed results using $n\in\{50,100\}$ diffusion steps.

**Table 2: Mean LPIPS for various linear and nonlinear imaging tasks on the FFHQ and ImageNet $256\times 256$ datasets with $\sigma_{y}=0.05$.**
|  | FFHQ | ImageNet |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Task | MGPS | DPS | PGDM | DDNM | DiffPir | RedDiff | MGPS | DPS | PGDM | DDNM | DiffPir | RedDiff |
| SR ($\times 4$) | 0.09 | 0.09 | 0.33 | 0.14 | 0.13 | 0.36 | 0.30 | 0.41 | 0.78 | 0.34 | 0.36 | 0.56 |
| SR ($\times 16$) | 0.26 | 0.24 | 0.44 | 0.30 | 0.28 | 0.51 | 0.53 | 0.50 | 0.60 | 0.70 | 0.63 | 0.83 |
| Box inpainting | 0.10 | 0.19 | 0.17 | 0.12 | 0.18 | 0.19 | 0.22 | 0.34 | 0.29 | 0.28 | 0.28 | 0.36 |
| Half mask | 0.20 | 0.24 | 0.26 | 0.22 | 0.23 | 0.28 | 0.29 | 0.44 | 0.38 | 0.38 | 0.35 | 0.44 |
| Gaussian Deblur | 0.15 | 0.16 | 0.87 | 0.19 | 0.12 | 0.23 | 0.32 | 0.35 | 1.00 | 0.45 | 0.29 | 0.54 |
| Motion Deblur | 0.13 | 0.16 | $-$ | $-$ | $-$ | 0.21 | 0.22 | 0.39 | $-$ | $-$ | $-$ | 0.40 |
| JPEG (QF = 2) | 0.16 | 0.39 | 1.10 | $-$ | $-$ | 0.32 | 0.42 | 0.63 | 1.31 | $-$ | $-$ | 0.51 |
| Phase retrieval | 0.11 | 0.46 | $-$ | $-$ | $-$ | 0.25 | 0.47 | 0.62 | $-$ | $-$ | $-$ | 0.60 |
| Nonlinear deblur | 0.23 | 0.52 | $-$ | $-$ | $-$ | 0.66 | 0.44 | 0.88 | $-$ | $-$ | $-$ | 0.67 |
| High dynamic range | 0.07 | 0.49 | $-$ | $-$ | $-$ | 0.20 | 0.10 | 0.85 | $-$ | $-$ | $-$ | 0.21 |

Figure: Figure 4: MGPS samples with LDM on FFHQ dataset.

Table 3: Mean LPIPS with LDM on FFHQ.
Task
MGPS
ReSample
PSLD
SR ($\times 4$)
0.11
0.20
0.22
SR ($\times 16$)
0.30
0.36
0.35
Box inpainting
0.16
0.22
0.26
Half mask
0.25
0.30
0.31
Gaussian Deblur
0.16
0.15
0.35
Motion Deblur
0.18
0.19
0.41
JPEG (QF = 2)
0.20
0.26
$-$
Phase retrieval
0.34
0.41
$-$
Nonlinear deblur
0.26
0.30
$-$
High dynamic range
0.15
0.15
$-$

##### ECG.

An electrocardiogram (ECG) is an electrical recording of the heart’s activity in which the signals generated by the heartbeats are recorded in order to diagnose various cardiac conditions such as cardiac arrhythmias and conduction abnormalities.
Unlike static images, ECGs constitute complex time-series data consisting of 12 electrical signals acquired using 10 electrodes, 4 of which are attached to the limbs and record the ’limb leads’, while the others record the ’precordial leads’ around the heart.
We study two conditional generation problems in ECGs. The first is a forecasting or missing-block (MB) reconstruction problem, where one half of the ECG is reconstructed from the other (see fig. 7). This task evaluates the algorithm’s ability to capture temporal information to predict a coherent signal. The second problem is an inpainting or missing-leads (ML) reconstruction, where the entire ECG is reconstructed from the lead I (see fig. 5). The question is whether we can capture the subtle information contained in lead I and reconstruct a coherent ECG with the same diagnosis as the real ECG. This task is challenging because lead I, being acquired with limb electrodes far from the heart, may contain very subtle features related to specific cardiac conditions.
We train a state-space diffusion model Goel et al. () to generate ECGs using the 20k training ECGs from the ptb-xl dataset Wagner et al. (), and benchmark the posterior sampling algorithm on the 2k test ECGs (see section D.5). We report the Mean Absolute Error (MAE) and the Root Mean Squared Error (RMSE) between ground-truth and reconstructions in Table 4. We demonstrate that a diffusion model trained to generate ECGs can serve as a prior to solve imputation tasks without additional training or fine-tuning, yielding superior results to a diffusion model trained conditionally on the MB task Alcaraz & Strodthoff (). The rationale for this result is discussed in Section D.5.2.
We report in Section 4 the balanced accuracy of diagnosing Left Bundle Branch Block (LBBB), Right Bundle Branch Block (RBBB), Atrial Fibrillation (AF), and Sinus Bradycardia (SB) using the downstream classifier proposed in Strodthoff et al. () (see section D.5.1), applied to both ground-truth and to reconstructed samples from lead I.
See section D.5.2 for sample figures. MGPS outperforms all other posterior sampling algorithms with just 50 diffusion steps.

**Table 4: MAE and RMSE for missing block task on the ptb-xl dataset.**
| Metric | MGPS | DPS | PGDM | DDNM | DiffPir | RedDiff | TrainedDiff |
| --- | --- | --- | --- | --- | --- | --- | --- |
| MAE | $\underline{0.111\pm 2\mathrm{e}\!{-3}}$ | $0.117\pm 4\mathrm{e}\!{-3}$ | $0.118\pm 2\mathrm{e}\!{-3}$ | $\mathbf{0.103\pm 2\mathrm{e}\!{-3}}$ | $0.115\pm 2\mathrm{e}\!{-3}$ | $0.171\pm 3\mathrm{e}\!{-3}$ | $0.116\pm 2\mathrm{e}\!{-3}$ |
| RMSE | $\mathbf{0.225\pm 4\mathrm{e}\!{-3}}$ | $0.232\pm 4\mathrm{e}\!{-3}$ | $0.233\pm 4\mathrm{e}\!{-3}$ | $\mathbf{0.224\pm 4\mathrm{e}\!{-3}}$ | $0.233\pm 4\mathrm{e}\!{-3}$ | $0.287\pm 5\mathrm{e}\!{-3}$ | $0.266\pm 3\mathrm{e}\!{-3}$ |

Figure: Figure 5: 10s ECG reconstruction from lead I. Ground-truth in blue, $10\%$–$90\%$ quantile range in green, random sample in orange.

Table 5: Balanced acc. downstream diagnosis from ECG reconstructed from lead I.
Method
RBBB
LBBB
AF
SB
MGPS50
0.81
0.92
0.94
0.66
MGPS300
0.90
0.93
0.92
0.66
DPS
0.54
0.84
0.79
0.50
PGDM
0.65
0.87
0.88
0.55
DDNM
0.71
0.83
0.86
0.59
DiffPir
0.57
0.80
0.77
0.53
RedDiff
0.73
0.86
0.88
0.60
Ground-truth
0.99
0.98
0.94
0.70

## 5 Conclusion

We have introduced MGPS, a novel posterior sampling algorithm designed to solve general inverse problems using both diffusion models and latent diffusion models as a prior. Our approach is based on trading off the approximation error in the prior backward dynamics for a more accurate approximation of the guidance term.
This strategy has proven to be effective in a variety of numerical experiments across various tasks.
The results show that MGPS consistently performs competitively and often even superior to state-of-the-art methods.

##### Limitations and future directions.

Although the proposed method is promising, it has certain limitations that invite further exploration. First, a detailed analysis of how the algorithm’s performance depends on the choice of the intermediate time steps $({\ell_{k}})_{k=1}^{n-1}$ remains a challenging but critical subject for future research. Although our image and ECG experiments suggest that setting ${\ell_{k}}=\lfloor k/2\rfloor$ leads to significant performance gains on most tasks, we have observed that using an adaptive sequence ${\ell_{k}}=\lfloor\eta_{k}k\rfloor$, where $\eta_{k}$ increases as $k$ decreases, further enhances results on certain tasks, such as Gaussian and motion deblurring, as well as when using latent diffusion models. A second direction for improvement lies in refining the approximation of $g_{k}$ which we believe could lead to overall algorithmic improvements. Specifically, devising an approximation $\hat{g}^{\theta}_{k}$ such that $\nabla\log\hat{g}^{\theta}_{k}$ does not require a vector-Jacobian product could significantly reduce the algorithm’s runtime. Lastly, we see potential in using the two-stage approximation introduced in our warm-start strategy (see Algorithm 3) at every diffusion step. Although this technique is promising, it currently leads to instability as $k$ decreases. Understanding and resolving this instability will be a key focus of our future work.

##### Acknowledgements.

The work of Y.J. and B.M. has been supported by Technology Innovation Institute (TII), project Fed2Learn. The work of Eric Moulines has been partly funded by the European Union (ERC-2022-SYG-OCEAN-101071601). Views and opinions expressed are however those of the author(s) only and do not necessarily reflect those of the European Union or the European Research Council Executive Agency. Neither the European Union nor the granting authority can be held responsible for them. We would like to thank IHU-LIRYC for the computing power made available to us.

## Appendix A Background on denoising diffusion models

### A.1 Denoising diffusion probabilistic models

In this section, we provide further background on DDMs based on the DDPM framework Ho et al. (); Dhariwal & Nichol (); Song et al. (). We rely on definitions provided in the main text.

DDPMs define generative models for $q$ relying only on parametric approximations $(\smash{m^{\theta}_{0|t}})_{t=1}^{T}$ of the denoisers $(\smash{m_{0|t}})_{t=1}^{T}$. These approximate denoisers are usually defined through the parameterization

$$ $\smash{m^{\theta}_{0|t}}(x_{t})=(x_{t}-\sqrt{1-\alpha_{t}}\smash{\epsilon^{ \theta}_{t}}(x_{t}))/\sqrt{\alpha_{t}}$ (A.1) $$

and trained by minimizing the denoising loss

$$ $\sum_{t=1}^{T}w_{t}\mathbb{E}\left[\|\epsilon_{t}-\smash{\epsilon^{\theta}_{t} }(\sqrt{\alpha_{t}}X_{0}+\sqrt{1-\alpha_{t}}\epsilon_{t})\|^{2}\right]\,,$ (A.2) $$

w.r.t. the neural network parameter $\theta$,
where $(\epsilon_{t})_{t=1}^{T}$ are i.i.d. standard normal vectors, $X_{0}\sim q$, and $(w_{t})_{t=1}^{T}$ are some nonnegative weights.
Having trained the denoisers, the generative model for $q$ is defined as follows. Let $(t_{k})_{k=0}^{n}$ be an increasing sequence of time steps in $\llbracket 0,T\rrbracket$ with $t_{0}=0$. We assume that $t_{n}$ is large enough so that $q_{t_{n}}(x_{t_{n}})=\int q(x_{0})q_{t_{n}|0}(x_{t_{n}}|x_{0})\,\mathrm{d}x_{0}$ is approximately the density of a multivariate standard normal distribution. For convenience we assign the index $k$ to any quantity depending on $t_{k}$; *e.g.*, we denote $q_{t_{k}}$ by $q_{k}$. Consider the backward decomposition

$$ $\displaystyle q_{0:n}(x_{0:n})=q(x_{0})\prod_{k=0}^{n-1}q_{k+1|k}(x_{k+1}|x_{k })=q_{n}(x_{n})\prod_{k=0}^{n-1}q_{k|k+1}(x_{k}|x_{k+1})$ $$

of the forward process initialized at $q$, where $q_{k|k+1}(x_{k}|x_{k+1})\propto q_{k}(x_{k})q_{k+1|k}(x_{k+1}|x_{k})$. Next, for $(j,\ell,k)\in\llbracket 0,n\rrbracket^{3}$ such that $j<\ell<k$, define

$$ $\displaystyle m_{\ell|j,k}(x_{j},x_{k})$ $\displaystyle\vcentcolon=\frac{\sqrt{\alpha_{\ell}{}/\alpha_{j}}(1-\alpha_{k}/ \alpha_{\ell})}{1-\alpha_{k}/\alpha_{j}}x_{j}+\frac{\sqrt{\alpha_{k}/\alpha_{ \ell}}(1-\alpha_{\ell}/\alpha_{j})}{1-\alpha_{k}/\alpha_{j}}x_{k}\,,$ (A.3) $\displaystyle v_{\ell|j,k}$ $\displaystyle\vcentcolon=\frac{(1-\alpha_{\ell}/\alpha_{j})(1-\alpha_{k}/ \alpha_{\ell})}{1-\alpha_{k}/\alpha_{j}}\,.$ (A.4) $$

Then the bridge kernel

$$ $q_{\ell|j,k}(x_{\ell}|x_{j},x_{k})=q_{\ell|j}(x_{\ell}|x_{j})q_{k|\ell}(x_{k}| x_{\ell})\big{/}q_{k|j}(x_{k}|x_{j})$ (A.5) $$

is the density of a Gaussian distribution with mean $m_{\ell|j,k}(x_{j},x_{k})$ and covariance $v_{\ell|j,k}\mathrm{I}_{d}$. Using the bridge kernel, $q_{0:n}$ is approximated using the variational approximation

$$ $p^{\theta}_{0:n}(x_{0:n})=p^{\theta}_{n}(x_{n})\prod_{k=0}^{n-1}p^{\theta}_{k| k+1}(x_{k}|x_{k+1})\,,$ $$

where $p^{\theta}_{0|1}(x_{0}|x_{1})=\mathcal{N}(x_{0};\smash{m^{\theta}_{0|1}}(x_{1}
),v_{0|1}\mathrm{I}_{d})$, $v_{0|1}$ being a tunable parameter, and

$$ $p^{\theta}_{k|k+1}(x_{k}|x_{k+1})\vcentcolon=q_{k|0,k+1}(x_{k}|\smash{m^{ \theta}_{0|k+1}}(x_{k+1}),x_{k+1})\,\quad k\in\llbracket 1,n-1\rrbracket\,.$ (A.6) $$

When $n=T$, the denoising objective (A.2)
corresponds to the KL divergence $\mathsf{KL}(q_{0:n}\parallel p^{\theta}_{0:n})$ for a specific choice of weights $(w_{t})_{t}$. In practice, a DDPM is trained using the objective (A.2) with large $T$ but at inference $n$ is usually much smaller.

### A.2 Latent diffusion models

Latent diffusion models (LDM) Rombach et al. () define a DDM in a latent space. Let $\mathcal{E}:\mathbb{R}^{d}\to\mathbb{R}^{p}$ be an encoder function and $\mathcal{D}:\mathbb{R}^{p}\to\mathbb{R}^{d}$ a decoder function. We assume that these functions satisfy $\mathcal{D}_{\sharp}\mathcal{E}_{\sharp}q\approx q$, where, for instance, $\mathcal{E}_{\sharp}q$ denotes the law of the random variable $\mathcal{E}(X)$ where $X\sim q$. A LDM approximating $q$ is given by $\mathcal{D}_{\sharp}p^{\theta}_{0}$, where $p^{\theta}_{0}$ is a diffusion model trained on samples from $\mathcal{E}_{\sharp}q$. Finally, when solving inverse problems with LDMs, we assume that the target distribution is instead

$$ $\pi(x)\propto g(x)\mathcal{D}_{\sharp}\mathcal{E}_{\sharp}q(x)\,.$ $$

Let $X\sim\pi$, $D\sim\mathcal{D}_{\sharp}\mathcal{E}_{\sharp}q$, $E\sim\mathcal{E}_{\sharp}q$, and $Z\sim\overline{\pi}$, where $\overline{\pi}(z)\propto g(\mathcal{D}_{\sharp}(z))\mathcal{E}_{\sharp}q(z)$. For any bounded function $f$ on $\mathbb{R}^{d}$, we have, following the definition of $\pi$, that

$$ $\displaystyle\mathbb{E}[f(X)]=\frac{\mathbb{E}[g(D)f(D)]}{\mathbb{E}[g(D)]}= \frac{\mathbb{E}[g(\mathcal{D}(E))f(\mathcal{D}(E))]}{\mathbb{E}[g(\mathcal{D} (E))]}=\mathbb{E}[f(\mathcal{D}(Z))]\,.$ $$

Hence $\mathsf{Law}(X)=\mathsf{Law}(\mathcal{D}(Z))$. As a result, to sample approximately from $\pi$, we first sample approximately from $\overline{\pi}$ using any diffusion posterior sampling algorithm with pre-trained DDM for $\mathcal{E}_{\sharp}q$, then decode the obtained samples using $\mathcal{D}$.

### A.3 Midpoint decomposition

Before we proceed with the midpoint decomposition, we first recall that under the joint distribution obtained by initializing the posterior $\pi$ with the forward process (see (3.1)), it holds that for all $i<j$,

$$ $\smash{\pi}_{i}(x_{i})q_{j|i}(x_{j}|x_{i})=\smash{\pi}_{j}(x_{j})\smash{\pi}_{ i|j}(x_{i}|x_{j})\,,$ (A.7) $$

where $\smash{\pi}_{i|j}(x_{i}|x_{j})\vcentcolon=\pi_{i}(x_{i})q_{j|i}(x_{j}|x_{i})/
\pi_{j}(x_{j})$ and integrates to one.

###### Proof of Lemma 3.1 .

Let $(\ell,k)\in\llbracket 0,n\rrbracket$ be such that $\ell<k$.
Applying repeatedly the definition of the bridge kernel (A.5) and the identity (A.7),

$$ $\displaystyle q_{k|\ell,k+1}(x_{k}|x_{\ell},x_{k+1})$ $\displaystyle=\frac{q_{k|\ell}(x_{k}|x_{\ell})q_{k+1|k}(x_{k+1}|x_{k})}{q_{k+1 |\ell}(x_{k+1}|x_{\ell})}$ $\displaystyle=\frac{\smash{\pi}_{\ell}(x_{\ell})q_{k|\ell}(x_{k}|x_{\ell})q_{k +1|k}(x_{k+1}|x_{k})}{\smash{\pi}_{\ell}(x_{\ell})q_{k+1|\ell}(x_{k+1}|x_{\ell })}$ $\displaystyle=\frac{\smash{\pi}_{\ell|k}(x_{\ell}|x_{k})\smash{\pi}_{k|k+1}(x_ {k}|x_{k+1})\smash{\pi}_{k+1}(x_{k+1})}{\smash{\pi}_{\ell|k+1}(x_{\ell}|x_{k+1 })\smash{\pi}_{k+1}(x_{k+1})}$ $\displaystyle=\frac{\smash{\pi}_{\ell|k}(x_{\ell}|x_{k})\smash{\pi}_{k|k+1}(x_ {k}|x_{k+1})}{\smash{\pi}_{\ell|k+1}(x_{\ell}|x_{k+1})}\enspace.$ $$

It then follows that

$$ $\displaystyle\smash{\pi}_{k|k+1}(x_{k}|x_{k+1})$ $\displaystyle=\int\smash{\pi}_{\ell|k}(x_{\ell}|x_{k})\smash{\pi}_{k|k+1}(x_{k }|x_{k+1})\,\mathrm{d}x_{\ell}$ $\displaystyle=\int q_{k|\ell,k+1}(x_{k}|x_{\ell},x_{k+1})\smash{\pi}_{\ell|k+1 }(x_{\ell}|x_{k+1})\,\mathrm{d}x_{\ell}\,.$ $$

∎

###### Proof of Lemma 3.1 .

## Appendix B The Gaussian case

### B.1 Derivation

In this section we derive the recursions verified by the first and second moments of the marginal distribution $\smash{\hat{\pi}^{\ell}_{0}}$ of the surrogate model (3.7) in the simplified setting of Example 3.2. We recall that in this specific example we assume that $q=\mathrm{N}(m,\Sigma)$ where $(m,\Sigma)\in\mathbb{R}^{d}\times\mathcal{S}^{++}_{d}$ and $g:x\mapsto\mathcal{N}(y;\mathrm{A}x,\sigma_{y}^{2}\mathrm{I}_{d_{y}})$.

##### Denoiser and DDPM transitions.

Since we are dealing with a Gaussian prior, the denoiser $\smash{m_{0|k}}$ can be computed in closed form for any $k\in\llbracket 1,n\rrbracket$. Using Bishop (, Eqn. 2.116), we have that

$$ $\displaystyle q_{0|k}(x_{0}|x_{k})$ $\displaystyle\propto q(x_{0})q_{k|0}(x_{k}|x_{0})$ $\displaystyle=\mathcal{N}\big{(}x_{0};\Sigma_{0|k}\big{(}(\sqrt{\alpha_{k}}/v_ {k})x_{k}+\Sigma^{-1}m\big{)},\Sigma_{0|k}\big{)},$ $$

where $\Sigma_{0|k}\vcentcolon=((\alpha_{k}/v_{k})\mathrm{I}+\Sigma^{-1})^{-1}$. Hence,

$$ $\smash{m_{0|k}}(x_{k})=\Sigma_{0|k}\big{(}(\sqrt{\alpha_{k}}/v_{k})x_{k}+ \Sigma^{-1}m\big{)},$ $$

and we assume in the remainder of this section that $\smash{m^{\theta}_{0|k}}=\smash{m_{0|k}}$. From the expression of $q_{0|k}(\cdot|x_{k})$ we can immediately derive the more general backward transitions $q_{\ell|k}(\cdot|x_{k})$ for $\ell\in\llbracket 1,k-1\rrbracket$ by noting that

$$ $q_{\ell|k}(x_{\ell}|x_{k})=\int q_{\ell|0,k}(x_{\ell}|x_{0},x_{k})q_{0|k}(x_{0 }|x_{k})\,\mathrm{d}x_{0}\,.$ $$

From this, (A.3), (A.4), and the law of total expectation and covariance, it follows that $q_{\ell|k}(\cdot|x_{k})=\mathrm{N}(m_{\ell|k}(x_{k}),\Sigma_{\ell|k}(x_{k}))$, where

$$ $\displaystyle m_{\ell|k}(x_{k})=m_{\ell|0,k}(\smash{m_{0|k}}(x_{k}),x_{k})\,, \quad\Sigma_{\ell|k}=\frac{\alpha_{\ell}(1-\alpha_{k}/\alpha_{\ell})^{2}}{(1- \alpha_{k})^{2}}\Sigma_{0|k}+v_{\ell|0,k}\,.$ $$

On the other hand, the DDPM transitions are

$$ $p^{\theta}_{\ell|k}(\cdot|x_{k})=q_{\ell|0,k}(\cdot|\smash{m_{0|k}}(x_{k}),x_{ k})=\mathrm{N}(m_{\ell|0,k}(\smash{m_{0|k}}(x_{k}),x_{k}),v_{\ell|0,k}\mathrm{ I}_{d}),$ (B.1) $$

which shows that in this case, the true transitions and approximate ones differ only by their covariance.

##### Moments recursion.

For $k\in\llbracket 0,n\rrbracket$ we let $\smash{\hat{\pi}^{\ell}_{k}}$ denote the $x_{k}$ marginal of the surrogate model (3.7). We remind the reader that $\smash{\hat{\pi}^{\ell}_{n}}=\mathrm{N}(\mathrm{0}_{d},\mathrm{I}_{d})$. The marginals satisfy the recursion

$$ $\smash{\hat{\pi}^{\ell}_{k}}(x_{k})=\int\smash{\hat{\pi}^{\ell}_{k|k+1}(x_{k}| x_{k+1})}\smash{\hat{\pi}^{\ell}_{k+1}}(x_{k+1})\,\mathrm{d}x_{k+1}\,,\quad k \in\llbracket 0,n-1\rrbracket\,.$ $$

Since $\hat{g}^{\theta}_{k}(x_{k})=\mathcal{N}(y;\mathrm{A}\smash{m^{\theta}_{0|k}}(x
_{k}),\sigma_{y}^{2}\mathrm{I}_{d_{y}})$ and $\smash{m^{\theta}_{0|k}}$ is linear in $x_{k}$, it is easily seen that $\smash{\hat{\pi}^{\theta}_{{\ell_{k}}|k+1}(\cdot|x_{k+1})}$ is the density of a Gaussian distribution. Consequently, by definition (3.5) and the definition (A.5) of the bridge kernel, this is also the case for $\smash{\hat{\pi}^{\ell}_{k|k+1}(\cdot|x_{k+1})}$. Now assume that

$$ $\displaystyle\smash{\hat{\pi}^{\ell}_{k+1}}(x_{k+1})$ $\displaystyle=\mathcal{N}(x_{k+1};\hat{\mu}^{\ell}_{k+1},\hat{\Sigma}^{\ell}_{ k+1})\,,$ (B.2) $\displaystyle\smash{\hat{\pi}^{\ell}_{k|k+1}(x_{k}|x_{k+1})}$ $\displaystyle=\mathcal{N}(x_{k};\mathrm{M}^{\ell}_{k|k+1}x_{k+1}+c^{\ell}_{k|k +1},\hat{\Sigma}^{\ell}_{k|k+1})\,,$ (B.3) $$

where $\mathrm{M}^{\ell}_{k|k+1}\in\mathbb{R}^{{d}\times{d}}$, $\hat{\Sigma}^{\ell}_{k|k+1}\in\mathcal{S}^{++}_{d}$, and $c^{\ell}_{k|k+1}\in\mathbb{R}^{{d}}$.
Using the definition (A.5) of the bridge kernel we find that $\smash{\hat{\pi}^{\ell}_{k}}=\mathrm{N}(\hat{\mu}_{k},\hat{\Sigma}_{k})$, where

$$ $\displaystyle\hat{\mu}^{\ell}_{k}$ $\displaystyle=\mathrm{M}^{\ell}_{k|k+1}\hat{\mu}^{\ell}_{k+1}+c^{\ell}_{k|k+1}\,,$ $\displaystyle\hat{\Sigma}^{\ell}_{k}$ $\displaystyle=\mathrm{M}^{\ell}_{k|k+1}\hat{\Sigma}^{\ell}_{k+1}\mathrm{M}^{ \intercal}_{k|k+1}+\hat{\Sigma}^{\ell}_{k|k+1}\,.$ $$

Iterating these updates until reaching $k=0$, starting from the initialization $\hat{\mu}^{\ell}_{n}=\mathrm{0}_{d}$ and $\hat{\Sigma}^{\ell}_{n}=\mathrm{I}_{d}$, yields the desired moments of the surrogate posterior $\smash{\pi}^{\ell}_{0}$. It now remains to show that the backward transition $\smash{\hat{\pi}^{\ell}_{k|k+1}(\cdot|x_{k+1})}$ writes in the form (B.3) and identify $\mathrm{M}^{\ell}_{k|k+1}$ and $c^{\ell}_{k|k+1}$.

First, we write the approximate potential in the form

$$ $\hat{g}^{\theta}_{k}(x_{k})=\mathcal{N}(y;\hat{\mathrm{A}}_{k}x_{k}+b_{k}, \sigma_{y}^{2}\mathrm{I}_{d_{y}}),$ $$

where

$$ $\hat{\mathrm{A}}_{k}=(\sqrt{\alpha_{k}}/v_{k})\mathrm{A}\Sigma_{0|k},\quad b_{ k}=\mathrm{A}\Sigma_{0|k}\Sigma^{-1}m\,.$ $$

We also denote by $m^{\theta}_{\ell|k}(x_{k})$ the mean of the Gaussian distribution with density given by the DDPM transition $p^{\theta}_{\ell|k}(\cdot|x_{k})$ in (B.1). We have that

$$ $m^{\theta}_{{\ell_{k}}|k+1}(x_{k+1})=\mathrm{H}_{{\ell_{k}}|k+1}x_{k+1}+h_{{ \ell_{k}}|k+1}\,,$ $$

where

$$ $\displaystyle\mathrm{H}_{{\ell_{k}}|k+1}$ $\displaystyle\vcentcolon=\frac{\alpha_{{\ell_{k}}}(1-\alpha_{k+1}/\alpha_{{ \ell_{k}}})}{v_{\ell_{k}}v_{k+1}}\Sigma_{0|k}+\frac{\sqrt{\alpha_{k+1}}(1- \alpha_{{\ell_{k}}})}{v_{k+1}}\mathrm{I}_{d}\,,$ (B.4) $\displaystyle h_{{\ell_{k}}|k+1}$ $\displaystyle\vcentcolon=\frac{\sqrt{\alpha_{{\ell_{k}}}}(1-\alpha_{k+1}/ \alpha_{{\ell_{k}}})}{v_{k+1}}\Sigma_{0|k}\Sigma^{-1}m\,.$ (B.5) $$

Then, applying (Bishop, , Eqn. 2.116), we get

$$ $\displaystyle\smash{\hat{\pi}^{\theta}_{{\ell_{k}}|k+1}(x_{{\ell_{k}}}|x_{k+1} )}=\mathcal{N}(x_{{\ell_{k}}};\widetilde{\mathrm{M}}^{\theta}_{{\ell_{k}}|k+1} x_{k+1}+\tilde{c}^{\theta}_{{\ell_{k}}|k+1},\Gamma_{{\ell_{k}}|k+1})\,,$ $$

where

$$ $\displaystyle\Gamma_{{\ell_{k}}|k+1}$ $\displaystyle\vcentcolon=\big{(}v^{-1}_{{\ell_{k}}|0,k+1}\mathrm{I}+\sigma_{y} ^{-2}\hat{\mathrm{A}}_{{\ell_{k}}}^{\top}\hat{\mathrm{A}}_{{\ell_{k}}}\big{)}^ {-1}\,,$ $\displaystyle\widetilde{\mathrm{M}}^{\theta}_{{\ell_{k}}|k+1}$ $\displaystyle\vcentcolon=v^{-1}_{{\ell_{k}}|0,k+1}\Gamma_{{\ell_{k}}|k+1} \mathrm{H}_{{\ell_{k}}|k+1}\,,$ $\displaystyle\tilde{c}^{\theta}_{{\ell_{k}}|k+1}$ $\displaystyle\vcentcolon=\Gamma_{{\ell_{k}}|k+1}\big{[}\sigma_{y}^{-2}\hat{ \mathrm{A}}_{{\ell_{k}}}^{\top}(y-b_{{\ell_{k}}})+v^{-1}_{{\ell_{k}}|0,k+1}h_{ {\ell_{k}}|k+1}\big{]}\,.$ $$

Finally, following (3.5) we find that

$$ $\displaystyle\mathrm{M}^{\ell}_{k|k+1}$ $\displaystyle=\frac{\sqrt{\alpha_{k}/\alpha_{{\ell_{k}}}}(1-\alpha_{k+1}/ \alpha_{k})}{1-\alpha_{k+1}/\alpha_{{\ell_{k}}}}\widetilde{\mathrm{M}}^{\theta }_{{\ell_{k}}|k+1}+\frac{\sqrt{\alpha_{k+1}/\alpha_{k}}(1-\alpha_{k}/\alpha_{{ \ell_{k}}})}{1-\alpha_{k+1}/\alpha_{{\ell_{k}}}}\mathrm{I}_{d}\,,$ $\displaystyle c^{\ell}_{k|k+1}$ $\displaystyle=\frac{\sqrt{\alpha_{k}/\alpha_{{\ell_{k}}}}(1-\alpha_{k+1}/ \alpha_{k})}{1-\alpha_{k+1}/\alpha_{{\ell_{k}}}}\tilde{c}^{\theta}_{{\ell_{k}} |k+1}\,.$ $$

### B.2 Experimental setup

In the experiment described in Example 3.2, we set ${d}=100$ and generate $500$ instances of inverse problems $(\mathrm{A},m,\Sigma)$.
For each instance, we compute the Wasserstein-2 distance between the resulting posterior distribution $\pi$ and $\smash{\hat{\pi}^{\ell(\eta)}_{0}}$.

##### Prior.

The mean of the prior is sampled from a standard Gaussian distribution.
To generate the covariance $\Sigma\in\mathcal{S}^{++}_{d}$, we first draw a matrix $\mathrm{G}\in\mathbb{R}^{{d}\times{d}}$ with i.i.d. entries sampled from $\mathrm{N}(0,1)$.
Then, for a better conditioning of $\Sigma$, we normalize the columns of $\mathrm{G}$ and set $\Sigma=\bar{\lambda}^{2}\mathrm{I}+\mathrm{G}\mathrm{G}^{\top}$, where $\bar{\lambda}^{2}$ is the mean of the squared singular values of $\mathrm{G}$.

##### Likelihood.

To sample an ill-posed problem we generate a rank-deficient matrix $\mathrm{A}\in\mathbb{R}^{{d}_{y}\times{d}}$ with ${d}_{y}\leq{d}$.
We sample uniformly ${d}_{y}$ from the interval $\llbracket{d}/10,{d}\rrbracket$ and draw the entries of $\mathrm{A}$ i.i.d. from $\mathrm{N}(0,1)$.
Regarding $\sigma_{y}$, we sample it uniformly from the interval $[0.1,0.5]$.

Finally, the resulting posterior is also Gaussian (Bishop, , Eqn. 2.116)

$$ $\pi(x)\propto g(x)p(x)\propto\mathcal{N}(x;m_{y},\Sigma_{y})\,,$ $$

where $\Sigma_{y}=(\Sigma^{-1}+(1/\sigma_{y}^{2})\mathrm{A}^{\top}\mathrm{A})^{-1}$ and $m_{y}=\Sigma_{y}\big{(}(1/\sigma_{y}^{2})\mathrm{A}^{\top}y+\Sigma^{-1}m\big{)}$.
The Wasserstein-2 distance between the true posterior and $\smash{\hat{\pi}^{\ell(\eta)}_{0}}$ is thus Olkin & Pukelsheim ()

$$ $W_{2}(\pi,\smash{\hat{\pi}^{\ell}_{0}})^{2}=\|m_{y}-\hat{\mu}^{\ell}_{0}\|^{2} +\operatorname{tr}\big{(}\Sigma_{y}+\hat{\Sigma}^{\ell}_{0}-(\Sigma^{\frac{1}{ 2}}_{y}\hat{\Sigma}^{\ell}_{0}\Sigma^{\frac{1}{2}}_{y})^{\frac{1}{2}}\big{)}\,.$ $$

## Appendix C More details on MGPS

### C.1 Loss gradient

Recall that the loss function we optimize is
$\varphi\mapsto\mathsf{KL}(\lambda^{\varphi}_{{\ell_{k}}|k+1}\parallel\smash{
\hat{\pi}^{\theta}_{{\ell_{k}}|k+1}(\cdot|x_{k+1})})$
for a given $x_{k+1}$. Its gradient writes

$$ $\displaystyle\nabla_{\varphi}\mathcal{L}_{k}(\varphi;X_{k+1})$ $\displaystyle=-\mathbb{E}\big{[}\nabla_{\varphi}\log\hat{g}^{\theta}_{{\ell_{k }}}(\hat{\mu}_{{\ell_{k}}}+\mathrm{diag}(\mathrm{e}^{\hat{\rho}_{\ell_{k}}})Z) \big{]}+\nabla_{\varphi}\mathsf{KL}(\lambda^{\varphi}_{{\ell_{k}}|k+1} \parallel p^{\theta}_{{\ell_{k}}|k+1}(\cdot|X_{k+1}))$ $$

and the second term can be expressed in a closed form since both distributions are Gaussians, *i.e.*,

$$ $\nabla_{\varphi}\mathsf{KL}(\lambda^{\varphi}_{{\ell_{k}}|k+1}\parallel p^{ \theta}_{{\ell_{k}}|k+1}(\cdot|x_{k+1}))\\ =\nabla_{\varphi}\left[-\sum_{j=1}^{d}\hat{\rho}_{{\ell_{k}},j}+\frac{\|\hat{ \mu}_{\ell_{k}}-m_{{\ell_{k}}|0,k+1}(\smash{m^{\theta}_{0|k+1}}(x_{k+1}),x_{k+ 1})\|^{2}+\sum_{j=1}^{d}\frac{\mathrm{e}^{2\hat{\rho}_{{\ell_{k}},j}}}{v_{{ \ell_{k}}|0,k+1}}}{2v_{{\ell_{k}}|0,k+1}}\right]\,.$ $$

### C.2 Warm start

In this section we describe the warm-start approach discussed in the main paper. The complete algorithm with the warm-start procedure is given in Algorithm 3. The original version presented in Algorithm 1 relies on first sampling approximately $X_{{\ell_{k}}}$, given $X_{k+1}$, from the surrogate transition $\smash{\hat{\pi}^{\theta}_{{\ell_{k}}|k+1}(\cdot|X_{k+1})}$ and then sampling $X_{k}$ from $q_{k|{\ell_{k}},k+1}(\cdot|X_{\ell_{k}},X_{k+1})$. In our warm-start approach, which we apply only during the first iterations of the algorithm (see the hyperparameter settings in Table 6), we draw inspiration from the decomposition

$$ $\smash{\pi}_{k|k+1}(x_{k}|x_{k+1})=\int q_{k|1,k+1}(x_{k}|x_{1},x_{k+1})\smash {\pi}_{1|{\ell_{k}}}(x_{1}|x_{\ell_{k}})\smash{\pi}_{{\ell_{k}}|k+1}(x_{\ell_{ k}}|x_{k+1})\,\mathrm{d}x_{1}\mathrm{d}x_{\ell_{k}}\,.$ (C.1) $$

It suggests introducing a second intermediary step that involves sampling approximately from the transition $\smash{X_{1}\sim\smash{\pi}_{1|{\ell_{k}}}(\cdot|X_{\ell_{k}})}$. $X_{k}$ is then sampled from the bridge $q_{k|1,k+1}(\cdot|X_{1},X_{k+1})$.

In order to draw approximate samples from $\smash{\pi}_{1|{\ell_{k}}}(\cdot|X_{{\ell_{k}}})$ we leverage again a Gaussian variational approximation $\smash{\lambda^{\psi}_{1|{\ell_{k}}}\vcentcolon=\mathrm{N}(\tilde{\mu}_{1},
\mathrm{diag}(\mathrm{e}^{2\tilde{\rho}_{1}}))}$, which we fit by minimizing a proxy of the KL divergence between $\lambda^{\psi}_{1|{\ell_{k}}}$ and $\smash{\pi}_{1|{\ell_{k}}}(\cdot|X_{\ell_{k}})$ of which the gradient w.r.t. $\psi\vcentcolon=(\tilde{\mu}_{1},\tilde{\rho}_{1})$ writes

$$ $\nabla_{\psi}\mathcal{L}^{\mathrm{ws}}_{1|{\ell_{k}}}(\psi;X_{\ell_{k}})_{| \psi_{0}}\vcentcolon=-\nabla_{\psi}\left[\sum_{j=1}^{d}\tilde{\rho}_{1,j}- \frac{\|X_{\ell_{k}}-(\alpha_{{\ell_{k}}}/\alpha_{1})^{1/2}\tilde{\mu}_{1}\|^{ 2}+(\alpha_{{\ell_{k}}}/\alpha_{1})\sum_{j=1}^{d}\mathrm{e}^{\tilde{\rho}^{2}_ {1,j}}}{2(1-\alpha_{{\ell_{k}}}/\alpha_{1})}\right]_{|\psi_{0}}\\ -\mathbb{E}\left[\nabla_{\psi}(\tilde{\mu}_{1}+\mathrm{diag}(\mathrm{e}^{ \tilde{\rho}_{1}})Z)^{\intercal}_{|\psi_{0}}(\nabla_{x}\log g+s^{\theta}_{1})( \tilde{\mu}_{1}+\mathrm{diag}(\mathrm{e}^{\tilde{\rho}_{1}})Z\,{}_{|\psi_{0}}) \right]\,.$ (C.2) $$

This expression corresponds to the gradient w.r.t. $\psi$ of $\mathsf{KL}(\lambda^{\psi}_{1|{\ell_{k}}}\parallel\smash{\pi}_{1|{\ell_{k}}}(
\cdot|X_{\ell_{k}}))$ combined with the score approximation $s^{\theta}_{1}$ of $\nabla_{x}\log q_{1}$ and the mild potential approximation $\nabla_{x}\log g_{1}$ of $\nabla_{x}\log g$. Indeed, we have that

$$ $\mathsf{KL}(\lambda^{\psi}_{1|{\ell_{k}}}\parallel\smash{\pi}_{1|{\ell_{k}}}( \cdot|X_{\ell_{k}}))=-\mathbb{E}_{\lambda^{\psi}_{1|{\ell_{k}}}}\big{[}\log g_ {1}(X_{1})+\log q_{1}(X_{1})\big{]}\\ +\int\log\frac{\lambda^{\psi}_{1|{\ell_{k}}}(x_{1})}{q_{{\ell_{k}}|1}(X_{\ell_ {k}}|x_{1})}\lambda^{\psi}_{1|{\ell_{k}}}(x_{1})\,\mathrm{d}x_{1}\,.$ $$

The second term can be computed in a closed form and its gradient corresponds to the first term in (C.2). As for the first term, under standard differentiability assumptions and applying the raparameterization trick, we find that

$$ $\displaystyle\nabla_{\psi}$ $\displaystyle\mathbb{E}_{\lambda^{\varphi}_{1|{\ell_{k}}}}\big{[}\log g_{1}(X_ {1})+\log q_{1}(X_{1})\big{]}_{|\psi_{0}}$ $\displaystyle=\nabla_{\psi}\mathbb{E}\left[\big{(}\log g_{1}+\log q_{1}\big{)} (\tilde{\mu}_{1}+\mathrm{diag}(\mathrm{e}^{\tilde{\rho}_{1}})Z)\right]_{|\psi_ {0}}$ $\displaystyle=\mathbb{E}\left[\nabla_{\psi}(\tilde{\mu}_{1}+\mathrm{diag}( \mathrm{e}^{\tilde{\rho}_{1}})Z)^{\intercal}_{|\psi_{0}}(\nabla_{x}\log g_{1}+ \nabla_{x}\log q_{1})(\tilde{\mu}_{1}+\mathrm{diag}(\mathrm{e}^{\tilde{\rho}_{ 1}})Z\,{}_{|\psi_{0}})\right]\,.$ $$

Plugging the previous approximations yields the second term in (C.2). The warm-start procedure is summarized in Algorithm 2. We also use a single sample Monte Carlo estimate to perform the optimization. Finally, at step $k$ and given $\hat{X}_{\ell_{k}}$, the initial parameter $\psi_{k}$ of the variational approximation is chosen so that

$$ $\lambda^{\psi_{k}}_{1|{\ell_{k}}}(x_{1})=q_{1|0,{\ell_{k}}}(x_{1}|\smash{m^{ \theta}_{0|{\ell_{k}}}}(\hat{X}_{\ell_{k}}),\hat{X}_{\ell_{k}})\,.$ $$

After having sampled $\tilde{X}_{1}\sim\lambda^{\psi^{*}_{k}}_{1|{\ell_{k}}}$, we use it to first sample $X_{k}\sim q_{k|1,k+1}(\cdot|\tilde{X}_{1},X_{k+1})$ and then initialize the next variational approximation $\lambda^{\varphi}_{{\ell_{k-1}}|k}$. The initial parameter $\varphi_{k-1}$ is set so that

$$ $\lambda^{\varphi_{k-1}}_{{\ell_{k-1}}|k}(x_{\ell_{k-1}})=q_{{\ell_{k-1}}|0,k}( x_{\ell_{k-1}}|\tilde{X}_{1},X_{k})\,,$ $$

see line 5 in Algorithm 3.

Figure: Algorithm 2 $\mathsf{WarmStart}$

Figure: Algorithm 3 MGPS with warm start strategy

## Appendix D Experiments

In this section we provide the implementation details on our algorithm as well as the algorithms we benchmark against.
We use the the hyperparameters recommended by the authors and tune them on each dataset if they are not provided.

### D.1 Implementation details and hyperparameters for MGPS

We implement Algorithm 1 with ${\ell_{k}}=\lfloor k/2\rfloor$ and use the Adam optimizer Kingma & Ba () with a learning rate of 0.03 for optimization. The number of gradient steps is adjusted based on the complexity of the task: posterior sampling with the ImageNet DDM prior or FFHQ LDM prior is more challenging and therefore requires additional gradient steps. Detailed hyperparameters are provided in Table 6.

The warm-start strategy outlined in Section C.2 improved reconstruction plausibility and eliminated potential artifacts. A similar effect was observed when performing multiple gradient steps ($M=20$) during the initial stages. For latent-space models, switching the intermediate step to ${\ell_{k}}=k$ for the second half of the diffusion process has been crucial and significantly enhanced reconstruction quality by mitigating the smoothing effect, which often removes important details. A similar strategy has been useful for the Gaussian deblurring and motion deblurring tasks on the ImageNet dataset.

**Table 6: The hyperparameters used in MGPS for the considered datasets.**
|  | Warm start | Threshold ($w$) | Diffusion steps | ${\ell_{k}}$ | Learning rate | Gradient steps |
| --- | --- | --- | --- | --- | --- | --- |
| FFHQ | ✓ | $\lfloor 3n/4\rfloor$ | $n\in\{50,100,300\}$ | $\lfloor k/2\rfloor$ | $0.03$ | $M_{k}=\begin{cases}20&\text{ if }k\geq n-5\\<br>20&\text{ if }k\mod 10=0\\<br>2&\text{ otherwise}\\<br>\end{cases}$ |
| FFHQ (LDM) | ✓ | $\lfloor 3n/4\rfloor$ | $n\in\{50,100,300\}$ | ${\ell_{k}}=\begin{cases}\lfloor k/2\rfloor&\text{ if }k>\lfloor n/2\rfloor\\<br>k&\text{ otherwise }\end{cases}$ | $0.03$ | $M_{k}=\begin{cases}20&\text{ if }k\geq n-5\\<br>10&\text{ if }k\mod 10=0\\<br>5&\text{ otherwise}\\<br>\end{cases}$ |
| ImageNet | ✓ | $\lfloor 3n/4\rfloor$ | $n\in\{50,100,300\}$ | $\lfloor k/2\rfloor$ | $0.03$ | $M_{k}=\begin{cases}20&\text{ if }k\geq n-5\\<br>10&\text{ if }k\mod 20=0\\<br>2&\text{ otherwise}\\<br>\end{cases}$ |
| Gaussian Mixture | $\times$ | – | $n=300$ | $\lfloor k/2\rfloor$ | $0.1$ | $M_{k}=\begin{cases}20&\text{ if }k\geq n-5\\<br>20&\text{ if }k\mod 10=0\\<br>2&\text{ otherwise}\\<br>\end{cases}$ |
| ptb-xl (ECG) | ✓ | $\lfloor 3n/4\rfloor$ | $n\in\{50,300\}$ | $\lfloor k/2\rfloor$ | $0.03$ | $M_{k}=\begin{cases}20&\text{ if }k\geq n-5\\<br>5&\text{ if }k\mod 20=0\\<br>5&\text{ otherwise}\\<br>\end{cases}$ |

### D.2 Implementation details of competitors

##### DPS.

We implemented Chung et al. (, Algorithm 1) and refer to Chung et al. (, App. D) for the values of its hyperparameters.
After tuning, we adopt $\gamma=0.2$ for JPEG $2\%$ and $\gamma=0.07$ for High Dynamic Range tasks.

##### DiffPIR.

We implemented Zhu et al. (, Algorithm 1) to make it compatible with our existing code base.
We used the hyperparameters recommended in the official, released version(^1^11[https://github.com/yuanzhi-zhu/DiffPIR](https://github.com/yuanzhi-zhu/DiffPIR)).
Unfortunately, we did not manage to make the algorithm converge for nonlinear problems.
While the authors give some guidelines to handle such problems (Zhu et al., , Eqn. (13)), examples are missing in the paper and the released code.
Similarly, we do not run it on the motion deblur task as the FFT-based solution provided in Zhu et al. () is only valid for circular convolution, yet we opted for the experimental setup of Chung et al. () which uses convolution with reflect padding.

##### DDNM.

We adapted the implementation in the released code(^2^22[https://github.com/wyhuai/DDNM](https://github.com/wyhuai/DDNM)) to our code base.
Since DDNM utilizes the pseudo-inverse of the degradation operator, we noticed that it is unstable for operators whose SVD are prone to numerical errors, such as Gaussian Blur with wide convolution kernel.

##### RedDiff.

We used the implementation of RedDiff available in the released code(^3^33[https://github.com/NVlabs/RED-diff](https://github.com/NVlabs/RED-diff)).
On nonlinear problems, for which the pseudo-inverse of the observation is not available, we initialized the variational optimization with a sample from the standard Gaussian distribution.

##### PGDM.

We relied on the implementation provided in RedDiff’s code as some of its authors are co-authors of PGDM.
The implementation features a subtle difference with Song et al. (, Algorithm 1): in the last line of the algorithm, the guidance term $g$ is multiplied by $\sqrt{\alpha}_{t}$ but in the implementation it is multiplied by $\sqrt{\alpha_{t-1}\alpha_{t}}$.
This modification stabilizes the algorithm on most tasks. For JPEG $2\%$ we found that it worsens the performance. In this case we simply multiply by $\sqrt{\alpha}_{t}$, as in the original algorithm.

##### PSLD

We implemented the PSLD algorithm provided in Rout et al. (, Algorithm (2)) and used the hyperparameters provided by the authors in the publicly available implementation(^4^44[https://github.com/LituRout/PSLD](https://github.com/LituRout/PSLD)).

##### ReSample

We used the original code provided by the authors(^5^55[https://github.com/soominkwon/resample](https://github.com/soominkwon/resample)) and modified it to expose several hyperparameters that were not directly accessible in the released version.
Specifically, we exposed the tolerance $\varepsilon$ and the maximum number of iterations $N$ for solving the optimization problems related to hard data consistency, as well as the scaling factor for the variance of the stochastic resampling distribution $\gamma$.
Our experiments revealed that the algorithm is highly sensitive to $\varepsilon$.
We found that setting it equal to the noise level of the inverse problem gave the best reconstruction across tasks and noise levels.
We set the maximum number of gradient iterations to $N=200$ to make the algorithm less computationally prohibitive.
Finally, we tuned $\gamma$ for each task but found it has less impact on the quality of the reconstruction compared to $\varepsilon$.

#### D.2.1 Runtime and memory

To get the runtime and GPU memory consumption of an algorithm on a dataset, we average these two metrics over both samples of the dataset and the considered tasks in Table 4.

Figure: Figure 6: Runtime and memory requirement of the considered algorithms on datasets: FFHQ latent space (1^st row), FFHQ pixel space (2^nd row), and ImageNet (3^rd row). The left axis displays the runtime in minutes, whereas the right axis the GPU memory requirement in Gigabytes (GB).
Refer to caption: x4.png

### D.3 Gaussian mixtures

In this section, we elaborate more on the Gaussian mixture experiment in Table 4, where we consider a linear inverse problem with a Gaussian mixture as prior.
Therefore, the potential is $g:x\mapsto\mathcal{N}(y;\mathrm{A}x,\sigma_{y}^{2}\mathrm{I}_{d_{y}})$.
Recall that a Gaussian mixture with $C\in\mathbb{N}$ components, whose weights, means, and covariances are $w_{i}>0$, $m_{i}\in\mathbb{R}^{{d}}$, and $\sigma^{2}_{i}\mathrm{I}_{{d}}$, respectively, has the density

$$ $q(x)=\sum_{i=1}^{C}w_{i}\mathcal{N}(x;m_{i},\sigma^{2}_{i}\mathrm{I}_{{d}})\,.$ $$

##### Denoiser.

The denoiser is obtained via Tweedie’s formula Robbins (),

$$ $\smash{m_{0|k}}(x_{k})=\big{(}x_{k}+(1-\alpha_{k})\nabla\log q_{k}(x_{k})\big{ )}/\sqrt{\alpha_{k}},$ $$

where the densities of the marginals $(q_{k})_{k=1}^{n}$ are straightforward,

$$ $\displaystyle q_{k}(x_{k})$ $\displaystyle=\int q(x_{0})q_{k|0}(x_{k}|x_{0})\ \mathrm{d}x_{0}=\sum_{i=1}^{C }w_{i}\int\mathcal{N}(x_{0};m_{i},\sigma^{2}_{i}\mathrm{I}_{{d}})q_{k|0}(x_{k} |x_{0})\ \mathrm{d}x_{0}$ $\displaystyle=\sum_{i=1}^{C}w_{i}\mathcal{N}(x_{k};\sqrt{\alpha_{k}}m_{i},( \alpha_{k}\sigma^{2}_{i}+v_{k})\mathrm{I}_{{d}})\,.$ $$

##### Posterior.

Furthermore, following Bishop (, Eqn. 2.116), the posterior can be shown to be a Gaussian mixture

$$ $\pi(x)\propto g(x)q(x)\propto\sum_{i=1}^{C}\bar{w}_{i}\mathcal{N}(x;\bar{m}_{i },\bar{\Sigma}_{i})\,,$ $$

with new weights, covariances, and means given by

$$ $\displaystyle\bar{w}_{i}$ $\displaystyle=w_{i}\mathcal{N}(y;\mathrm{A}m_{i},\sigma_{y}^{2}\mathrm{I}_{{d} _{y}}+\sigma^{2}_{i}\mathrm{A}\mathrm{A}^{\top})\,,$ $\displaystyle\bar{\Sigma}_{i}$ $\displaystyle=\big{(}(1/\sigma^{2}_{i})\mathrm{I}_{{d}}+(1/\sigma_{y}^{2}) \mathrm{A}^{\top}\mathrm{A}\big{)}^{-1}\,,$ $\displaystyle\bar{m}_{i}$ $\displaystyle=\bar{\Sigma}_{i}\big{(}(1/\sigma_{y}^{2})A^{\top}y+(1/\sigma^{2} _{i})m_{i}\big{)}\,,$ $$

where the weights are un-normalized.

##### Experimental setup.

We consider two setups where ${d}\in\{20,200\}$ and generate Gaussian mixtures of $C=25$ components with means $(m_{i})_{i=1}^{C}=\{(8i,8j,\ldots,8i,8j)\in\mathbb{R}^{{d}}:(i,j)\in\llbracket
-2,2\rrbracket^{2}\}$, unit covariances, and weights $w_{i}$ drawn from a uniform distribution on $[0,1]$ then normalized to sum to $1$.
The linear inverse problem $(y,\mathrm{A})$ is generated by first drawing the matrix $\mathrm{A}\in\mathbb{R}^{1\times{d}}$ with entries i.i.d. from $\mathrm{N}(0,1)$, then sampling $x^{*}$ from the prior $q$ and finally computing $y=\mathrm{A}x^{*}+\sigma_{y}\varepsilon$ with $\varepsilon\in\mathrm{N}(0,1)$ to get the observation.
The standard deviation of the inverse problem is fixed to $\sigma_{y}=0.05$ across all problem instances.

To assess the performance of each algorithm, we draw $2000$ samples and compare against $2000$ samples from the true posterior distribution using the Sliced Wasserstein (SW) distance by averaging over $10^{4}$ slices.
In Section 4 we report the average SW and the 95% confidence interval over $100$ replicates.

### D.4 Extended image experiments

We extend, in Tables 7 and 8, the results in Table 2 by including MGPS with $n=50$ and $n=100$. In the setting $n=100$, the runtime of MGPS is twice lower than that of DPS, PGDM and comparable to that of DiffPir, DDNM and RedDiff, see Figure 6. It is also outperforming them on most of the tasks, especially the nonlinear ones, as is seen in the table below. In the setting $n=50$, MGPS has the lowest runtime among all the methods while maintaining competitive performance.

Finally, in Section D.6, we present sample reconstructions on the various tasks. For each algorithm, we generate five samples and select the four most visually appealing ones. Completely black or white images correspond to failure cases of DPS.

**Table 7: Mean LPIPS values for various linear and nonlinear imaging tasks on the FFHQ $256\times 256$ dataset.**
| Task | MGPS50 | MGPS100 | MGPS300 | DPS | PGDM | DDNM | DiffPir | RedDiff |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| SR ($\times 4$) | 0.13 | 0.10 | 0.09 | 0.09 | 0.33 | 0.14 | 0.13 | 0.36 |
| SR ($\times 16$) | 0.27 | 0.26 | 0.26 | 0.24 | 0.44 | 0.30 | 0.28 | 0.51 |
| Box inpainting | 0.16 | 0.12 | 0.10 | 0.19 | 0.17 | 0.12 | 0.18 | 0.19 |
| Half mask | 0.24 | 0.22 | 0.20 | 0.24 | 0.26 | 0.22 | 0.23 | 0.28 |
| Gaussian Deblur | 0.21 | 0.18 | 0.15 | 0.16 | 0.87 | 0.19 | 0.12 | 0.26 |
| Motion Deblur | 0.19 | 0.15 | 0.13 | 0.16 | $-$ | $-$ | $-$ | 0.21 |
| JPEG (QF = 2) | 0.20 | 0.17 | 0.16 | 0.39 | 1.10 | $-$ | $-$ | 0.32 |
| Phase retrieval | 0.20 | 0.14 | 0.11 | 0.46 | $-$ | $-$ | $-$ | 0.25 |
| Nonlinear deblur | 0.23 | 0.23 | 0.23 | 0.52 | $-$ | $-$ | $-$ | 0.66 |
| High dynamic range | 0.13 | 0.09 | 0.07 | 0.49 | $-$ | $-$ | $-$ | 0.20 |

**Table 8: Mean LPIPS values for various linear and nonlinear imaging tasks on the ImageNet $256\times 256$ dataset.**
| Task | MGPS50 | MGPS100 | MGPS300 | DPS | PGDM | DDNM | DiffPir | RedDiff |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| SR ($\times 4$) | 0.36 | 0.33 | 0.30 | 0.41 | 0.78 | 0.34 | 0.36 | 0.56 |
| SR ($\times 16$) | 0.58 | 0.55 | 0.53 | 0.50 | 0.60 | 0.70 | 0.63 | 0.83 |
| Box inpainting | 0.31 | 0.26 | 0.22 | 0.34 | 0.29 | 0.28 | 0.28 | 0.36 |
| Half mask | 0.39 | 0.34 | 0.29 | 0.44 | 0.38 | 0.38 | 0.35 | 0.44 |
| Gaussian Deblur | 0.35 | 0.29 | 0.32 | 0.35 | 1.00 | 0.45 | 0.29 | 0.52 |
| Motion Deblur | 0.35 | 0.25 | 0.22 | 0.39 | $-$ | $-$ | $-$ | 0.40 |
| JPEG (QF = 2) | 0.50 | 0.46 | 0.42 | 0.63 | 1.31 | $-$ | $-$ | 0.51 |
| Phase retrieval | 0.54 | 0.52 | 0.47 | 0.62 | $-$ | $-$ | $-$ | 0.60 |
| Nonlinear deblur | 0.49 | 0.47 | 0.44 | 0.88 | $-$ | $-$ | $-$ | 0.67 |
| High dynamic range | 0.22 | 0.15 | 0.10 | 0.85 | $-$ | $-$ | $-$ | 0.21 |

### D.5 ECG experiments

#### D.5.1 Implementation details

##### ptb-xl dataset

We use 12-lead ECGs at a sampling frequency of 100 Hz from the ptb-xl dataset Wagner et al. (). For both training and generation, we do not use the augmented limb leads aVL, aVR and aVF as they can be obtained with the following relations: aVL=(I-III)/2, aVF=(II+III)/2 and aVR=-(I+II)/2. This leads to inputs of shape $T\times 9$ instead of $T\times 12$ (where $T$ is the length of the signal). For evaluation metrics, we reconstruct the augmented limb leads.
For the MB task, following Alcaraz & Strodthoff (), we use 2.56-second ECG random crops, hence the inputs are of shape $256\times 9$. For the ML task, we use the full 10-second ECGs and pad them into a tensor of shape $1024\times 9$. Table 9 summarizes the input shapes for models and algorithms for each task. Table 10 summarizes the distribution of train/val/test splits and the number of cardiac conditions (RBBB, LBBB, AF, SB) per split. Note that for posterior sampling, we only use the test set, and for both training and sampling we never use the cardiac condition.

**Table 9: ECG size and input-shape per task.**
| Task | ECG size (seconds) | Total leads | Used leads | Input shape |
| --- | --- | --- | --- | --- |
| MB | 2,56 | 12 | I, II, III, V1–6 | $256\times 9$ |
| ML | 10 | 12 | I, II, III, V1–6 | $1024\times 9$ |

**Table 10: ptb-xl dataset description.**
| Split | All | RBBB | LBBB | AF | SB |
| --- | --- | --- | --- | --- | --- |
| Train | 17,403 | 432 | 428 | 1211 | 503 |
| Val | 2,183 | 55 | 54 | 151 | 64 |
| Test | 2,203 | 54 | 54 | 152 | 64 |

##### Diffusion model

Following Alcaraz & Strodthoff (; ) we used Structured State Space Diffusion (SSSD) models Gu et al. (); Goel et al. ().
SSSD generates a sequence $u(t)$ with the following differential equation denpending on the input seuqnece $x(t)$ and a hidden state $h(t)$:

$$ $\displaystyle h^{\prime}(t)$ $\displaystyle=Ah(t)+Bx(t)$ (D.1) $\displaystyle u(t)$ $\displaystyle=Ch(t)+Dx(t)\,,$ (D.2) $$

where $A,B,C,D$ are transition matrices. We use the SSSD^SA Alcaraz & Strodthoff () (also denoted as Sashimi Goel et al. ()) architecture publicly available(^6^66[https://github.com/AI4HealthUOL/SSSD](https://github.com/AI4HealthUOL/SSSD))(^7^77[https://github.com/albertfgu/diffwave-sashimi/tree/master](https://github.com/albertfgu/diffwave-sashimi/tree/master)). We parametrize the matrix $A=\Lambda-pp^{*}$ where $p$ is a vector and $\Lambda$ is a diagonal matrix. This parametrization with facilitates the control of the spectrum of $A$ to enforce stability (negative real part of eigenvalues).
We use a multi-scale architecture, with a stack of residual S4 blocks. The top tier processes the raw audio signal at its original sampling rate, while lower tiers process downsampled versions of the input signal. The outputs of the lower tiers are upsampled with up-pooling and down-pooling operations and combined with the input to the upper tier to provide a stronger conditioning signal. Hyper-parameters are described in table 11. The selected model is the model of the last epoch.
Training time is 30 hours for 2.5 seconds ECGs, 42 hours for 10 seconds ECGs.

**Table 11: ECG diffusion generative model hyper-parameters.**
| Hyper-parameter | Value |
| --- | --- |
| Residual layers | 4 |
| Pooling factor | [1, 2, 2] |
| Feature expansion | 2 |
| Diffusion embedding dim. 1 | 128 |
| Diffusion embedding dim. 2 | 512 |
| Diffusion embedding dim. 3 | 512 |
| Diffusion steps | 1000 |
| Optimizer | Adam |
| Number of iterations | 150k |
| Loss function | MSE |
| Learning rate | 0.002 |
| Batch size | 128 |
| Number of parameters | 16 millions |

##### Algorithm parameters

We use the same parameters as described in table 6 for MGPS and the implementation described in section D.2 for competitors. Since we had to run the experiments over 2k samples, we set the number of diffusion steps such that the runtime approximates 30 seconds for generating 10 samples. This leads to the parameters shown in table 12. For the ML experiment, we also tested MGPS with 300 diffusion steps, which resulted in a runtime of 5 minutes and 30 seconds per batch of 10 samples. The results improved, but we already outperform competitors with just 50 diffusion steps and a 30-second runtime.

**Table 12: Number of diffusion steps used in posterior sampling algorithms for ECG tasks.**
| MGPS50 | MGPS300 | DPS | PGDM | DDNM | DiffPir | RedDiff |
| --- | --- | --- | --- | --- | --- | --- |
| 50 | 300 | 400 | 200 | 200 | 500 | 400 |

##### Evaluation metrics

For the MB task, for each observation, we generate 10 samples - instead of 100 as in in Alcaraz & Strodthoff (). We then compute the Mean Absolute Error (MAE) and Root Mean Squared Error (RMSE) between each generated sample and the ground-truth. The final score is the average of these errors over the 10 generated samples. We report the confidence intervals over the test set in table 4 and table 14.

For the ML task, we use a classifier trained to detect four cardiac conditions: Right Bundle Branch Block (RBBB), Left Bundle Branch Block (LBBB), Atrial Fibrillation (AF), and Sinus Bradycardia (SB) on the ptb-xl dataset. We follow the XResNet1d50 described in Strodthoff et al. () with hyper-parameters reported in table 13.
We apply the classifier to the ground-truth ECG of the test set and to the samples generated from lead I. As in the MB task, for each observation, 10 samples are generated. The output of the classifier is averaged over these 10 samples.
For each cardiac condition, we compute balanced accuracy to account for class imbalance (see table 10). The classification threshold is selected using ptb-xl validation-set.

**Table 13: XResNet1d50 downstream classifier hyper-parameters.**
| Hyper-parameter | Value |
| --- | --- |
| Blocks | 4 |
| Layers per block | [3, 4, 6, 3] |
| Expansion | 4 |
| Stride | 1 |
| Optimizer | Adam |
| Learning rate | $0.001$ |
| Batch size | $0.001$ |
| Epochs | 100 |

#### D.5.2 Additional Results

##### Discussion

In table 4, we demonstrated that most posterior sampling algorithms outperform the trained diffusion model for the missing block reconstruction task. This result is particularly interesting as it suggests that training a diffusion model (which takes several days) is not necessary for this task.
However, when Alcaraz & Strodthoff () trained the model on the reconstruction of random missing blocks (RMB), where 50% of each lead is independently removed, the model outperformed all posterior sampling algorithms on the MB task. We report in table 14 the results of the top two algorithms, as well as the model trained on the MB task and the RMB task. The significant improvement between RMB and MB can be seen as an enhancement due to data augmentation.

**Table 14: MAE and RMSE for missing block task on the ptb-xl dataset.**
| Metric | MGPS | DDNM | TrainedDiff-MB | TrainedDiff-RMB |
| --- | --- | --- | --- | --- |
| MAE | $0.111\pm 2\mathrm{e}\!{-3}$ | $0.103\pm 2\mathrm{e}\!{-3}$ | $0.116\pm 2\mathrm{e}\!{-3}$ | $\mathbf{0.0879\pm 2\mathrm{e}\!{-3}}$ |
| RMSE | $0.225\pm 4\mathrm{e}\!{-3}$ | $0.224\pm 4\mathrm{e}\!{-3}$ | $0.266\pm 3\mathrm{e}\!{-3}$ | $\mathbf{0.217\pm 6\mathrm{e}\!{-3}}$ |

##### Generated samples

Figure: Figure 7: Missing block imputation with MGPS on 2.56s 12-lead ECG. Ground-truth in blue, $10\%$–$90\%$ quantile range in green, random sample in orange.
Refer to caption: x7.png

Figure: Figure 8: Missing lead reconstruction from lead I on 10s 12-lead Normal Sinus Rythm (NSR) ECGs. Ground-truth in blue, $10\%$–$90\%$ quantile range in green, random sample in orange.
Refer to caption: x8.png

Figure: Figure 9: Missing lead reconstruction from lead I on 10s 12-lead Right Bundle Branch Block (RBBB) ECGs. Ground-truth in blue, $10\%$–$90\%$ quantile range in green, random sample in orange.
Refer to caption: x14.png

Figure: Figure 10: Missing lead reconstruction from lead I on 10s 12-lead Left Bundle Branch Block (LBBB) ECGs. Ground-truth in blue, $10\%$–$90\%$ quantile range in green, random sample in orange.
Refer to caption: x20.png

Figure: Figure 11: Missing lead reconstruction from lead I on 10s 12-lead Atrial Fibrillation (AF) ECGs. Ground-truth in blue, $10\%$–$90\%$ quantile range in green, random sample in orange.
Refer to caption: x26.png

Figure: Figure 12: Missing lead reconstruction from lead I on 10s 12-lead Sinus Bradycardia (SB) ECGs. Ground-truth in blue, $10\%$–$90\%$ quantile range in green, random sample in orange.
Refer to caption: x32.png

### D.6 Sample images

Figure: Figure 13: Sample reconstructions with box mask on ImageNet dataset.
Refer to caption: extracted/5901560/figures/imagenet_ims/inpainting_center/30.jpeg

Figure: Figure 14: Sample reconstructions with half mask on FFHQ dataset.
Refer to caption: extracted/5901560/figures/ffhq_ims/outpainting_half/13.jpeg

Figure: Figure 15: Sample reconstructions with half mask on ImageNet dataset.
Refer to caption: extracted/5901560/figures/imagenet_ims/outpainting_half/45.jpeg

Figure: Figure 16: Sample reconstructions on motion deblurring task.
Refer to caption: extracted/5901560/figures/imagenet_ims/motion_blur/122.jpeg

Figure: Figure 17: Sample reconstructions on Gaussian deblurring task.
Refer to caption: extracted/5901560/figures/imagenet_ims/blur_svd/92.jpeg

Figure: Figure 18: Sample reconstructions on phase retrieval task.
Refer to caption: extracted/5901560/figures/ffhq_ims/phase_retrieval/30.jpeg

Figure: Figure 19: Sample reconstructions on JPEG 2 task.
Refer to caption: extracted/5901560/figures/imagenet_ims/jpeg2/82.jpeg

Figure: Figure 20: Sample reconstructions on high dynamic range task.
Refer to caption: extracted/5901560/figures/ffhq_ims/high_dynamic_range/99.jpeg

Figure: Figure 21: Sample reconstructions on SR (16$\times$) task.
Refer to caption: extracted/5901560/figures/ffhq_ims/sr16/6.jpeg

Figure: Figure 22: More sample reconstructions on SR (16$\times$) task.
Refer to caption: extracted/5901560/figures/imagenet_ims/sr16/50.jpeg

Figure: (a) JPEG 2
Refer to caption: extracted/5901560/figures/ldm_ims/jpeg2/77.jpeg

Figure: Figure 24: More sample reconstructions with MGPS on the expand task.
Refer to caption: extracted/5901560/figures/ffhq_ims/outpainting_expand/3.jpeg

Figure: Figure 25: More sample reconstructions with MGPS and half mask.
Refer to caption: extracted/5901560/figures/ffhq_ims/outpainting_half/25.jpeg

Figure: Figure 26: More sample reconstructions with MGPS and half mask.
Refer to caption: extracted/5901560/figures/imagenet_ims/outpainting_half/41.jpeg

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: 1. Would there be an optimal mid-point schedule for $l_k$, which allows fewer diffusion model evaluations and comparable performance? Or better technique can be used to find the optimal schedule? Or can we suggest a criteria for evaluating the trade-off between performance and computational cost?

2. What is the key reason that makes mid-point guidance perform better? When the same prior is applied, what difference does the midpoint sequence make to sample algorithm as difference stages of the posterior distributions?
soundness: 3
strengths: 1. The approach introduces a midpoint guidance mechanism that provides a novel trade-off for guidance and complexity of the learned transition term in diffusion models. This make it different from other diffusion-based posterior sampling methods.

2. The method is well-justified with mathematical rigor. The decomposition of the backward transition and the use of Gaussian variational approximations.

3. MGPS is extensively evaluated across both synthetic and real-world tasks, such as Gaussian mixture sampling, image super-resolution, and ECG imputation. Experimental results show significant improvements in reconstruction quality over baseline methods.

4. The paper is generally well-organized, providing clear problem definitions, methodological details, and experimental results. Detailed explanations and algorithmic steps (Algorithm 1) support reproducibility.
summary: This paper introduces a novel diffusion-based posterior sampling method called Midpoint Guidance Posterior Sampling (MGPS) to address Bayesian inverse problems. In cases where denoising diffusion models (DDMs) are used as priors, MGPS aims to approximate the posterior while balancing the complexity between guidance and prior transition terms. The method leverages an intermediate midpoint state to improve posterior approximation and incorporates a Gaussian variational approximation for additional flexibility. MGPS is validated on linear and nonlinear inverse problems across various domains, including image and ECG signal reconstruction, and outperforms several state-of-the-art methods.
weaknesses: 1. More exploration of $\eta$ and the midpoint sequence's impact on different task types and data complexities would clarify MGPS's adaptability across scenarios. The authors showed the effect of $\eta$ in the Gaussian toy example. It would be interesting to see $\eta$'s influence in other tasks.

2. The trade-off introduced by the midpoint state could be theoretically explored further. The authors mention the need for tuning this midpoint sequence but provide limited theoretical insights on why this works well. For instance, the bounds on the approximation errors may be further derived, or the balance between the prior and guidance terms could be analyzed with the midpoint sequence.

3. Concerns about the correct implementation of other compared methods remain for me, as those methods perform in Figures 13 and 15 worse than those in the original publications related to some image domain tasks. Or are those methods undertuned?

---

confidence: 5
contribution: 3
initial_rating: 6
presentation: 3
questions: Following up on the weaknesses:
- Why was a different notation used compared to existing works on posterior sampling for inverse problems?
- Why was only 50 randomly selected evaluation points used for FFHQ and ImageNet? How was this random selection done?
- Why was only LPIPS used for evaluation, and other standard metrics like PSNR/SSIM were not reported?
- Why was DPS run for 300 steps instead of the more standard 1000?
soundness: 4
strengths: - The idea of midpoint guidance is novel, and intuitively appealing to solve the issues associated with guidance issues especially at the early stages of the diffusion process.
- Very thorough evaluation.
- Results are good, showing improvement over SOTA methods, including both DDMs and LDMs (that are commonly used in these applications).
- Multiple nonlinear inverse problems are studied, an area where other posterior sampling methods have issues. A good level of improvement is shown in these applications.
- Performance also evaluated across NFEs.
- Good sample variety is shown.
summary: The paper proposes a novel midpoint guidance strategy for posterior sampling of diffusion models. The main idea is to first move to a mid-point state l_k that is a function of k for guidance, and then noise back to obtain X_k unconditionally. This seems to solve issues associated with other SOTA approaches that perform the guidance based on variations of Tweedie's formula. Extensive experiments are provided, along with comparison to SOTA methods. Improvements are shown in almost all inverse problems, and notably for nonlinear inverse problems. 

I think the paper has good merit, but its exposition is overly complicated and notation does not match the rest of the literature. These dampened my enthusiasm, but I'm happy to reevaluate after the authors respond.
weaknesses: - Unfortunately, the exposition is overly complicated. The idea can be explained much more clearly, but also partly due to non-standard notation, the gist does not come across easily. Section 3 would benefit from a substantial rewrite that changes the notation (please see below), and highlights the main ideas, and potentially even including a figure to show the midpoint guidance idea.
- Similarly, the notation does not match rest of the literature on posterior diffusion sampling for inverse problems. For instance, the posterior is denoted by \pi(x), which is written as a marginal, though this should depend on the observations y. Similarly g(.) in (2.1) should be conditional on y, but this is not done either. This propagates throughout the paper, and makes it hard to appreciate the contributions. Please use standard notation to match other existing works.
- The 50 randomly selected evaluation points for FFHQ and ImageNet is a bit unconvincing. There are standard validation sets that are publicly available for both databases, and these are commonly used in other papers (e.g. DPS, PGDM). Please report your results over a larger set (& also explain how random selection was done).
- Unclear why only LPIPS is provided. PSNR/SSIM must be provided as well. I understand improvement may not be uniform for those metrics, but this should be up to the reader to figure out.
- DPS is typically run with NFE = 1000, but it was used for N = 300 in this paper for comparison. This may further close the LPIPS gap.

The following points are not really weaknesses, but I wanted to note them:
- The ECG application comes out of nowhere. Without enough motivation and knowing the difficulty of the task, the contribution here is a bit hard to appreciate.
- There are also two additional references that may be of interest:
a) arXiv:2402.02149, which is similar to Boys et al in spirit
b) arXiv:2407.11288 (ECCV 2024), which learns w_k in (2.5) and also uses a diagonal approximation to circumvent vector-Jacobian calculations

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
questions: See weaknesses.
soundness: 4
strengths: S1) The approach is well-motivated and has strong theoretical backing, with a novel theoretical result included in Appendix A.3.

S2) The authors evaluate many different problems on multiple datasets, achieving relatively strong performance across the board.

S3) The paper is well-written and reasonably easy to follow (although I have some gripes with notation, outlined in the 'Weaknesses' section below.

S4) Many great visuals, especially in the appendices.

S5) The authors provide very detailed and explicit implementation details for their competitors. I think that this is of vital importance and appreciate the efforts made by the authors to share these details.
summary: The authors propose 'midpoint guidance posterior sampling,' where the reverse diffusion process is decomposed into two steps:

1. Denoising the current diffusion measurements into a 'midpoint' state.
2. Renoising the midpoint measurements to obtain the next diffusion iterate.

The denoising approach consists of approximating a Gaussian variational approximation in conjunction with the DPS guidance proposed by Chung et al to sample an estimated image. The variational approximation is learned at each diffusion timestep.

The renoising stage appears to be the typical DDIM/DDPM computation of the next diffusion iterate.

This approach of midpoint guidance achieves very strong empirical performance on a wide variety of problems/datasets and has strong theoretical foundations.
weaknesses: At a high level, I quite like this work. However, as described below, I feel that the experimental results are incomplete.

W1) Even though the paper is well-written, and all mathematical elements check out (at least to me), the notation makes all of the math in Section 3. In particular, the differentiation between scalar and vector quantities is not sufficient. I would suggest that the authors make vector quantities boldface (e.g., $\boldsymbol{x}$). This would make things much easier to read.

W2) While I appreciate the robust slate of experiments, using test sets of size 50 (at least for FFHQ, ImageNet) is not sufficient for two reasons:
1. I would argue that the standard test set size for any diffusion inverse solver is 1k images. This seems to be the accepted number, and large enough to truly understand model performance. A test of size 50 is just not sufficient for acceptance to a conference like ICLR. I feel that the results would be more convincing if there were fewer experiments with a larger test set. Note that I am explicitly speaking on the image datasets, where samples are plentiful. There are plenty of problems where there is not enough data to have a test set comprised of 1k samples and that is fine. To summarize: I think that fewer experiments are not a bad thing if it means more robust testing.
2. Only using 50 test samples does not enable reliable computation of FID (a metric that is noticeably absent from the paper). FID is a standard metric when evaluating diffusion inverse solvers. With it absent, the experimental results are incomplete.

W3) I also think that the experimental results are incomplete due to the lack of a pixel-space quality metric. The authors argue against PSNR/SSIM in Section 4, but I disagree. I think that LPIPS is a great choice, but that PSNR is still a necessity when testing because similarity in the pixel-space matters too. If the samples don't respect the measurements, that is a problem and LPIPS may not catch it since it is a feature-based metric.

To summarize W2 and W3, I think that the authors need to reconsider their evaluation of problems which use the image datasets. I would suggest:
1. 1k test samples.
2. Computing PSNR and FID in addition to LPIPS.

These changes are critical to fully understanding model performance. Without this sufficient experimental evaluation, I cannot recommend this paper for acceptance.

---

confidence: 4
contribution: 4
initial_rating: 8
presentation: 3
questions: 1. In Tab. 2, only the LPIPS values are reported. I understand that this is probably to save space, but I would recommend also including other standard metrics such as PSNR, SSIM, and FID.

2. MGPS is based on DDPM sampling. Would there be any way to incorporate faster solvers into this?

I am willing to further raise the score after clarification.
soundness: 4
strengths: 1. The paper is well-written and relatively straightforward to follow.

2. When an approximation is used, the authors do a good job of explaining the rationale, by either showing it theoretically or experimentally.

3. MGPS is a good balance between a theoretically grounded solution and a practical solution, not requiring too much computation.

4. The results of image restoration tasks are clearly SOTA.

5. Experiments are complete. Numerical experiments on toy data, image restoration experiments, and ECG completion all indicate the superiority of the method.
summary: The paper proposes MGPS, a variational inference approach to diffusion model-based inverse problem solving (DIS) in the standard guiding reverse diffusion setup. Two clever ideas are used:

1. Instead of using the DPS approx after denoising $k + 1 \rightarrow k$, denoise it up to $0 < \ell_k < k$, and use the posterior mean of $x_{\ell_k}$ instead of the posterior mean of $x_k$ to compute the DPS gradient. This way, you can control the trade-off arising in the approximation error for the denoising kernel and the likelihood computation.

2. For every reverse diffusion timestep $k$, use $j$ iterations of stochastic optimization for fitting a variational distribution for the reverse posterior kernel. The authors propose to use a Gaussian kernel but also optimize over the diagonal elements of covariance. This is different from the previous works where typically, an isotropic variance is used, probably with the exception of [1].

Overall, the paper is well-written, has a clear theory, and has great results. The reviewer especially appreciates the efforts that the authors made to make the experiments as fair as possible, carefully addressing the details that are often ignored. I do have some questions and concerns on the practical implementation of the method, and in some places, how to derive it. Nevertheless, I think the paper should be a clear accept.


**References**

[1] Peng, Xinyu, et al. "Improving Diffusion Models for Inverse Problems Using Optimal Posterior Covariance." ICML 2024.
weaknesses: 1. Some parts of the derivation are unclear. The authors propose a variational distribution $\lambda^\varphi$ for $\hat{\pi}^\theta$ in (3.6), which, to my understanding, is already tractable. Since $\hat{g}^\theta$ is Gaussian from DPS, and $p_{\ell_k}^\theta$ is also Gaussian, then isn't $\hat{\pi}^\theta$ already a Gaussian? Why would one need an additional variational distribution to approximate this?

2. Adding on 1, I don't quite understand why including the diagonal terms of the covariance for additional optimization would induce better fitting, when $\hat{\pi}^\theta$ would have istropic covariance. Both these points maybe from my misunderstanding, but it should be clarified.

3. Balancing the denoising and likelihood approximation errors is interesting. Is it safe to say that the mainstream previous works that use something similar to (2.5) be considered as simply taking $\ell_k = k$? Or is it not directly comparable? I believe the authors could spend some more effort on clarifying the difference between the existing methods. There is a related works section, but the connections seem a bit vague. An additional appendix section could be nice.

4. The upper bound of the variational inference problem is defined and used, but it is not explained how this is derived.

Meta-review: The paper proposes a diffusion-based method for posterior sampling in diffusion models. The four reviewers all indicated that the paper is clearly above threshold for acceptance. They found the paper to be written well and present an idea that could be found widely useful. The empirical evaluation was sufficient to demonstrate its improvement over other SOTA models. During a detailed discussion phase, the authors provided additional information that caused three reviewers to improve their scores. It is important for this additional information to be incorporated into the final draft of the paper.

Decision: Accept (Oral)