# Optimizing Sampling Patterns for Compressed Sensing MRI with Diffusion
Generative Models

Sriram Ravula
University of Texas at Austin
Electrical and Computer Engineering
<EMAIL_ADDRESS>
&Brett Levac
University of Texas at Austin
Electrical and Computer Engineering
<EMAIL_ADDRESS>
&Ajil Jalal
University of California, Berkeley
Electrical Engineering and Computer Sciences
<EMAIL_ADDRESS>
&Jonathan I. Tamir
University of Texas at Austin
Electrical and Computer Engineering
<EMAIL_ADDRESS>
&Alexandros G. Dimakis
University of Texas at Austin
Electrical and Computer Engineering
<EMAIL_ADDRESS>

###### Abstract

Diffusion-based generative models have been used as powerful priors for
magnetic resonance imaging (MRI) reconstruction. We present a learning method
to optimize sub-sampling patterns for compressed sensing multi-coil MRI that
leverages pre-trained diffusion generative models. Crucially, during training
we use a single-step reconstruction based on the posterior mean estimate given
by the diffusion model and the MRI measurement process. Experiments across
varying anatomies, acceleration factors, and pattern types show that sampling
operators learned with our method lead to competitive, and in the case of 2D
patterns, improved reconstructions compared to baseline patterns. Our method
requires as few as five training images to learn effective sampling patterns.

††Code available at https://github.com/Sriram-Ravula/MRI_Sampling_Diffusion.

##  1 Introduction

Compressed sensing (CS) [11, 15] is a powerful framework for recovering a
structured signal from a limited number of measurements, with applications in
many imaging inverse problems. In particular, CS has been used to accelerate
magnetic resonance imaging (MRI) beyond the Nyquist rate by sampling a pseudo-
random subset of Fourier coefficients in k-space and imposing a sparse prior
on the image [33]. In the multi-coil setting, CS-MRI can be combined with
parallel imaging [40, 36, 17] to exploit additional structure. While the
incoherence condition in CS-MRI is typically realized through pseudo-random
sampling, optimal recovery for parallel imaging is based on linear
predictability from local neighborhoods in k-space [18]. Thus, there is a
tradeoff in balancing random and structured sampling.

More recently, deep learning has been used as a powerful tool to solve ill-
posed inverse problems beyond the capabilities of sparsity priors, through
either supervised or unsupervised deep neural networks [34]. While these
techniques have demonstrated unprecedented reconstruction quality for
accelerated multi-coil MRI, the sub-sampling pattern is still typically chosen
based on heuristics from CS or parallel imaging. Several approaches have been
introduced to optimize the sampling pattern, either separately or jointly with
the reconstruction network, through end-to-end learning [7, 2, 49, 39, 56, 3].
However, these approaches are only suitable when the gradient of the sampling
operator can be calculated through the full reconstruction network. Thus,
while these techniques have been very successful for optimizing the sampling
pattern for end-to-end reconstruction networks such as the U-Net [54], Var-Net
[19, 45], and other unrolled methods [1], it is unclear how to extend it to
unsupervised methods such as CS with generative models (CSGM) proposed in [9]
and applied to MRI in [24].

In this work we optimize the collected samples for CSGM MRI. We use the
recently proposed unsupervised method of posterior sampling with diffusion
models [24, 13, 32, 43], since it has been shown to be robust to changes in
imaging anatomy, acceleration factor, and sub-sampling patterns, without
requiring retraining. As these reconstructions are not amenable to algorithm
unrolling, we propose a simple and effective alternative to the full gradient
of the reconstruction network that can be used to optimize the selection of
which samples to keep. We apply our approach to multi-coil Cartesian MRI and
demonstrate that optimized sampling patterns can be used to reduce the
reconstruction error for a given acceleration factor.

###  1.1 Contributions

  * •

We propose a novel training objective for selecting which samples to obtain
for MRI reconstruction using diffusion-based generative models. Our method is
based on a theoretical result which provides a closed-form solution for
calculating the mean of the posterior distribution.

  * •

We use pre-trained diffusion models as priors without requiring further
tuning, making our method modular. Our method is also independent of the
choice of posterior sampling algorithm.

  * •

We empirically demonstrate that our approach can learn competitive sampling
patterns using as few as five training samples. Qualitative results show that
patterns learned with our method lead to reconstructions with fewer artifacts
than baseline equispaced and Poisson disc masks.

###  1.2 Related Work

#### Compressed Sensing with Generative Models

The theory of compressed sensing [11, 15] demonstrates that a signal can be
accurately recovered from measurements taken below the Nyquist rate, given
that the signal is sparse in an appropriately-chosen basis. The authors of
CSGM [9] extended this finding by showing that the range of a generative model
is a more powerful prior than sparsity. Generative models have been used as
priors for inverse problems in deblurring [53], phase retrieval [20], 1-bit
compressed sensing [31, 37], and many more. Various reconstruction algorithms
have been proposed to leverage generative models for these reconstruction
tasks [42, 25, 6, 28, 44], including methods which use untrained models [47,
21] or models trained entirely on compressive measurements [10]. See the
survey in [34] for taxonomy of image reconstruction algorithms and
applications.

#### MRI Sampling Pattern Selection

Previous works have looked at optimal sampling pattern selection for
accelerated MRI [7, 55, 49, 50, 2, 30]. In each approach, the reconstruction
network is trained in an end-to-end fashion jointly with the sampling mask.
The learned sampling masks in these approaches were either treated as a
probability density function to select points on the Cartesian grid, or as
non-Cartesian coordinates in the non-uniform Fast Fourier Transform (NUFFT)
which were subsequently treated as learnable parameters. Other works have
taken an adaptive approach to selecting sampling patterns for individual
patients [51].

(a)  (b)

Figure 1:  Reconstructions using baseline masks vs masks learned with our
method. We perform posterior sampling using a diffusion model to reconstruct
MRI scans using varying anatomies, pattern types, and acceleration factors
(R). Artifacts in the reconstructions are highlighted with red arrows and
insets (zoom may be needed). Masks learned with our method produce
reconstructions with fewer artifacts.

##  2 Background

###  2.1 Diffusion-Based Generative Models

Diffusion-based models learn to generate new signals by reversing a corruption
process. The two main classes of diffusion models, Score-Based Models (SBMs)
[41] and Denoising Diffusion Probabilistic Models (DDPMs) [22], specifically
are trained to remove Gaussian noise. The noise magnitude is determined by the
variance $\sigma^{2}_{t}$ and indexed by time step $t\in[0,T]$. Typically,
$\sigma^{2}_{0}=0$ and corresponds to clean signals, and $\sigma^{2}_{T}$ is
chosen to be very large so that the distribution of noisy signals at time
$t=T$ is indistinguishable from pure Gaussian noise.

Song et al. [44] unify the discrete-time noise processes of SBMs and DDPMs
under the framework of continuous-time Stochastic Differential Equations
(SDEs). They cast SBMs as reversing _Variance Exploding_ (VE) SDEs, with noisy
image distributions having increasing variance as $t$ increases, and DDPMs as
reversing _Variance Preserving_ (VP) SDEs, whose noisy image distributions
typically maintain unit variance throughout the diffusion process. In this
paper, we focus on the VE-SDE framework.

The noisy signal at time $t$ in the forward diffusion process is given by
${\mathbf{x}}_{t}\in{\mathbb{R}}^{n}$. The diffusion process is modeled as the
solution to an Itô SDE of the form

$d{\mathbf{x}}={\mathbf{f}}({\mathbf{x}},t)dt+g(t)d{\mathbf{w}}.$ (1)

Here, ${\mathbf{w}}$ is the standard $n$-dimensional Wiener process, and in
the VE case, ${\mathbf{f}}({\mathbf{x}},t)=\mathbf{0}$ and
$g(t)=\sqrt{\frac{d\sigma_{t}^{2}}{dt}}$. The variance $\sigma^{2}_{t}$ is a
monotonically increasing function that defines the distribution of the
diffused signal, with the property that at time $t=0$ we recover the data
distribution: ${\mathbf{x}}_{0}\sim p_{0}=p_{data}$.

The goal of Diffusion-based models is to start from samples
${\mathbf{x}}_{T}\sim p_{T}$ consisting of pure Gaussian noise and reverse the
forward diffusion given by Eq. (1) to arrive at samples ${\mathbf{x}}_{0}\sim
p_{0}$. Conveniently, the reverse of the forward SDE is also an SDE [4], with
the form:

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

where $dt$ is now a negative time step and $\bar{{\mathbf{w}}}$ is the
standard Wiener process when time flows backward. The reverse SDE depends on
the _score function_ $\nabla_{{\mathbf{x}}_{t}}\log p_{t}({\mathbf{x}}_{t})$
of the marginal distribution at time $t$.

A result from Vincent [48] states that we can learn the _unconditional_ score
function in Eq(1) through denoising score matching (DSM). DSM only requires
the score of the _conditional_ noise distribution at time $t$, which is known
in closed-form for additive Gaussian noise. The training objective is

$\mathbf{\theta}^{*}=\operatorname*{argmin\,}_{\mathbf{\theta}}{\mathbb{E}}_{t\sim
U[0,T]}\big{[}\lambda_{t}{\mathbb{E}}_{({\mathbf{x}}_{0},{\mathbf{x}}_{t})\sim
p_{0}({\mathbf{x}}_{0})p_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{0})}\lVert{\mathbf{s}}_{\mathbf{\theta}}({\mathbf{x}}_{t},t)-\nabla_{{\mathbf{x}}_{t}}\log
p_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{0})\rVert^{2}\big{]},$ (3)

where $\lambda_{t}$ is a time-dependent positive weighting function and
${\mathbf{s}}_{\mathbf{\theta}}({\mathbf{x}}_{t},t)$ is the _score network_.
Given enough data and model capacity, the score network learns to approximate
the unconditional score function:
${\mathbf{s}}_{\mathbf{\theta}}({\mathbf{x}}_{t},t)\simeq\nabla_{{\mathbf{x}}_{t}}\log
p_{t}({\mathbf{x}}_{t})$. Subsequently, sampling with the score network is a
matter of substituting ${\mathbf{s}}_{\mathbf{\theta}}({\mathbf{x}}_{t},t)$
for the score function in Eq(2), then solving the reverse SDE.

###  2.2 Inverse Problems

In the inverse problem setting, we are given measurements of the form
$\mathbf{y}={\cal A}({\mathbf{x}}_{0})+\epsilon$, where ${\cal
A}:{\mathbb{R}}^{n}\rightarrow{\mathbb{R}}^{m}$ is a known forward operator,
$\epsilon\in\mathbb{R}^{m}$ is some additive noise, and
${\mathbf{x}}_{0}\in\mathbb{R}^{n}$ is a signal we want to recover. The
measurements are usually compressive, i.e. $m<n$, which means that there are
infinitely many solutions that match the observed measurements. We turn to a
Bayesian approach to find a solution.

Bayes’ rule gives the relationship between the _prior_ distribution
$p({\mathbf{x}})$ and the _posterior_ distribution
$p({\mathbf{x}}|{\mathbf{y}})$ as
$p({\mathbf{x}}|{\mathbf{y}})=p({\mathbf{x}})p({\mathbf{y}}|{\mathbf{x}})/p({\mathbf{y}})$.
If we take the gradient w.r.t. ${\mathbf{x}}$ of the log of both sides, we
find that

$\nabla_{{\mathbf{x}}}\log
p({\mathbf{x}}|{\mathbf{y}})=\nabla_{{\mathbf{x}}}\log
p({\mathbf{x}})+\nabla_{{\mathbf{x}}}\log p({\mathbf{y}}|{\mathbf{x}}).$ (4)

Returning to the SDE framework, we can sample from the posterior distribution
by using Eq.(4) in place of the score in Eq.(2):

$d{\mathbf{x}}=[{\mathbf{f}}({\mathbf{x}},t)-g(t)^{2}(\nabla_{{\mathbf{x}}_{t}}\log
p_{t}({\mathbf{x}}_{t})+\nabla_{{\mathbf{x}}_{t}}\log
p_{t}({\mathbf{y}}|{\mathbf{x}}_{t}))]dt+g(t)d\bar{{\mathbf{w}}}.$ (5)

Therefore, we need to know the score of the prior,
$\nabla_{{\mathbf{x}}_{t}}\log p_{t}({\mathbf{x}}_{t})$, and the likelihood,
$\nabla_{{\mathbf{x}}_{t}}\log p({\mathbf{y}}|{\mathbf{x}}_{t}))$. To find the
score of the prior distribution, we can train a diffusion model as described
in Sec2.1. Unfortunately, the time-dependent likelihood
$p_{t}({\mathbf{y}}|{\mathbf{x}}_{t})$ is not easy to obtain. In graphical
terms, ${\mathbf{x}}_{0}\rightarrow{\mathbf{y}}$ and
${\mathbf{x}}_{0}\rightarrow{\mathbf{x}}_{t}$, but there is no closed-form
dependency between the measurements ${\mathbf{y}}$ and the noisy signal
${\mathbf{x}}_{t}$.

The authors of Diffusion Posterior Sampling (DPS) [12] propose to approximate
$\nabla_{{\mathbf{x}}_{t}}\log p_{t}({\mathbf{y}}|{\mathbf{x}}_{t})$ by
exploiting a result from Tweedie [16]. Tweedie’s formula gives a closed-form
expression for the mean of $p_{0}({\mathbf{x}}_{0}|{\mathbf{x}}_{t})$ when
$p_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{0})$ belongs to an exponential family
distribution. For the case of VE diffusions, we have that
$p_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{0})={\cal
N}({\mathbf{x}}_{t};{\mathbf{x}}_{0},\sigma_{t}^{2}{\mathbf{I}})$, with the
posterior mean given by Tweedie as

$\hat{{\mathbf{x}}}_{0}\coloneqq{\mathbb{E}}_{{\mathbf{x}}_{t}\sim
p_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{0})}[{\mathbf{x}}_{0}|{\mathbf{x}}_{t}]={\mathbf{x}}_{t}+\sigma_{t}^{2}\nabla_{{\mathbf{x}}_{t}}\log
p_{t}({\mathbf{x}}_{t}).$ (6)

This formula allows us to get a “one-step” denoised estimate of the clean
signal ${\mathbf{x}}_{0}$ using only the noisy signal ${\mathbf{x}}_{t}$ and
the score of the unconditional distribution, $\nabla_{{\mathbf{x}}_{t}}\log
p_{t}({\mathbf{x}}_{t})$. We can replace the unconditional score in Eq.(6)
using our trained score network
${\mathbf{s}}_{\mathbf{\theta}}({\mathbf{x}}_{t},t)$ to approximate the
posterior mean as

$\hat{{\mathbf{x}}}_{0}\coloneqq{\mathbb{E}}_{{\mathbf{x}}_{t}\sim
p_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{0})}[{\mathbf{x}}_{0}|{\mathbf{x}}_{t}]={\mathbf{x}}_{t}+\sigma_{t}^{2}{\mathbf{s}}_{\mathbf{\theta}}({\mathbf{x}}_{t},t).$
(7)

Assuming the forward operator ${\cal A}$ is differentiable, DPS approximates
the time-dependent score of the likelihood as

$\nabla_{{\mathbf{x}}_{t}}\log
p_{t}({\mathbf{y}}|{\mathbf{x}}_{t})\simeq\nabla_{{\mathbf{x}}_{t}}\log
p_{t}({\mathbf{y}}|\hat{{\mathbf{x}}}_{0}).$ (8)

For the case of Gaussian measurement noise, ${\mathbf{y}}\sim{\cal
N}({\mathbf{y}};{\cal
A}({\mathbf{x}}_{0}),\sigma^{2}_{{\mathbf{y}}}{\mathbf{I}})$, the
approximation in Eq.(8) gives us

$\nabla_{{\mathbf{x}}_{t}}\log
p_{t}({\mathbf{y}}|{\mathbf{x}}_{t})\simeq-\frac{1}{\sigma_{{\mathbf{y}}}^{2}}\nabla_{{\mathbf{x}}_{t}}\lVert{\cal
A}(\hat{{\mathbf{x}}}_{0})-{\mathbf{y}}\rVert_{2}^{2},$ (9)

where $\hat{{\mathbf{x}}}_{0}$ is calculated as in Eq.(7). Finally, using the
approximation from Eq.(9), we can sample from the posterior using the reverse
SDE from Eq.(5). In practice, calculating the gradient in Eq.(9) requires
backpropagating through the score network with respect to the noisy input
${\mathbf{x}}_{t}$.

###  2.3 MRI Reconstruction

MRI is a popular medical imaging modality due to its superior soft tissue
contrast and lack of ionizing radiation. Unfortunately, MRI suffers from long
acquisition times. A common approach for accelerating scans is to collect
fewer measurements and solve the resulting ill-posed inverse problem. The
measurement process for multi-coil MRI can be written in the form

$\mathbf{y}_{i}=\mathbf{PFS}_{i}{\mathbf{x}}_{0}+\mathbf{\epsilon},$ (10)

where $\mathbf{y}_{i}\in{\mathbb{C}}^{m}$ are the measurements in the spatial
frequency domain (or _k-space_) for the $i^{th}$ coil,
${\mathbf{x}}_{0}\in{\mathbb{C}}^{n}$ is the image of interest,
$\mathbf{S}_{i}\in{\mathbb{C}}^{n\times n}$ is the coil sensitivity map for
the $i^{th}$ coil ($c$ coils in total), $\mathbf{F}\in{\mathbb{C}}^{n\times
n}$ is the Fourier transform matrix, $\mathbf{P}\in{\mathbb{C}}^{m\times n}$
is a sub-sampling operator whose rows are a subset of the rows of the $n\times
n$ identity matrix, and $\mathbf{\epsilon}\in{\mathbb{C}}^{m}$ is i.i.d
Gaussian noise. We also define the _acceleration factor_ $R\coloneqq m/n$ as
the under-sampling ratio. Note that even if the number of coils is large and
$mc\geq n$ the problem may still be ill-posed as coil sensitivity maps have
linear dependence due to spatial correlations.

Many early approaches for MRI reconstruction relied on sparsity-based or low-
rank priors [33, 18, 38]. More recently, however, there has been a push to
leverage advances in deep learning to assist in under-sampled image
reconstruction for MRI. These methods include unrolled techniques [1, 5, 19],
as well as approaches leveraging generative models [21, 24, 13, 32].

(a) 2D, Brain (b) 3D, Brain (c) 2D, Knee (d) 3D, Knee

Figure 2:  Mean test SSIM [52] for 2D and 3D patterns on Brains and Knees. We
compare reconstructions with masks learned using our method to those with
fixed baseline masks across a range of acceleration factors. The shaded areas
indicate a 95% confidence interval. Masks learned with our method consistently
lead to better reconstructions than baseline masks for 2D patterns. For 3D
patterns, our masks offer competitive performance with baselines.

##  3 Methods

We would like to learn a distribution $p_{\theta_{{\cal A}}}({\cal A})$,
parameterized by weights $\theta_{{\cal A}}$, over forward operators ${\cal
A}$ that produce measurements ${\mathbf{y}}={\cal
A}({\mathbf{x}}_{0})+\epsilon$. Our goal is to minimize the reconstruction
error between the true signal ${\mathbf{x}}_{0}$ and the estimates
$\tilde{{\mathbf{x}}}_{0}\sim p_{0}({\mathbf{x}}_{0}|{\mathbf{y}})$ produced
by sampling from the posterior using a diffusion-based generative model. We
can write the problem as an optimization that can be solved using gradient
descent-based methods:

$\theta_{{\cal A}}^{*}=\operatorname*{argmin\,}_{\theta_{{\cal
A}}}{\mathbb{E}}_{{\mathbf{x}}_{0}\sim
p_{0}({\mathbf{x}}_{0}),\>\tilde{{\mathbf{x}}}_{0}\sim
p_{0}({\mathbf{x}}_{0}|{\mathbf{y}}),\>{\mathbf{y}}\sim
p({\mathbf{y}}|{\mathbf{x}}_{0}),\>{\cal A}\sim p_{\theta_{{\cal A}}}({\cal
A})}\lVert{\mathbf{x}}_{0}-\tilde{{\mathbf{x}}}_{0}\rVert^{2}_{2}.$ (11)

In practice, however, sampling from the posterior to get
$\tilde{{\mathbf{x}}}_{0}$ involves an iterative application of an SDE solver
or ancestral sampling. This makes differentiating with respect to
$\theta_{{\cal A}}$ challenging, as naïvely backpropagating through the
sampling procedure is infeasible due to memory constraints. In the following
sections, we describe our approach to solve this problem.

(a)  (b)  (c)

Figure 3: Test results for varying training and sampling parameters. We vary
(a) the number of training samples for our method, (b) the number of
deterministic sampling steps, and (c) the stochasticity of sampling for 1000
steps (determined by S_churn as in Algorithm 1), and present the mean SSIM on
the test set. The shaded areas indicate a 95% confidence interval. All
experiments are performed on brain scans using 2D patterns with an
acceleration of R=10.

###  3.1 Learning Forward Operators

To find a solution, we return to Tweedie’s formula from Section 2.2. For the
case of VE SDEs, recall that the noisy signal at time $t$ is given by
$p_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{0})={\cal
N}({\mathbf{x}}_{t};{\mathbf{x}}_{0},\sigma_{t}^{2}{\mathbf{I}})$, allowing us
to use Eq. (6) to get a one-step approximation of the denoised posterior mean
${\mathbb{E}}_{{\mathbf{x}}_{t}\sim
p_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{0})}[{\mathbf{x}}_{0}|{\mathbf{x}}_{t}]$.
We show in 3.1 that we can extend Tweedie’s formula to include measurements
${\mathbf{y}}\sim p({\mathbf{y}}|{\mathbf{x}}_{0})$.

######  Proposition 3.1 (Tweedie’s formula with additional measurements).

Let ${\mathbf{x}}_{0}\sim p_{0}({\mathbf{x}}_{0})$ be an unknown signal,
${\mathbf{x}}_{t}\sim p_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{0})={\cal
N}({\mathbf{x}}_{t};{\mathbf{x}}_{0},\sigma^{2}_{t}{\mathbf{I}})$ a version of
${\mathbf{x}}_{0}$ corrupted by additive Gaussian noise, and ${\mathbf{y}}\sim
p({\mathbf{y}}|{\mathbf{x}}_{0})$ some additional measurements of
${\mathbf{x}}_{0}$. Furthermore, let ${\mathbf{x}}_{t}$ and ${\mathbf{y}}$ be
conditionally independent given ${\mathbf{x}}_{0}$:
$p_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{0},{\mathbf{y}})=p_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{0})$.
Finally, assume that $p_{t}({\mathbf{x}}_{t}|{\mathbf{y}})$ is supported
everywhere. Then, the posterior mean of ${\mathbf{x}}_{0}$ conditioned on
${\mathbf{x}}_{t}$ and ${\mathbf{y}}$ is given by

${\mathbb{E}}[{\mathbf{x}}_{0}|{\mathbf{x}}_{t},{\mathbf{y}}]={\mathbf{x}}_{t}+\sigma^{2}_{t}\nabla_{{\mathbf{x}}_{t}}\log
p_{t}({\mathbf{x}}_{t}|{\mathbf{y}}).$ (12)

We give the proof in Appendix A. Recalling that $\nabla_{{\mathbf{x}}_{t}}\log
p_{t}({\mathbf{x}}_{t}|{\mathbf{y}})=\nabla_{{\mathbf{x}}_{t}}\log
p_{t}({\mathbf{x}}_{t})+\nabla_{{\mathbf{x}}_{t}}\log
p_{t}({\mathbf{y}}|{\mathbf{x}}_{t})$, our result essentially allows us to
leverage the score of the likelihood _in addition to_ the prior to obtain a
finer estimate of ${\mathbf{x}}_{0}$ than using the prior alone. The
assumption of conditional independence of ${\mathbf{x}}_{t}$ and
${\mathbf{y}}$ given ${\mathbf{x}}_{0}$ is satisfied in the inverse problem
setting, as we have that ${\mathbf{x}}_{0}\rightarrow{\mathbf{y}}$ and
${\mathbf{x}}_{0}\rightarrow{\mathbf{x}}_{t}$ with no other dependencies.

In reality, we only have access to a score network and a surrogate form of the
score of the time-dependent likelihood from Eq. (8). Therefore, a tractable
approximation of the expectation in Eq. (12) is

${\mathbb{E}}[{\mathbf{x}}_{0}|{\mathbf{x}}_{t},{\mathbf{y}}]\simeq{\mathbf{x}}_{t}+\sigma^{2}_{t}[{\mathbf{s}}_{\mathbf{\theta}}({\mathbf{x}}_{t},t)+\nabla_{{\mathbf{x}}_{t}}\log
p_{t}({\mathbf{y}}|\hat{{\mathbf{x}}}_{0})],$ (13)

where $\hat{{\mathbf{x}}}_{0}$ is given by Tweedie’s formula as in Eq. (6). In
practice, we assume that $p_{t}({\mathbf{y}}|\hat{{\mathbf{x}}}_{0})\sim{\cal
N}({\mathbf{y}};{\cal A}(\hat{{\mathbf{x}}}_{0}),\sigma^{2}_{t}{\mathbf{I}})$
and calculate the posterior mean as

${\mathbb{E}}[{\mathbf{x}}_{0}|{\mathbf{x}}_{t},{\mathbf{y}}]\simeq{\mathbf{x}}_{t}+\sigma^{2}_{t}{\mathbf{s}}_{\mathbf{\theta}}({\mathbf{x}}_{t},t)-\gamma\nabla_{{\mathbf{x}}_{t}}\lVert{\cal
A}(\hat{{\mathbf{x}}}_{0})-{\mathbf{y}}\rVert_{2}^{2},$ (14)

where $\gamma$ is a likelihood step size.

Using the result from 3.1, we finally present our training objective:

$\begin{gathered}\theta_{{\cal A}}^{*}=\operatorname*{argmin\,}_{\theta_{{\cal
A}}}{\mathbb{E}}_{{\mathbf{x}}_{0}\sim
p_{0}({\mathbf{x}}_{0})}\lVert{\mathbf{x}}_{0}-\tilde{{\mathbf{x}}}_{0}\rVert^{2}_{2}\\\
\tilde{{\mathbf{x}}}_{0}={\mathbb{E}}_{{\cal A}\sim p_{\theta_{{\cal
A}}}({\cal A}),\>{\mathbf{y}}\sim{\cal N}({\cal
A}({\mathbf{x}}_{0}),\sigma_{{\mathbf{y}}}^{2}{\mathbf{I}}),\>{\mathbf{x}}_{t}\sim
p_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{0}),\>t\sim
q_{t}(t)}[{\mathbf{x}}_{0}|{\mathbf{x}}_{t},{\mathbf{y}}],\end{gathered}$ (15)

where $q_{t}(t)$ is a distribution over time steps and we use the
approximation from Eq. (14) to calculate the posterior mean estimate
$\tilde{{\mathbf{x}}}_{0}$. During training, we calculate the gradient w.r.t
$\theta_{{\cal A}}$ with automatic differentiation (e.g. using PyTorch [35]),
which involves backpropagating through $\tilde{{\mathbf{x}}}_{0}$ and
therefore through the score network.

Starting with a noisy signal at time $t$, samplers for diffusion models
generally follow the approach of: get the denoised estimate
$\hat{{\mathbf{x}}}_{0}$ at time $t$ using the score network, then add back
the appropriate amount of noise for time $t-1$ and repeat. Therefore, by
learning forward operators that produce good one-step posterior denoised
estimates, our method facilitates good final reconstructions. We are able to
use powerful generative models “off-the-shelf” without further training.
Furthermore, since we decouple our training objective from the posterior
sampling scheme, forward operators learned with our method are flexible and
agnostic to the choice of sampler.

###  3.2 Parameterizing the Sampling Pattern

2

We wish to learn a distribution over optimal sampling pattern for
reconstructing MRI images for some fixed acceleration $R$. This amounts to
learning ${\cal A}_{i}({\mathbf{x}})=\mathbf{PFS}_{i}{\mathbf{x}}$, for
$i\in[c]$ with $\mathbf{F}$ and $\mathbf{S}_{i}$ fixed. Therefore, we only
need to learn a distribution $p_{\theta_{\mathbf{P}}}$ parameterized by
weights $\theta_{\mathbf{P}}$ over the sub-sampling operator $\mathbf{P}$. We
base our parameterization and sampling scheme on that of LOUPE [55], with some
changes. For $n$-dimensional images, we learn parameters
$\theta_{\mathbf{P}}\in\mathbb{R}^{n}$. Each entry
$\theta_{\mathbf{P},i},i\in[n],$ defines an independent Bernouli random
variable $\mathcal{B}(p(\theta_{\mathbf{P},i}))$ at each location in the
k-space of an image, where $p(\theta_{\mathbf{P},i})$ is the probability that
the variable takes the value 1. The sampling mask distribution is defined as
$p_{\theta_{\mathbf{P}}}=\prod_{i=1}^{n}\mathcal{B}(p(\theta_{\mathbf{P},i}))$.

Algorithm 1 Posterior sampling

1:${\mathbf{s}}_{\mathbf{\theta}}({\mathbf{x}}_{t},t),\>\sigma_{t\in\\{t_{N},\dots,t_{0}\\}},\>S_{churn},\rho_{dps},{\mathbf{y}}$

2:sample ${\mathbf{x}}_{N}\sim{\cal
N}(\mathbf{0},\sigma_{t_{N}}^{2}{\mathbf{I}})$

3:for $i\in\\{N,\dots,1\\}$ do

4: sample $\mathbf{z}_{i}\sim{\cal N}(\mathbf{0},{\mathbf{I}})$

5: $\alpha_{i}\leftarrow min(S_{churn}/N,\sqrt{2}-1)$

6: $\hat{\sigma}_{t_{i}}\leftarrow\sigma_{t_{i}}+\alpha_{i}\sigma_{t_{i}}$

7:
$\hat{{\mathbf{x}}}_{i}\leftarrow{\mathbf{x}}_{i}+\sqrt{\hat{\sigma}_{t_{i}}^{2}-\sigma_{t_{i}}^{2}}\mathbf{z}_{i}$

8:
$\hat{{\mathbf{x}}}_{0}\leftarrow\hat{{\mathbf{x}}}_{i}+\hat{\sigma}_{t_{i}}^{2}{\mathbf{s}}_{\mathbf{\theta}}(\hat{{\mathbf{x}}}_{i},\hat{t}_{i})$

9:
$\hat{{\mathbf{x}}}_{i}^{\prime}\leftarrow\hat{{\mathbf{x}}}_{i}+(\hat{\sigma}_{t_{i}}-\sigma_{t_{i-1}}){\mathbf{s}}_{\mathbf{\theta}}(\hat{{\mathbf{x}}}_{i},\hat{t}_{i})$

10:
${\mathbf{x}}_{i-1}\leftarrow\hat{{\mathbf{x}}}_{i}^{\prime}-\rho_{dps}\nabla_{\hat{{\mathbf{x}}}_{i}}\lVert{\cal
A}(\hat{{\mathbf{x}}}_{0})-{\mathbf{y}}\rVert_{2}^{2}$

11:return ${\mathbf{x}}_{0}$

Following LOUPE, we re-scale the values $\sigma(\theta_{\mathbf{P},i})$, where
$\sigma(x)=1/(1+e^{-x})$ is the standard sigmoid function, so that they have
mean $1/R$. Since the Bernouli random variables are independent, sample
patterns drawn from the re-scaled distribution will have an average
acceleration factor of $R$. Given the unnormalized mean
$\bar{p}=\frac{1}{n}\sum_{i=1}^{n}\sigma(\theta_{\mathbf{P},i})$, we define
$p(\theta_{\mathbf{P},i})$ as

$p(\theta_{\mathbf{P},i})=\begin{cases}\frac{1}{\bar{p}R}\sigma(\theta_{\mathbf{P},i})&\text{if
}\bar{p}\geq\frac{1}{R}\\\
1-\frac{R-1}{R-\bar{p}R}(1-\sigma(\theta_{\mathbf{P},i})&\text{otherwise
}\end{cases},$ (16)

which outputs values in the range $[0,1]$ and has the desired mean
$\frac{1}{n}\sum_{i=1}^{n}p(\theta_{\mathbf{P},i})=1/R$.

Next, we need to sample from $p_{\theta_{\mathbf{P}}}$ in a way that is
differentiable with respect to $\theta_{\mathbf{P}}$. We use the Gumbel
Straight-Through estimator [26], which shows superior performance to the
vanilla Straight-Through estimator [8] used by LOUPE. For samples $g_{1}$ and
$g_{2}$ drawn i.i.d. from the Gumbel(0, 1) distribution and temperature
$\tau>0$, the Gumbel Straight-Through estimator generates a sample
$z_{i}\in\\{0,1\\}$ for k-space location $i\in[n]$ as

$z_{i}=\mathbf{1}_{\geq
0.5}(y_{i}),\>\>\>y_{i}=\frac{u_{\tau}(p(\theta_{\mathbf{P},i}),g_{1})}{u_{\tau}(p(\theta_{\mathbf{P},i}),g_{1})+u_{\tau}(1-p(\theta_{\mathbf{P},i}),g_{2})},\>\>\>u_{\tau}(p,g)=\exp\bigg{(}\frac{\log
p+g}{\tau}\bigg{)},$ (17)

where $\mathbf{1}_{\geq 0.5}(\cdot)$ is the indicator function that takes the
value 1 if its argument is $\geq 0.5$ and 0 otherwise. As $\tau\rightarrow 0$,
we have that $\mathbb{E}[y_{i}]\rightarrow p(\theta_{\mathbf{P},i})$. The
Gumbel Straight-Through estimator replaces $\nabla_{\theta_{\mathbf{P}}}z_{i}$
with $\nabla_{\theta_{\mathbf{P}}}y_{i}$ during backpropagation. This trick
allows us to draw realistic binary samples $z_{i}$ while enjoying well-defined
gradients from the smooth softmax sample $y_{i}$.

##  4 Experiments

We experiment on the fastMRI multi-coil brain and knee datasets [54]. We
create single-coil minimum variance unbiased estimator (MVUE) images from the
fully-sampled k-space data and sensitivity maps (calculated using ESPIRiT
[46]) to use as our ground truth reference. For training and testing sub-
sampling patterns, we create a subset of 200 training and 50 validation scans
from the fastMRI training set and 100 test scans from the fastMRI validation
set. We perform retrospective sub-sampling in our sampling pattern
experiments, meaning that we artificially apply sampling masks to fully-
sampled k-space data.

We train score-based networks using the ADM architecture [14] with the EDM
repo [27]. We use the EDM [27] default pre-conditioning and training
parameters, with a learning rate of $1\times 10^{-4}$, batch size of 15, and
an exponential moving average (ema) half-life of 100K images, and train the
models for 140 epochs. For posterior sampling, we modify the the stochastic
sampling algorithm from EDM to have no second-order correction and add a log-
likelihood step. We present our posterior sampler in Algorithm 1. Following
the implementation of DPS [12], we set the sampling likelihood step size
parameter as $\rho_{dps}=\rho/\lVert{\mathbf{y}}-{\cal
A}(\hat{{\mathbf{x}}}_{0})\rVert$, where $\rho$ is a tuneable hyperparameter.
Unless otherwise stated, we use 100 sampling steps with $\rho=10$ and the
stochasticity parameter $S_{churn}=0$.

When learning sub-sampling patterns, we use the noise schedule from EDM with
$\sigma_{t}=t$ and noise distribution $q_{t}(t)=q_{\sigma_{t}}(\sigma_{t})$
such that $\ln\sigma_{t}\sim{\cal N}(P_{mean},P_{std}^{2})$, where we use the
default values $P_{mean}=-1.2$ and $P_{std}=1.2$ [27]. We set the measurement
noise $\sigma_{{\mathbf{y}}}=0$, training likelihood step size to $\gamma=1$,
training batch size to $1$, and Gumbel temperature to $\tau=1$. We use the
Adam optimizer [29] with a learning rate of $1\times 10^{-2}$ and optimize for
10 epochs over all 200 training samples, unless stated otherwise. We fix a
central 16-pixel wide region that is always fully sampled for calibration
purposes on both our masks and baselines, and initialize pattern weights as
$\theta_{\mathbf{P}}=\mathbf{0}$. We perform all experiments on a machine with
$4\times$ Nvidia A100 GPUs.

Figure 4: Learned masks and reconstructions for fixed noise level during
training. We fix the value of $\sigma$ during training instead of sampling it
randomly to investigate the effects of noise level on learning. Very low
($\sigma=0.05,\>0.5$) and very high ($\sigma=50$) noise levels lead to useful
masks while moderate levels ($\sigma=5$) do not converge. We colorize the
learned probability distribution and scale the residual images by a factor of
$4$ for easier viewing. All patterns are 3D with $R=10$

###  4.1 Comparison to Baseline Patterns

First, we train 2D and 3D sub-sampling pattern distributions for brain and
knee scans across various accelerations. As a baseline, we compare to
posterior sampling reconstructions using fixed equispaced masks for the 2D
case and Poisson disc masks for 3D. We display qualitative results in Figure
1. Our masks can reduce small- and large-scale artifacts compared to baseline
masks, as shown by the areas indicated with red arrows and insets in Figure 1.
We note that our brain reconstruction at $R=20$ displays a shading artifact in
the top-left region, and both the baseline and our knee reconstructions
display shading artifacts. While our method displays good performance,
reconstructions still show limited quality for 2D patterns and high
accelerations regardless of training.

We evaluate the performance of each method using the structural similarity
index measure (SSIM) [52] between the reference MVUE and the reconstruction
and plot the results in Figure 2. Our method achieves better reconstruction
error for 2D imaging in both anatomies at acceleration values $R>4$, as
indicated in Figure 2(a) and (c). Interestingly Figures 2(b) and (d) show that
both our masks and the Poisson disc masks perform similarly in the 3D case,
with the only significant difference being that our masks lead to slightly
worse SSIM for knee scans for $R=8$ and $R=10$. We note that training a
diffusion model end-to-end or fine-tuning a pre-trained model with the
objective in Eq. (15) may offer better performance. However, these approaches
would be more computationally expensive and couple the diffusion model with
the forward operator, decreasing generalization ability and reducing the
usefulness of the generator for other tasks.

###  4.2 Varying Parameters

Here, we test how changing various training and sampling parameters changes
the quality of reconstructions. We fix the test set to brain scans and
evaluate 2D patterns with $R=10$ to easily compare results between the
experiments. We show the results in Figure 3.

First, we vary the number of training samples for our method. We train
sampling patterns for a total of 2000 iterations with a batch size of 1
regardless of the number of samples (e.g. for 5 samples we train for 400
epochs and for 25 samples we train for 80 epochs) and evaluate on the same 100
test images. The results in Figure 3(a) show that our method learns a pattern
that can outperform a fixed, equispaced baseline with as few as five samples.
Next, we vary the number of deterministic sampling steps, keeping the learned
sampling pattern distribution fixed. As shown in Figure 3(b), the relative
performance between the two methods remains invariant to changing number of
steps. Third, we introduce randomness into sampling by varying $S_{churn}$,
fixing the number of sampling steps to 1000 and changing $\rho$ to 2. Figure 3
indicates that the relative performance between our masks and the equispaced
masks does not change. The modularity of our method allows us to train
sampling operators that are uncoupled from the sampler choice at inference
time.

###  4.3 Effects of Noise Level on Training

In this experiment, we fix the anatomy to knees and the pattern to 3D with
$R=10$. During training, instead of drawing the noise level $\sigma_{t}$
randomly, we fix it to a single value $\sigma$. We show the visual
reconstruction results in Figure 4. At lower noise levels ($\sigma=0.05,0.5$),
the patterns capture more high-frequency portions of k-space, while at the
highest noise level ($\sigma=50$) the pattern densely sampled the central
region of k-space. Intuitively, lower levels of additive noise during training
means that the score network can more accurately predict the denoised image,
with the measurements aiding by adding fine details to the blurry denoised
image. For $\sigma=50$, the signal-to-noise ratio is almost zero, so the
posterior mean estimate relies heavily on the measurements to provide low-
frequency information. We note that the pattern for $\sigma=5$ did not
converge to a useful mask. This indicates that there may be a “dead zone” of
poor values for $\sigma$ during training.

##  5 Limitations and Broader Impact

#### Limitations.

The main limitation of our method is that it has a high memory complexity. Our
training objective requires backpropagating through the posterior mean
estimate, which itself relies on backpropagating through the score network to
obtain the approximation of the log-likelihood gradient. This process
calculates and stores second-order derivatives, which are very large for high-
dimensional problems. Another limitation is that we only consider Cartesian
sampling patterns, while non-Cartesian patterns may offer improved
reconstruction quality and robustness to patient motion during scanning. We
emphasize that our work focuses on the novel methodology and additional
clinical evaluation would be needed before it can be usable for MR protocol
designs.

#### Broader Impact.

Learning sampling patterns that offer the same reconstruction quality as
variable density or fixed patterns at higher accelerations would allow
patients to be scanned in a shorter time. This would benefit patients by
decreasing costs. However, a potential negative impact of our method is that
the training data may bias the sampling patterns toward a certain demographic.
For example, if the training data are all from patients under the age of 20,
the learned sampling patterns may offer poor reconstruction quality for older
patients. Care must be taken to construct the training and evaluation sets to
mitigate any such bias.

##  6 Conclusion

We introduced a method to learn sub-sampling operators for reconstructing MRI
images using posterior sampling with a diffusion-based generative model. We
rely on an extension of Tweedie’s formula for inverse problems, using pre-
trained generators with no need for re-training. Further, our method is
agnostic to the choice of sampling scheme at inference time. Experiments
demonstrate better performance than equispaced masks for 2D imaging and
similar performance to Poisson Discs for 3D imaging. Our method can learn
competitive patterns with as few as five training samples.

## References

  * [1] Hemant K Aggarwal, Merry P Mani, and Mathews Jacob.  Modl: Model-based deep learning architecture for inverse problems.  IEEE transactions on medical imaging, 38(2):394–405, 2018.
  * [2] Hemant Kumar Aggarwal and Mathews Jacob.  J-MoDL: Joint model-based deep learning for optimized sampling and reconstruction.  IEEE Journal of Selected Topics in Signal Processing, 14(6):1151–1162, oct 2020.
  * [3] Cagan Alkan, Morteza Mardani, Shreyas Vasanawala, and John M. Pauly.  Learning to sample MRI via variational information maximization.  In NeurIPS 2020 Workshop on Deep Learning and Inverse Problems, 2020\.
  * [4] Brian D.O. Anderson.  Reverse-time diffusion equation models.  Stochastic Processes and their Applications, 12(3):313–326, 1982\.
  * [5] Marius Arvinte, Sriram Vishwanath, Ahmed H. Tewfik, and Jonathan I. Tamir.  Deep j-sense: Accelerated mri reconstruction via unrolled alternating optimization, 2021.
  * [6] Muhammad Asim, Ali Ahmed, and Paul Hand.  Invertible generative models for inverse problems: mitigating representation error and dataset bias.  ArXiv, abs/1905.11672, 2020.
  * [7] Cagla D. Bahadir, Alan Q. Wang, Adrian V. Dalca, and Mert R. Sabuncu.  Deep-learning-based optimization of the under-sampling pattern in mri.  2019\.
  * [8] Yoshua Bengio, Nicholas Léonard, and Aaron C. Courville.  Estimating or propagating gradients through stochastic neurons for conditional computation.  ArXiv, abs/1308.3432, 2013.
  * [9] Ashish Bora, Ajil Jalal, Eric Price, and Alexandros G. Dimakis.  Compressed sensing using generative models.  In ICML, 2017.
  * [10] Ashish Bora, Eric Price, and Alexandros G. Dimakis.  AmbientGAN: Generative models from lossy measurements.  In International Conference on Learning Representations, 2018.
  * [11] Emmanuel J. Candès, Justin K. Romberg, and Terence Tao.  Stable signal recovery from incomplete and inaccurate measurements.  Communications on Pure and Applied Mathematics, 59:1207–1223, 2005\.
  * [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.
  * [13] Hyungjin Chung and Jong Chul Ye.  Score-based diffusion models for accelerated mri.  Medical Image Analysis, 80:102479, 2022.
  * [14] Prafulla Dhariwal and Alexander Nichol.  Diffusion models beat gans on image synthesis.  Advances in Neural Information Processing Systems, 34:8780–8794, 2021.
  * [15] D.L. Donoho.  Compressed sensing.  IEEE Transactions on Information Theory, 52(4):1289–1306, 2006\.
  * [16] Bradley Efron.  Tweedie’s formula and selection bias.  Journal of the American Statistical Association, 106:1602–1614, 12 2011.
  * [17] Mark A Griswold, Peter M Jakob, Robin M Heidemann, Mathias Nittka, Vladimir Jellus, Jianmin Wang, Berthold Kiefer, and Axel Haase.  Generalized autocalibrating partially parallel acquisitions (grappa).  Magnetic Resonance in Medicine: An Official Journal of the International Society for Magnetic Resonance in Medicine, 47(6):1202–1210, 2002\.
  * [18] Justin P. Haldar and Kawin Setsompop.  Linear predictability in magnetic resonance imaging reconstruction: Leveraging shift-invariant fourier structure for faster and better imaging.  IEEE Signal Processing Magazine, 37(1):69–82, 2020.
  * [19] 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.  Magnetic Resonance in Medicine, 79(6):3055–3071, 2018.
  * [20] Paul Hand, Oscar Leong, and Vlad Voroninski.  Phase retrieval under a generative prior.  In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
  * [21] Reinhard Heckel and Paul Hand.  Deep decoder: Concise image representations from untrained non-convolutional networks.  ArXiv, abs/1810.03982, 2019.
  * [22] Jonathan Ho, Ajay Jain, and Pieter Abbeel.  Denoising diffusion probabilistic models.  Advances in Neural Information Processing Systems, 33:6840–6851, 2020.
  * [23] Jonathan Ho and Tim Salimans.  Classifier-free diffusion guidance.  arXiv preprint arXiv:2207.12598, 2022.
  * [24] Ajil Jalal, Marius Arvinte, Giannis Daras, Eric Price, Alexandros G Dimakis, and Jon Tamir.  Robust compressed sensing mri with deep generative priors.  In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 14938–14954. Curran Associates, Inc., 2021.
  * [25] Ajil Jalal, Sushrut Karmalkar, Alexandros G. Dimakis, and Eric Price.  Instance-optimal compressed sensing via posterior sampling, 2021.
  * [26] Eric Jang, Shixiang Gu, and Ben Poole.  Categorical reparameterization with gumbel-softmax.  In International Conference on Learning Representations, 2017.
  * [27] Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine.  Elucidating the design space of diffusion-based generative models.  In Proc. NeurIPS, 2022.
  * [28] Bahjat Kawar, Michael Elad, Stefano Ermon, and Jiaming Song.  Denoising diffusion restoration models.  arXiv preprint arXiv:2201.11793, 2022.
  * [29] Diederik P. Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  In Yoshua Bengio and Yann LeCun, editors, 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015.
  * [30] Carole Lazarus, Pierre Weiss, Nicolas Chauffert, Franck Mauconduit, Loubna El Gueddari, Christophe Destrieux, Ilyess Zemmoura, Alexandre Vignaud, and Philippe Ciuciu.  Sparkling: variable-density k-space filling curves for accelerated t2*-weighted mri.  Magnetic Resonance in Medicine, 81(6):3643–3661, 2019.
  * [31] Zhaoqiang Liu, Selwyn Gomes, Avtansh Tiwari, and Jonathan Scarlett.  Sample complexity bounds for 1-bit compressive sensing and binary stable embeddings with generative priors, 2020.
  * [32] Guanxiong Luo, Martin Heide, and Martin Uecker.  Mri reconstruction via data driven markov chain with joint uncertainty estimation.  arXiv preprint arXiv:2202.01479, 2022.
  * [33] Michael Lustig, David Donoho, and John M Pauly.  Sparse mri: The application of compressed sensing for rapid mr imaging.  Magnetic Resonance in Medicine: An Official Journal of the International Society for Magnetic Resonance in Medicine, 58(6):1182–1195, 2007\.
  * [34] Gregory Ongie, Ajil Jalal, Christopher A Metzler, Richard G Baraniuk, Alexandros G Dimakis, and Rebecca Willett.  Deep learning techniques for inverse problems in imaging.  IEEE Journal on Selected Areas in Information Theory, 1(1):39–56, 2020.
  * [35] Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer.  Automatic differentiation in pytorch.  2017\.
  * [36] Klaas P Pruessmann, Markus Weiger, Markus B Scheidegger, and Peter Boesiger.  Sense: sensitivity encoding for fast mri.  Magnetic Resonance in Medicine: An Official Journal of the International Society for Magnetic Resonance in Medicine, 42(5):952–962, 1999\.
  * [37] Shuang Qiu, Xiaohan Wei, and Zhuoran Yang.  Robust one-bit recovery via relu generative networks: Improved statistical rates and global landscape analysis.  arXiv preprint arXiv:1908.05368, 2019.
  * [38] Saiprasad Ravishankar and Yoram Bresler.  Mr image reconstruction from highly undersampled k-space data by dictionary learning.  IEEE Transactions on Medical Imaging, 30(5):1028–1041, 2011.
  * [39] Ferdia Sherry, Martin Benning, Juan Carlos De los Reyes, Martin J Graves, Georg Maierhofer, Guy Williams, Carola-Bibiane Schönlieb, and Matthias J Ehrhardt.  Learning the sampling pattern for mri.  IEEE Transactions on Medical Imaging, 39(12):4310–4321, 2020.
  * [40] Daniel K Sodickson and Warren J Manning.  Simultaneous acquisition of spatial harmonics (smash): fast imaging with radiofrequency coil arrays.  Magnetic resonance in medicine, 38(4):591–603, 1997.
  * [41] Yang Song and Stefano Ermon.  Generative modeling by estimating gradients of the data distribution.  Advances in neural information processing systems, 32, 2019.
  * [42] Yang Song and Stefano Ermon.  Improved techniques for training score-based generative models.  arXiv preprint arXiv:2006.09011, 2020.
  * [43] Yang Song, Liyue Shen, Lei Xing, and Stefano Ermon.  Solving inverse problems in medical imaging with score-based generative models.  In International Conference on Learning Representations, 2022.
  * [44] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole.  Score-based generative modeling through stochastic differential equations.  In International Conference on Learning Representations, 2021.
  * [45] Anuroop Sriram, Jure Zbontar, Tullie Murrell, Aaron Defazio, C Lawrence Zitnick, Nafissa Yakubova, Florian Knoll, and Patricia Johnson.  End-to-end variational networks for accelerated mri reconstruction.  In Medical Image Computing and Computer Assisted Intervention–MICCAI 2020: 23rd International Conference, Lima, Peru, October 4–8, 2020, Proceedings, Part II 23, pages 64–73. Springer, 2020.
  * [46] 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.  Magnetic resonance in medicine, 71(3):990–1001, 2014.
  * [47] Dmitry Ulyanov, Andrea Vedaldi, and Victor S. Lempitsky.  Deep image prior.  2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 9446–9454, 2018.
  * [48] Pascal Vincent.  A connection between score matching and denoising autoencoders.  Neural computation, 23(7):1661–1674, 2011.
  * [49] Guanhua Wang, Tianrui Luo, Jon-Fredrik Nielsen, Douglas C Noll, and Jeffrey A Fessler.  B-spline parameterized joint optimization of reconstruction and k-space trajectories (bjork) for accelerated 2d mri.  IEEE Transactions on Medical Imaging, 41(9):2318–2330, 2022.
  * [50] Guanhua Wang, Jon-Fredrik Nielsen, Jeffrey A. Fessler, and Douglas C. Noll.  Stochastic optimization of 3d non-cartesian sampling trajectory (snopy), 2022.
  * [51] Guanhua Wang, Douglas C. Noll, and Jeffrey A. Fessler.  Adaptive sampling for linear sensing systems via langevin dynamics, 2023\.
  * [52] Zhou Wang, A.C. Bovik, H.R. Sheikh, and E.P. Simoncelli.  Image quality assessment: from error visibility to structural similarity.  IEEE Transactions on Image Processing, 13(4):600–612, 2004.
  * [53] Jay Whang, Mauricio Delbracio, Hossein Talebi, Chitwan Saharia, Alexandros G Dimakis, and Peyman Milanfar.  Deblurring via stochastic refinement.  arXiv preprint arXiv:2112.02475, 2021.
  * [54] Jure Zbontar, Florian Knoll, Anuroop Sriram, Tullie Murrell, Zhengnan Huang, Matthew J Muckley, Aaron Defazio, Ruben Stern, Patricia Johnson, Mary Bruno, et al.  fastmri: An open dataset and benchmarks for accelerated mri.  arXiv preprint arXiv:1811.08839, 2018.
  * [55] Jinwei Zhang, Hang Zhang, Alan Wang, Qihao Zhang, Mert Sabuncu, Pascal Spincemaille, Thanh D. Nguyen, and Yi Wang.  Extending loupe for k-space under-sampling pattern optimization in multi-coil mri, 2020.
  * [56] Marcelo VW Zibetti, Gabor T Herman, and Ravinder R Regatte.  Fast data-driven learning of parallel mri sampling patterns for large scale problems.  Scientific Reports, 11(1):19312, 2021.

##  Appendix A Proof of 3.1

See 3.1

###### Proof.

We begin by representing the distribution
$p_{t}({\mathbf{x}}_{t}|{\mathbf{y}})$ as marginalizing out ${\mathbf{x}}_{0}$
conditioned on ${\mathbf{y}}$:

$\displaystyle
p_{t}({\mathbf{x}}_{t}|{\mathbf{y}})=\int_{{\mathbf{x}}_{0}}p_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{0},{\mathbf{y}})p_{0}({\mathbf{x}}_{0}|{\mathbf{y}})d{\mathbf{x}}_{0}.$

Next, we take the gradient w.r.t. ${\mathbf{x}}_{t}$ on both sides:

$\displaystyle\nabla_{{\mathbf{x}}_{t}}p_{t}({\mathbf{x}}_{t}|{\mathbf{y}})$
$\displaystyle=\nabla_{{\mathbf{x}}_{t}}\int_{{\mathbf{x}}_{0}}p_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{0},{\mathbf{y}})p_{0}({\mathbf{x}}_{0}|{\mathbf{y}})d{\mathbf{x}}_{0}$
$\displaystyle=\int_{{\mathbf{x}}_{0}}p_{0}({\mathbf{x}}_{0}|{\mathbf{y}})\nabla_{{\mathbf{x}}_{t}}p_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{0},{\mathbf{y}})d{\mathbf{x}}_{0}$
$\displaystyle=\int_{{\mathbf{x}}_{0}}p_{0}({\mathbf{x}}_{0}|{\mathbf{y}})p_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{0},{\mathbf{y}})\nabla_{{\mathbf{x}}_{t}}\log
p_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{0},{\mathbf{y}})d{\mathbf{x}}_{0}.$

On the last line, we use the identity $\nabla_{x}\log
f(x)=\nabla_{x}f(x)/f(x)$. We note that since
$p_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{0},{\mathbf{y}})=p_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{0})$,
and $p_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{0})$ is Gaussian,
$p_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{0},{\mathbf{y}})$ has non-zero value
everywhere and we avoid singularities from the denominator.

Continuing, we use the conditional independence of ${\mathbf{x}}_{t}$ and
${\mathbf{y}}$ given ${\mathbf{x}}_{0}$ to replace
$\nabla_{{\mathbf{x}}_{t}}\log
p_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{0},{\mathbf{y}})$ on the right-hand side
with $\nabla_{{\mathbf{x}}_{t}}\log p_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{0})$
and obtain:

$\displaystyle\nabla_{{\mathbf{x}}_{t}}p_{t}({\mathbf{x}}_{t}|{\mathbf{y}})$
$\displaystyle=\int_{{\mathbf{x}}_{0}}p_{0}({\mathbf{x}}_{0}|{\mathbf{y}})p_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{0},{\mathbf{y}})\nabla_{{\mathbf{x}}_{t}}\log
p_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{0})d{\mathbf{x}}_{0}$
$\displaystyle=\int_{{\mathbf{x}}_{0}}p_{0}({\mathbf{x}}_{0}|{\mathbf{y}})p_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{0},{\mathbf{y}})\bigg{(}\frac{{\mathbf{x}}_{0}-{\mathbf{x}}_{t}}{\sigma^{2}_{t}}\bigg{)}d{\mathbf{x}}_{0}.$

Here, we use the fact that $p_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{0})={\cal
N}({\mathbf{x}}_{t};{\mathbf{x}}_{0},\sigma^{2}_{t}{\mathbf{I}}_{n})$ is a
Gaussian and replace the score function $\nabla_{{\mathbf{x}}_{t}}\log
p_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{0})$ by its exact value,
$({\mathbf{x}}_{0}-{\mathbf{x}}_{t})/\sigma^{2}_{t}$.

Expanding the right-hand-side, we get:

$\displaystyle\nabla_{{\mathbf{x}}_{t}}p_{t}({\mathbf{x}}_{t}|{\mathbf{y}})$
$\displaystyle=\frac{1}{\sigma^{2}_{t}}\bigg{[}\int_{{\mathbf{x}}_{0}}p_{0}({\mathbf{x}}_{0}|{\mathbf{y}})p_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{0},{\mathbf{y}}){\mathbf{x}}_{0}d{\mathbf{x}}_{0}-\int_{{\mathbf{x}}_{0}}p_{0}({\mathbf{x}}_{0}|{\mathbf{y}})p_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{0},{\mathbf{y}}){\mathbf{x}}_{t}d{\mathbf{x}}_{0}\bigg{]}$
$\displaystyle=\frac{1}{\sigma^{2}_{t}}\bigg{[}\int_{{\mathbf{x}}_{0}}p_{0}({\mathbf{x}}_{0}|{\mathbf{y}})p_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{0},{\mathbf{y}}){\mathbf{x}}_{0}d{\mathbf{x}}_{0}-{\mathbf{x}}_{t}\int_{{\mathbf{x}}_{0}}p_{0}({\mathbf{x}}_{0}|{\mathbf{y}})p_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{0},{\mathbf{y}})d{\mathbf{x}}_{0}\bigg{]}$
$\displaystyle=\frac{1}{\sigma^{2}_{t}}\bigg{[}\int_{{\mathbf{x}}_{0}}p_{0}({\mathbf{x}}_{0}|{\mathbf{y}})p_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{0},{\mathbf{y}}){\mathbf{x}}_{0}d{\mathbf{x}}_{0}-{\mathbf{x}}_{t}p_{t}({\mathbf{x}}_{t}|{\mathbf{y}})\bigg{]}.$

In the previous line, we marginalize out ${\mathbf{x}}_{0}$ conditioned on
${\mathbf{y}}$ as in the first line of the proof to recover
$p_{t}({\mathbf{x}}_{t}|{\mathbf{y}})$.

Next, we observe that Bayes’ rule tells us
$p_{0}({\mathbf{x}}_{0}|{\mathbf{y}})p_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{0},{\mathbf{y}})=p_{t}({\mathbf{x}}_{t}|{\mathbf{y}})p_{0}({\mathbf{x}}_{0}|{\mathbf{x}}_{t},{\mathbf{y}})$
and replace the former quantity by the latter on the right-hand side:

$\displaystyle\nabla_{{\mathbf{x}}_{t}}p_{t}({\mathbf{x}}_{t}|{\mathbf{y}})$
$\displaystyle=\frac{1}{\sigma^{2}_{t}}\bigg{[}\int_{{\mathbf{x}}_{0}}p_{t}({\mathbf{x}}_{t}|{\mathbf{y}})p_{0}({\mathbf{x}}_{0}|{\mathbf{x}}_{t},{\mathbf{y}}){\mathbf{x}}_{0}d{\mathbf{x}}_{0}-{\mathbf{x}}_{t}p_{t}({\mathbf{x}}_{t}|{\mathbf{y}})\bigg{]}$
$\displaystyle=\frac{1}{\sigma^{2}_{t}}\bigg{[}p_{t}({\mathbf{x}}_{t}|{\mathbf{y}}){\mathbb{E}}[{\mathbf{x}}_{0}|{\mathbf{x}}_{t},{\mathbf{y}}]-{\mathbf{x}}_{t}p_{t}({\mathbf{x}}_{t}|{\mathbf{y}})\bigg{]}$
$\displaystyle=\frac{p_{t}({\mathbf{x}}_{t}|{\mathbf{y}})}{\sigma^{2}_{t}}\bigg{[}{\mathbb{E}}[{\mathbf{x}}_{0}|{\mathbf{x}}_{t},{\mathbf{y}}]-{\mathbf{x}}_{t}\bigg{]}$
$\displaystyle\frac{\nabla_{{\mathbf{x}}_{t}}p_{t}({\mathbf{x}}_{t}|{\mathbf{y}})}{p_{t}({\mathbf{x}}_{t}|{\mathbf{y}})}$
$\displaystyle=\frac{1}{\sigma^{2}_{t}}\bigg{[}{\mathbb{E}}[{\mathbf{x}}_{0}|{\mathbf{x}}_{t},{\mathbf{y}}]-{\mathbf{x}}_{t}\bigg{]}.$

We note that from our assumption, $p_{t}({\mathbf{x}}_{t}|{\mathbf{y}})$ is
fully supported everywhere, so we avoid singularities when dividing by this
quantity.

Finally, we again invoke the identity $\nabla_{x}\log
f(x)=\nabla_{x}f(x)/f(x)$ to rewrite the left-hand side and rearrange to
obtain the desired result:

$\displaystyle\nabla_{{\mathbf{x}}_{t}}\log
p_{t}({\mathbf{x}}_{t}|{\mathbf{y}})$
$\displaystyle=\frac{1}{\sigma^{2}_{t}}\bigg{[}{\mathbb{E}}[{\mathbf{x}}_{0}|{\mathbf{x}}_{t},{\mathbf{y}}]-{\mathbf{x}}_{t}\bigg{]}$
$\displaystyle{\mathbb{E}}[{\mathbf{x}}_{0}|{\mathbf{x}}_{t},{\mathbf{y}}]$
$\displaystyle={\mathbf{x}}_{t}+\sigma^{2}_{t}\nabla_{{\mathbf{x}}_{t}}\log
p_{t}({\mathbf{x}}_{t}|{\mathbf{y}}).$

∎

##  Appendix B Extended Experimental details

###  B.1 Dataset

All MRI data are initially stored as complex-valued multi-coil k-space
measurements. We appropriately crop the raw k-space data and MVUE images for
brain scans to $384\times 384$ pixels and knee scans to $320\times 320$ pixels
for all experiments.

For training the score network on brain scans, we take volumes from the
fastMRI multi-coil brain training set and remove the last two noisy slices
from each volume for a total of 57,297 scans. For training on knee scans, we
keep all data in the fastMRI multi-coil knee training set for a total of
28,102 scans. We use brain and knee scans from all available contrasts and
field strengths for training the score networks. To make the data compatible
with real-valued network weights, we represent the complex-valued MVUE images
as two-channel, real-valued images.

Since the data exhibit a wide dynamic range of pixel values, we linearly scale
images to the range [-1,1] when training the score networks using the minimum
and maximum pixel values from the two-channel fully-sampled MVUE images. When
learning sampling patterns and performing posterior sampling, we use the
minimum and maximum pixel values from the MVUE calculated from the
retrospectively-undersampled k-space data to perform this scaling.

For training and evaluating sampling patterns, we use T2-weighted brain scans
and proton density-weighted knee scans without fat suppression, all with a
field strength of 3 Teslas. We remove the first ten slices from knee volumes
and last five slices from brain volumes when creating the training,
validation, and test sets for learning sampling patterns with our method.

###  B.2 Training Diffusion Models

We train score-based networks using the ADM architecture [14] and the default
parameters for ImageNet 256 with some changes. We use 128 base channels
instead of 256 and self-attention in the two smallest resolution scales
instead of three smallest. The models are trained with classifier-free
diffusion guidance [23] with a label dropout probability of 0.1, treating each
unique (contrast, field strength) pair as a different class. We train one
network for brain scans and one for knees.

###  B.3 Training Sampling Patterns

When training the sampling patterns using Eq. (15), we use a single forward
operator ${\cal A}({\mathbf{x}})=\mathbf{PF}{\mathbf{x}}$ instead of using
${\cal A}_{i}({\mathbf{x}})=\mathbf{PFS}_{i}{\mathbf{x}}$ for each of the $c$
coils with $i\in[c]$. In other words, we do not use the coil sensitivity maps
$\mathbf{C}_{i}$ when retrospectively sub-sampling training data to create
measurements. We find that this method accelerates training by avoiding memory
and computational costs for multiple coils. We also find that this method
improves test performance and leads to more stable convergence of the learned
sampling patterns. During validation and testing, we still perform
reconstructions using undersampled multi-coil k-space data as measurements.

During validation, we perform posterior sampling to reconstruct each image in
the validation set and track the error between the reconstructions and fully-
sampled images. Before testing, we restore the weights of the sampling
operator from the iteration with the lowest mean validation error.

###  B.4 Hyperparameters

We tune the following hyperparameters: the training likelihood step size
$\gamma$, Gumbel temperature $\tau$, learning rate for sampling patterns,
sampling likelihood step size $\rho$, and sampling stochasticity parameter
$S_{churn}$.

For the general experiments in Figures 1 and 2, we fix the number of sampling
steps to 100 and search for the values of
$\rho\in\\{0.1,0.2,\dots,1,2,\dots,10\\}$ and
$S_{churn}\in\\{0,10,\dots,50\\}$. We choose the pair of values with the
smallest reconstruction error on a hold-out set of 30 scans. Once we find the
optimal values $\rho=10$ and $S_{churn}=0$, we fix them and search for values
of $\gamma\in\\{0.1,0.5,1,5,10\\}$ and $\tau\in\\{0.1,0.5,1\\}$. We train
sampling patterns for 2 epochs on our training set using a learning rate of
$0.1$ and choose the pair $\gamma=1$ and $\tau=1$ that gives the best
reconstruction error on the 30 hold-out scans using the learned pattern.
Finally, we fix all previous tuned hyperparameters and search for the learning
rate in $\\{10^{-4},10^{-3},10^{-2},10^{-1}\\}$ by again training for 2 epochs
and reconstructing hold-out images, finding $10^{-2}$ to be optimal. We tune
all listed hyperparameters using 2D patterns with $R=4$ on brain scans and use
the same values across all anatomies, pattern types, and accelerations.

For the experiment in Figure 3 (b), we set $S_{churn}=0$ and set $\rho=20$ for
25 steps, $\rho=10$ for 100 steps, $\rho=8$ for 200 steps, $\rho=6$ for 400
steps, $\rho=4$ for 600 steps, $\rho=3$ for 800 steps, and $\rho=2$ for 1000
steps. To tune these values, we start from 25 steps and search for
$\rho\in\\{1,2,\dots,25\\}$ and notice that the optimal value $\rho=20$ is
larger than for 100 steps ($\rho=10$). Continuing, as we find the optimal
$\rho$ for each step value, we increase the number of steps and restrict the
search space to smaller and smaller subsets, with the maximum search value
being the optimal value of $\rho$ for the previous step value. For example, we
search $\rho\in\\{1,2,\dots,10\\}$ for 200 steps and
$\rho\in\\{1,2,\dots,8\\}$ for 400 steps. For the experiment in Figure 3 (c),
we re-use the optimal value of $\rho=2$ for 1000 steps. We fix the same value
$\rho=2$ for all values of $S_{churn}$ in this experiment. We tune the
hyperparameters for these experiments on a hold-out set of 30 brain scans with
2D patterns and $R=10$.

##  Appendix C PSNR Plots

In this section we present PSNR plots for the experiments from Figures 2 and
3.

(a) 2D, Brain (b) 3D, Brain (c) 2D, Knee (d) 3D, Knee

Figure 5:  Mean test PSNR for 2D and 3D patterns on Brains and Knees. We
compare reconstructions with masks learned using our method to those with
fixed baseline masks across a range of acceleration factors. The shaded areas
indicate a 95% confidence interval.

(a)  (b)  (c)

Figure 6: Test results for varying training and sampling parameters. We vary
(a) the number of training samples for our method, (b) the number of
deterministic sampling steps, and (c) the stochasticity of sampling for 1000
steps (determined by S_churn as in Algorithm 1), and present the mean PSNR on
the test set. The shaded areas indicate a 95% confidence interval. All
experiments are performed on brain scans using 2D patterns with an
acceleration of R=10.

##  Appendix D Additional Plots of Learned Masks and Reconstructions

Figure 7: Learned 3D Sampling Patterns and Reconstructions for Brains. We
perform posterior sampling using a diffusion model to reconstruct brain scans
for varying acceleration factors (R) using 3D patterns. We colorize the
probability distributions of our learned patterns for easier viewing. Figure
8: Learned 3D Sampling Patterns and Reconstructions for Knees. We perform
posterior sampling using a diffusion model to reconstruct knee scans for
varying acceleration factors (R) using 3D patterns. We colorize the
probability distributions of our learned patterns for easier viewing. Figure
9: Learned 2D Sampling Patterns and Reconstructions for Brains. We perform
posterior sampling using a diffusion model to reconstruct brain scans for
varying acceleration factors (R) using 2D patterns. We colorize the
probability distributions of our learned patterns for easier viewing. Figure
10: Learned 2D Sampling Patterns and Reconstructions for Knees. We perform
posterior sampling using a diffusion model to reconstruct knee scans for
varying acceleration factors (R) using 2D patterns. We colorize the
probability distributions of our learned patterns for easier viewing.