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

Asad Aali UT Austin<EMAIL_ADDRESS>Giannis Daras∗ UT Austin
<EMAIL_ADDRESS>Brett Levac UT Austin<EMAIL_ADDRESS>
Sidharth Kumar UT Austin<EMAIL_ADDRESS>Alexandros G. Dimakis UT Austin
<EMAIL_ADDRESS>Jonathan I. Tamir UT Austin<EMAIL_ADDRESS>Asad
Aali and Giannis Daras contributed equally to this work.

###### 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$). 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 [15, 20, 47] 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 [17, 8, 37]. Prior works have shown how to train Generative
Adversarial Networks (GANs) [7, 14], flow models [29] and restoration models
[32, 31, 47, 35] with corrupted training data. More recently, there has been a
shift towards training diffusion generative models given corrupted data [17,
1, 28, 16, 30]. 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 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}}$.

We propose the first framework to solve inverse problems with diffusion models
learned from linearly corrupted data, for example as in Ambient Diffusion
[17]. 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 [12] to Ambient Models and
fully unlocks the potential of generative models trained on corrupted data for
image restoration.

Our contributions:

  1. 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. 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. 3.

We extend the Ambient Diffusion framework to train models directly on multi-
coil Fourier subsampled MRI.

  4. 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. 5.

We open-source our code and models to accelerate research in this area:
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 [39, 3]:

$\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 [12, 27, 25, 38, 13, 19, 21]. One of the simplest and
most effective approximations is Diffusion Posterior Sampling (DPS) [12]. 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.

[17] 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$. [17] 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 [43, 18, 42]. In settings
such as 3D+time dynamic contrast-enhanced MRI [49] it is impossible to collect
fully sampled data due to the time-varying dynamics of the contrast agent
[50].

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 [1, 28] 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 [2, 22]. 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 [47].
Other approaches leverage structure in the MRI acquisition to learn from
limited-resolution data [46]. 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 [40, 36, 9, 41].

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 [25, 48].

###  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
[17]: 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 [17]).

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}]$.

#### 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 [17]. 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

(a) LPIPS per Number of Measurements.

(b) MSE per Number of Measurements.

Figure 2: Compressed Sensing results, AFHQ: performance metric and standard
deviation. As shown, the model trained with clean data ($p=0.0$) only
outperforms the models trained with corrupted data for more than $1000$
measurements, in both LPIPS and MSE.

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 [17] 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 [4] 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 [10, 11]. 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.

(a) LPIPS per downsampling factor.

(b) MSE per downsampling factor.

Figure 3: Super-resolution results, AFHQ: Performance metric and standard
deviation. The model trained with clean data ($p=0.0$) performs worse, except
at downscaling factor 2.

(a) Compressed Sensing with $4000$ measurements per Number of Function
Evaluations (NFEs).

(b) $2\times$ Super-Resolution per Number of Function Evaluations (NFEs).

Figure 4: Speed performance plots for AFHQ.

#### 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 [44]. 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 [26]. 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 [25]. As an
additional baseline, we also train models with EDM after L1-Wavelet compressed
sensing reconstruction of the training set at each acceleration factor [33].
We used the BART implementation [45, 6] 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 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$). 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.

#### 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 [17], 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
[5] 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 [2] or VarNet [22] 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 [2]. 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 7: NRMSE performance metrics at $R=2,4,6,8$ for all methods. Figure 8:
PSNR performance metrics at $R=2,4,6,8$ for all methods. Figure 9: SSIM
performance metrics at $R=2,4,6,8$ for all methods.

#### 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 [2]. 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) [12] 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 [25]. 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.

## References

  * [1] Asad Aali, Marius Arvinte, Sidharth Kumar and Jonathan I Tamir  “Solving Inverse Problems with Score-Based Generative Priors learned from Noisy Data”  In _arXiv preprint arXiv:2305.01166_ , 2023
  * [2] Hemant K. Aggarwal, Merry P. Mani and Mathews Jacob  “MoDL: Model-Based Deep Learning Architecture for Inverse Problems”  In _IEEE Transactions on Medical Imaging_ 38.2, 2019, pp. 394–405  DOI: 10.1109/TMI.2018.2865356
  * [3] Brian D.O. Anderson  “Reverse-time diffusion equation models”  In _Stochastic Processes and their Applications_ 12.3 Elsevier, 1982, pp. 313–326
  * [4] Richard G Baraniuk  “Compressive sensing [lecture notes]”  In _IEEE signal processing magazine_ 24.4 IEEE, 2007, pp. 118–121
  * [5] Matthew Bendel, Rizwan Ahmad and Philip Schniter  “A Regularized Conditional GAN for Posterior Sampling in Image Recovery Problems”  In _arXiv e-prints_ , 2022, pp. arXiv:2210.13389  DOI: 10.48550/arXiv.2210.13389
  * [6] Moritz Blumenthal, Christian Holme, Volkert Roeloffs, Sebastian Rosenzweig, Philip Schaten, Nick Scholand, Jon Tamir, Xiaoqing Wang and Martin Uecker  “mrirecon/bart: version 0.8.00”  Zenodo, 2022  DOI: 10.5281/zenodo.7110562
  * [7] Ashish Bora, Eric Price and Alexandros G Dimakis  “AmbientGAN: Generative models from lossy measurements”  In _International conference on learning representations_ , 2018
  * [8] Nicolas Carlini, Jamie Hayes, Milad Nasr, Matthew Jagielski, Vikash Sehwag, Florian Tramer, Borja Balle, Daphne Ippolito and Eric Wallace  “Extracting training data from diffusion models”  In _32nd USENIX Security Symposium (USENIX Security 23)_ , 2023, pp. 5253–5270
  * [9] Dongdong Chen, Julián Tachella and Mike E. Davies  “Robust Equivariant Imaging: A Fully Unsupervised Framework for Learning To Image From Noisy and Partial Measurements”  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , 2022, pp. 5647–5656
  * [10] Sitan Chen, Sinho Chewi, Jerry Li, Yuanzhi Li, Adil Salim and Anru R Zhang  “Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions”  In _arXiv preprint arXiv:2209.11215_ , 2022
  * [11] Sitan Chen, Giannis Daras and Alex Dimakis  “Restoration-degradation beyond linear diffusions: A non-asymptotic analysis for ddim-type samplers”  In _International Conference on Machine Learning_ , 2023, pp. 4462–4484  PMLR
  * [12] Hyungjin Chung, Jeongsol Kim, Michael Thompson Mccann, Marc Louis Klasky and Jong Chul Ye  “Diffusion Posterior Sampling for General Noisy Inverse Problems”  In _The Eleventh International Conference on Learning Representations_ , 2023  URL: https://openreview.net/forum?id=OnD9zGAGT0k
  * [13] Hyungjin Chung, Byeongsu Sim, Dohoon Ryu and Jong Chul Ye  “Improving diffusion models for inverse problems using manifold constraints”  In _Advances in Neural Information Processing Systems_ 35, 2022, pp. 25683–25696
  * [14] Elizabeth K. Cole, Frank Ong, Shreyas S. Vasanawala and John M. Pauly  “Fast Unsupervised MRI Reconstruction Without Fully-Sampled Ground Truth Data Using Generative Adversarial Networks”  In _2021 IEEE/CVF International Conference on Computer Vision Workshops (ICCVW)_ , 2021, pp. 3971–3980  DOI: 10.1109/ICCVW54120.2021.00444
  * [15] The Event Horizon Telescope Collaboration et al.  “First M87 Event Horizon Telescope Results. IV. Imaging the Central Supermassive Black Hole”  In _The Astrophysical Journal Letters_ 875.1 The American Astronomical Society, 2019, pp. L4  DOI: 10.3847/2041-8213/ab0e85
  * [16] Zhuo-Xu Cui, Chentao Cao, Shaonan Liu, Qingyong Zhu, Jing Cheng, Haifeng Wang, Yanjie Zhu and Dong Liang  “Self-score: Self-supervised learning on score-based models for mri reconstruction”  In _arXiv preprint arXiv:2209.00835_ , 2022
  * [17] Giannis Daras, Kulin Shah, Yuval Dagan, Aravind Gollakota, Alexandros G Dimakis and Adam Klivans  “Ambient Diffusion: Learning Clean Distributions from Corrupted Data”  In _arXiv preprint arXiv:2305.19256_ , 2023
  * [18] Arjun D Desai, Andrew M Schmidt, Elka B Rubin, Christopher M Sandino, Marianne S Black, Valentina Mazzoli, Kathryn J Stevens, Robert Boutin, Christopher Ré, Garry E Gold, Brian A Hargreaves and Akshay S Chaudhari  “SKM-TEA: A Dataset for Accelerated MRI Reconstruction with Dense Image Labels for Quantitative Clinical Evaluation”, 2022  arXiv:2203.06823 [eess.IV]
  * [19] Berthy T Feng, Jamie Smith, Michael Rubinstein, Huiwen Chang, Katherine L Bouman and William T Freeman  “Score-Based diffusion models as principled priors for inverse imaging”  In _arXiv preprint arXiv:2304.11751_ , 2023
  * [20] Angela F Gao, Oscar Leong, He Sun and Katherine L Bouman  “Image Reconstruction without Explicit Priors”  In _ICASSP 2023-2023 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_ , 2023, pp. 1–5  IEEE
  * [21] Alexandros Graikos, Nikolay Malkin, Nebojsa Jojic and Dimitris Samaras  “Diffusion models as plug-and-play priors”  In _Advances in Neural Information Processing Systems_ 35, 2022, pp. 14715–14728
  * [22] Kerstin Hammernik, Teresa Klatzer, Erich Kobler, Michael P. Recht, Daniel K. Sodickson, Thomas Pock and Florian Knoll  “Learning a variational network for reconstruction of accelerated MRI data”  In _Magnetic Resonance in Medicine_ 79.6, 2018, pp. 3055–3071  DOI: https://doi.org/10.1002/mrm.26977
  * [23] Reinhard Heckel and Paul Hand  “Deep Decoder: Concise Image Representations from Untrained Non-convolutional Networks”, 2019  arXiv:1810.03982 [cs.CV]
  * [24] Yuyang Hu, Mauricio Delbracio, Peyman Milanfar and Ulugbek S. Kamilov  “A Restoration Network as an Implicit Prior” arXiv:2310.01391, 2023
  * [25] Ajil Jalal, Marius Arvinte, Giannis Daras, Eric Price, Alexandros G Dimakis and Jon Tamir  “Robust compressed sensing mri with deep generative priors”  In _Advances in Neural Information Processing Systems_ 34, 2021, pp. 14938–14954
  * [26] Tero Karras, Miika Aittala, Timo Aila and Samuli Laine  “Elucidating the design space of diffusion-based generative models”  In _arXiv preprint arXiv:2206.00364_ , 2022
  * [27] Bahjat Kawar, Michael Elad, Stefano Ermon and Jiaming Song  “Denoising Diffusion Restoration Models”  In _Advances in Neural Information Processing Systems_
  * [28] Bahjat Kawar, Noam Elata, Tomer Michaeli and Michael Elad  “GSURE-Based Diffusion Model Training with Corrupted Data”  In _arXiv preprint arXiv:2305.13128_ , 2023
  * [29] Varun A Kelkar, Rucha Deshpande, Arindam Banerjee and Mark A Anastasio  “AmbientFlow: Invertible generative models from incomplete, noisy measurements”  In _arXiv preprint arXiv:2309.04856_ , 2023
  * [30] Kwanyoung Kim and Jong Chul Ye  “Noise2score: tweedie’s approach to self-supervised image denoising without clean images”  In _Advances in Neural Information Processing Systems_ 34, 2021, pp. 864–874
  * [31] Alexander Krull, Tim-Oliver Buchholz and Florian Jug  “Noise2void-learning denoising from single noisy images”  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , 2019, pp. 2129–2137
  * [32] Jaakko Lehtinen, Jacob Munkberg, Jon Hasselgren, Samuli Laine, Tero Karras, Miika Aittala and Timo Aila  “Noise2Noise: Learning image restoration without clean data”  In _arXiv preprint arXiv:1803.04189_ , 2018
  * [33] Michael Lustig, David Donoho and John M. Pauly  “Sparse MRI: The application of compressed sensing for rapid MR imaging”  In _Magnetic Resonance in Medicine_ 58.6, 2007, pp. 1182–1195  DOI: https://doi.org/10.1002/mrm.21391
  * [34] Christopher A Metzler, Arian Maleki and Richard G Baraniuk  “From denoising to compressed sensing”  In _IEEE Transactions on Information Theory_ 62.9 IEEE, 2016, pp. 5117–5144
  * [35] Charles Millard and Mark Chiew  “A Theoretical Framework for Self-Supervised MR Image Reconstruction Using Sub-Sampling via Variable Density Noisier2Noise”  In _IEEE Transactions on Computational Imaging_ 9, 2023, pp. 707–720  DOI: 10.1109/TCI.2023.3299212
  * [36] Jérémy Scanvic, Mike Davies, Patrice Abry and Julián Tachella  “Self-Supervised Learning for Image Super-Resolution and Deblurring”, 2023  arXiv:2312.11232 [eess.IV]
  * [37] Gowthami Somepalli, Vasu Singla, Micah Goldblum, Jonas Geiping and Tom Goldstein  “Diffusion Art or Digital Forgery? Investigating Data Replication in Diffusion Models”  In _arXiv preprint arXiv:2212.03860_ , 2022
  * [38] Yang Song, Liyue Shen, Lei Xing and Stefano Ermon  “Solving inverse problems in medical imaging with score-based generative models”  In _arXiv preprint arXiv:2111.08005_ , 2021
  * [39] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon and Ben Poole  “Score-based generative modeling through stochastic differential equations”  In _arXiv preprint arXiv:2011.13456_ , 2020
  * [40] Julián Tachella, Dongdong Chen and Mike Davies  “Sensing Theorems for Unsupervised Learning in Linear Inverse Problems”, 2022  arXiv:2203.12513 [stat.ML]
  * [41] Julián Tachella, Dongdong Chen and Mike Davies  “Unsupervised Learning From Incomplete Measurements for Inverse Problems”  In _arXiv preprint arXiv:2201.12151_ , 2022
  * [42] U Tariq, P Lai, M Lustig, M Alley, M Zhang, G Gold and Vasanawala Shreyas S  “MRI Data: Undersampled Knees”  In _Undersampled Knees | MRI Data_ URL: http://old.mridata.org/undersampled/knees
  * [43] Radhika Tibrewala, Tarun Dutt, Angela Tong, Luke Ginocchio, Mahesh B Keerthivasan, Steven H Baete, Sumit Chopra, Yvonne W Lui, Daniel K Sodickson, Hersh Chandarana and Patricia M Johnson  “FastMRI Prostate: A Publicly Available, Biparametric MRI Dataset to Advance Machine Learning for Prostate Cancer Imaging”, 2023  arXiv:2304.09254 [physics.med-ph]
  * [44] Martin Uecker, Peng Lai, Mark J. Murphy, Patrick Virtue, Michael Elad, John M. Pauly, Shreyas S. Vasanawala and Michael Lustig  “ESPIRiT—an eigenvalue approach to autocalibrating parallel MRI: Where SENSE meets GRAPPA”  In _Magnetic Resonance in Medicine_ 71.3, 2014, pp. 990–1001  DOI: https://doi.org/10.1002/mrm.24751
  * [45] Martin Uecker, Jonathan I Tamir, Frank Ong and Michael Lustig  “The BART toolbox for computational magnetic resonance imaging”  In _Proc Intl Soc Magn Reson Med_ 24, 2016
  * [46] Frederic Wang, Han Qi, Alfredo De Goyeneche, Reinhard Heckel, Michael Lustig and Efrat Shimron  “K-band: Self-supervised MRI Reconstruction via Stochastic Gradient Descent over K-space Subsets”  In _arXiv e-prints_ , 2023, pp. arXiv:2308.02958  DOI: 10.48550/arXiv.2308.02958
  * [47] Burhaneddin Yaman, Seyed Amir Hossein Hosseini, Steen Moeller, Jutta Ellermann, Kâmil Uğurbil and Mehmet Akçakaya  “Self-supervised learning of physics-guided reconstruction neural networks without fully sampled reference data”  In _Magnetic Resonance in Medicine_ 84.6, 2020, pp. 3172–3191  DOI: https://doi.org/10.1002/mrm.28378
  * [48] Martin Zach, Florian Knoll and Thomas Pock  “Stable Deep MRI Reconstruction using Generative Priors”  In _arXiv e-prints_ , 2022, pp. arXiv:2210.13834  DOI: 10.48550/arXiv.2210.13834
  * [49] Tao Zhang, Joseph Y. Cheng, Aaron G. Potnick, Richard A. Barth, Marcus T. Alley, Martin Uecker, Michael Lustig, John M. Pauly and Shreyas S. Vasanawala  “Fast pediatric 3D free-breathing abdominal dynamic contrast enhanced MRI with high spatiotemporal resolution”  In _Journal of Magnetic Resonance Imaging_ 41.2, 2015, pp. 460–473  DOI: https://doi.org/10.1002/jmri.24551
  * [50] Tao Zhang, John Pauly, Shreyas Vasanawala and Michael Lustig  “MRI Data: Undersampled Abdomens”  In _Undersampled Abdomens | MRI Data_ URL: http://old.mridata.org/undersampled/abdomens
  * [51] Hongkai Zheng, Weili Nie, Arash Vahdat and Anima Anandkumar  “Fast Training of Diffusion Models with Masked Transformers”  In _arXiv preprint arXiv:2306.09305_ , 2023

##  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$ 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 [17]. 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.

∎

##  Appendix B Additional Performance Results

Here we include additional performance results.

(a) LPIPS per Number of Measurements.

(b) MSE per Number of Measurements.

Figure 10: Compressed Sensing Results for FFHQ.

(a) LPIPS per Downscaling Factor.

(b) MSE per Downscaling Factor.

Figure 11: Downscaling Results for FFHQ.

(a) LPIPS per Number of Measurements.

(b) MSE per Number of Measurements.

Figure 12: Compressed Sensing Results for Celeb-A.

(a) LPIPS per Downscaling Factor.

(b) MSE per Downscaling Factor.

Figure 13: Downscaling Results for Celeb-A.

##  Appendix C Additional Speed Results

We further include additional speed results.

(a) Compressed Sensing with $4000$ measurements per Number of Function
Evaluations (NFEs).

(b) $2\times$ Super-Resolution per Number of Function Evaluations (NFEs).

Figure 14: Speed LPIPS performance plots for AFHQ.

(a) Compressed Sensing with $4000$ measurements per Number of Function
Evaluations (NFEs).

(b) $2\times$ Super-Resolution per Number of Function Evaluations (NFEs).

Figure 15: Speed MSE performance plots for FFHQ.

(a) Compressed Sensing with $4000$ measurements per Number of Function
Evaluations (NFEs).

(b) $2\times$ Super-Resolution per Number of Function Evaluations (NFEs).

Figure 16: Speed LPIPS performance plots for FFHQ.

(a) Compressed Sensing with $4000$ measurements per Number of Function
Evaluations (NFEs).

(b) $2\times$ Super-Resolution per Number of Function Evaluations (NFEs).

Figure 17: Speed MSE performance plots for Celeb-A.

(a) Compressed Sensing with $4000$ measurements per Number of Function
Evaluations (NFEs).

(b) $2\times$ Super-Resolution per Number of Function Evaluations (NFEs).

Figure 18: Speed LPIPS performance plots for Celeb-A.

##  Appendix D Additional MRI Results

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$). Figure 20: Ambient diffusion one step performance
metrics at $R=2,4,6,8$ for models trained at $R=2,4,6,8$. 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$. 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) Figure 23: Fully supervised MoDL
performance metrics at $R=2,4,6,8$ for models trained at $R=2,4,6,8$. Figure
24: SSDU MoDL performance metrics at $R=2,4,6,8$ for models trained at
$R=2,4,6,8$.

##  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 [17] paper.