Ambient Diffusion Posterior Sampling: Solving Inverse Problems with Diffusion Models Trained on Corrupted Data

## Abstract

Abstract We provide a framework for solving inverse problems with diffusion models learned from linearly corrupted data. Our method, Ambient Diffusion Posterior Sampling (A-DPS), leverages a generative model pre-trained on one type of corruption (e.g. image inpainting) to perform posterior sampling conditioned on measurements from a potentially different forward process (e.g. image blurring).
We test the efficacy of our approach on standard natural image datasets (CelebA, FFHQ, and AFHQ) and we show that A-DPS can sometimes outperform models trained on clean data for several image restoration tasks in both speed and performance. We further extend the Ambient Diffusion framework to train MRI models with access only to Fourier subsampled multi-coil MRI measurements at various acceleration factors (R = 2 , 4 , 6 , 8 absent 2 4 6 8 =2,4,6,8 = 2 , 4 , 6 , 8 ). We again observe that models trained on highly subsampled data are better priors for solving inverse problems in the high acceleration regime than models trained on fully sampled data. We open-source our code and the trained Ambient Diffusion MRI models: github.com/utcsilab/ambient-diffusion-mri .

## 1 Introduction

For some applications, it is expensive or impossible to acquire fully observed or uncorrupted data [, , ] but possible to acquire partially observed samples.
Further, in some cases, it may be desirable to train generative models with
noisy or corrupted data since that reduces memorization of the training set [, , ]. Prior works have shown how to train Generative Adversarial Networks (GANs) [, ], flow models [] and restoration models  [, , , ] with corrupted training data. More recently, there has been a shift towards training diffusion generative models given corrupted data  [, , , , ]. What remains unexplored is how to use generative models trained on a certain type of corruption (e.g. inpainted data) to solve inverse problems that arise from a different forward process (e.g. blurring). Thus far, it is not known whether models trained on corrupted data are useful priors for solving inverse problems.

Figure: Figure 1: Illustration of Ambient Diffusion Posterior Sampling (Ambient DPS). During training, we only have access to linearly corrupted data from a forward operator $A_{\mathrm{train}}$. We use this data and the Ambient Diffusion framework to learn a generative model, $G_{\mathrm{ambient}}$, for the uncorrupted distribution, $p({\bm{x}}_{0})$. At inference time, we use the learned generator to sample from the posterior distribution $p({\bm{x}}_{0}|{\bm{y}}_{\mathrm{A_{\mathrm{inf}}}})$, for measurements ${\bm{y}}_{\mathrm{inf}}$ coming from a different forward operator, $A_{\mathrm{inf}}$.
Refer to caption: extracted/5468845/assets/figs/framework_fig.png

We propose the first framework to solve inverse problems with diffusion models learned from linearly corrupted data, for example as in Ambient Diffusion []. Ambient Diffusion models estimate the ambient score, i.e. how to best reconstruct given an input with a corrupted linear forward operator corrupted noisy input. We show how to use these models for solving linear inverse problems outside of their training distribution. Our experiments on standard datasets of natural images and multi-coil MRI show something surprising: Ambient Models can outperform (in the high corruption regime) models trained on clean data. Further, they do so while being substantially faster. Our algorithm extends Diffusion Posterior Sampling [] to Ambient Models and fully unlocks the potential of generative models trained on corrupted data for image restoration.

Our contributions:

- 1.
We propose Ambient Diffusion Posterior Sampling (Ambient DPS), an algorithm that uses diffusion models trained on linearly corrupted data as priors for solving inverse problems with arbitrary linear measurement models.
- 2.
We use pre-trained Ambient Diffusion models to solve inverse problems (compressed sensing, super-resolution) on several natural image datasets (CelebA, FFHQ, AFHQ) and show that they can even outperform models trained on clean data in the high corruption regime.
- 3.
We extend the Ambient Diffusion framework to train models directly on multi-coil Fourier subsampled MRI.
- 4.
We train Ambient Diffusion MRI models on FastMRI at various retrospective acceleration factors (R=$2,4,6,8$); we again observe that models trained on subsampled data are better priors for solving inverse problems in the high acceleration regime.
- 5.
We open-source our code and models to accelerate research in this area: [github.com/utcsilab/ambient-diffusion-mri](https://github.com/utcsilab/ambient-diffusion-mri).

## 2 Method

### 2.1 Background and Notation

#### Diffusion Posterior Sampling (DPS).

Diffusion models are typically trained (up to network reparametrizations) to reconstruct a clean image ${\bm{x}}_{0}\sim p_{0}({\bm{x}}_{0})$ from a noisy observation ${\bm{x}}_{t}={\bm{x}}_{0}+\sigma_{t}\bm{\eta},\quad\eta\sim\mathcal{N}(\mathbf
{0},I)$. Despite the simplicity of the training objective, diffusion models can approximately sample from $p({\bm{x}})$ by running a discretized version of the Stochastic Differential Equation [, ]:

$$ $\displaystyle\mathrm{d}{\bm{x}}=-2\dot{\sigma}_{t}(\mathbb{E}[{\bm{x}}_{0}|{ \bm{x}}_{t}]-{\bm{x}}_{t})\mathrm{d}t+g(t)\mathrm{d}{\bm{w}},$ (2.1) $$

where ${\bm{w}}$ is the standard Wiener process and $\mathbb{E}[{\bm{x}}_{0}|{\bm{x}}_{t}]$ is estimated by the trained neural network. Given a measurement ${\bm{y}}_{\mathrm{inf}}=A_{\mathrm{inf}}{\bm{x}}_{0}$, one can sample from the posterior distribution $p({\bm{x}}_{0}|{\bm{y}}_{\mathrm{inf}})$ by running the process:

$$ $\displaystyle\mathrm{d}{\bm{x}}=-2\dot{\sigma}_{t}\sigma_{t}\left(\frac{ \mathbb{E}[{\bm{x}}_{0}|{\bm{x}}_{t}]-{\bm{x}}_{t}}{\sigma_{t}}+\underbrace{ \nabla\log p({\bm{y}}_{\mathrm{inf}}|{\bm{x}}_{t})}_{\mathrm{likelihood\ term} }\right)\mathrm{d}t+g(t)\mathrm{d}{\bm{w}}.$ (2.2) $$

For most forward operators it is intractable to write the likelihood in closed form. Hence, several approximations have been proposed to use diffusion models for inverse problems [, , , , , , ]. One of the simplest and most effective approximations is Diffusion Posterior Sampling (DPS) []. DPS estimates ${\bm{x}}_{0}$ using ${\bm{x}}_{t}$ and uses the conditional likelihood $p({\bm{y}}_{\mathrm{inf}}|\hat{\bm{x}}_{0})$ instead of the intractable term, i.e. DPS approximates $p({\bm{y}}_{\mathrm{inf}}|{\bm{x}}_{t})$ with $p({\bm{y}}_{\mathrm{inf}}|{\bm{x}}_{0}=\mathbb{E}[{\bm{x}}_{0}|{\bm{x}}_{t}])$. The update rule becomes:

$$ $\displaystyle\mathrm{d}{\bm{x}}=-2\dot{\sigma}_{t}\sigma_{t}\left(\frac{ \mathbb{E}[{\bm{x}}_{0}|{\bm{x}}_{t}]-{\bm{x}}_{t}}{\sigma_{t}}+\gamma_{t} \nabla_{{\bm{x}}_{t}}\log p({\bm{y}}_{\mathrm{inf}}|{\bm{x}}_{0}=\mathbb{E}[{ \bm{x}}_{0}|{\bm{x}}_{t}])\right)\mathrm{d}t+g(t)\mathrm{d}{\bm{w}},$ (2.3) $$

where $\gamma_{t}$ is a tunable guidance parameter.

#### Ambient Diffusion.

In some settings we do not have a large training set of clean data but we have access to a large set of lossy measurements that we would like to leverage to train a diffusion model for the clean distribution.

[, ] considered the setting of having access to linearly corrupted data $\{{\bm{y}}_{0}=A_{\mathrm{train}}x_{0},A_{\mathrm{train}}\}$, where the distribution of $A_{\mathrm{train}}$, denoted as $p(A_{\mathrm{train}})$, is assumed to be known. The ultimate goal is to learn the best restoration model given a linearly corrupted noisy input ${\bm{y}}_{t,\mathrm{train}}=A_{\mathrm{train}}({\bm{x}}_{0}+\sigma_{t}\bm{\eta})$, at all noise levels $t$. [, ] form further corrupted iterates $\tilde{\bm{y}}_{t,\mathrm{train}}=\tilde{A}_{\mathrm{train}}({\bm{x}}_{0}+
\sigma_{t}\bm{\eta})$ and train with the following objective:

$$ $\displaystyle J(\theta)=\mathbb{E}_{{\bm{x}}_{0},{\bm{x}}_{t},A_{\mathrm{train }},\tilde{A}_{\mathrm{train}}}\left[\left|\left|A_{\mathrm{train}}{\bm{h}}_{ \theta}(\tilde{\bm{y}}_{0,\mathrm{train}},\tilde{A}_{\mathrm{train}})-{\bm{y}} _{t,\mathrm{train}}\right|\right|\right],$ (2.4) $$

that provably learns $\mathbb{E}[{\bm{x}}_{0}|\tilde{A}_{\mathrm{train}},\tilde{\bm{y}}_{t,\mathrm{
train}}]$ as long as the matrix $\mathbb{E}[A_{\mathrm{train}}^{T}A_{\mathrm{train}}|\tilde{A}_{\mathrm{train}}]$ is full-rank. In certain cases, it is possible to introduce minimal additional corruption and satisfy this condition. For example, if $A_{\mathrm{train}}$ is a random inpainting matrix (i.e. $A_{ij}\sim\mathrm{Be}(1-p))$, then $\tilde{A}_{\mathrm{train}}$ can be formed by taking $A_{\mathrm{train}}$ and erasing additional pixels with any non-zero probability $\delta>0$.

#### Multi-coil Magnetic Resonance Imaging.

MRI is a prototypical use case for a framework that can learn generative models from linearly corrupted data, as in many cases it is not feasible to collect a large training set of fully sampled data [, , ]. In settings such as 3D+time dynamic contrast-enhanced MRI [] it is impossible to collect fully sampled data due to the time-varying dynamics of the contrast agent [].

In the multi-coil MRI setting, the acquisition involves collecting measurements of an image directly in the spatial frequency, known as $k$-space, from a set of spatially localized coils. Mathematically, there are $N_{\mathrm{c}}$ coils, each of which gives measurements:

$$ $\displaystyle{\bm{z}}_{{\bm{x}},i}=P\mathcal{F}S_{i}{\bm{x}}+{\bm{w}}_{i}, \quad i\in[N_{\mathrm{c}}],$ (2.5) $$

where ${\bm{x}}$ is the (complex-valued) image of interest, $S_{i}$ represents the coil-sensitivity profile of the $i^{th}$ coil, $\mathcal{F}$ is the (2-D or 3-D) Fourier transform, $P$ represents the Fourier subsampling operator and ${\bm{w}}_{i}$ is complex-valued Gaussian i.i.d noise. For simplicity in this work, we assume that we work in the noiseless case and we point the reader to the works of [, , ] for approaches that handle the noisy case.

For the discrete approximation of the continuous signal as an image ${\bm{x}}$ $\in$ $C^{n}$, the composition of $P,\mathcal{F},S_{i}$ can be written as a matrix $A_{i}\in C^{m\times n}$, where the number of measurements, $m$, depends on the subsampling operator $P$. It is common to denote with $R$ the ratio $\frac{n}{m}$, which is known as the acceleration factor. At inference time (i.e., for a new patient), we typically want to acquire data with a high acceleration factor because this reduces scan time and patient discomfort.

Due to the time required to collect $k$-space measurements, it is often not possible to acquire fully sampled $k$-space for an entire dataset. This means that training cannot rely on a fully sampled image to guide reconstruction quality as is done in the fully supervised setting [, ]. To overcome this limitation, some end-to-end techniques use a loss on the measurement domain by partitioning the training measurements for each sample into (1) measurements for reconstruction, and (2) measurements for applying the loss function []. Other approaches leverage structure in the MRI acquisition to learn from limited-resolution data []. More recently, works have begun leveraging signal set properties such as group invariance to assist in learning from subsampled measurements for a variety of inverse problems [, , , ].

All these approaches learn a restoration model without access to fully sampled reference data. However, as they are inherently end-to-end methods, their performance on out-of-distribution tasks (e.g., due to different acquisition trajectories) is known to degrade [, ].

### 2.2 Ambient Diffusion for MRI

As explained above, the MRI acquisition process results in linearly corrupted measurements and thus it should be possible to use the Ambient Diffusion framework to learn diffusion models from the corrupted observations. Yet, we identified three important changes that differ from the setting studied by [, ]: i) the inpainting happens in the Fourier Domain, ii) the inpainting has structure, i.e. whole vertical lines in the spatial Fourier are either observed or not observed (instead of masking random pixels) and, iii) the image is measured from an array of spatially varying receive coils. In what follows, we show how to account for these factors and apply the Ambient Diffusion framework to MRI data.

First, as in Ambient Diffusion, we will corrupt further the given measurements, but in our case, the additional corruption will happen in the Fourier space. We create further corrupted measurements for each coil by decreasing the Fourier subsampling ratio, i.e. we create iterates:

$$ $\displaystyle\tilde{\bm{z}}_{{\bm{x}},i}=\tilde{P}\mathcal{F}S_{i}{\bm{x}}.$ (2.6) $$

Note that this is always possible to do since we can just subsample the available data ${\bm{z}}_{{\bm{x}},i}$. We also underline that to further corrupt the given measurements we keep the inpainting structure that there is in the measurements, i.e. we delete (at random) 2-D lines in the Fourier space (instead of independent pixel masking as in [, ]).

In that regards the multiplicity of the coils, we combine their measurements to form a crude estimator of ${\bm{x}}$ by taking the adjoint of $\tilde{A}_{\mathrm{train}}$:

$$ $\displaystyle\tilde{\bm{y}}_{\mathrm{train}}=\sum_{i}S_{i}^{H}\mathcal{F}^{-1} (\tilde{\bm{z}}_{{\bm{x}},i}).$ (2.7) $$

Notice that for a discrete signal ${\bm{x}}\in\mathbb{C}^{n}$, all these operations are linear and hence $\tilde{\bm{y}}$ can be written as:

$$ $\displaystyle\tilde{\bm{y}}_{\mathrm{train}}=\underbrace{\left(\sum_{i}S_{i}^{ H}\mathcal{F}^{-1}\tilde{P}\mathcal{F}S_{i}\right)}_{\tilde{A}_{\mathrm{train} }}{\bm{x}}.$ (2.8) $$

As in Ambient Diffusion, we will train the network to predict the corresponding signal before the additional corruption, i.e. our target will be:

$$ $\displaystyle{\bm{y}}_{\mathrm{train}}=\underbrace{\left(\sum_{i}S_{i}^{H} \mathcal{F}^{-1}P\mathcal{F}S_{i}\right)}_{A_{\mathrm{train}}}{\bm{x}}.$ (2.9) $$

We are now ready to state our main Theorem.

###### Theorem 2.1 ((Informal)) .

Let ${\bm{x}}_{t},{\bm{y}}_{t,\mathrm{train}},\tilde{\bm{y}}_{t,\mathrm{train}}$ represent the noisy versions of ${\bm{x}},{\bm{y}}_{\mathrm{train}},\tilde{\bm{y}}_{\mathrm{train}}$ respectively, i.e.:

$$ $\displaystyle\begin{cases}{\bm{x}}_{t}={\bm{x}}+\sigma_{t}\bm{\eta}\\ {\bm{y}}_{t,\mathrm{train}}=A_{\mathrm{train}}\left({\bm{y}}_{t,\mathrm{train} }+\sigma_{t}\bm{\eta}\right)\\ \tilde{\bm{y}}_{t,\mathrm{train}}=\tilde{A}_{\mathrm{train}}\left(\tilde{\bm{y }}_{t,\mathrm{train}}+\sigma_{t}\bm{\eta}\right)\end{cases}.$ (2.10) $$

Then, the minimizer of the objective:

$$ $\displaystyle J(\theta)=\mathbb{E}_{{\bm{y}}_{0,\mathrm{train}},\tilde{\bm{y}} _{t,\mathrm{train}},A_{\mathrm{train}},\tilde{P}}\left[\left|\left|A_{\mathrm{ train}}{\bm{h}}_{\theta}(\tilde{\bm{y}}_{t,\mathrm{train}},\tilde{P})-{\bm{y}} _{0,\mathrm{train}}\right|\right|^{2}\right],$ (2.11) $$

is: ${\bm{h}}_{\theta}(\tilde{\bm{y}}_{t,\mathrm{train}},\tilde{P})=\mathbb{E}[{\bm
{x}}_{0}|\tilde{\bm{y}}_{t,\mathrm{train}},\tilde{P}]$.

###### Theorem 2.1 ((Informal)) .

#### Proof overview.

The formal proof of this Theorem is given in the Appendix. We provide a sketch of the proof here. The techniques are based on Theorem 4.2 of Ambient Diffusion []. In Theorem 4.2, the condition that needs to be satisfied is that
$\mathbb{E}[A_{\mathrm{train}}^{T}A_{\mathrm{train}}|\tilde{A}_{\mathrm{train}}]$ is full-rank. Our objective is slightly different because the network doesn’t see the whole forward operator $\tilde{A}_{\mathrm{train}}$, but only part of it, i.e. the matrix $\tilde{P}$. By following the steps of the Ambient Diffusion proof, we arrive at the necessary condition for our case, which is to prove that $\mathbb{E}[A_{\mathrm{train}}|\tilde{P}]$ is full-rank.

We start the argument by noting that $\mathbb{E}[P|\tilde{P}]$ is full-rank. Intuitively, this is because there is a non-zero probability of observing any Fourier coefficient and hence any deleted Fourier co-efficient could be due to the extra corruption introduced by $\tilde{P}$. Since $\mathcal{F}$ is an invertible matrix, then also $\mathbb{E}[\mathcal{F}^{-1}P\mathcal{F}|\tilde{P}]$ is full-rank. Finally, we use the properties of the sensitivity masks to show that
$\mathbb{E}\left[\sum_{i}S_{i}^{H}\mathcal{F}^{-1}P\mathcal{F}S_{i}|\tilde{P}\right]$ is also full-rank and this completes the proof.

#### Note on alternative Ambient Diffusion designs.

Our approach to extending Ambient Diffusion to the Fourier subsampled MRI case is by no means unique. For example, we could have used other linear operators to aggregate the measurements from the different coils (e.g. we could have used the pseudoinverse), we could further condition on the sensitivity maps, or we could have directly fed all the coil measurements as input to the network without aggregating them. Such approaches could potentially lead to improved performance since they fully leverage the coil information. In this work, we opted for a simple and fast method (aggregation through the adjoint) and we leave such investigations for future work.

### 2.3 Ambient Diffusion Posterior Sampling

DPS requires access to $\mathbb{E}[{\bm{x}}_{0}|{\bm{x}}_{t}]$ to approximately sample from $p({\bm{x}}_{0}|{\bm{y}}_{\mathrm{inf}})$. Since Ambient Diffusion models can only work with corrupted inputs, we propose the following update rule instead:

$$ $\displaystyle\mathrm{d}{\bm{x}}=-2\dot{\sigma}_{t}\sigma_{t}\left(\underbrace{ \frac{\mathbb{E}[{\bm{x}}_{0}|{\bm{y}}_{t,\mathrm{train}},A_{\mathrm{train}}]- {\bm{x}}_{t}}{\sigma_{t}}}_{\mathrm{Ambient\ Score}}+\gamma_{t}\nabla_{{\bm{x} }_{t}}\log p({\bm{y}}_{\mathrm{inf}}|{\bm{x}}_{0}=\mathbb{E}[{\bm{x}}_{0}|{\bm {y}}_{t,\mathrm{train}},A_{\mathrm{train}}])\right)\mathrm{d}t+g(t)\mathrm{d}{ \bm{w}},$ (2.12) $$

for a fixed $A_{\mathrm{train}}\sim p(A_{\mathrm{train}})$. Comparing this to the DPS update rule (E.q. 2.3), all the $\mathbb{E}[{\bm{x}}_{0}|{\bm{x}}_{t}]$ terms have been replaced with their ambient counterparts, i.e. with $\mathbb{E}[{\bm{x}}_{0}|{\bm{y}}_{t,\mathrm{train}},A_{\mathrm{train}}]$. We remark that similar to DPS, the proposed algorithm is an approximation to sampling from the true posterior distribution $\mathbb{E}[{\bm{x}}_{0}|{\bm{y}}_{\inf}]$. We term our approximate sampling algorithm for solving inverse problems with diffusion models learned from corrupted data Ambient DPS (A-DPS).

### 2.4 Ambient Diffusion for in-domain reconstructions.

Ambient DPS can be used to solve any inverse problem for which the forward operator is known, using a diffusion model trained on linearly corrupted data of some form. It is important to underline though that if the inverse problem that we want to solve at inference time has the same forward operator as the one that was used for the training measurements, then we can use the Ambient Diffusion as a supervised restoration model. This is because Ambient Diffusion models are trained to estimate $\mathbb{E}[{\bm{x}}_{0}|\tilde{\bm{y}}_{\mathrm{train}},\tilde{A}_{\mathrm{
train}}]$, and hence if $A_{\mathrm{inf}}$ comes from the same distribution as $A_{\mathrm{train}}$, then, the one-step prediction of the model, Ambient One Step (A-OS), is the Mean Squared Error (MSE) minimizer. Similarly, if we are interested in unconditional generation, we can simply run Ambient DPS without the likelihood term.

## 3 Experiments

### 3.1 Ambient Diffusion with Pre-trained Models on Datasets of Natural Images

Figure: (a) LPIPS per Number of Measurements.
Refer to caption: x1.png

In this section, we evaluate the performance of Ambient DPS, which uses diffusion models trained on corrupted data, and we compare it to DPS, which uses diffusion models trained on clean data. For our experiments, we use the models from the Ambient Diffusion [] that are trained on randomly inpainted data with different erasure probabilities. Specifically, for AFHQ we use the Ambient Models with erasure probability $p\in\{0.2,0.4,0.6,0.8\}$ and for Celeb-A we use the pre-trained models with $p\in\{0.6,0.8,0.9\}$.

We underline that all the Ambient Models have worse performance for unconditional generation compared to the models trained with clean data (i.e. the models trained with $p=0.0$). The goal of this work is to explore the conditional generation performance of Ambient Models, where the conditioning is in the measurements ${\bm{y}}_{\mathrm{inf}}$, and compare it with models trained on uncorrupted data. To ensure that Ambient Models do not have an unfair advantage, we test only on restoration tasks that are different from the ones encountered in their training. Specifically, we use models trained on random inpainting and we evaluate on Gaussian Compressed Sensing [] and Super Resolution.

#### Hyperparameters.

The only tunable parameters for DPS (Eq. 2.3) and Ambient DPS (Eq. 2.12) are in the scheduling of the magnitude of the measurements likelihood term. In all the experiments in the DPS paper, this term is kept constant throughout the diffusion sampling trajectory and the authors recommend selecting a value in the range between $[0.1,10]$. We follow this recommendation and we keep this term constant. The value of the step size for each model is selected with a hyperparameter search in the recommended range. For all our experiments, we follow exactly the DPS implementation provided in the official code repository of the paper. The other parameter that impacts performance is the number of steps we are going to run each algorithm for, i.e. the discretization level of the SDEs of Equations 2.3, 2.12. Typically, the higher the number of steps the better the performance since the discretization error decreases [, ]. For the performance results, we run each method for several steps $\in\{50,100,150,200,250,300\}$, and we report the best result among them.

Figure: (a) LPIPS per downsampling factor.
Refer to caption: x3.png

Figure: (a) Compressed Sensing with $4000$ measurements per Number of Function Evaluations (NFEs).
Refer to caption: x5.png

#### Results.

Figure 2 presents Gaussian Compressed Sensing reconstruction results (i.e. reconstructing a signal from Gaussian random projections).
We show MSE and LPIPS performance metrics for the AFHQ dataset as we vary the number of measurements. The results are given for models that are trained with inpainted images at different levels of corruption, indicated by the erasure probability $p$. As shown in the Figure, the model trained with clean data outperforms the models trained with corrupted data when the number of measurements is high. However, as we reduce the number of measurements,
Ambient Models outperform the models trained with clean data in the very low measurements regime. To the best of our knowledge, there is no known theoretical argument that explains this performance cross-over, and understanding this further is an interesting research direction. Similar results are presented in Figure 3 for the task of super-resolution at AFHQ. The model trained on clean data ($p=0$) slightly outperforms the Ambient Models in both LPIPS and MSE for reconstructing a $2\times$ downsampled image, as expected. Yet, as the resolution decreases, there is again a cross-over in performance and models trained on corrupted data start to outperform the models trained on uncorrupted data. We include results for LPIPS and MSE for Compressed Sensing and Downsampling in FFHQ and Celeb-A in the Appendix (Figs 10, 11, 12, 13).

Finally, we examine how the number of sampling steps affects the performance. The MSE results for Compressed Sensing with $4000$ measurements on AFHQ are shown in Figure 4. As shown, the higher the erasure probability $p$ during training, the better the Compressed Sensing performance of the model for low Number of Function Evaluations (NFEs). Models trained with higher corruption are faster since they require fewer steps for the same performance. For increased NFEs, the models that are trained on clean(er) data finally outperform. This result is consistent across different datasets (AFHQ, FFHQ, CelebA), reconstruction tasks (Compressed Sensing, Downsampling), and metrics (MSE, LPIPS) (Figures 14, 15, 16, 17, 18 in the Appendix).

### 3.2 Ambient MRI Diffusion Models

In the previous section, we used pre-trained diffusion models on datasets of natural images. In this section, we extend Ambient Diffusion to the multi-coil Fourier subsampled MRI setting (as detailed in Section 2.2) and we train our own MRI models from scratch.

**Table 1: Unconditional Sampling FID Scores for models trained on MRI data at different acceleration factors.**
| Inference/Training | Training Data R | Reconstruction Method | FID |
| --- | --- | --- | --- |
| EDM | $1$ | MVUE | $10.4082$ |
| EDM | $2$ | L1-Wavelet | $18.5534$ |
| EDM | $4$ | L1-Wavelet | $27.6405$ |
| EDM | $6$ | L1-Wavelet | $51.4272$ |
| EDM | $8$ | L1-Wavelet | $102.984$ |
| NCSNV2 | $1$ | MVUE | $13.2006$ |
| Ambient | $2$ | Adjoint | $30.3425$ |
| Ambient | $4$ | Adjoint | $32.3101$ |
| Ambient | $6$ | Adjoint | $31.5033$ |
| Ambient | $8$ | Adjoint | $48.1525$ |

#### Dataset.

For the dataset preparation, we follow standard practices in the MRI literature. Specifically, we download the FastMRI dataset, randomly select 2,000 T2-Weighted scans, and pick 5 center slices from each scan.
This gives us a dataset of 10,000 T2-weighted brain scans ($k$-space measurements).
We pass each multi-channel $k$-space sample through a noise pre-whitening filter to transform the noise into standard white Gaussian.
Then, we normalize each $k$-space sample by the absolute value of $99$-th percentile of the root sum-of-squares reconstruction of the autocalibration region (24x24 pixels from the center).
Given the pre-processed fully sampled $k$-space samples, we estimate sensitivity maps using the ESPIRiT calibration method []. These measurements serve as a fully sampled reference dataset (i.e., $R=1$).

#### Sampling masks.

We retrospectively subsample the $k$-space training data by applying randomly subsampled masks that correspond to acceleration factors $R\in\{2,4,6,8\}$. The sampling masks include fully sampled vertical (readout) lines corresponding to the observed Fourier coefficients (phase encodes). We always sample the central $20$ lines for autocalibration. To form further corrupted measurements, we randomly remove additional lines such that we create measurements at acceleration factor $R+1$. For example, we take data at $R=2$, we create further corrupted measurements at $R=3$ and we train the model to predict the clean image by measuring its error with the available data at $R=2$ given its prediction for the $R=3$ corrupted input.

We note that, based on the Ambient Diffusion theory, the extra corruption is a free parameter that can be chosen to be minimal, and perhaps, we could have chosen to corrupt the available data less aggressively. Due to the lack of computational resources, we did not ablate the degree of further corruption and we leave this as future work.

#### Comparison Models.

We train one model for each acceleration factor $R\in\{1,2,4,6,8\}$. The $R=1$ model is trained directly on clean data (no extra corruption) based on the EDM approach []. The $R>1$ models are trained with our modified Ambient Diffusion framework. While we focus on EDM for fully sampled training, we also train an NCSNV2 model for $R=1$ following the approach in []. As an additional baseline, we also train models with EDM after L1-Wavelet compressed sensing reconstruction of the training set at each acceleration factor []. We used the BART implementation [, ] and searched for the best regularization parameter over the training set. Specifically, we used a regularization weighting of 0.001 with 100 iterations. We point the interested reader to Section E in the Appendix (and to our open-sourced code) for all the training hyperparameters.

Figure: Figure 5: Prior samples from Ambient Diffusion trained with under sampled data at $R=2,4,6,8$ (columns $1-4$), and EDM trained with fully sampled data (column $5$).
Refer to caption: extracted/5468845/assets/figs/mri/prior.png

Figure: Figure 6: Posterior reconstructions using A-DPS (columns $2-5$) and FS-DPS (column $6$). Rows $1$ and $3$ show reconstructions at $R=4$ and $R=8$, respectively, while rows $2$ and $4$ display their difference images on a $10\times$ scale.
Refer to caption: extracted/5468845/assets/figs/mri/posterior.png

#### Unconditional generation evaluation.

We first evaluate the unconditional generation performance of each model. We show prior samples for the Ambient Diffusion models and EDM at $R=1$ in Figure 5. Additional samples for all trained models are included in the Appendix, Figure 19. Similar to the results in [, ], we observe that the Ambient Diffusion samples are qualitatively similar to the EDM models at low acceleration, and become slightly blurry at higher accelerations. Notably, the samples generated by Ambient Diffusion at $R=8$ have no residual aliasing artifacts, in contrast to the samples generated from L1-Wavelet reconstructions (as seen in the Appendix).

To further quantify unconditional sample quality, we follow the approach in [] and calculate FID scores from 100 samples using a pre-trained VGG network. Table 1, shows the FID scores for the different diffusion models. While there is an increase in FID for the Ambient models, we see that those trained at higher acceleration factors outperform the models trained on L1-Wavelet reconstructions.

### 3.3 MRI Reconstruction

After verifying that our trained models have learned the underlying distribution reasonably well, we want to evaluate how they perform in the task of accelerated image reconstruction. We start by listing the details of our models and the considered baselines.

#### Ambient Diffusion Posterior Sampling (A-DPS)

We implemented our method by following Eq. 2.12 for $500$ steps with a likelihood weighting $\gamma_{t}=\frac{1}{\left\|{\bm{y}}-A\mathbb{E}[{\bm{x}}_{0}|{\bm{y}}_{t,
\mathrm{train}},A_{\mathrm{train}}]\right\|_{2}}$.

#### Ambient One Step (A-OS)

As mentioned previously, our method also admits a one-step solution. This is outlined in section 2.4 by noting that we train our model to estimate $\mathbb{E}[{\bm{x}}_{0}|\tilde{\bm{y}}_{\mathrm{train}},\tilde{A}_{\mathrm{
train}}]$. Thus we can use our models in a one-step fashion. The performance of the one-step prediction should only be expected to be good in-distribution, i.e. when the model is evaluated at the same acceleration factor as the one used during training.

#### L1-Wavelet Compressed Sensing

We use the same L1-Wavelet reconstruction described previously as a standalone non-deep learning comparison.

#### Self-Supervised learning via Data Under-sampling (SSDU) [ 47 ]

In a supervised setting, end-to-end image restoration networks can be trained by passing measurements through a network like MoDL [] or VarNet [] and taking a loss directly on the output of the network with the ground truth image. In the self-supervised setting, however, this is not possible. SSDU, trains the end-to-end restoration network by splitting the available measurement set ($\Omega$) of each sample $y^{\Omega}_{i}$ into training sets ($\Theta$) and loss sets ($\Lambda$) or measurements $y^{\Theta}_{i}$ and $y^{\Lambda}_{i}$ respectively where $\Omega=\Theta\cup\Lambda$. Where the reconstruction network is given access to measurements in the training set to obtain an estimated image $\hat{{\bm{x}}}^{\Theta}_{i}={\bm{h}}_{\theta}(y^{\Theta}_{i})$ and the loss is then defined over the loss measurements as

$$ $L({\bm{y}}^{\Lambda}_{i},\hat{{\bm{x}}}^{\Theta}_{i})=\frac{\left\|{\bm{y}}^{ \Lambda}_{i}-A_{\Lambda}\hat{{\bm{x}}}^{\Theta}_{i}\right\|_{1}}{\left\|{\bm{y }}^{\Lambda}_{i}\right\|_{1}}+\frac{\left\|{\bm{y}}^{\Lambda}_{i}-A_{\Lambda} \hat{{\bm{x}}}^{\Theta}_{i}\right\|_{2}}{\left\|{\bm{y}}^{\Lambda}_{i}\right\| _{2}}$ (3.1) $$

where $A_{\Lambda}$ is the forward operator with ones in the selection mask $P$ at only the measurement locations in the loss set and zeros elsewhere. In this work, we selected ${\bm{h}}_{\theta}(\cdot)$ to have the MoDL architecture []. We trained individual models at four different acceleration levels ($R=2,4,6,8$). Each model was trained for ten epochs using the same training set ($10,000$ slices) as previously described. We used uniform random sampling to separate the training and loss measurement groups for each sample and found that a split of $\rho=\frac{|\Lambda|}{|\Omega|}=0.2$ provided the best performance and thus was used for reporting all subsequent metrics.

Figure: Figure 7: NRMSE performance metrics at $R=2,4,6,8$ for all methods.
Refer to caption: extracted/5468845/assets/figs/mri/NRMSE.png

Figure: Figure 8: PSNR performance metrics at $R=2,4,6,8$ for all methods.
Refer to caption: extracted/5468845/assets/figs/mri/PSNR.png

Figure: Figure 9: SSIM performance metrics at $R=2,4,6,8$ for all methods.
Refer to caption: extracted/5468845/assets/figs/mri/SSIM.png

#### MoDL Supervised (MoDL-Sup)

To provide an upper bound on performance for supervised end-to-end methods we also trained the MoDL architecture in a supervised fashion []. Again, we trained individual models at four different acceleration levels ($R=2,4,6,8$). Each model was trained using the same $10,000$ slices as above. However, for these models, we used the normalized root mean squared error (NRMSE) as the loss function:

$$ $L({\bm{x}}_{i},\hat{{\bm{x}}}_{i})=\frac{\left\|{\bm{x}}_{i}-\hat{{\bm{x}}}_{i }\right\|_{2}}{\left\|{\bm{x}}_{i}\right\|_{2}}$ (3.2) $$

where ${\bm{x}}_{i}$ are the ground truth images from our dataset and $\hat{{\bm{x}}}_{i}={\bm{h}}_{\theta}({\bm{y}}_{i},A_{i})$ are the reconstructions provided by our network based on under-sampled measurements ${\bm{y}}_{i}$.

#### Diffusion Posterior Sampling (DPS)

For models trained on clean data, we can use the Diffusion Posterior Sampling (DPS) [] algorithm. We implemented DPS with a variety of trained diffusion models. Specifically, we implemented the DPS algorithm with ($\mathbf{1}$) a Diffusion model trained with fully sampled ground truth images (FS-DPS), ($\mathbf{2-5}$) four different diffusion models trained on the L1-Wavelet reconstructions of data from measurement sets at $R=2,4,6,8$ (L1-DPS). Inference was run using the update rule of Eq. 2.3, for a total of $500$ steps and with $\gamma_{t}=\frac{1}{\left\|{\bm{y}}-A\mathbb{E}[{\bm{x}}_{0}|{\bm{x}}_{t}]
\right\|_{2}}$.

#### Annealed Langevin Dynamics (ALD)

Another popular method for solving inverse problems via score-based priors is with ALD []. Here we trained a score-based model using the same fully sampled $10,000$ samples as above and ran inference for $1,300$ steps. This method acts as another fully sampled comparison to our technique and thus serves a similar purpose as FS-DPS.

#### Results

Quantitative metrics for the listed methods are included in Figs. 7, 8, 9 under the NRMSE, PSNR and SSIM metrics respectively. We observe that at low acceleration levels ($R=2,4$) FS-DPS outperforms most methods including all A-DPS models. This is not unexpected, since FS-DPS uses models trained on clean data. However, as the validation acceleration increases to higher ratios $R$, A-DPS begins to outperform all models, including the ones trained on fully sampled data such as FS-DPS and ALD. This trend is consistent with the results that we observed in Section 3.1 for Ambient Models trained on natural images.

To get a better sense of the performance of the different methods, we visualize reconstructions for a subset of the methods listed at various accelerations in Fig. 6. Here, we see that at lower accelerations ($R=4$), just as the metrics suggest, FS-DPS outperforms A-DPS. However, at higher acceleration ($R=8$) we can see that FS-DPS introduces significant artifacts into the reconstruction while our A-DPS reconstructions maintain a high degree of visual fidelity.

## 4 Conclusions

We presented a simple framework based on DPS for solving inverse problems with Ambient Diffusion models. We showed that diffusion models trained on missing data are state-of-the-art inverse problem solvers for high corruption levels even if the forward process at inference time is different from the one used during training. Our framework fully unlocks the potential of Ambient Diffusion models that are critical in applications where access to full data is impossible or undesirable.

## 5 Acknowledgements

This research has been supported by NSF Grants CCF 1763702, NSF CCF-2239687 (CAREER),
AF 1901292, CNS 2148141, Tripods CCF 1934932, NSF AI Institute for Foundations of Machine Learning (IFML) 2019844, Oracle for Research Fellowship, Google Research Scholars Award, the Texas Advanced Computing Center (TACC) and research gifts by Western Digital, WNCG IAP, UT Austin Machine Learning Lab (MLL), Cisco and the Archie Straiton Endowed Faculty Fellowship. Giannis Daras has been supported by the Onassis Fellowship (Scholarship ID: F ZS 012-1/2022-2023), the Bodossaki Fellowship and the Leventis Fellowship.

## Appendix A Theorem Proof

In this section, we provide the proof of Theorem 2.1 that was stated in the main paper. We begin by giving the formal statement of the theorem.

###### Theorem A.1 (Formal Statement of Theorem 2.1 ) .

Let ${\bm{x}}\in\mathbb{C}^{n}$ be an unknown signal from a distribution that admits density $p({\bm{x}})$. Let $\{S_{i}\}$ be a random set of diagonal matrices $\in\mathbb{C}^{n}$ that satisfies: $\sum_{i}S_{i}^{H}S_{i}=I_{n\times n}$. Assume sample access to linearly corrupted measurements of ${\bm{x}}$ given by:

$$ $\displaystyle{\bm{y}}_{\mathrm{train}}=\underbrace{\left(\sum_{i}S_{i}^{H} \mathcal{F}^{-1}P\mathcal{F}S_{i}\right)}_{A_{\mathrm{train}}}{\bm{x}},$ (A.1) $$

where $\mathcal{F}$ is the discrete Fourier matrix and $P$ is a random inpainting matrix, i.e. a diagonal matrix with either zeros or ones in the diagonal, such that $\Pr[P_{ii}=1]>0$, for all entries $i$ of the diagonal. Define:

$$ $\displaystyle\tilde{\bm{y}}_{\mathrm{train}}=\underbrace{\left(\sum_{i}S_{i}^{ H}\mathcal{F}^{-1}\tilde{P}\mathcal{F}S_{i}\right)}_{\tilde{A}_{\mathrm{train} }}{\bm{x}},$ (A.2) $$

where $\tilde{P}$ is a further corrupted version of $P$ in the sense that with some non-zero probability $p$, a diagonal element that was $1$ in $P$ becomes 0 0 in $\tilde{P}$.
Define also ${\bm{x}}_{t},{\bm{y}}_{t,\mathrm{train}},\tilde{\bm{y}}_{t,\mathrm{train}}$ the noisy versions of ${\bm{x}},{\bm{y}}_{\mathrm{train}},\tilde{\bm{y}}_{\mathrm{train}}$ respectively, as in:

$$ $\displaystyle{\bm{x}}_{t}={\bm{x}}+\sigma_{t}\bm{\eta},\quad{\bm{y}}_{t, \mathrm{train}}=A_{\mathrm{train}}\left({\bm{y}}_{t,\mathrm{train}}+\sigma_{t} \bm{\eta}\right),\quad\tilde{\bm{y}}_{t,\mathrm{train}}=\tilde{A}_{\mathrm{ train}}\left(\tilde{\bm{y}}_{t,\mathrm{train}}+\sigma_{t}\bm{\eta}\right).$ (A.3) $$

Then, the minimizer of the objective:

$$ $\displaystyle J(\theta)=\mathbb{E}_{{\bm{y}}_{0,\mathrm{train}},\tilde{\bm{y}} _{t,\mathrm{train}},A_{\mathrm{train}},\tilde{P}}\left[\left|\left|A_{\mathrm{ train}}{\bm{h}}_{\theta}(\tilde{\bm{y}}_{t,\mathrm{train}},\tilde{P})-{\bm{y}} _{0,\mathrm{train}}\right|\right|^{2}\right],$ (A.4) $$

is: ${\bm{h}}_{\theta}(\tilde{\bm{y}}_{t,\mathrm{train}},\tilde{P})=\mathbb{E}[{\bm
{x}}_{0}|\tilde{\bm{y}}_{t,\mathrm{train}},\tilde{P}]$.

###### Proof.

We adapt the proof Theorem 4.1 in Ambient Diffusion []. To avoid notation clutter, we will be omitting the subscript “train”, when necessary.

Let ${\bm{h}}_{\theta^{*}}(\tilde{\bm{y}}_{t},\tilde{P})=\mathbb{E}[{\bm{x}}_{0}|
\tilde{\bm{y}}_{t},\tilde{P}]+{\bm{f}}(\tilde{\bm{y}}_{t},\tilde{P})$ be the optimal solution. The value of the objective for the optimal solution becomes:

$$ $\displaystyle J(\theta^{*})=\mathbb{E}_{{\bm{y}}_{0},\tilde{\bm{y}}_{t},A, \tilde{P}}\left[\left|\left|A\mathbb{E}[{\bm{x}}_{0}|\tilde{\bm{y}}_{t},\tilde {P}]+A{\bm{f}}(\tilde{\bm{y}}_{t},\tilde{P})-{\bm{y}}_{0}\right|\right|^{2}\right]$ (A.5) $\displaystyle=\mathbb{E}_{{\bm{y}}_{0},\tilde{\bm{y}}_{t},A,\tilde{P}}\left[ \underbrace{\left|\left|A\mathbb{E}[{\bm{x}}_{0}|\tilde{\bm{y}}_{t},\tilde{P}] -{\bm{y}}_{0}\right|\right|^{2}}_{\mathrm{irreducible\ error}}+{\bm{f}}(\tilde {\bm{y}}_{t},\tilde{P})^{T}A^{T}A{\bm{f}}(\tilde{\bm{y}}_{t},\tilde{P})-2\left (\mathbb{E}[{\bm{x}}_{0}|\tilde{\bm{y}}_{t},\tilde{P}]-{\bm{x}}_{0}\right)^{T} A^{T}A{\bm{f}}(\tilde{\bm{y}}_{t},\tilde{P})\right].$ (A.6) $$

We will now work with the last term.

$$ $\displaystyle\mathbb{E}_{{\bm{x}}_{0},\tilde{\bm{y}}_{t},A,\tilde{P}}\left[ \left(\mathbb{E}[{\bm{x}}_{0}|\tilde{\bm{y}}_{t},\tilde{P}]-{\bm{x}}_{0}\right )^{T}A^{T}A{\bm{f}}(\tilde{\bm{y}}_{t},\tilde{P})\right]$ (A.7) $\displaystyle=\mathbb{E}_{\tilde{\bm{y}}_{t},A,\tilde{P}}\left[\mathbb{E}_{{ \bm{x}}_{0}|\tilde{\bm{y}}_{t},A,\tilde{P}}\left[\left(\mathbb{E}[{\bm{x}}_{0} |\tilde{\bm{y}}_{t},\tilde{P}]-{\bm{x}}_{0}\right)^{T}A^{T}A{\bm{f}}(\tilde{ \bm{y}}_{t},\tilde{P})\right]\right]$ (A.8) $\displaystyle=\mathbb{E}_{\tilde{\bm{y}}_{t},A,\tilde{P}}\left[\left(\mathbb{E }[{\bm{x}}_{0}|\tilde{\bm{y}}_{t},\tilde{P}]-\mathbb{E}[{\bm{x}}_{0}|\tilde{ \bm{y}}_{t},A,\tilde{P}]\right)^{T}A^{T}A{\bm{f}}(\tilde{\bm{y}}_{t},\tilde{P} )\right]$ (A.9) $\displaystyle=0,$ (A.10) $$

since $\mathbb{E}[{\bm{x}}_{0}|\tilde{\bm{y}}_{t},\tilde{P}]=\mathbb{E}[{\bm{x}}_{0}|
\tilde{\bm{y}}_{t},A,\tilde{P}]$, i.e. the value of ${\bm{x}}_{0}$ does not depend on $\tilde{A}$ given $\tilde{P}$ and $\tilde{\bm{y}}_{t}$.

We will now work with the middle term. We have that:

$$ $\displaystyle\mathbb{E}_{\tilde{\bm{y}}_{t},A,\tilde{P}}\left[{\bm{f}}(\tilde{ \bm{y}}_{t},\tilde{P})^{T}A^{T}A{\bm{f}}(\tilde{\bm{y}}_{t},\tilde{P})\right]$ (A.11) $\displaystyle=\mathbb{E}_{\tilde{\bm{y}}_{t},\tilde{P}}\left[\mathbb{E}_{A| \tilde{\bm{y}}_{t},\tilde{P}}\left[{\bm{f}}(\tilde{\bm{y}}_{t},\tilde{P})^{T}A ^{T}A{\bm{f}}(\tilde{\bm{y}}_{t},\tilde{P})\right]\right]$ (A.12) $\displaystyle=\mathbb{E}_{\tilde{\bm{y}}_{t},\tilde{P}}\left[{\bm{f}}(\tilde{ \bm{y}}_{t},\tilde{P})^{T}\mathbb{E}[A^{T}A|\tilde{P}]{\bm{f}}(\tilde{\bm{y}}_ {t},\tilde{P})\right].$ (A.13) $$

From the last equation, it is evident that if $\mathbb{E}[A^{T}A|\tilde{P}]$ is full-rank, and hence the minimizer is $\mathbb{E}[{\bm{x}}_{0}|\tilde{\bm{y}}_{t},\tilde{P}]$ as needed.

Since our $A$ matrix is square, it suffices to show that $\mathbb{E}[A|\tilde{P}]$ is full-rank. By Corollary A.1. in Ambient Diffusion, we have that: $\mathbb{E}[P|\tilde{P}]$ is full-rank.
We will now show that $\mathcal{F}^{-1}\mathbb{E}[P|\tilde{P}]\mathcal{F}$
is also full-rank. This can be easily proved with contradiction. Assume that $\mathcal{F}^{-1}\mathbb{E}[P|\tilde{P}]\mathcal{F}$ is not full-rank. Then, there exists a vector ${\bm{w}}\neq 0$ such that:

$$ $\displaystyle\mathcal{F}^{-1}\mathbb{E}[P|\tilde{P}]\mathcal{F}{\bm{w}}=0\iff$ (A.14) $\displaystyle\mathbb{E}[P|\tilde{P}]\underbrace{\mathcal{F}{\bm{w}}}_{{\bm{z}} \neq\bm{0}}=0\iff$ (A.15) $\displaystyle\mathbb{E}[P|\tilde{P}]{\bm{z}}=0,$ (A.16) $$

which is a contradiction, since ${\bm{z}}\neq\bm{0}$ and $\mathbb{E}[P|\tilde{P}]$ is full-rank.

So far, we have established that $\mathcal{F}^{-1}\mathbb{E}[P|\tilde{P}]\mathcal{F}$ is full-rank.
By linearity of expectation, we further have that:

$$ $\displaystyle\mathcal{F}^{-1}\mathbb{E}[P|\tilde{P}]\mathcal{F}=\mathbb{E}[ \mathcal{F}^{-1}P\mathcal{F}|\tilde{P}],$ (A.17) $$

and hence, the latter is also full-rank. Finally, since $\sum_{i}S_{i}^{H}S_{i}=I$, for any full-rank matrix $C$, we have that $\sum_{i}S_{i}^{H}CS_{i}$ is also full-rank. This is true for any given set $S_{i}$ and hence it is also true for the expectation over the sets $\{S_{i}\}$. Putting everything together, we have that:

$$ $\displaystyle\mathbb{E}_{\{S_{i}\},P}\left[\sum_{i}S_{i}^{H}\mathcal{F}^{-1} \mathbb{E}[P|\tilde{P}]\mathcal{F}S_{i}\bigg{|}\tilde{P}\right]=\mathbb{E}[A| \tilde{P}],$ (A.18) $$

is full-rank and the proof is complete.

∎

###### Theorem A.1 (Formal Statement of Theorem 2.1 ) .

###### Proof.

## Appendix B Additional Performance Results

Here we include additional performance results.

Figure: (a) LPIPS per Number of Measurements.
Refer to caption: x7.png

Figure: (a) LPIPS per Downscaling Factor.
Refer to caption: x9.png

Figure: (a) LPIPS per Number of Measurements.
Refer to caption: x11.png

Figure: (a) LPIPS per Downscaling Factor.
Refer to caption: x13.png

## Appendix C Additional Speed Results

We further include additional speed results.

Figure: (a) Compressed Sensing with $4000$ measurements per Number of Function Evaluations (NFEs).
Refer to caption: x15.png

Figure: (a) Compressed Sensing with $4000$ measurements per Number of Function Evaluations (NFEs).
Refer to caption: x17.png

Figure: (a) Compressed Sensing with $4000$ measurements per Number of Function Evaluations (NFEs).
Refer to caption: x19.png

Figure: (a) Compressed Sensing with $4000$ measurements per Number of Function Evaluations (NFEs).
Refer to caption: x21.png

Figure: (a) Compressed Sensing with $4000$ measurements per Number of Function Evaluations (NFEs).
Refer to caption: x23.png

## Appendix D Additional MRI Results

Figure: Figure 19: Prior samples from Ambient Diffusion trained with under sampled data at $R=2,4,6,8$ (columns $1-4$), EDM trained with L1-wavelet reconstructions of subsampled data at $R=2,4,6,8$ (columns $5-8$), NCSNV2 trained with fully sampled data (column $9$) and EDM trained with fully sampled data (column $10$).
Refer to caption: extracted/5468845/assets/figs/mri/prior-2.png

Figure: Figure 20: Ambient diffusion one step performance metrics at $R=2,4,6,8$ for models trained at $R=2,4,6,8$.
Refer to caption: extracted/5468845/assets/figs/mri/A-OS.png

Figure: Figure 21: Ambient diffusion posterior sampling multi-step performance metrics at $R=2,4,6,8$ for models trained at $R=2,4,6,8$.
Refer to caption: extracted/5468845/assets/figs/mri/A-DPS.png

Figure: Figure 22: Performance metrics at $R=2,4,6,8$ for L1-wavelet compressed sensing, diffusion/score-based models trained with fully sampled data (FS-DPS, ALPS), diffusion models trained with L1 reconstructions of subsampled data (L1-DPS)
Refer to caption: extracted/5468845/assets/figs/mri/DPS-L1-ALPS-baselines.png

Figure: Figure 23: Fully supervised MoDL performance metrics at $R=2,4,6,8$ for models trained at $R=2,4,6,8$.
Refer to caption: extracted/5468845/assets/figs/mri/MoDL_sup.png

Figure: Figure 24: SSDU MoDL performance metrics at $R=2,4,6,8$ for models trained at $R=2,4,6,8$.
Refer to caption: extracted/5468845/assets/figs/mri/SSDU.png

## Appendix E Training Hyperparameters for MRI models

The EDM model trained on the fully-sampled ($R=1$) dataset, as well as the four EDM models trained on L1-Wavelet reconstructions at each acceleration factor $R\in\{2,4,6,8\}$ were all trained with 65 million parameters. The Ambient Diffusion models on the other hand were trained with 36 million parameters, for faster training. All the Ambient Diffusion models were trained for 250 epochs. For the Ambient Diffusion training experiments, the further corrupted sampling masks are given as an input to the model by concatenating with the measurements along the channel dimension, as in the original Ambient Diffusion [] paper.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 4
questions: Please take a look at the comments under weaknesses.
soundness: 2
strengths: 1. The paper is motivated by the challenge of solving inverse problems without access to clean images, a common problem in real-world imaging applications.
2. Experimental results show the proposed framework is effective for severely ill-posed problems, sometimes outperforming priors trained on clean images.
summary: This paper uses ambient diffusion models as a prior for solving ill-posed inverse problems. The ambient diffusion is trained on corrupted data acquired from a linear ill-posed operator. The authors also extend the ambient diffusions to use MRI k-space subsampled measurements as training data. The authors show the proposed framework can outperform DPS trained on clean images when dealing with severely ill-posed inverse problems.
weaknesses: 1. Despite promising experiments, the novelty of the paper seems limited; the authors primarily extend ambient diffusion to MRI and adapt it within the DPS framework for solving inverse problems.

2. Below Eq. 2.11, the authors state, "We remark that similar to DPS, the proposed algorithm is an approximation to sampling from the true posterior distribution". However, the accuracy of this approximation should be rigorously compared to that of the original DPS, as it relies on estimating the score function from corrupted data.

3. The authors should consider including additional, recent self-supervised learning approaches as baselines, such as the following paper:
Gan, Weijie, et al. "Self-supervised deep equilibrium models with theoretical guarantees and applications to MRI reconstruction." IEEE Transactions on Computational Imaging (2023).

4. The paper needs to be further polished, for example, Eq 2.9 seems incorrect.

---

confidence: 5
contribution: 2
initial_rating: 5
presentation: 3
questions: 1) The further corruption ratio is simply R+1. How does this further corruption ratio impact the performance, says R+2, R+3, etc?
2) Does the supervised DPS share the same network architecture with the proposed method?
3) How does the coil sensitivity map in this training dataset and the L1-wavelet baseline be estimated?
4) What is the time complexity of the methods?
5) What is the actual under sampled ratio considering the sampling mask for different R has the same 20 lines of the auto calibration signal?
soundness: 3
strengths: 1) The paper represents an early investigation into the application of self-supervised learning for diffusion models in the context of inverse problems.
2) The paper is well-organized, and the presentation is easy to follow.
summary: This paper applies the recent advancement in self-supervised learning for diffusion models, named ambient diffusion, to inverse problems. The effectiveness of this approach is shown through its application to MRI reconstruction and image inpainting.
weaknesses: 1) Although this paper introduces the first self-supervised diffusion method for inverse problems, the overall novelty is somewhat limited. It appears to be a straightforward integration of existing ambient diffusion with well-known diffusion posterior sampling, with application to MRI reconstruction. While the investigation holds value from a research perspective, the new insight/concept might be limited.

2) In the training phase (as described in Equation 2.10), the network takes \tilde{A} as input, considering one type of imaging artifact. However, during testing (Equation 2.11), the network uses A as input, which features a different imaging artifact. This inconsistency raises concerns about potential suboptimal performance. Could the authors clarify this aspect?

3) There is insufficient discussion regarding why the proposed A-DPS method outperforms FS-DPS, especially at high undersampling rates. Given that FS-DPS uses a diffusion model trained on fully sampled data, the results seem unexpectedly favorable.

4) The manuscript lacks discussions with recent theoretical developments in self-supervised learning for MRI [1,2], which share similar proof strategy with Theorem 1 in this paper.

5) The numerical validation could be improved, with certain experimental setups are not fully explored (see specific questions below). Additionally, the baseline methods for MRI reconstruction cited are well-known but somewhat dated (2019, 2020). Given the extensive literature on MRI reconstruction, these baselines may not represent the current state of the field. Comparing recent alternatives, such as [1-3] and self-supervised or supervised methods in survey [3] and [4], respectively, could improve this study. There is also no comparison between A-OS and other end-to-end methods. Despite that, I fully understand the time constraints of the rebuttal period and do not intend to suggest that the authors include comparisons with new baselines.

I am willing to increase my score if the authors could address my concerns.

[1] Millard, Charles, and Mark Chiew. "A theoretical framework for self-supervised MR image reconstruction using sub-sampling via variable density Noisier2Noise." IEEE transactions on computational imaging (2023).

[2] Gan, Weijie, et al. "Self-supervised deep equilibrium models with theoretical guarantees and applications to MRI reconstruction." IEEE Transactions on Computational Imaging (2023).

[3] Sriram, Anuroop, et al. "End-to-end variational networks for accelerated MRI reconstruction." Medical Image Computing and Computer Assisted Intervention–MICCAI 2020: 23rd International Conference, Lima, Peru, October 4–8, 2020, Proceedings, Part II 23. Springer International Publishing, 2020.

[4] Akçakaya, Mehmet, et al. "Unsupervised deep learning methods for biological image reconstruction and enhancement: An overview from a signal processing perspective." IEEE Signal Processing Magazine 39.2 (2022): 28-44.

[5] Heckel, Reinhard, et al. "Deep learning for accelerated and robust MRI reconstruction: a review." arXiv preprint arXiv:2404.15692 (2024).

---

confidence: 5
contribution: 3
initial_rating: 8
presentation: 4
questions: The questions have been covered in the weakness part.
soundness: 3
strengths: (1) This paper presents a novel framework for learning a ‘diffusion model’ solely from corrupted data, which is especially valuable in the field of medical imaging reconstruction, where only undersampled data is often available. 
(2) A solid theoretical proof for learning MMSE restoration estimator from corrupted data under suitable assumption.
(3) Based on this diffusion model, the authors demonstrate its effectiveness in both generative tasks and solving inverse problems, highlighting the efficiency and potential impact of this approach.
summary: This paper proposes a novel method for learning a diffusion model from corrupted data, supported by a theoretical analysis of training the MMSE restoration estimator under such conditions. Empirical results indicate that the learned MMSE restoration estimator can operate similarly to an MMSE denoiser, enabling effective score-based generation and restoration tasks. Overall, this is a highly interesting paper with the potential to open many new research directions. However, certain aspects of the analysis and discussion could benefit from further improvement.
weaknesses: (1) This work introduces an intriguing approach to learning a diffusion model from corrupted data. However, the authors’ analysis assumes that the learned MMSE restoration estimator can approximate a learned MMSE denoising estimator to derive a score function for sampling—an assumption that could benefit from further clarification. In approaches like DPS, posterior sampling depends on the score derived from Tweedie’s formula, which is specifically suited for denoisers. In this study, however, the corresponding score function is not discussed in sufficient detail. A discussion of recent approaches such as DRP[1] and SHARP[2], which utilize the MMSE restoration estimator with a modified score function, could enhance the analysis and provide a more fitting framework for this model.

(2) The theoretical approach in this work is compelling and well-suited for training with randomly inpainted images. However, in the MRI setting used here, different training samples may share the auto-calibration signal (ACS) region, potentially leading to an over-representation of corresponding k-space regions in the training data. A more solid approach might involve introducing a diagonal weighting matrix in the loss function to account for these oversampled regions. A similar work discussing this point for MRI reconstruction is [3].

(3) Further explanation to how the FID value is calculated for MRI data will be helpful.

(4) Some minor suggestions regarding the paper's presentation: (1) In Table 3, it would be beneficial to standardize the significant digits for the FID values to improve readability. (2) For Figures 7 and 8, adjusting the color bar could enhance visual clarity and make the figures easier to interpret.

Reference: 
[1] Y. Hu, M. Delbracio, P. Milanfar, and U. S. Kamilov, “A Restoration Network as an Implicit Prior,” Proc. Int. Conf. Learn. Represent. (ICLR 2024) (Vienna, Austria, May 7-11).
[2] Y. Hu, A. Peng, W. Gan, P. Milanfar, M. Delbracio, and U. S. Kamilov, “Stochastic Deep Restoration Priors for Imaging Inverse Problems.”  [https://arxiv.org/abs/2410.02057]
[3] C. Millard and M. Chiew, "A Theoretical Framework for Self-Supervised MR Image Reconstruction Using Sub-Sampling via Variable Density Noisier2Noise," in IEEE Transactions on Computational Imaging, vol. 9, pp. 707-720,

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
questions: - Would it work better to train a single network on multiple acceleration rates rather than a separate network for each acceleration rate?
- Were the evaluations done on a single posterior sample for each method? Or do you generate many posterior samples and compute a posterior mean?
soundness: 2
strengths: - The paper addresses a significant problem in solving inverse problems when ground truth data is unavailable. By only requiring corrupted measurements for training, the method presents an avenue to solve a wider set of inverse problems.
- The framework is novel as it incorporates the benefits of unconditional Ambient Diffusion with the conditional flexibility of DPS. 
- The experiments are adequately extensive and effectively demonstrate the benefits (and shortcomings) of their approach compared to several existing methods.
summary: The paper proposes a posterior sampling method for inverse problems called Ambient Diffusion Posterior Sampling (A-DPS), which utilizes a diffusion model trained only on corrupted images. They modify the existing Ambient Diffusion framework to work for the accelerated MRI problem, allowing the training of an unconditional diffusion model with only subsampled MRI measurements. Then, they integrate the Ambient Diffusion models into the DPS framework to form a posterior sampling approach. The experiments on accelerated MRI and natural image datasets consistently show better recovery performance than the state-of-the-art DPS when the conditioning is highly corrupted but not when the corruption level is low.
weaknesses: While the general methodology is explained well, the description of particular details could be modified to improve clarity and soundness
- Eq: 2.9: Are these equations correct? It seems to me that it should be $y_{t,\text{train}} = A_{\text{train}}(y_{\text{train}} + \sigma_{t}\eta)$ if $A_{\text{train}}$ is an orthogonal projection matrix or $y_{t,\text{train}} = A_{\text{train}}(x_{\text{train}} + \sigma_{t}\eta)$ if not
- What is the input to the diffusion model $h_\theta$? In Eq 2.10, it takes the additional subsampling matrix $\tilde{P}$. In Eq 2.11, it takes the full forward operator $A_\text{train}$. 
- An intuitive explanation of $h_{\theta}(\tilde{y}\_{t, t_{\text{train}}}, \tilde{P})$ would be helpful for understanding the overall algorithm. Specifically, what does the training accomplish and why can this be directly plugged into Eq. 2.11?
- What is the reason or intuition for why an A-DPS trained on one corruption can generate posterior samples for a different corruption? A good hypothesis could better sell this property of the method.

Minor Points:
- L93-94: This sentence is confusing because it sounds like the Ambient Diffusion framework is "the first framework to solve inverse problems with diffusion models learned from linearly corrupted data". I would reword this to avoid confusion.
- Eq 2.4: It was not clear what $\tilde{A}_\text{train}$ represents. I would explain this matrix first before the equation, so the equation is more interpretable
- L255: Typo with "In that regards the multiplicity..."
- L364: Typo with "Similar to the results in The authors of ..."
- The construction of Fig 3 is confusing with the labeling of "Recon-FS". Since it's vertically displaced, it is easy to misread. I would suggest changing this to "Error" or rotate it vertically so "Recon-FS" is on a single line

---

confidence: 4
contribution: 4
initial_rating: 8
presentation: 3
questions: See weaknesses. Although, I do have one actual question:

Q1) Did you consider training your MRI diffusion model (ambient or other) with the multi-coil data? To be clear: I am not asking you to do this, as it is a fundamental departure from what you have here. That said, I wonder if you tried this. I know that in, e.g., fastMRI, the number of coils is inconsistent, but you can get around this issue with virtual coils [1]. I suspect that the overall performance could be further improved with a multi-coil diffusion process, rather than a coil-combined one.

[1] T. Zhang, J. M. Pauly, S. S. Vasanawala, and M. Lustig, “Coil compression for accelerated imaging with
Cartesian sampling,” 2013
soundness: 4
strengths: S1) The paper is very well written and provides an appropriate amount of detail. Similarly, all analysis is easy to follow.

S2) The authors include a nice theoretical result in Appendix A.

S3) The quantitative performance of A-DPS is impressive, and the selected baselines offer an easy way to contextualize it (particularly for MRI).

S4) In my opinion, and based on the results, this is the first diffusion-MRI method to truly push the state-of-the-art forward since [1].

[1] Ajil Jalal, Marius Arvinte, Giannis Daras, Eric Price, Alexandros G Dimakis, and Jon Tamir, "Robust
compressed sensing mri with deep generative priors," 2021
summary: In this work the authors provide a framework for solving inverse problems with diffusion models that are trained on corrupted data. The authors extend the existing ambient diffusion framework to enable Fourier domain-corrupted measurements and propose ambient diffusion posterior sampling, an algorithm that leverages generative models pre-trained on one type of corruption to solve other inverse problems. The authors show the performance of their method on accelerated MRI reconstruction and on various problems for natural images.
weaknesses: I find this paper to be very impactful (hence my score). That said, it is very important to me that the following two weaknesses are addressed:

W1) While the experimental performance of A-DPS is impressive, I think that the MRI results suffer from insufficient evaluation with all competitors. NRMSE being the only metric reported in Table 2 is a concern. It would be worthwhile to add some feature-based metrics like LPIPS/DISTS for all methods. You use VGG to compute FID, you can also use VGG for those metrics (since VGG agrees with radiologists' perceptions [1]). Alternatively, you could also use AlexNet as your feature extractor [1]. Either way, I think that it is important to see these results for your method and competitors. The lack of additional metrics holds this paper back. I understand the space concern given the size of Table 2, but even putting these results in an appendix would be useful to provide readers full context and make a convincing statement about A-DPS' superiority.

W2) I am confused by the bolding in Table 2. For methods trained with clean data, no values are bolded. For the R=2 and R=4 sections, A-DPS has equal performance to A-ALD in some metrics. The A-ALD metrics that are equally performant should be bolded for fairness. Please revise this table so that it is consistent.

[1] P. M. Adamson, A. D. Desai, J. Dominic, C. Bluethgen, J. P. Wood, A. B. Syed, R. D. Boutin, K. J. Stevens,
S. Vasanawala, J. M. Pauly, A. S. Chaudhari, and B. Gunel, “Using deep feature distances for evaluating
MR image reconstruction quality,” 2023

Meta-review: This paper combines the recently proposed ambient diffusion with diffusion posterior sampling (DPS) to make ambient diffusion posterior sampling (A-DPS). They adapt ambient diffusion to allow training on corrupted (say subsampled) Fourier measurements. This results in a self-supervised learning scheme that uses diffusion model. The authors apply it to MRI at different acceleration factors and find that often at high acceleration factors it is more favorable to train the models on subsampled than on clean data. The results are overall strong. Acknowledging experiments on compressed sensing, the writing and the experiments are really geared towards MRI which could have been featured in the title.

I think that this is a nice paper, but that the novelty is weak since the method combines ambient diffusion and DPS in a straightforward way and applies them to MRI. Extending ambient diffusion to Fourier is nice but technically straightforward. I agree with the authors that the finding about training on subsampled data is nice, but it alone does not immediately warrant acceptance. Given the clear writing, clear delineation of contributions, the fact that the ideas and the results are good, the relative consensus between the reviewers about the merits of this work, and the fact that it's the first work that does self-supervised diffusion for inverse problems, my recommendation is ultimately to accept.

Decision: Accept (Poster)