Amortized Posterior Sampling with Diffusion Prior Distillation

## Abstract

Abstract We propose a variational inference approach to sample from the posterior distribution for solving inverse problems. From a pre-trained diffusion model, our approach trains a conditional flow model to minimize the divergence between the proposal variational distribution and the posterior distribution implicitly defined through the diffusion model. Once trained, the flow model is capable of sampling from the posterior distribution with a single NFE, amortized with respect to the measurement. The proposed method paves a new path for distilling a diffusion prior for efficient posterior sampling. We show that our method is applicable to standard signals in Euclidean space, as well as signals on manifold.

## 1 Introduction

We consider the following inverse problem

$$ $\displaystyle{\bm{y}}={\mathcal{A}}({\bm{x}})+{\bm{n}},\quad{\bm{y}}\in{ \mathbb{R}}^{m},\,{\bm{x}}\in{\mathbb{R}}^{n},\,{\bm{A}}:{\mathbb{R}}^{n} \mapsto{\mathbb{R}}^{m},\,{\bm{n}}\sim{\mathcal{N}}(0,\sigma_{y}^{2}{\bm{I}}),$ (1) $$

where the goal is to infer an unknown signal ${\bm{x}}$ from the degraded measurement ${\bm{y}}$ obtained through some forward operator ${\mathcal{A}}$, leveraging the information contained in the measurement and the prior $p({\bm{x}})$. A powerful modern way to define the prior is through diffusion models [ho2020denoising, song2020score], where we train a parametrized model ${\bm{s}}_{\theta}$ to estimate the gradient of the log prior $\nabla_{\bm{x}}\log p({\bm{x}})$.

Solving inverse problems with the diffusion model can be achieved through posterior sampling with Bayesian inference. Arguably the standard way to achieve this is through modifying the reverse diffusion process of diffusion models. This adjustment shifts the focus from sampling from the trained prior distribution $p_{\theta}({\bm{x}}_{0})$ to sampling from the posterior distribution $p_{\theta}({\bm{x}}_{0}|{\bm{y}})$. This transition is facilitated by employing iterative projections to the measurement subspace [kadkhodaie2020solving, song2020score, chung2022come, wang2023zeroshot], guiding the samples through gradients pointing towards measurement consistency [chung2023diffusion, song2023pseudoinverseguided].

It should be noted that diffusion models learn the gradient of the prior and diffusion samplers [song2020denoising, lu2022dpm, song2020score] are methods that numerically solve the probability-flow ODE (PF-ODE) that defines the reverse diffusion sampling trajectory. Consequently, regardless of the specifics of the methods, standard diffusion model-based inverse problem solvers (DIS), even those that are considered fast, take at least a few tens of neural function evaluation (NFE), making them less effective for time-critical applications such as medical imaging and computational photography.

Figure: Figure 1: Diverse inverse problem solving can be done with a single NFE, with the same network for all the different measurements. (row 1) Denoising on celebA, (row 2) inpainting MNIST on the bunny manifold, (row 3) imputation of ERA5 on the spherical manifold.
Refer to caption: extracted/5754552/figs/main.png

Figure: Figure 2: Concept of the proposed method. (a) Training can be performed in an unsupervised fashion with a dataset consisting of degraded measurements ${\bm{y}}$ to train a conditional normalizing flow $G_{\phi}$ with the diffusion prior ${\bm{s}}_{\theta}$. (b) Once trained, one can achieve multiple posterior samples by inputting different noise vectors ${\bm{z}}\sim{\mathcal{N}}(0,{\bm{I}})$ concatenated with the condition ${\bm{y}}$ with a single NFE, generalizable across any measurement ${\bm{y}}$.
Refer to caption: extracted/5754552/figs/concept.png

Another class of methods [feng2023score, feng2023efficient] introduces the use of variational inference (VI) to train a new proposal distribution $q_{\phi}^{\bm{y}}({\bm{x}})$ to distill the prior learned through the pre-trained diffusion model. The problem is defined as the following optimization problem

$$ $\displaystyle\min_{\phi}D_{KL}(q_{\phi}^{\bm{y}}({\bm{x}}_{0})||p_{\theta}({ \bm{x}}_{0}|{\bm{y}})),$ (2) $$

where the superscript ${\bm{y}}$ emphasizes that the proposal distribution is specific for a single measurement ${\bm{y}}$. For tractable optimization, $q$ is often taken to be a normalizing flow [rezende2015variational, dinh2016density] (NF) so that the computation of likelihood can be done instantly, and one can sample multiple reconstructions from the posterior samples by plugging in different noise values from the reference distribution. Posing the problem this way yields a method that can achieve posterior samples with just a single NFE. Nevertheless, it is still impractical as training a measurement-specific variational distribution takes hours of training. It is often unrealistic to train a whole new model from scratch every time when a new measurement is taken.

In this work, we take a step towards a practical VI-based posterior sampler by distilling a diffusion model prior. To this end, we propose a conditional normalizing flow $q_{\phi}({\bm{x}}_{0}|{\bm{y}})$ as our variational distribution and amortize the optimization problem in Eq. (2) over the conditions ${\bm{y}}$. By using a network that additionally takes in the condition ${\bm{y}}$ as the input, we can train a single model that generalizes across the whole dataset without the need for cumbersome re-training for specific measurements. (See Fig. 2 for the conceptual illustration of the proposed method, as well as representative results presented in Fig. 1.) Interestingly, we find that the speed of optimization is not hampered with such amortization, and the proposed method achieves comparable performance against the measurement-specific flow model [feng2023efficient, feng2023score]. Furthermore, we extend the theory to consider inverse problems on the Riemannian manifold, showing that the proposed idea is generalizable even when the signal is not one the Euclidean manifold. In summary, our contributions and key takeaways are as follows

- 1.
We propose an amortized variational inference framework to enable 1-step posterior sampling constructed implicitly from the pre-trained diffusion prior $p_{\theta}({\bm{x}})$ for any measurement ${\bm{y}}$.
- 2.
To the best of our knowledge, our method is the first diffusion prior distillation approach for solving inverse problems that are unsupervised (i.e. does not require any ground-truth data ${\bm{x}}$), as opposed to standard conditional NFs [lugmayr2020srflow] that required supervised paired data.
- 3.
Experimentally, we show that the proposed method easily scales to signals that lie on the standard Euclidean manifold (e.g. images) as well as signals that lie on the Riemannian manifold, achieving strong performance regardless of the representation.

## 2 Preliminaries

### 2.1 Score-based diffusion models

We adopt the standard framework for constructing a continuous diffusion process ${\bm{x}}(t)$, where $t\in[0,T]$ and ${\bm{x}}(t)\in\mathbb{R}^{d}$, as outlined by Song *et al*. [song2020score]. Specifically, our goal is to initialize ${\bm{x}}(0)$ from a distribution $p_{0}({\bm{x}})=p_{\text{data}}$, and evolve ${\bm{x}}(t)$ towards a reference distribution $p_{T}$ at time $T$, which is easily samplable.

The evolution of ${\bm{x}}(t)$ is governed by the Itô stochastic differential equation (SDE):

$$ $\displaystyle d{\bm{x}}={\bm{f}}({\bm{x}},t)dt+g(t)d{\bm{w}},$ (3) $$

where ${\bm{f}}:\mathbb{R}^{d}\times\mathbb{R}\rightarrow\mathbb{R}^{d}$ represents the drift function, and $g:\mathbb{R}\rightarrow\mathbb{R}^{d}$ denotes the diffusion coefficient. These coefficients are designed to drive ${\bm{x}}(t)$ towards a spherical Gaussian distribution as $t$ approaches $T$.
When the drift function ${\bm{f}}({\bm{x}},t)$ is affine, the perturbation kernel $p_{0t}({\bm{x}}(t)|{\bm{x}}(0))$ is Gaussian, allowing for the parameters to be determined analytically. This facilitates data perturbation via $p_{0t}({\bm{x}}(t)|{\bm{x}}(0))$ efficiently, without necessitating computations through a neural network.

Furthermore, corresponding to the forward SDE, there exists a reverse-time SDE:

$$ $\displaystyle d{\bm{x}}$ $\displaystyle=[{\bm{f}}({\bm{x}},t)-g(t)^{2}\nabla_{\bm{x}}\log p_{t}({\bm{x}} )]dt+g(t)d\bar{{\bm{w}}},$ (4) $$

where $dt$ represents an infinitesimal negative time step, and $\bar{{\bm{w}}}$ denotes the backward standard Brownian motion. While the trajectory of Eq. (4) is stochastic, there also exists a corresponding probability-flow ODE (PF-ODE) that recovers the same law $p_{t}({\bm{x}})$ as the time progresses [song2020score, song2020denoising]

$$ $\displaystyle d{\bm{x}}$ $\displaystyle=[{\bm{f}}({\bm{x}},t)-\frac{g(t)^{2}}{2}\nabla_{\bm{x}}\log p_{t }({\bm{x}})]dt.$ (5) $$

This allows a deterministic mapping between the reference and the target distribution, and hence diffusion models can also be seen as a neural ODE [chen2018neural].

A neural network can be trained to approximate the true score function $\nabla_{\bm{x}}\log p_{t}({\bm{x}})$ through score matching techniques, as demonstrated in previous works [song2019generative, song2020score]. This approximation, denoted ${\bm{s}}_{\theta}({\bm{x}},t)\approx\nabla_{\bm{x}}\log p_{t}({\bm{x}})$, is then utilized to numerically integrate the reverse-time SDE. To effectively train the score function, denoising score matching (DSM) is often employed [hyvarinen2005estimation]

$$ $\displaystyle\theta^{*}=\operatorname*{arg\,min}_{\theta}{\mathbb{E}}_{t\sim U (\varepsilon,1),{\bm{x}}(t),{\bm{x}}(0)}\left[\|{\bm{s}}_{\theta}({\bm{x}}(t), t)-\nabla_{{\bm{x}}_{t}}\log p({\bm{x}}(t)|{\bm{x}}(0))\|_{2}^{2}\right],$ (6) $$

Interestingly, the posterior mean, or the so-called denoised estimate can be computed via Tweedie’s formula [efron2011tweedie]. Specifically, for $p({\bm{x}}_{t}|{\bm{x}}_{0})={\mathcal{N}}({\bm{x}}_{t};\alpha_{t}{\bm{x}}_{0}
,\beta_{t}^{2}{\bm{I}})$,

$$ $\displaystyle\hat{\bm{x}}_{0|t}^{\theta}:={\mathbb{E}}_{p({\bm{x}}_{0}|{\bm{x} }_{t})}[{\bm{x}}_{0}|{\bm{x}}_{t}]=\frac{1}{\alpha_{t}}({\bm{x}}_{t}+\beta_{t} ^{2}\nabla_{{\bm{x}}_{t}}\log p({\bm{x}}_{t})).$ (7) $$

### 2.2 Diffusion models for inverse problems (DIS)

Solving the reverse SDE in Eq. (4) or the PF-ODE in Eq. (5) results in sampling from the prior distribution $p_{\theta}({\bm{x}}_{0})$, with the subscript emphasizing the time variable in the diffusion model context ${\bm{x}}_{0}\equiv{\bm{x}}$. When solving an inverse problem as posed in Eq. (1), our goal is to sample from the posterior $p_{\theta}({\bm{x}}_{0}|{\bm{y}})\propto p_{\theta}({\bm{x}})p({\bm{y}}|{\bm{x
}}_{0})$. Using Bayes rule for a general timestep $t$ yields

$$ $\displaystyle\nabla_{{\bm{x}}_{t}}\log p_{\theta}({\bm{x}}_{t}|{\bm{y}})= \nabla_{{\bm{x}}_{t}}\log p_{\theta}({\bm{x}}_{t})+\nabla_{{\bm{x}}_{t}}\log p ({\bm{y}}|{\bm{x}}_{t}).$ (8) $$

While the former term can be replaced with a pre-trained diffusion model, the latter term is intractable and needs some form of approximation. Existing DIS [kawar2022denoising, chung2023diffusion, wang2023zeroshot] propose different approximations for $\nabla_{{\bm{x}}_{t}}\log p({\bm{y}}|{\bm{x}}_{t})$, which yields sampling from slightly different posteriors $\nabla_{{\bm{x}}_{t}}\log p_{\theta}({\bm{x}}_{t}|{\bm{y}})$.

Algorithmically, the posterior samplers are often implemented so that the original numerical solver for sampling from the prior distribution remains intact, while modifying the Tweedie denoised estimate at each time $\hat{\bm{x}}_{0|t}^{\theta}$ to satisfy the measurement condition given as Eq. (1). From Tweedie’s formula, we can see that this corresponds to approximating the conditional posterior mean ${\mathbb{E}}[{\bm{x}}_{0}|{\bm{x}}_{t},{\bm{y}}]$ in the place of the unconditional counterpart ${\mathbb{E}}[{\bm{x}}_{0}|{\bm{x}}_{t}]$. The algorithms are inherently iterative, and the modern solvers [chung2023diffusion, wang2023zeroshot, zhu2023denoising] require at least 50 NFE to yield a high-quality sample. Moreover, as existing methods can be interpreted as approximating the reverse distribution $p({\bm{x}}_{0}|{\bm{x}}_{t})$ with a simplistic Gaussian distribution $q({\bm{x}}_{0}|{\bm{x}}_{t})={\mathcal{N}}({\bm{x}}_{0};\hat{\bm{x}}_{0|t}^{
\theta},s_{t}^{2}{\bm{I}})$ [peng2024improving], it typically yields a large approximation error, especially in the earlier steps of the reverse diffusion.

**Table 1: Methods that leverage diffusion priors for solving inverse problems according to their class, and their characteristics.**
| Class | Score-based | Variational Inference |  |  |  |
| --- | --- | --- | --- | --- | --- |
| Methods | DIS | Noise2Score [kim2021noise2score] | RED-Diff | Feng *et al*. [feng2023score, feng2023efficient] | Ours |
| One-step inference | ✗ | ✗ | ✓ | ✓ | ✓ |
| Tackles general<br>inverse problems | ✓ | ✓ | ✗ | ✓ | ✓ |
| Exact likelihood<br>computation | ✗ | ✗ | ✗ | ✓ | ✓ |
| Amortized across ${\bm{y}}$ | ✗ | ✗ | ✗ | ✗ | ✓ |
| Generalizable<br>across dataset | ✗ | ✗ | ✗ | ✗ | ✓ |

## 3 Related works

### 3.1 Variational inference in DIS

Standard DIS discussed in Sec. 2.2 sample from the posterior distribution by following the reverse diffusion trajectory. Another less studied approach uses VI to use a new proposal distribution, where the problem is cast as an optimization problem in Eq. (2).
RED-diff [mardani2023variational] places a unimodal Gaussian distribution as the proposal distribution $q_{\phi}^{\bm{y}}({\bm{x}})$, and the KL minimization is done in a coarse-to-fine manner, similar to standard DIS, starting from high noise level to low noise level. While motivated differently, RED-diff and standard DIS have similar downsides of requiring at least a few tens of NFEs, as well as placing a simplistic proposal distribution. Furthermore, one can achieve only a single sample per optimization.

Recently, Feng *et al*. [feng2023score, feng2023efficient] uses an NF model for the proposal distribution while solving the same VI problem. The optimization problem involves computing the diffusion prior log likelihood $\log p_{\theta}({\bm{x}})$. It was shown that it can be exactly computed by solving the PF-ODE [feng2023score, song2020score], but numerically solving the PF-ODE per every optimization step is extremely computationally heavy, and hence does not scale well. To circumvent this issue, it was proposed to use a lower bound [feng2023efficient, song2021maximum]. Once trained, the NF model can be given different noise inputs ${\bm{z}}\sim{\mathcal{N}}(0,{\bm{I}})$ to generate diverse posterior samples with a single forward pass through the network. However, the training should be performed with respect to all the different measurements, not being able to generalize across the dataset. Our work follows along this path to overcome the current drawback
as shown in Table 1.

### 3.2 Distillation of the diffusion prior

Our method involves distillation of the diffusion prior into a student deep neural network, in our case an NF model. Particularly, it involves evaluating the output of the model by checking the denoising loss gradients from the pre-trained diffusion model. This idea is closely related to variants of score distillation sampling (SDS) [poole2023dreamfusion, wang2024prolificdreamer], where the gradient from the denoising loss is used to distill the diffusion prior by discarding the score Jacobian. Possibly a closely related work is Diff-instruct [luo2024diff], where the authors propose to train a one-step generative model similar to GANs [goodfellow2014generative] by distillation of the diffusion prior with VI. By proposing an integral KL divergence (IKL) by considering KL minimization across multiple noise levels across the diffusion, it was shown that SDS-like gradients can be used to effectively train a new generative model. While having similarities, our method directly minimizes the KL divergence and does not require dropping the score Jacobian.

Orthogonal to the score distillation approaches, there have been recent efforts to train a student network to emulate the PF-ODE trajectory itself [song2023consistency, gu2023boot] with a single NFE, one of the most prominent directions being consistency distillation (CD) [song2023consistency]. While promising, the performance of CM is upper-bounded by the teacher PF-ODE. Thus, in order to leverage CD-type approaches for diffusion posterior sampling, one has to choose one of the approximations of DIS as its teacher model. In this regard, applying CD for diffusion inverse problem solving is inherently limited.

## 4 Proposed Method

### 4.1 Conditional NF for amortized score prior

The goal is to use a variational distribution that is conditioned on ${\bm{y}}$, such that the resulting distilled conditional NF model $G_{\phi}$ generalizes to any condition ${\bm{y}}$. To this end, inspired from the choices of [sun2021deep, feng2023score] we modify the objective in Eq. (2) to

$$ $\displaystyle\min_{\phi}D_{KL}(q_{\phi}({\bm{x}}_{0}|{\bm{y}})||p_{\theta}({ \bm{x}}_{0}|{\bm{y}}))$ (9) $\displaystyle=$ $\displaystyle\min_{\phi}\int q_{\phi}({\bm{x}}_{0}|{\bm{y}})[-\log p({\bm{y}}| {\bm{x}}_{0})-\log p_{\theta}({\bm{x}}_{0})+\log q_{\phi}({\bm{x}}_{0}|{\bm{y} })]$ (10) $\displaystyle=$ $\displaystyle\min_{\phi}{\mathbb{E}}_{{\bm{z}}}\left[-\underbrace{\log p({\bm{ y}}|G_{\phi}({\bm{z}},{\bm{y}}))}_{\rm fidelity}-\underbrace{\log p_{\theta}(G _{\phi}({\bm{z}},{\bm{y}}))}_{\rm prior}+\underbrace{\log\pi({\bm{z}})-\log \left|\det\frac{dG_{\phi}({\bm{z}},{\bm{y}})}{d{\bm{z}}}\right|}_{\rm induced \>entropy}\right],$ (11) $$

where the second equality is the result of choosing a conditional NF as our proposal distribution, and now the expectation is over random noise ${\bm{z}}\sim{\mathcal{N}}(0,{\bm{I}})$. Notice that our network takes in both a random noise ${\bm{z}}$ and the condition ${\bm{y}}$ as an input to the network.

Under the Gaussian measurement model in Eq. (1), the fidelity loss reads

$$ $\displaystyle-{\mathbb{E}}_{\bm{z}}[\log p({\bm{y}}|G_{\phi}({\bm{z}},{\bm{y}} ))]$ $\displaystyle=-{\mathbb{E}}_{\bm{z}}\left[\frac{\|{\bm{y}}-{\mathcal{A}}(G_{ \phi}({\bm{z}},{\bm{y}}))\|_{2}^{2}}{2\sigma_{y}^{2}}\right].$ (12) $$

Moreover, the induced entropy can be easily computed as it is an NF

$$ $\displaystyle{\mathbb{E}}_{q_{\phi}({\bm{x}})}[\log q_{\phi}({\bm{x}})]={ \mathbb{E}}_{{\bm{z}}}\left[\log\pi({\bm{z}})-\log\left|\det\frac{dG_{\phi}({ \bm{z}},{\bm{y}})}{d{\bm{z}}}\right|\right]$ (13) $$

where $\pi(z)$, in our case, is the reference Gaussian distribution ${\mathcal{N}}$. For simplicity, let us denote $\hat{\bm{x}}_{0}:=G_{\phi}({\bm{z}},{\bm{y}})$.

Computation of $\log p_{\theta}(\hat{\bm{x}}_{0})$ is more involved: to exactly compute the value, we would have to solve the PF-ODE, which is compute-heavy [song2020score, feng2023score]. To circumvent this burden, we leverage the evidence lower bound (ELBO) [song2021maximum, feng2023efficient] $b_{\theta}(\hat{\bm{x}}_{0})\leq\log p_{\theta}(\hat{\bm{x}}_{0})$:

$$ $\displaystyle b_{\theta}(\hat{\bm{x}}_{0})={\mathbb{E}}_{p(\hat{\bm{x}}_{T}| \hat{\bm{x}}_{0})}\left[\log\pi(\hat{\bm{x}}_{T})\right]-\frac{1}{2}\int_{0}^{ T}g(t)^{2}h(t)\,dt$ (14) $$

where

$$ $\displaystyle h(t):={\mathbb{E}}_{p(\hat{\bm{x}}_{t}|\hat{\bm{x}}_{0})}$ $\displaystyle\Bigg{[}\underbrace{\|{\bm{s}}_{\theta}(\hat{\bm{x}}_{t})-\nabla_ {\hat{\bm{x}}_{t}}\log p(\hat{\bm{x}}_{t}|\hat{\bm{x}}_{0})\|_{2}^{2}}_{\rm DSM (Eq.~{}(\ref{eq:dsm}))}$ $\displaystyle-\|\nabla_{\hat{\bm{x}}_{t}}\log p(\hat{\bm{x}}_{t}|\hat{\bm{x}}_ {0})\|_{2}^{2}-\frac{2}{g(t)^{2}}\nabla_{\hat{\bm{x}}_{t}}\cdot{\bm{f}}(\hat{ \bm{x}}_{t},t).\Bigg{]}$ (15) $$

When we have $p({\bm{x}}_{t}|{\bm{x}}_{0})={\mathcal{N}}({\bm{x}}_{t};\alpha_{t}{\bm{x}}_{0}
,\beta_{t}^{2}{\bm{I}})$ and a standard diffusion model with a linear SDE ${\bm{f}}({\bm{x}}_{t},t)=f(t){\bm{x}}_{t}$,

$$ $\displaystyle\|\nabla_{{\bm{x}}_{t}}\log p({\bm{x}}_{t}|{\bm{x}}_{0})\|_{2}^{2 }=\frac{1}{\beta(t)^{2}}\|{\bm{\epsilon}}\|_{2}^{2},\quad\frac{2}{g(t)^{2}} \nabla_{{\bm{x}}_{t}}\cdot{\bm{f}}({\bm{x}}_{t},t)=\frac{2d\beta(t)}{g(t)^{2}},$ (16) $$

where $d$ is the dimensionality of ${\bm{x}}_{t}$, and both terms are independent of $\phi$ and $\theta$. Intuitively, the DSM term evaluates the probability of ${\bm{x}}_{0}$ by measuring how easy it is to denoise the given ${\bm{x}}_{0}$. When the network easily denoises the given image, then it will assign a high probability. When not, a low probability is assigned.
We can now define an equivalent ELBO $b^{\prime}_{\theta}({\bm{x}}_{0})$ in terms of optimization, which reads

$$ $\displaystyle b^{\prime}_{\theta}({\bm{x}}_{0})={\mathbb{E}}_{p_{0T}}[\log\pi( {\bm{x}}_{T})]-\frac{1}{2}\int_{0}^{T}g(t)^{2}\|{\bm{s}}_{\theta}({\bm{x}}_{t} )-\nabla_{{\bm{x}}_{t}}\log p({\bm{x}}_{t}|{\bm{x}}_{0})\|_{2}^{2}\,dt$ (17) $$

Plugging $b^{\prime}_{\theta}(G_{\phi}({\bm{z}},{\bm{y}}))$ of Eq. (17) in the place of $\log p_{\theta}(G_{\phi}({\bm{z}},{\bm{y}}))$ in Eq. (11), we can efficiently update $\phi$ by distilling the prior information contained in the diffusion model.

### 4.2 Architecture

It has been demonstrated in [lugmayr2020srflow] that Conditional NFs are capable of learning distributions on the ambient space that are constrained on measurement. To achieve an architecture with invertible transformations, we extend RealNVP [dinh2016density] architecture to the conditional settings by borrowing insight from [sun2021deep]. In its plain form, RealNVP architecture mainly consists of Flow steps, each containing two Affine Coupling layers. In each affine coupling layer, input signal ${\bm{x}}$ is split into two parts: ${\bm{x}}_{a}$ which stays unchanged, and ${\bm{x}}_{b}$ which is fed into the neural network. In order to invoke the condition, we simply concatenate the conditioning input ${\bm{y}}$ to the ${\bm{x}}_{b}$ as these layers serve as the main and basic building blocks of entire invertible architecture. This seemingly simple integration led to very promising results in both Euclidean and non-Euclidean geometries as will be depicted in Section 5.

### 4.3 Manifold

Manifolds exhibit more complex or curved structures compared to flat Euclidean spaces. In many real-world scenarios such as in environmental science, data naturally resides on these non-Euclidean geometries, and removing artifacts or denoising over these shapes becomes both necessary and hard. Traditional DIS methods usually rely on spatial domain and requires for off-the-shelf projection methods to get plausible results. To mitigate this, our work extends the application of conditional NFs to distributions lying on non-Euclidean manifolds and solve intricate inverse problems right on the surface without the need for additional renderings. By leveraging the expressive power of conditional NFs, we aim to capture the intrinsic geometry and structure of manifold data while enabling efficient inference and sampling. We treat data as a point cloud of size $V\times C$, in which $V$ and $C$ denotes the number of vertices on the discretization of mesh and dimension of signal features (e.g. $3$ in case of rgb), respectively. Several complex geometries and practical data have been considered and CNF is optimized to perform noisy inpainting and imputation tasks. Findings suggest that our framework can capture the distribution even in severe masking levels across different geometries (see Section 5).

## 5 Experiments

We validate our approach through various experiments, including (i) Denoising, Super Resolution, and Deblurring with CelebA face image data [liu2015faceattributes]; (ii) Inpainting on Stanford Bunny Manifold with MNIST data; and (iii) Imputation on Sphere with ERA5 [hersbach2020era5] temperature data. (i) Denoising is performed on the Euclidean in the image domain. In contrast, noisy (ii) inpainting and (iii) imputation are solved directly on the bunny and sphere manifolds. Throughout all the experiments, we use $24$ flow steps and we set the batch size to the $64$. We conduct all the training and optimization experiments on a single RTX3090 GPU instance. Our code is implemented in the JAX framework [jax2018github].

### 5.1 Experimental Settings

Inverse Problems on CelebA. We follow the usual formulation and adapt $32\times 32$ resolution of facial images. Data is normalized into [0, 1] range and measurement is acquired by the appropriate choices of forward operator depending on the task (See Appendix A for details). $G_{\phi}$ is optimized over the $19,962$ test images by using the forward operator and prior from diffusion models. We optimize CNF for $1$M iterations (convergency was observed earlier but continued for potential refinement).

Inpainting on Bunny MNIST. In order to demonstrate the geometric awareness of our model, we conduct experiment on Stanford Bunny Manifold. We choose the mesh resolution of $1889$ vertices and then project the [0, 1] normalized MNIST digits onto the bunny manifold [turk1994zippered]. In order to ensure the dimensionality compatibility for the models, we use $1888$ vertices and zero mask the last vertex throughout the experiments. We obtain the measurement by occluding $30\%$ of vertices randomly and adding some Gaussian noise, i.e. ${\mathcal{A}}$ is the random masking operator and $\sigma_{y}=0.1$ in (1). Our proposed CNF is trained on the test chunk of $10,000$ digit examples for $1$M iterations.

Imputation on ERA5. To show the essence and practical importance of our pipeline, we further conduct experiments on ERA5 temperature dataset. Even though data is available in a rectangular format, due to the spherical shape of Earth, it inherits some geometric information. We use $4^{\circ}$ resolution dataset with $4140$ vertices borrowed from [dupontetal5gasp] with only temperature channel as it is quite popular to analyze in the domain of generative AI. Again, due to the dimensionality, we add $20$ more vertices with a signal value of zero, and the data is [0, 1] normalized. In contrast to Bunny MNIST, we use more severe occlusion of $60\%$ random masking with additional Gaussian noise of $\sigma_{y}=0.05$. We perform the optimization of our CNF model on the test part of the dataset with $2420$ examples for $315$k iterations.

Score Networks. For all the diffusion priors VPSDE formulation has been adapted. In the case of image domain CelebA, we borrow the same score checkpoint used in the [feng2023score, feng2023efficient], which uses NCSN++ [song2021maximum] architecture and has been trained for $1$M iterations. For Bunny MNIST, we adapt the $1$D formulation of DDPM [ho2020denoising] and train the score network for $500$k, at which the convergence was clearly observed through the generated samples. In the case of spherical weather data, we followed the same strategy as Bunny MNIST but achieved convergence of score network earlier at $360$k iterations.

**Table 2: Quantitative results on our 3 main experiments. Best, second best**
| Geometry | Euclidean | Riemannian |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Dataset & Task | celebA (denoising) | Bunny-MNIST (inpainting) | ERA5 (imputation) |  |  |  |  |  |  |
| Metric | Time[s]$\downarrow$ | PSNR$\uparrow$ | SSIM$\uparrow$ | Time[s]$\downarrow$ | PSNR$\uparrow$ | MSE$\downarrow$ | Time[s]$\downarrow$ | PSNR$\uparrow$ | SSIM$\uparrow$ |
| MCG [chung2022improving] | - | - | - | 19.85 | 26.69 | 0.0024 | 16.16 | 27.52 | 0.871 |
| Noise2Score [kim2021noise2score] | 0.0172 | 24.36 | 0.871 | - | - | - | - | - | - |
| DPS [chung2023diffusion] | 16.95 | 27.93 | 0.932 | 19.39 | 28.03 | 0.0017 | 15.36 | 28.95 | 0.953 |
| Proposed ($N=1$) | 0.0021 | 23.37 | 0.836 | 0.0021 | 24.04 | 0.0041 | 0.0012 | 33.17 | 0.883 |
| Proposed ($N=128$) | 0.0035 | 25.82 | 0.901 | 0.0035 | 24.88 | 0.0034 | 0.0018 | 34.61 | 0.959 |

### 5.2 Results

In this section, we provide the general results of each different task described above. We compare our model with the various baselines including, DPS [chung2023diffusion], MCG [chung2022improving], and Noise2Score [kim2021noise2score].
It should be noted that MCG and DPS are identical for denoising, and Noise2Score is only applicable to denoising. In such cases, we do not report the metrics.
We also demonstrate the comparisons and results against Feng *et al*. [feng2023efficient]. Finally, we experimentally confirm the robustness of our model across different unseen data or datasets. For evaluation purposes, we use peak signal-to-noise ratio (PSNR) and structural-similarity-index-measure (SSIM) which are widely used to assess the performance of inverse solvers with the ground truth and reconstructed signals. We further evaluate Fréchet Inception Distance (FID) to showcase the perceptual quality of generated samples. As the proposed method can sample multiple different posterior samples with a single forward pass, and this process is easily parallelizable, we report two different types for the proposed method. One by taking a single posterior sample ($N=1$), and another by taking 128 posterior samples and taking the mean ($N=128$).

#### 5.2.1 General Results.

In general, our approach achieves competitive quantitative and qualitative results across different datasets on Euclidean and non-Euclidean geometries. We observe significant time improvements due to the single-step generation ability of our framework. Tab. 2 and Fig. 1 depict competitive quantitative and qualitative results confirming discussions along with instant time generations. It should be noted that the boosted version ($N=128$) of the proposed method only marginally increases the compute time, as we can sample multiple reconstructions in parallel. To demonstrate that our method can be applied to more general inverse problems, similar to [chung2023diffusion], we conducted 2$\times$ Super Resolution and Gaussian Deblurring experiments on the celebA data as shown in Tab. 3 & Fig. 4, where we see that the perceptual quality of the proposed method is better while being $\sim\times 1000$ faster, and the difference in the distortion metrics are small. Interestingly, while Noise2Score approximates the posterior mean, and the boosted version of the proposed method also approximates the posterior mean by taking the average of the posterior samples, our method outperforms Noise2Score by more than 1 db, showcasing the superiority of the proposed method.

#### 5.2.2 Comparing with DIS and Noise2Score.

Both DPS [chung2023diffusion] and MCG [chung2022improving] leverage the pre-trained diffusion model to sample from the posterior distribution. However, these methods require thousands of NFEs to achieve stable performance. The required time for DPS and MCG is reported in Tab. 2, 3. When decreasing the NFE as shown in Fig. 3 (a), PSNR heavily degrades and eventually diverges when we take an NFE value of less than 30. Our method achieves competitive performance even with a single NFE.
Moreover, it is shown in Fig. 3 (b) that even slightly incorrectly choosing the step size parameter leads to a large degradation in performance, whereas our method is free from such cumbersome hyperparameter tuning. Finally, it is shown in Fig. 3 (c) that DPS collapses to the mean of the prior distribution, altering the content of the measurement heavily when we take a smaller number of NFEs.

Figure: Figure 3: Comparison of our method against DPS [chung2023diffusion] on celebA denoising. (a) NFE vs. PSNR plot, (b) step size (used in DPS only) vs. PSNR plot, (c) representative results by varying the NFE.
Refer to caption: extracted/5754552/figs/comparison_dps_stability.png

It is worth mentioning that Noise2Score [kim2021noise2score] is applicable for one-step denoising of the measurements by leveraging the Tweedie’s formula. However, as discussed in Sec. 6, our method is generally applicable to a wide class of inverse problems, whereas the applicability of Noise2Score is limited.

Figure: Table 3: Quantitative results
Refer to caption: x1.png

#### 5.2.3 Comparing with score prior method.

Compared to the exact score prior [feng2023score], surrogate counterpart [feng2023efficient] presents $100$ times faster approach along with competitive or slightly better results in terms of quality. Despite being fast in terms of optimization of NF, [feng2023efficient] still requires training the network for a considerable amount of time for every single measurement. We observed that under same conditions, conditional NF does not increase the complexity and training stage takes 0.15 seconds which is 0.14 seconds in case of unconditional version. We further sample a random point from test data of celebA and optimize unconditional NF with the same configurations as ours on this single measurement. NF trained solely on this data reaches 23.75dB in PSNR score, which is almost same as our result of 23.43dB on this measurement. All these confirms that under same conditions, our approach of introducing conditions can simply achieve best results being also amortized for plenty of measurements.

#### 5.2.4 Generalizability across datasets and Blind Inverse Problems.

We further observe that our optimized framework can be used on unseen data as well. Table 4 and Figure 7 depicts that we achieve similar quantitative and qualitative reconstruction results when we sample from unseen celebA or ERA5 validation datasets. Note that score network is trained on train signals and CNF optimization has been conducted on test signals, i.e. validation is totally hidden to both teacher and student models. More strongly, our approach can also be leveraged on various datasets as pre-optimized inverse solver. To this end, we use our celebA optimized CNF model to perform Denoising task on the FFHQ [karras2019style] dataset. Same Table and Figure show that our model can remove the noise artifacts with a similar performance as it does on original data, confirming the generalizibility feature.

We further observed that our amortized inference approach can work in the absence of forward operator. In other words, we can perform blind inverse problems through our amortized posterior sampling. We used various imputation levels between 30% to 60% for ERA5 dataset, and conducted experiments with random choice of imputation in a blind manner. As a result, Fig. 5shows that results as good as the original inverse solver with the known forward operation (all levels with more than 33 PSNR value).

Figure: Figure 5: Blind inverse problem solving with varying imputation levels.
Refer to caption: x2.png

**Table 4: Our approach is robust against unseen data samples and even generalizable accross different datasets once it is optimized. For the first $2$ columns, we use validation part of datasets and feed-forward our CNF on this totally unseen data. For the third column, we even show that our pre-optimized CNF can be leveraged to restore back the noised data samples from across various datasets. All the quantitative results align with their counterparts in Table 2 that confirms robustness and generalization ability of our pipeline which was not possible before.**
| Dataset & Task | celebA val (denoising) | ERA5 val (imputation) | FFHQ (denoising) |
| --- | --- | --- | --- |
| PSNR$\uparrow$ | 23.26 | 33.12 | 21.92 |
| SSIM$\uparrow$ | 0.831 | 0.882 | 0.822 |

Figure: Figure 6: Figure demonstrating the visual results of robustness and generalization ability of our approach. First and second rows show the results on unseen validation data, and third row depicts generalization to another dataset. Corresponding quantitative analysis can be found in Table 4.
Refer to caption: extracted/5754552/figs/ood.png

## 6 Discussion

We show a first proof of concept that we can construct a one-step posterior sampler that generalizes across any measurements in an unsupervised fashion (only having access to the measurements ${\bm{y}}$). Notably, our method extends to wide use cases with minimal constraints: 1) the operator ${\mathcal{A}}$ can be arbitrarily complex and non-linear, as in DPS [chung2023diffusion], unlike many recent DIS that requires linearity of the operator [wang2023zeroshot, chung2024decomposed, zhu2023denoising]; 2) Training of the sampler can be done without any strict conditions on the measurement, unlike recent unsupervised score training methods that require i.i.d. measurement conditions with the same randomized forward operator [daras2023ambient, kawar2023gsure]. We opted for simplicity in the architecture design of $G_{\phi}$, and avoided introducing inductive bias of spatial information by taking a vectorized input, potentially explaining the slight background noise in the reconstructions. Further optimization in the choice of network architecture is left as a future direction of study.

## 7 Conclusion

In this work, we propose to use a conditional NF for a VI-based optimization strategy to train a one-step posterior sampler, which implicitly samples from the posterior distribution defined from the pre-trained diffusion prior. We show that our method is highly generalizable, being able to reconstruct samples that are not seen during training, applicable to diverse forward measurements, and types of data, encompassing standard Euclidean geometry as well as data on general Riemannian manifolds. We believe that our work can act as a cornerstone for developing a fast, practical posterior sampler that distills the diffusion prior.

## A Appendix

### A.1 Reproducibility and Details of Parameters

We further provide all the necessary details to replicate the results with our proposed method. Table 5 demonstrates the details of tasks and datasets along with the parameter choices for both prior score network training and Conditional NF optimization. Note that, we have validated our approach through $3$ different inverse problems on the image dataset (celebA), where noise level was set to 0.1 for denoising and 0.01 for Super-Resolution and Gaussian Deblurring. In addition to this, we also submit all the codes as a supplementary material to reproduce the results shown throughout the paper.

**Table 5: Different configurations of hyperparameter choices for varying datasets and manifolds learned by our model. (S) and (C) denotes the parameter choices for score network and CNF optimization, respectively.**
| Parameter | CelebA | Bunny | Sphere |
| --- | --- | --- | --- |
| resolution (#vertices) | 32 $\times$ 32 | 1889 | 4140 |
| distribution on manifold | - | MNIST | ERA5 |
| task | varying | Inpainting | Imputation |
| mask level | - | 30% | 60% |
| noise level | varying | 0.1 | 0.05 |
| #channels | 3 | 1 | 1 |
| normalized range | [0, 1] | [0, 1] | [0, 1] |
| (S) #train data | 162,770 | 60,000 | 8,510 |
| (S) batch size | 128 | 64 | 64 |
| (S) learning rate | 2e-4 | 2e-4 | 2e-4 |
| (S) #training iters | 1M | 500k | 360k |
| (C) #test data | 19,962 | 10,000 | 2,420 |
| (C) batch size | 64 | 64 | 64 |
| (C) learning rate | 1e-5 | 1e-5 | 1e-5 |
| (C) #optimization iters | 925k | 1M | 315k |

### A.2 Qualitative comparisons with baselines

Figure: Figure 7: Comparisons with different baselines for CelebA denoising, ERA5 imputation, and Bunny MNIST inpainting tasks. Note that second row shows the results of Noise2Score for the CelebA denoising task and MCG for inpainting and imputation of manifold data.
Refer to caption: extracted/5754552/figs/app_baselines.png

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: - The method seems to require a rather large test set to learn a proper approximation of the conditional normalizing flow. How does the size of the test set affect the performance of APS?  
- Do you think it is mainly the architecture of $G\_{\\phi}(z,y)$ that resulted in the reduced sample quality or is there a more fundamental reason related perhaps to the ELBO on $\\log p\_{\\theta}(\\hat{x}\_0)$?   
- Did you benchmark the diversity of the samples from APS? For example, against the baseline Score-Prior by Feng et al. (2023)? Or even against “standard” zero-shot diffusion posterior samplers such as DPS/DDNM? 
- Does APS lose the great advantage of zero-shot DIS methods being adaptable to any operator at test time without additional training?
- Is there an inherent limitation with respect to the resolution of images that APS can handle? If so, it is recommended to mention this explicitly in your discussion. If not, then perhaps APS can be applied to larger images (e.g. 128x128) for better benchmarking?  
- The distortion (e.g. PSNR/SSIM) of APS can be improved by sampling N=128 posterior samples and averaging them to approximate the posterior mean. However, shouldn’t a fair comparison do the same for the baselines?
soundness: 3
strengths: - The paper is very well written and mathematically sound.  
- The results cover diverse settings and show clear favorable runtime for APS.  
- Related work is addressed well and the paper’s contribution is put in proper context.
summary: The paper proposes to accelerate inverse problem-solving with diffusion priors. This is achieved by distilling the unconditional diffusion prior for a given set of inverse problems into a conditional normalizing flow model $q_{\phi}(\boldsymbol{x}|\boldsymbol{y})$ that can (approximately) sample from the posterior distribution $p(\boldsymbol{x}|\boldsymbol{y})$ given any input $\boldsymbol{y}$ with a single NFE. Training is performed with a principled variational objective that is mathematically sound, and experiments are performed on multiple data domains (euclidean/non-euclidean) showing $\approx$comparable results with significantly faster inference time.
weaknesses: - The method requires training at test time (e.g. on the test set). This stands in contradiction to the speed claim of the method being a fast sampler. The question then becomes: How large should the test set be to compensate for this prolonged training?   
- The method is claimed to produce high-quality and diverse posterior samples. Nonetheless, the resulting samples look rather noisy (e.g. CelebA denoising in Fig. 1, row 1), and they seem to have little to no diversity.  
- The datasets used in the experiments are rather unorthodox and make it hard to properly appreciate the performance of the method. For example, 32x32 CelebA is tiny \- it is hard to notice any meaningful differences at this resolution in the resulting samples. Similarly, the non-euclidean datasets seem to be relatively simple and quite constrained.

The main reason for my negative score is the relatively limited experiments and my current lack of understanding of the test scenario where APS could be advantageous. Nonetheless, the paper is well structured and the overall idea is interesting. I'm willing to raise my score if I'm provided convincing answers to these two critical issues. 

A few caught minor typos:
- L093 \- **one** the Euclidean manifold \-\> **on** the Euclidean manifold  
- L356 \- we set the batch size to **the** 64 \-\> we set the batch size to 64  
- L364 \- **convergency**  \-\> **convergence**  
- L377 \- “due to dimensionality”?  
- L467 \- leveraging **the** Tweedie’s formula \-\> leveraging Tweedie’s formula  
- L523 \- unlike recent DIS methods require \-\> unlike recent DIS methods **which** require

---

confidence: 5
contribution: 1
initial_rating: 3
presentation: 2
questions: - Please clarify the points raised in the weaknesses section.

- Is the neural network for the variational distribution optimized separately for each task, or can a single network generalize across multiple inverse problem tasks?

- The ability of APS to generalize to out-of-distribution datasets (e.g., applying a CelebA-trained model to the FFHQ dataset) is an appealing property. Can the authors compare these results to relevant baselines? Recent work, such as [1], has demonstrated similar generalization capabilities, and a direct comparison would provide useful insights.
soundness: 3
strengths: - Clarity: The paper is clearly written and well-structured, making the methodology and key contributions easy to understand.

- Significance: Previous approaches to solving inverse problems often require multiple neural network evaluations, leading to high computational costs. This paper’s proposed method addresses this limitation by using an amortized variational inference framework, enabling efficient, single-step sampling. This improvement in computational efficiency represents a meaningful step forward for practical applications of inverse problem-solving.
summary: This paper introduces Amortized Posterior Sampling (APS), a variational inference method designed for efficient posterior sampling in inverse problems. APS trains a conditional flow model to align the variational distribution with the posterior distribution defined by a diffusion model, enabling diverse posterior samples with a single neural function evaluation. The authors demonstrate APS’s effectiveness across tasks like image restoration, manifold signal reconstruction, and climate data imputation, showing it significantly improves computational efficiency while maintaining competitive reconstruction quality for real-time applications across various fields.
weaknesses: - Limited novelty: The primary contribution of the paper is the use of amortization and diffusion model distillation to enable efficient, single-step sampling. However, this concept has already been explored in recent work [1] with a very similar approach, also leveraging **unsupervised diffusion prior distillation** with **amortized variational inference** for inverse problems. This challenges the paper's claim of novelty (i.e., "the first diffusion prior distillation" in line 101 or "the first proof of concept" in line 515). This highly relevant work is neither discussed in the paper nor compared experimentally. Thorough discussion and performance comparison with this prior work to highlight meaningful distinctions would strengthen the methodological contribution.

- Several existing frameworks, such as RED-Diff, also utilize variational inference objectives for inverse problem-solving. The primary distinction between RED-Diff and APS seems to be the replacement of the posterior distribution with a neural network conditioned on y; the diffusion prior loss remains similar. While the amortization of neural networks does improve efficiency, both methods are based on similar principles. Thus, an experimental comparison between RED-Diff and APS could highlight potential trade-offs in accuracy and generalization, providing a clearer picture of APS’s advantages.

- Limited performance and baselines: The paper lacks experimental comparisons with key baseline methods [1-5] that are also designed to solve inverse problems. For instance, while RED-Diff [2] is mentioned in Table 1 as a related framework, there is no direct comparison of its performance with APS. This gap in baseline evaluations weakens the empirical foundation of the paper, making it challenging to assess the benefits of APS. Additionally, the experimental results of APS do not seem strong in table 2 and qualitatively in Figure 6; artifacts and noise are still noticeable, even with limited baselines.

[1] Diffusion Prior-Based Amortized Variational Inference for Noisy Inverse Problems, ECCV24 \
[2] A variational perspective on solving inverse problems with diffusion models, ICLR23 \
[3] Pseudoinverse-guided diffusion models for inverse problems, ICLR23 \
[4] Zero-shot image restoration using denoising diffusion null-space model, ICLR23 \
[5] Denoising diffusion models for plug-and-play image restoration, CVPRW23

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 3
questions: Please see the weakness for some of my concerns.

1. Can you clarify how many test images you use to obtain the metrics like PSNR and SSIM?

2. The use of varying metrics is a bit confusing; sometimes you report MSE, and other times FID or SSIM. A unified framework of metrics would be helpful.
soundness: 3
strengths: The paper is well-written and easy to follow. It employs a conditional normalizing flow model to achieve one-step posterior inference. Experiments on both Euclidean and non-Euclidean datasets across various tasks—including inpainting, denoising, Gaussian deblurring, and super-resolution—demonstrate the proposed distillation method's fast speed and competitive performance.
summary: This paper introduces Amortized Posterior Sampling (APS), a novel unsupervised variational inference approach for efficient posterior sampling in inverse problems using a conditional flow model. APS enables fast, diverse posterior sampling with a single neural evaluation and adapts to both Euclidean and non-Euclidean data. It achieves superior computational efficiency across tasks like image restoration and climate data imputation, providing real-time, high-quality solutions to various inverse problems.
weaknesses: 1. Though the proposed method enables one-step inference, training such a model using equation (11) in the main text is time-consuming. Additionally, it requires training on test images, whereas previous methods like DPS perform zero-shot inference on test images and achieve strong performance. Given the strong performance in the field of diffusion-based inverse problems, I find APS's performance underwhelming. Furthermore, including more challenging tasks on higher-resolution images would be beneficial.

2. Although the paper highlights its novelty in variational inference, it lacks comparisons with relevant baselines. Methods like RED-DIFF and the score prior have not been included, likely because they are not distillation methods. More distillation-based approaches for posterior sampling should be considered.

3. The method appears to combine the reconstruction step and the normalizing-flow step in the paper, score prior,  into a single step, which limits its novelty.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: Have the authors considered testing APS on broader inverse problem domains, beyond image and environmental data?

Could the authors clarify the performance of APS in terms of generalizability to entirely unseen problem classes?

Given APS’s unsupervised nature, how does it perform on ill-posed inverse problems where the prior distribution diverges significantly from the true posterior?

Is the observed background noise in reconstructions specific to the choice of vectorized input representation, or might this indicate a broader limitation of APS?
soundness: 3
strengths: The method allows one-step inference for posterior sampling, offering substantial time savings over traditional DIS methods

APS is tested on both Euclidean and non-Euclidean data, showcasing its potential in various fields.

The integration of conditional NFs with diffusion priors, particularly in unsupervised settings, is promising.
summary: This paper presents APS, an efficient variational inference-based posterior sampling method that combines conditional normalizing flows with diffusion models to achieve one-step posterior samples. APS is unsupervised, does not require paired training data, and is capable of tackling inverse problems in both Euclidean and non-Euclidean spaces. Experimentally, APS is evaluated on image denoising, inpainting, and data imputation tasks, demonstrating speed improvements while maintaining competitive accuracy.
weaknesses: Limited to specific models, such as MCG, DPS, and Noise2Score, but does not benchmark against other efficient variational methods recently introduced for inverse problems. Also benchmarking against similar methods would help demonstrate the behavior of this approach to others: DDNM, Pi-GDM, FPS, FPS-SMC, ect. For example, it’s unclear how the model can fail and what the failure cases may look like.