In image recovery problems, one seeks to reconstruct an image from distorted, incomplete, and/or noise-corrupted measurements.
Such problems arise in magnetic resonance imaging (MRI), computed tomography, deblurring, super-resolution, inpainting, phase retrieval, image-to-image translation, and other applications.
Given a training set of signal/measurement pairs, we design a method to generate posterior samples rapidly and accurately.
In particular, we propose a regularized conditional Wasserstein GAN that generates dozens of high-quality posterior samples per second.
Our regularization comprises an $\ell_1$ penalty and an adaptively weighted standard-deviation reward.
Using quantitative evaluation metrics like conditional Fréchet inception distance, we demonstrate that our method produces state-of-the-art posterior samples in both multicoil MRI and large-scale inpainting applications.

§ INTRODUCTION

We consider image recovery, where one observes a measurement $\vec{y}=\mc{M}(\vec{x})$ of the true image $\vec{x}$ through a function $\mc{M}(\cdot)$ that masks, distorts, and/or corrupts $\vec{x}$ with noise.
This includes linear inverse problems arising in, e.g., deblurring, super-resolution, inpainting, colorization, computed tomography (CT), and magnetic resonance imaging (MRI), where $\vec{y}=\vec{Ax}+\vec{w}$ with known linear operator $\vec{A}$ and noise $\vec{w}$.
But it also includes non-linear inverse problems like those arising in phase-retrieval and dequantization, as well as image-to-image translation problems.

In all cases, it is impossible to perfectly infer $\vec{x}$ from $\vec{y}$.

Image recovery is often posed as finding the single “best” recovery $\hvec{x}$, which is known as a point estimate of $\vec{x}$ <cit.>.
However, point estimation is problematic due to the perception-distortion tradeoff <cit.>, which establishes a fundamental tradeoff between distortion (defined as some distance between $\hvec{x}$ and $\vec{x}$) and perceptual quality (defined as some distance between $\hvec{x}$ and the closest clean image).
For example, the minimum mean-squared error (MMSE) recovery $\hvec{x}\mmse$ is optimal in terms of $\ell_2$ distortion, but can be too smooth.
Although one could instead compute an approximation of the maximum a posteriori (MAP) estimate <cit.> or minimize some combination of perceptual and distortion losses such as $\ell_1$ plus SSIM <cit.>, it's unclear which criterion is most appropriate.

Another major limitation with point estimation is that it's unclear how certain one can be about the recovered $\hvec{x}$.
For example, with deep-learning-based recovery, it's possible to hallucinate a nice-looking $\hvec{x}$, but is it correct?
Quantifying the uncertainty in $\hvec{x}$ is especially important in medical applications such as MRI, where a diagnosis must be made based on the measurements $\vec{y}$.
Rather than simply reporting our best guess of whether a pathology is present or absent based on $\hvec{x}$, we'd like to report the probability that the pathology is present (given all available data).

Yet another problem with point estimation is that the estimated $\hvec{x}$ could pose issues with fairness <cit.>.
For example, say we are inpainting a face within an image.
With a racially heterogeneous prior distribution, $\hvec{x}\mmse$ (being the posterior mean) will be biased towards the most predominant race.
The same will be true of many other point estimates $\hvec{x}$.

To address the aforementioned limitations of point estimation, we focus on generating samples from the posterior distribution $\pxgy(\cdot|\vec{y})$, which represents the complete state-of-knowledge about $\vec{x}$ given the measurements $\vec{y}$.
The posterior correctly fuses prior knowledge with measurement knowledge, thereby alleviating any concerns about fairness (assuming the data used to represent the prior is fairly chosen <cit.>).
Furthermore, the posterior directly facilitates uncertainty quantification via, e.g., pixel-wise standard deviations or pathology detection probabilities (see detection).
Also, if it was important to report a single “good” recovery, then posterior sampling leads to an easy navigation of the perception-distortion tradeoff.
For example, averaging $P\geq 1$ posterior samples gives a close approximation to the less-distorted-but-oversmooth $\hvec{x}\mmse$ with large $P$ and sharp-but-more-distorted $\hvec{x}$ with small $P$.
Additionally, posterior sampling unlocks other important capabilities such as adaptive acquisition <cit.> and counterfactual diagnosis <cit.>.

Concretely, given a training dataset of image/measurement pairs $\{(\vec{x}_t,\vec{y}_t)\}_{t=1}^T$,

our goal is to learn a generating function $G_{\vec{\theta}}$ that, for a new $\vec{y}$, maps random code vectors $\vec{z}\sim\mc{N}(\vec{0},\vec{I})$ to posterior samples $\hvec{x}=G_{\vec{\theta}}(\vec{z},\vec{y})\sim\pxgy(\cdot|\vec{y})$.

There exist several well-known approaches to this task, with recent literature focusing on

conditional generative adversarial networks (cGANs)
conditional variational autoencoders (cVAEs)
conditional normalizing flows (cNFs)
and score/diffusion/Langevin-based generative models
Despite it being a longstanding problem, posterior image sampling remains challenging.
Although score/diffusion/Langevin approaches have dominated the recent literature due to advances in accuracy and diversity, their sample-generation speeds remain orders-of-magnitude behind those of cGANs, cVAEs, and cNFs.

We focus on cGANs, which are typically regarded as generating samples of high quality but low diversity.
Our proposed cGAN tackles the lack-of-diversity issue using a novel regularization that consists of supervised-$\ell_1$ loss plus an adaptively weighted standard-deviation (SD) reward.
This is not a heuristic choice; we prove that our regularization enforces consistency with the true posterior mean and covariance under certain conditions.

Experimentally, we demonstrate our regularized cGAN on accelerated MRI and large-scale face completion/inpainting.
We consider these applications for three main reasons.
First, uncertainty quantification in MRI, and fairness in face-generation, are both of paramount importance.
Second, posterior-sampling has been well studied for both applications, and fine-tuned cGANs <cit.> and score/Langevin-based approaches <cit.> are readily available.
Third, the linear operator “$\vec{A}$” manifests very differently in these two applications,[In MRI, the forward operator acts locally in the frequency domain but globally in the pixel domain, while in inpainting, the operator acts locally in the pixel domain but globally in the frequency domain.]
which illustrates the versatility of our approach.
To quantify performance, we focus on conditional Fréchet inception distance (CFID) <cit.>,
which is a principled way to quantify the difference between two high-dimensional posterior distributions,
although we also report other metrics.

Our results show the proposed regularized cGAN (rcGAN) outperforming existing cGANs <cit.> and the score/diffusion/Langevin approaches from <cit.> and <cit.> in all tested metrics, while generating samples $\sim 10^4$ times faster than <cit.>.

§ PROBLEM FORMULATION AND BACKGROUND

We build on the Wasserstein cGAN framework from <cit.>.

The goal is to design a generator network $G_{\vec{\theta}}:\mc{Z}\times\mc{Y}\rightarrow\mc{X}$ such that, for fixed $\vec{y}$, the random variable $\hvec{x}=G_{\vec{\theta}}(\vec{z},\vec{y})$ induced by $\vec{z}\sim\pz$ has a distribution that best matches the posterior $\pxgy(\cdot|\vec{y})$ in Wasserstein-1 distance.
Here, $\vec{z}$ is drawn independently of $\vec{y}$, and $\mc{X}$, $\mc{Y}$, and $\mc{Z}$ denote the sets of $\vec{x}$, $\vec{y}$, and $\vec{z}$, respectively.

The Wasserstein-1 distance can be expressed as

\begin{align}
W_1(\pxgy(\cdot,\vec{y}),\pxhgy(\cdot,\vec{y})) = \sup_{D\in L_1} \Exgy\{D(\vec{x},\vec{y})\} - \Exhgy\{D(\hvec{x},\vec{y})\}
\label{eq:W1},
\end{align}

where $L_1$ denotes functions that are 1-Lipschitz with respect to their first argument and $D:\mc{X}\times\mc{Y}\rightarrow \Real$ is a “critic” or “discriminator” that tries to distinguish between true $\vec{x}$ and generated $\hvec{x}$ given $\vec{y}$.
Since we want the method to work for typical values of $\vec{y}$, we define a loss by taking an expectation of (<ref>) over $\vec{y}\sim\py$.
Since the expectation commutes with the supremum in (<ref>),
we have <cit.>

\begin{align}
\Ey\{W_1(\pxgy(\cdot,\vec{y}),\pxhgy(\cdot,\vec{y}))\} &= \sup_{D\in L_1} \Exy\{D(\vec{x},\vec{y})\}  - \Egy\{D(\hvec{x},\vec{y})\} \\
&= \sup_{D\in L_1} \Exzy\{D(\vec{x},\vec{y})  - D(G_{\vec{\theta}}(\vec{z},\vec{y}),\vec{y})\}
\label{eq:EW1} .
\end{align}

In practice, $D_{\vec{\phi}}$ is implemented by a neural network and
$(\vec{\theta},\vec{\phi})$ are trained by alternately minimizing

\begin{align}
\Ladv(\vec{\theta},\vec{\phi})
\defn \Exzy\{D_{\vec{\phi}}(\vec{x},\vec{y}) - D_{\vec{\phi}}(G_{\vec{\theta}}(\vec{z},\vec{y}),\vec{y})\}
\label{eq:Ladv}
\end{align}

with respect to $\vec{\theta}$ and minimizing $-\Ladv(\vec{\theta},\vec{\phi}) + \Lgp(\vec{\phi})$ with respect to $\vec{\phi}$, where $\Lgp(\vec{\phi})$ is a gradient penalty that is used to encourage $D_{\vec{\phi}}\in L_1$ <cit.>.
Also, the expectation over $\vec{x}$ and $\vec{y}$ in (<ref>) is replaced by a sample average over the training examples $\{(\vec{x}_t,\vec{y}_t)\}_{t=1}^T$.

One of the main challenges with the cGAN framework in image recovery problems is that, for each training measurement example $\vec{y}_t$, there is only a single image example $\vec{x}_t$.
Thus, with the previously described training methodology, there is no incentive for the generator to produce diverse samples $G(\vec{z},\vec{y})|_{\vec{z}\sim\pz}$ for a fixed $\vec{y}$.
This can lead to the generator learning to ignore the code vector $\vec{z}$, which causes a form of “mode collapse.”

Although issues with stability and mode collapse are also present in unconditional GANs (uGANs),

the causes of mode-collapse in uGANs or discretely conditioned cGANs <cit.> are fundamentally different than in cGANs like ours.
In the latter case, there is only one example of a valid $\vec{x}_t$ for each given $\vec{y}_t$, whereas in the former there are either many $\vec{x}_t$, or many $\vec{x}_t$ for each conditioning class.
As a result, most strategies used to combat mode-collapse in uGANs <cit.> are not well suited to cGANs.
For example, mini-batch discrimination strategies like MBSD <cit.>

, where the discriminator aims to distinguish a mini-batch of true samples $\{\vec{x}_t\}$ from a mini-batch of generated samples $\{\hvec{x}_t\}$,
don't work with cGANs because the posterior statistics are very different than the prior statistics.

To combat mode collapse in cGANs, Adler & Öktem <cit.> proposed to use a three-input discriminator $\Dadler_{\vec{\phi}}:\mc{X}\times\mc{X}\times\mc{Y}\rightarrow \Real$ and replace $\Ladv$ from (<ref>) with the loss

\begin{align}
\Ladler(\vec{\theta},\vec{\phi})
&\defn \Exzzy\big\{
\tfrac{1}{2}\Dadler_{\vec{\phi}}(\vec{x},G_{\vec{\theta}}(\vec{z}_1,\vec{y}),\vec{y}) + \tfrac{1}{2}\Dadler_{\vec{\phi}}(G_{\vec{\theta}}(\vec{z}_2,\vec{y}),\vec{x},\vec{y})
\nonumber\\&\quad
- \Dadler_{\vec{\phi}}(G_{\vec{\theta}}(\vec{z}_1,\vec{y}),G_{\vec{\theta}}(\vec{z}_2,\vec{y}),\vec{y})
\big\}
\label{eq:Ladler},
\end{align}

which rewards variation between the first and second inputs to $\Dadler_{\vec{\phi}}$.
They then proved that minimizing $\Ladler$ in place of $\Ladv$ does not compromise the Wasserstein cGAN objective, i.e., $\arg\min_{\vec{\theta}} \Ladler(\vec{\theta},\vec{\phi}) = \arg\min_{\vec{\theta}} \Ladv(\vec{\theta},\vec{\phi})$.
As we show in experiments, this approach does avoid complete mode collapse, but it leaves much room for improvement.

§ PROPOSED METHOD

§.§ Proposed regularization: supervised-L1 plus SD reward

We now propose a novel cGAN regularization framework.
To train the generator, we propose to solve

\begin{align}\textstyle
\arg\min_{\vec{\theta}} \{\betaadv\Ladv(\vec{\theta},\vec{\phi})+\LonestdPt(\vec{\theta},\betastd)\}
\end{align}

with appropriately chosen $\betaadv,\betastd>0$ and $P\train\geq 2$, where the regularizer

\begin{align}\textstyle
\LonestdPt(\vec{\theta},\betastd) \defn \LonePt(\vec{\theta}) - \betastd\LstdPt(\vec{\theta})
\label{eq:LonestdP}
\end{align}

is constructed from the $P\train$-sample supervised-$\ell_1$ loss and standard-deviation (SD) reward terms

\begin{align}
\LonePt(\vec{\theta})
&\defn \textstyle \ExzPy\big\{ \|\vec{x}-\hvec{x}\avgPt\|_1\big\}
\label{eq:LoneP} \\
\LstdPt(\vec{\theta})
&\defn \textstyle \sqrt{\frac{\pi}{2P\train(P\train-1)}} \sum_{i=1}^{P\train} \EzPy\big\{ \|\hvec{x}_i-\hvec{x}\avgPt\|_1 \big\}
\label{eq:LstdP} ,
\end{align}

and where $\{\hvec{x}_i\}$ denote the generated samples and $\hvec{x}\avgP$ their $P$-sample average:

\begin{align}
\hvec{x}_i \defn G_{\vec{\theta}}(\vec{z}_i,\vec{y}),
\qquad
\hvec{x}\avgP \defn \textstyle \frac{1}{P} \sum_{i=1}^P \hvec{x}_i
\label{eq:xhat}.
\end{align}

The use of supervised-$\ell_1$ loss and SD reward is not heuristic.
As shown in L1std, it encourages the samples $\{\hvec{x}_i\}$ to match the true posterior in both mean and covariance.

Say $P\train\geq 2$ and $\vec{\theta}$ has complete control over the $\vec{y}$-conditional mean and covariance of $\hvec{x}_i$.
Then the parameters
$\vec{\theta}_* = \arg\min_{\vec{\theta}} \LonestdPt(\vec{\theta},\betastdG)$ with

\begin{align}
\textstyle \betastdG \defn \sqrt{\frac{2}{\pi P\train(P\train+1)}}
\label{eq:betastdG}
\end{align}

yield generated statistics

\begin{align}
\Ezigy\{\hvec{x}_i(\vec{\theta}_*)|\vec{y}\} &= \Exgy\{\vec{x}|\vec{y}\} = \hvec{x}\mmse \\
\Covzigy\{\hvec{x}_i(\vec{\theta}_*)|\vec{y}\} &= \Covxgy\{\vec{x}|\vec{y}\}
\end{align}

when the elements of $\hvec{x}_i$ and $\vec{x}$ are independent Gaussian conditioned on $\vec{y}$.
Thus, minimizing $\LonestdPt$ encourages the $\vec{y}$-conditional mean and covariance of $\hvec{x}_i$ to match those of the true $\vec{x}$.

See L1std for a proof.
In imaging applications, $\hvec{x}_i$ and $\vec{x}$ may not be independent Gaussian conditioned on $\vec{y}$, and so the $\betastd$ in (<ref>) may not be appropriate.
Thus we propose a method to automatically tune $\betastd$ in betastd.

illustration shows a toy example with
parameters $\vec{\theta}=[\mu,\sigma]\tran$,
generator $G_{\vec{\theta}}(z,y)=\mu+\sigma z$,
and $z\sim\mc{N}(0,1)$, giving generated posterior
Assuming true $\pxgy(x|y)=\mc{N}(x;\mu_0,\sigma_0^2)$,
illustration(a)-(b) show that, by minimizing the proposed $\LonestdPt(\vec{\theta},\betastdG)$ regularization over $\vec{\theta}=[\mu,\sigma]\tran$ for any $P\train\geq 2$, we recover the true $\vec{\theta}_0=[\mu_0,\sigma_0]\tran$.
They also the cost function steepens as $P\train$ decreases, with agrees with our finding that $P\train=2$ works best in practice.

The contours show the regularizer value versus $\vec{\theta}\!=\![\mu,\sigma]\tran$ for four different regularizers:
(a) supervised-$\ell_1$ plus SD reward with $\betastd\!=\!\betastdG$ at $P\train\!=\!2$,
(b) supervised-$\ell_1$ plus SD reward with $\betastd\!=\!\betastdG$ at $P\train\!=\!8$,
(c) supervised-$\ell_2$ at $P\train\!=\!8$,
(d) supervised-$\ell_2$ plus variance reward at $P\train\!=\!8$.
The red star shows the true posterior parameters $[\mu_0,\sigma_0]\tran$.

We note that regularizing a cGAN with supervised-$\ell_1$ loss alone is not new; see, e.g., <cit.>.

In fact, the use of supervised-$\ell_1$ loss is often preferred over $\ell_2$ in image recovery because it results in sharper, more visually pleasing results <cit.>.
But regularizing a cGAN using supervised-$\ell_1$ loss alone can push the generator towards mode collapse, for reasons described below.

For example, in <cit.>, $\ell_1$-induced mode collapse led the authors to use dropout to induce generator variation, instead of random codes.

Why not supervised-L2 regularization?
One may wonder: Why regularize using supervised-$\ell_1$ loss plus an SD reward in (<ref>) and not a more conventional choice like supervised-$\ell_2$ loss plus a variance reward, or even supervised-$\ell_2$ loss alone?
We start by discussing the latter.

The use of supervised-$\ell_2$ regularization in a cGAN can be found in <cit.>.
In this case, to train the generator, one would like to solve $\arg\min_{\vec{\theta}} \{\Ladv(\vec{\theta},\vec{\phi}) + \lambda \Ltwo(\vec{\theta})\}$ with some $\lambda>0$ and

\begin{align}
\Ltwo(\vec{\theta})
&\defn \Exy\big\{ \|\vec{x}-\Ez\{G_{\vec{\theta}}(\vec{z},\vec{y})\}\|^2_2 \big\}
\label{eq:Ltwo} .
\end{align}

Ohayon et al. <cit.> revived this idea for the explicit purpose of fighting mode collapse.

In practice, however, the $\Ez$ term in (<ref>) must be implemented by a finite-sample, giving

\begin{align}
\LtwoPt(\vec{\theta})
&\defn \textstyle \ExzPy\big\{ \big\|\vec{x}-\frac{1}{P\train}\sum_{i=1}^{P\train} G_{\vec{\theta}}(\vec{z}_i,\vec{y}) \big\|^2_2 \big\}
\label{eq:LtwoP}
\end{align}

for some $P\train\geq 2$.
For example, Ohayon's implementation <cit.> used $P\train=8$.
As we show in L2P, $\LtwoPt$ induces mode collapse rather than prevents it.

Say $P\train$ is finite and $\vec{\theta}$ has complete control over the $\vec{y}$-conditional mean and covariance of $\hvec{x}_i$.
Then the parameters
$\vec{\theta}_* = \arg\min_{\vec{\theta}} \LtwoPt(\vec{\theta})$
yield generated statistics

\begin{align}
\Ezigy\{\hvec{x}_i(\vec{\theta}_*)|\vec{y}\} &= \Exgy\{\vec{x}|\vec{y}\} = \hvec{x}\mmse \\
\Covzigy\{\hvec{x}_i(\vec{\theta}_*)|\vec{y}\} &= \vec{0} .
\end{align}

Thus, minimizing $\LtwoPt$ encourages mode collapse.

The proof (see L2P) follows from the bias-variance decomposition of (<ref>), i.e.,

\begin{align}
&= \textstyle \Ey\big\{
\|\hvec{x}\mmse-\Ezigy\{\hvec{x}_i(\vec{\theta})|\vec{y}\}\|_2^2
+ \tfrac{1}{P\train}
\tr [ \Covzigy\{\hvec{x}_i(\vec{\theta}) | \vec{y}\} ]
+ \Exgy\{\|\vec{e}\mmse\|_2^2 | \vec{y}\}
\big\}
\label{eq:LtwoPerrt} ,
\end{align}

where $\vec{e}\mmse\defn \vec{x}-\vec{x}\mmse$ is the MMSE error.
illustration(c) shows that $\LtwoPt$ regularization causes mode collapse in the toy example, and MRIexperiments shows that it causes mode collapse in MRI.

Why not supervised $\ell_2$ plus a variance reward?
To mitigate $\LtwoPt$'s incentive for mode-collapse, the second term in (<ref>) could be canceled using a variance reward, giving

\begin{align}
% HARDWIRING \betavar = 1/P\train
\LtwovarPt(\vec{\theta})
&\defn \textstyle \LtwoPt(\vec{\theta}) - \frac{1}{P\train}\LvarPt(\vec{\theta})
\label{eq:LtwovarP} \\
\text{with~}\LvarPt(\vec{\theta})
&\defn \textstyle \frac{1}{P\train-1}\sum_{i=1}^{P\train} \EzPy\{\|\hvec{x}_i(\vec{\theta})-\hvec{x}\avgP(\vec{\theta})\|_2^2\}
\label{eq:LvarP} .
\end{align}

since LvarPerr shows that $\LvarPt(\vec{\theta})$ is an unbiased estimator of the posterior trace-covariance:

\begin{align}
\LvarPt(\vec{\theta})
%&= \textstyle \Ey\big\{ \Ezigy\{\|\vec{d}_i\|_2^2 | \vec{y}\} \big\}
&= \textstyle \Ey \{ \tr [ \Covzigy\{\hvec{x}_i(\vec{\theta}) | \vec{y}\} ] \}
\text{~~for any~} P\train\geq 2
\label{eq:LvarPerr}.
\end{align}

However, the resulting $\LtwovarPt$ regularizer in (<ref>) does not encourage the generated covariance to match the true posterior covariance, unlike the proposed $\LonestdPt$ regularizer in (<ref>) (recall L1std).
For the toy example, this behavior is confirmed in illustration(d).

§.§ Auto-tuning of SD reward weight betastd

$\betastd=0$     $\betastd=\betastdG$     $\betastd=1.2\betastdG$     $\betastd=1.4\betastdG$     $\betastd=1.6\betastdG$     $\betastd=1.8\betastdG$

Number of averaged outputs, $P$, on a log scale
Example PSNR of $\hvec{x}\avgP$ versus $P$, the number of averaged outputs, for several values of $\betastd$ used during training. Also shown is the theoretical behavior for true-posterior samples.

We now propose a method to auto-tune $\betastd$ in (<ref>) for a given training dataset.
Our approach is based on the observation that larger values of $\betastd$ tend to yield samples $\hvec{x}_i$ with more variation.
But more variation is not necessarily better; we want samples with the correct amount of variation.
To assess the correct amount of variation, we compare the expected $\ell_2$ error of the $P$-sample average $\hvec{x}\avgP$ to that of $\hvec{x}\avgone$.

When $\{\hvec{x}_i\}$ are true posterior samples, these errors follow a particular relationship, as established by ratio below (see ratio for a proof).

Say $\hvec{x}_i\sim \pxgy(\cdot|\vec{y})$ are independent samples of the true posterior and, for any $P\geq 1$, their $P$-sample average is $\hvec{x}\avgP\defn \frac{1}{P}\sum_{i=1}^P\hvec{x}_i$.

\begin{align}
\mc{E}_P
&\defn \E\{\|\hvec{x}\avgP-\vec{x}\|_2^2 | \vec{y}\}
= \tfrac{P+1}{P} \mc{E}\mmse
\text{,\quad which implies\quad}
\tfrac{\mc{E}_1}{\mc{E}_P} = \tfrac{2P}{P+1}
\label{eq:Eratio}.
\end{align}

Experimentally we find that $\mc{E}_1/\mc{E}_P$ grows with the SD reward weight $\betastd$.
(See betastd_ablation.)
Thus, we propose to adjust $\betastd$ so that the observed SNR-gain ratio $\mc{E}_1/\mc{E}_{P\val}$ matches the correct ratio $(2P\val)/(P\val+1)$ from (<ref>), for some $P\val\geq 2$ that does not need to match $P\train$.
(We use $P\val=8$ in experiments.)
In particular, at each training epoch $\tau$, we approximate $\mc{E}_{P\val}$ and $\mc{E}_1$

as follows:

\begin{align}
\hat{\mc{E}}_{P\val,\tau}
&\defn \textstyle \frac{1}{V}\sum_{v=1}^V \|\frac{1}{P\val}\sum_{i=1}^{P\val} G_{\vec{\theta}_\tau}(\vec{z}_{i,v},\vec{y}_v)-\vec{x}_v\|_2^2 \\
\hat{\mc{E}}_{1,\tau}
&\defn \textstyle \frac{1}{V}\sum_{v=1}^V \|G_{\vec{\theta}_\tau}(\vec{z}_{1,v},\vec{y}_v)-\vec{x}_v\|_2^2 ,
\end{align}

with validation set $\{(\vec{x}_v,\vec{y}_v)\}_{v=1}^V$ and
i.i.d. codes $\{\vec{z}_{i,v}\}_{i=1}^{P\val}$.

We update $\betastd$ using gradient descent:

\begin{align}
\beta_{\mathsf{SD},\tau+1}
&= \textstyle \beta_{\mathsf{SD},\tau} - \mustd \cdot \big(
%- \big[\frac{2P\val}{P\val+1}\big]_{\text{dB}}
- [2P\val/(P\val+1)]_{\text{dB}}
\big) \betastdG
\label{eq:betastd}
\end{align}

$\beta_{\mathsf{SD},0} = \betastdG$,
some $\mustd>0$,
and $[x]_{\text{dB}}\defn 10\log_{10}(x)$.

§ NUMERICAL EXPERIMENTS

§.§ Conditional Fréchet Inception Distance

As previously stated, our goal is to train a generator $G_{\vec{\theta}}$ so that, for typical fixed values of $\vec{y}$, the generated distribution $\pxhgy(\cdot|\vec{y})$ matches the true posterior $\pxgy(\cdot|\vec{y})$.
It is essential to have a quantitative metric for evaluating performance with respect to this goal.
For example, it is not enough that the generated samples are “accurate” in the sense of being close to the ground truth, nor is it enough that they are “diverse” according to some heuristically chosen metric.

We quantify posterior-approximation quality using the conditional Fréchet inception distance (CFID) <cit.>, a computationally efficient approximation to the conditional Wasserstein distance

\begin{align}
\CWD \defn \Ey\{W_2(\pxgy(\cdot,\vec{y}),\pxhgy(\cdot,\vec{y}))\}
\label{eq:CWD} .
\end{align}

In (<ref>), $W_2(\pa,\pb)$ denotes the Wasserstein-2 distance between distributions $\pa$ and $\pb$, defined as

\begin{align}
\defn \min_{\pab\in \Pi(\pa,\pb)} \Eab\{\|\vec{a}-\vec{b}\|^2_2\}
\label{eq:Pi} ,
\end{align}

where $\Pi(\pa,\pb) \defn \textstyle \big\{\pab : \pa=\int \pab\deriv \vec{b} \text{~and~} \pb=\int \pab \deriv \vec{a} \big\} $ denotes the set of joint distributions $\pab$ with prescribed marginals $\pa$ and $\pb$.
Similar to how FID <cit.>—a popular uGAN metric—is computed, CFID approximates CWD (<ref>) as follows:
i) the random vectors $\vec{x}$, $\hvec{x}$, and $\vec{y}$ are replaced by low-dimensional embeddings $\uvec{x}$, $\huvec{x}$, and $\uvec{y}$, generated by the convolutional layers of a deep network, and
ii) the embedding distributions $\pxugyu$ and $\pxuhgyu$ are approximated by multivariate Gaussians.
More details on CFID are given in cfid.

§.§ MRI Experiments

We consider parallel MRI recovery, where the goal is to recover a complex-valued multicoil image $\vec{x}$ from zero-filled measurements $\vec{y}$ (see MRI for details).

For training data $\{\vec{x}_t\}$, we used the first 8 slices of all fastMRI <cit.> T2 brain training volumes with at least 8 coils, cropping to $384 \times 384$ pixels and compressing to 8 virtual coils <cit.>, yielding 12 200 training images.
Using the same procedure, 2 376 testing and 784 validation images were obtained from the fastMRI T2 brain testing volumes.
From the 2 376 testing images, 72 were randomly selected to evaluate the Langevin technique <cit.>, in order to limit its sample generation to 6 days.

To create the measurement $\vec{y}_t$, we transformed $\vec{x}_t$ to the Fourier domain, sampled using the pseudo-random GRO pattern <cit.> at acceleration $R=8$, and Fourier-transformed the zero-filled k-space measurements back to the (complex, multicoil) image domain.

We use a UNet <cit.> for our generator and a standard CNN for our generator, along with data-consistency as in consistency.
Architecture and training details are given in implementation.

We compare our cGAN to the Adler and Öktem's cGAN <cit.>, Ohayon et al.'s cGAN <cit.>, Jalal et al.'s fastMRI Langevin approach <cit.>, and the UNet trained to minimize the weighted combination of supervised $\ell_1$ loss and negative SSIM from <cit.>.
The cGAN from <cit.> uses generator loss $\betaadv\Ladler(\vec{\theta},\vec{\phi})$ and discriminator loss $-\Ladler(\vec{\theta},\vec{\phi})+ \alpha_1 \Lgp(\vec{\phi}) + \alpha_2 \Ldrift(\vec{\phi})$,
while the cGAN from <cit.> uses generator loss $\betaadv\Ladv(\vec{\theta},\vec{\phi}) + \LtwoP(\vec{\theta})$ and discriminator loss $-\Ladv(\vec{\theta},\vec{\phi})+ \alpha_1 \Lgp(\vec{\phi}) + \alpha_2 \Ldrift(\vec{\phi})$.
Each used the value of $\betaadv$ specified in the original paper.
All cGANs used the same generator and discriminator architectures, except that <cit.> used extra discriminator input channels to facilitate the 3-input loss (<ref>).
For the fastMRI Langevin approach <cit.>, we did not modify the authors' implementation in <cit.> except to use the GRO sampling mask.

To evaluate performance, we converted the multicoil outputs $\hvec{x}_i$ to complex-valued images using SENSE-based coil combining <cit.> with ESPIRiT-estimated <cit.> coil sensitivity maps, as described in MRI.

Magnitude images were used to compute performance metrics.

Average results for $R\!=\!8$ MRI.

CFID$^1$, FID, and APSD used $72$ test samples and $P\!=\!32$,
CFID$^2$ used $2\,376$ test samples and $P\!=\!8$, and
CFID$^3$ used all $14\,576$ samples and $P\!=\!1$

Model     CFID$^1\!\!\downarrow$     CFID$^2\!\!\downarrow$     CFID$^3\!\!\downarrow$     FID$\downarrow$

APSD     Time (4)$\downarrow$

UNet     6.30     4.68     4.37     11.53

0     217 ms

Langevin (Jalal et al. <cit.>)     7.34     -     -     14.32

7.6e-6     14 min

cGAN (Adler & Öktem <cit.>)     10.10     6.30     5.72     10.77

7.7e-6     217 ms

cGAN (Ohayon et al. <cit.>)     6.04     4.59     4.27     11.05

7.7e-7     217 ms

cGAN (Ours)     4.87     2.23     1.79     7.72

7.6e-6     217 ms

mriresults shows CFID, FID, APSD $\defn(\frac{1}{NP}\sum_{i=1}^{P} \|\hvec{x}\avgP-\hvec{x}_i\|^2)^{1/2}$, and 4-sample generation time.
(C)FID was computed using VGG-16 (not Inception-v3) to better align with radiologists' perceptions <cit.>.

As previously described, we evaluated the Langevin method using only 72 test images.
But, because CFID is biased at small sample sizes <cit.>, we re-evaluated the other methods using all $2\,376$ test images, and again using all $14\,576$ training and test images.
mriresults shows that our approach gave significantly better CFID and FID than the competitors.

Furthermore, the APSD of Ohayon et al.'s cGAN was an order-of-magnitude smaller than the others, indicating mode collapse.

The UNet and cGANs generated samples $3\,800$ times faster than the Langevin approach from <cit.>.
CFID and FID results at acceleration $R=4$, presented in additionalresults, show similar trends.

versusP shows PSNR, SSIM, LPIPS <cit.>, and DISTS <cit.> for the $P$-sample average $\hvec{x}\avgP$ at $P\in\{1,2,4,8,16,32\}$.
The best performance is achieved by the proposed cGAN with $P=32$ for PSNR, $P=8$ for SSIM, and $P=2$ for both LPIPS and DISTS.
These results can be explained by the perception-distortion tradeoff <cit.>: as $P$ increases, $\hvec{x}\avgP$ transitions from better perceptual quality to lower $\ell_2$ distortion.
PSNR favors $P\rightarrow\infty$ (e.g., $\ell_2$ optimality) while the other metrics favor particular combinations of perceptual quality and distortion.
The plots in mriplots show zoomed-in versions of $\hvec{x}\avgP$ that visually demonstrate the perception-distortion tradeoff at $P\in\{1,2,4,32\}$: smaller $P$ yield sharper images with more variability, while larger $P$ yield smoother reconstructions.

Average PSNR, SSIM, LPIPS, and DISTS of $\hvec{x}\avgP$ versus $P$ for $R=8$ MRI

(lr)2-7 (lr)8-13
Model     $P\!=$1     $P\!=$2     $P\!=$4     $P\!=$8     $P\!=$16     $P\!=$32
$P\!=$1     $P\!=$2     $P\!=$4     $P\!=$8     $P\!=$16     $P\!=$32

UNet     34.77     34.77     34.77     34.77     34.77     34.77
0.9228     0.9228     0.9228     0.9228     0.9228     0.9228

Langevin (Jalal et al. <cit.>)
32.17     32.83     33.45     33.74     33.83     33.90
0.8725     0.8919     0.9031     0.9091     0.9120     0.9137

cGAN (Adler & Öktem <cit.>)
31.31     32.31     32.92     33.26     33.42     33.51
0.8865     0.9045     0.9103     0.9111     0.9102     0.9095

cGAN (Ohayon et al. <cit.>)
34.89     34.90     34.90     34.90     34.91     34.92
0.9222     0.9217     0.9213     0.9211     0.9211     0.9210

cGAN (Ours)
32.32     33.67     34.53     35.01     35.27     35.42
0.9030     0.9199     0.9252     0.9257     0.9251     0.9246

(lr)2-7 (lr)8-13
Model     $P\!=$1     $P\!=$2     $P\!=$4     $P\!=$8     $P\!=$16     $P\!=$32
$P\!=$1     $P\!=$2     $P\!=$4     $P\!=$8     $P\!=$16     $P\!=$32

UNet     0.0586     0.0586     0.0586     0.0586     0.0586     0.0586
0.1120     0.1120     0.1120     0.1120     0.1120     0.1120

Langevin (Jalal et al. <cit.>)
0.0769     0.0619     0.0579     0.0589     0.0611     0.0611
0.1341     0.1136     0.1086     0.1119     0.1175     0.1212

cGAN (Adler & Öktem <cit.>)
0.0698     0.0614     0.0623     0.0667     0.0704     0.0727
0.1407     0.1262     0.1252     0.1291     0.1334     0.1361

cGAN (Ohayon et al. <cit.>)
0.0532     0.0536     0.0539     0.0540     0.0534     0.0540
0.1128     0.1143     0.1151     0.1155     0.1157     0.1158

cGAN (Ours)
0.0418     0.0379     0.0421     0.0476     0.0516     0.0539
0.0906     0.0877     0.0965     0.1063     0.1135     0.1177

avgsampsamp_19 shows zoomed-in versions of two posterior samples and $\hvec{x}\avgP$ at $P=32$.
The posterior samples show meaningful variations for the proposed method, essentially no variation for Ohayon et al.'s cGAN, and vertical or horizontal artifacts for Adler & Öktem's cGAN and the Langevin method, respectively.
The $\hvec{x}\avgP$ plots show that these artifacts are mostly suppressed by sample averaging.

avgerrstd_19 shows examples of $\hvec{x}\avgP$, along with the corresponding pixel-wise absolute errors $|\hvec{x}\avgP-\vec{x}|$ and pixel-wise SD $(\frac{1}{P}\sum_{i=1}^P (\hvec{x}\avgP-\hvec{x}_i)^2)^{1/2}$, for $P=32$.
The absolute-error image for the Langevin technique looks more diffuse than those of the other methods in the brain region.
The fact that it is brighter in the air region (i.e., near the edges) is a consequence of minor differences in sensitivity-map estimation.
The pixel-wise SD maps show a lack of variability for the UNet, which does not generate posterior samples, as well as Ohayon et al.'s cGAN, due to mode collapse.
The Langevin pixel-wise SD maps show localized hot-spots that appear to be reconstruction artifacts.
inpaintingplots shows other example MRI recoveries with zoomed pixel-wise SD maps.
In mriapp_11 and mriapp_9, Langevin recovery also shows strong artifacts, highlighted by red arrows.

§.§ Inpainting Experiments

Here our goal is to complete a large missing square in a face image.

We used $256 \times 256$ CelebA-HQ face images <cit.> and a centered $128\times 128$ missing square.
We randomly split the dataset, yielding $27\,000$ training, $2\,000$ validation, and $1\,000$ testing images.

For our cGAN, we use the CoModGAN generator and discriminator from <cit.> with our proposed $\LonestdPt$ regularization. Unlike <cit.>, we do not use MBSD <cit.> at the discriminator.

Average results for inpainting:
FID was
computed from $1\,000$ test images with $P\!=\!32$,
while CFID was computed from all 30 000 images with $P\!=\!1$

Model     CFID$\downarrow$     FID$\downarrow$

Time (128)$\downarrow$

Score-SDE (Song et al. <cit.>)     5.11     7.92

48 min

CoModGAN (Zhao et al. <cit.>)     5.29     8.50

217 ms

cGAN (ours)     4.69     7.45

217 ms

We use the same general training and testing procedure described in MRIexperiments, but with $\betaadv=$ 5e-3, $n\batch=100$, and 110 epochs of cGAN training.
Running PyTorch on a server with 4 Tesla A100 GPUs, each with 82 GB of memory, the cGAN training takes approximately 2 days.
FID was evaluated on $1\,000$ test images using $P\!=\!32$ samples per measurement.
To avoid the bias that would result from evaluating CFID on only $1\,000$ images (see CFID), CFID was evaluated on all $30\,000$ images with $P=1$.

We compare with the state-of-the-art CoModGAN <cit.>
and Score-based SDE <cit.> approaches.
For CoModGAN, we use the PyTorch implementation from <cit.>.
CoModGAN differs from our cGAN only in its use of MBSD and lack of $\LonestdPt$ regularization.
For Song et al.'s SDE, we use the authors' implementation from <cit.> with their pretrained weights and the settings they suggested for the $256\times256$ CelebA-HQ dataset.

inpainting256 shows test CFID, FID, and 128-sample generation time.
The table shows that our approach gave the best CFID and FID, and that the cGANs generated samples 13 000 times faster than the score-based method.
inpainting_22 shows an example of five generated samples for the three methods under test.
The samples are all quite good, although a few generated by CoModGAN and the score-based technique have minor artifacts.
Some samples generated by our technique show almond-shaped eyes, demonstrating fairness.
Additional examples are given in inpaintingplots.

§ CONCLUSION

We propose a novel regularization framework for image-recovery cGANs that consists of supervised-$\ell_1$ loss plus an appropriately weighted standard-deviation reward.

For the case of an independent Gaussian posterior, we proved that our regularization yields generated samples that agree with the true-posterior samples in both mean and covariance.
We also established limitations for alternatives based on supervised-$\ell_2$ regularization with or without a variance reward.
For practical datasets, we proposed a method to auto-tune our standard-deviation reward weight.

Experiments on parallel MRI and large-scale face inpainting showed our proposed method outperforming all cGAN and score-based competitors in CFID, which measures posterior-approximation quality, as well as other metrics like FID, PSNR, SSIM, LPIPS, and DISTS.
Furthermore, it generates samples thousands of times faster than Langevin/score-based approaches.

Limitations.  In our current implementation, our cGAN is trained for a specific image-recovery task and may not generalize to new tasks.
In ongoing work, we are extending our approach so that it can be trained to handle a wide range of recovery tasks, such as MRI with a generic acceleration and sampling mask, or inpainting with a generic mask.

We are also extending our approach to other applications, such as super-resolution, deblurring, compressive sensing, denoising, and phase retrieval.

The authors are funded in part by the National Institutes of Health under grant R01-EB029957 and the National Science Foundation under grant CCF-1955587.

Supplementary Materials

§ POSTERIOR SAMPLES FACILITATE CALIBRATED/FAIR DETECTION

Say that $z\in\{1,0\}$ denotes the presence or absence of a particular pathology (e.g., brain tumor), and say that we train a soft classifier $c(\cdot)$ on ground-truth images $\vec{x}$ and calibrate it such that

\begin{align}
&= \Pr\{z=1|\vec{x}\}
\label{eq:classifier}.
\end{align}

Now say that we observe a distorted/corrupted/incomplete measurement $\vec{y}=\mc{M}(\vec{x})$.
We would like to infer the probability that the pathology is present given $\vec{y}$, i.e., compute $\Pr\{z=1|\vec{y}\}$.
Note that

\begin{align}
\Pr\{z=1|\vec{y}\}
&= \int \Pr\{z=1,\vec{x}|\vec{y}\} \deriv\vec{x}
= \int \Pr\{z=1|\vec{x},\vec{y}\} p(\vec{x}|\vec{y}) \deriv\vec{x} \\
&= \int \Pr\{z=1|\vec{x}\} p(\vec{x}|\vec{y}) \deriv\vec{x}
= \int c(\vec{x}) p(\vec{x}|\vec{y}) \deriv\vec{x}
= \E\{c(\vec{x})|\vec{y}\} \\
&= \lim_{P\rightarrow\infty} \frac{1}{P}\sum_{i=1}^P c(\hvec{x}_i)
\text{~~for~~} \hvec{x}_i\sim\text{i.i.d.~} p(\vec{x}|\vec{y})
\label{eq:Pz|y},
\end{align}

where the last equality follows from the law of large numbers.
So, given access to many independent posterior samples $\{\hvec{x}_i\}$, equation (<ref>) says that we can simply plug them into our calibrated classifier $c(\cdot)$ and average the result to compute $\Pr\{z=1|\vec{y}\}$.

Conversely, if we have access to only the posterior mean $\hvec{x}\mmse = \E\{\vec{x}|\vec{y}\}$, then because

\begin{align}
\Pr\{z=1|\vec{y}\}
= \E\{c(\vec{x})|\vec{y}\}
\neq c(\E\{\vec{x}|\vec{y}\})
= c(\hvec{x}\mmse)
\end{align}

for any non-linear $c(\cdot)$, the plug-in probability estimate will be incorrect.
In fact, there exists no point estimate $\hvec{x}$ that gives the correct $\Pr\{z=1|\vec{y}\}$ for general $c(\cdot)$.

Although above we defined $c(\cdot)$ as a (soft) binary pathology classifier, the same results hold if we define $c(\cdot)$ as a K-ary classifier of any protected attribute, such as race, gender, etc.
This implies that, if we have a machine-learning system that has been calibrated to classify fairly on clean ground-truth data $\vec{x}$, then the use of posterior samples $\{\hvec{x}_i\}$ enables it to classify fairly on distorted/corrupted/incomplete measurements $\vec{y}=\mc{M}(\vec{x})$, whereas the use of generic point-estimates $\hvec{x}$ does not.

§ PROOF OF L1STD

Here we prove L1std.
To begin, for an $N$-pixel image, we rewrite (<ref>)-(<ref>) as

\begin{align}
\LoneP(\vec{\theta})
&= \textstyle \sum_{j=1}^N \Ey\big\{
\ExzPgy\big\{ |x_j-\frac{1}{P}\sum_{i=1}^P \hat{x}_{ij}|~\big|\vec{y}\big\}
\big\}
\label{eq:LonePB}\\
\LstdP(\vec{\theta})
&= \textstyle \sum_{j=1}^N \Ey\big\{
\EzPgy\big\{ \frac{\gamma_P}{P}\sum_{i=1}^P |\hat{x}_{ij}-\frac{1}{P}\sum_{k=1}^P \hat{x}_{kj}| ~\big|\vec{y}\big\}
\big\}
\label{eq:LstdPB} ,
\end{align}

where $x_j\defn[\vec{x}]_j$, $\hat{x}_{ij}\defn[\hvec{x}_i]_j$, and

\begin{align}\textstyle
\gamma_P \defn \sqrt{\frac{\pi P}{2(P-1)}}
\label{eq:gammaP} .
\end{align}

To simplify the notation in the sequel, we will consider an arbitrary fixed value of $j$ and use the abbreviations

\begin{align}
x_j &\rightarrow X,
&\hat{x}_{ij} \rightarrow \hat{X}_i
\label{eq:abbrev} .
\end{align}

Recall that $\vec{x}$ and $\{\hvec{x}_i\}$ are mutually independent when conditioned on $\vec{y}$ because the code vectors $\{\vec{z}_i\}$ are generated independently of both $\vec{x}$ and $\vec{y}$.
In the context of L1std, we also assume that the vector elements $x_j$ and $\hat{x}_{ij}$ are independent Gaussian when conditioned on $\vec{y}$.
This implies that we can make the notational shift

\begin{align}
\pxgy(x_j|\vec{y}) &\rightarrow \mc{N}(X;\mu_0,\sigma_0^2) ,
&\pxhgy(\hat{x}_{ij}|\vec{y}) \rightarrow \mc{N}(\hat{X}_i;\mu,\sigma^2)
\label{eq:gauss} ,
\end{align}

where $X$ and $\{\hat{X}_i\}$ are mutually independent.
With this simplified notation, we note that $[\hvec{x}\mmse]_j\rightarrow \mu_0$, and that mode collapse corresponds to $\sigma=0$.

Furthermore, if $\vec{\theta}$ can completely control $(\mu,\sigma)$, then (<ref>) can be rewritten as

\begin{align} %\textstyle
= \arg\min_{\mu,\sigma} \big\{ \LoneP(\mu,\sigma) - \betastd \LstdP(\mu,\sigma) \big\}
\begin{cases}
{\mu}_*=\mu_0 \\
\end{cases}
\label{eq:goalB}
\end{align}

\begin{align}
\LoneP(\mu,\sigma)
&= \textstyle \ExxP\{|X-\frac{1}{P}\sum_{i=1}^P \hat{X}_i|\}
\label{eq:LonePC}\\
\LstdP(\mu,\sigma)
&= \textstyle \ExP\{\gamma_P \sum_{i=1}^P |\hat{X}_i-\frac{1}{P}\sum_{k=1}^P \hat{X}_k|\}
\label{eq:LstdPC} .
\end{align}

Although $\sigma_*$ must be positive, it turns out that we do not need to enforce this in the optimization (<ref>) because it will arise naturally.

To further analyze (<ref>) and (<ref>), we define

\begin{align}
\hat{\mu} &\defn \textstyle \frac{1}{P}\sum_{i=1}^P \hat{X}_i\\
\hat{\sigma} &\defn \textstyle \frac{\gamma_P}{P} \sum_{i=1}^P |\hat{X}_i - \hat{\mu}|
\label{eq:sigmahat} .
\end{align}

The quantity $\hat{\mu}$ can be recognized as the unbiased estimate of the mean $\mu$ of $\hat{X}_i$, and we now show that $\hat{\sigma}$ is an unbiased estimate of the SD $\sigma$ of $\hat{X}_i$ in the case that $\hat{X}_i$ is Gaussian.
To see this, first observe that the i.i.d. $\mc{N}(\mu,\sigma^2)$ property of $\{\hat{X}_i\}$ implies that $\hat{X}_i-\hat{\mu}=(1-\frac{1}{P})\hat{X}_i-\frac{1}{P}\sum_{k\neq i}\hat{X}_k$ is Gaussian with mean zero and variance $(1-\frac{1}{P})^2\sigma^2 + \frac{P-1}{P^2}\sigma^2 = \frac{P-1}{P}\sigma^2$.
The variable $|\hat{X}_i-\hat{\mu}|$ is thus half-normal distributed with mean $\sqrt{\frac{2(P-1)}{\pi P}\sigma^2}$ <cit.>.
Because $\{\hat{X}_i\}$ are i.i.d., the variable $\frac{1}{P}\sum_{i=1}^P|\hat{X}_i-\hat{\mu}|$ has the same mean as $|\hat{X}_i-\hat{\mu}|$.
Finally, multiplying $\frac{1}{P}\sum_{i=1}^P|\hat{X}_i-\hat{\mu}|$ by $\gamma_P$ yields $\hat{\sigma}$ from (<ref>), and multiplying its mean using the expression for $\gamma_P$ from (<ref>) implies

\begin{align}
\E\{\hat{\sigma}\} = \sigma ,
\end{align}

and so $\hat{\sigma}$ is an unbiased estimator of $\sigma$, the SD of $\hat{X}_i$.

With the above definitions of $\hat{\mu}$ and $\hat{\sigma}$, the optimization cost in (<ref>) can be written as

\begin{align}
\LoneP(\mu,\sigma) - \betastd \LstdP(\mu,\sigma)
&= \ExxP\big\{|X-\hat{\mu}|\big\} - \betastd\ExP\big\{\hat{\sigma}\big\} \\
&= \ExxP\big\{|X-\hat{\mu}|\big\} - \betastd\sigma
\label{eq:JA},
\end{align}

where in the last step we exploited the unbiased property of $\hat{\sigma}$.
To proceed further, we note that the i.i.d. Gaussian property of $\{\hat{X}_i\}$ implies $\hat{\mu}\sim\mc{N}(\mu,\sigma^2/P)$, after which the mutual independence of $\{\hat{X}_i\}$ and $X$ yields

\begin{align}
\sim \mc{N}(\mu_0-\mu,\sigma_0^2 + \sigma^2/P)
\label{eq:prefolded}.
\end{align}

Taking the absolute value of a Gaussian random yields a folded-normal random variable <cit.>.
Using the mean and variance in (<ref>), the expressions in <cit.> yield

\begin{align}
\ExxP\big\{ |X-\hat{\mu}| \big\}
&= \sqrt{\frac{2(\sigma_0^2+\sigma^2/P)}{\pi}}
\exp\Big(-\frac{(\mu_0 - \mu)^2}{2(\sigma_0^2+\sigma^2/P)}\Big)
\nonumber\\&\quad
+ (\mu_0 - \mu)\erf\Big(\frac{\mu_0 - \mu}{\sqrt{2(\sigma_0^2 + \sigma^2/P})}\Big) .
\end{align}

Thus the optimization cost (<ref>) can be written as

\begin{align}
&\defn \sqrt{\frac{2(\sigma_0^2+\sigma^2/P)}{\pi}}
\exp\Big(-\frac{(\mu-\mu_0)^2}{2(\sigma_0^2+\sigma^2/P)}\Big)
\nonumber\\&\quad
+ (\mu-\mu_0)\erf\Big(\frac{\mu-\mu_0}{\sqrt{2(\sigma_0^2 + \sigma^2/P})}\Big)
- \betastd \sigma
\label{eq:JB}.
\end{align}

Since $J(\cdot,\cdot)$ is convex, the minimizer $(\mu_*,\sigma_*)=\arg\min_{\mu,\sigma} J(\mu,\sigma)$ satisfies $\nabla J(\mu_*,\sigma_*)=(0,0)$.
To streamline the derivation, we define

\begin{align}
c &\defn \sqrt{2(\sigma_0^2 + \sigma^2/P)/\pi},
&s \defn \sqrt{\sigma_0^2 + \sigma^2/P}
\end{align}

so that

\begin{align}
&= c \exp\Big(-\frac{(\mu-\mu_0)^2}{2s^2}\Big)
+ (\mu-\mu_0)\erf\Big(\frac{\mu-\mu_0}{\sqrt{2s^2}}\Big)
- \betastd \sigma .
\end{align}

Because $c$ and $s$ are invariant to $\mu$, we get

\begin{align}
\frac{\partial J(\mu,\sigma)}{\partial \mu}
&= -c \exp\Big(\!-\!\frac{(\mu-\mu_0)^2}{2s^2}\Big) \frac{\mu-\mu_0}{s^2}
+ \erf\Big(\frac{\mu-\mu_0}{\sqrt{2s^2}}\Big)
+ (\mu-\mu_0)\frac{2}{\sqrt{\pi}}\exp\Big(\!-\!\frac{(\mu-\mu_0)^2}{2s^2}\Big),
\end{align}

which equals zero if and only if $\mu=\mu_0$.
Thus we have determined that $\mu_*=\mu_0$.
Plugging $\mu_*=\mu_0$ into (<ref>), we find

\begin{align}
&= \sqrt{2(\sigma_0^2+\sigma^2/P)/\pi} - \betastd \sigma
\label{eq:JC}.
\end{align}

Taking the derivative with respect to $\sigma$, we get

\begin{align}
\frac{\partial J(\mu_*,\sigma)}{\partial \sigma}
%&= \textstyle \frac{1}{2}\frac{4\sigma}{\pi P} \Big( \frac{2(\sigma_0^2+\sigma^2/P)}{\pi} \Big)^{-1/2} - \betastd \\
%&= \textstyle \frac{2\sigma}{\pi P} \Big( \frac{2(\sigma_0^2+\sigma^2/P)}{\pi} \Big)^{-1/2} - \betastd\\
%&= \textstyle \Big( \frac{\pi P(P\sigma_0^2+\sigma^2)}{2\sigma^2} \Big)^{-1/2} - \betastd \\
%&= \textstyle \sqrt{ \frac{2\sigma^2}{\pi P(P\sigma_0^2+\sigma^2)} } - \betastd \\
&= \sqrt{ \frac{2}{\pi P(P\sigma_0^2/\sigma^2+1)} } - \betastd \\
&= \sqrt{ \frac{2}{\pi P(P\sigma_0^2/\sigma^2+1)} } - \sqrt{\frac{2}{\pi P (P+1)}} ,
\end{align}

where in the last step we applied the value of $\betastd$ from (<ref>).
It can now be seen that $\frac{\partial J(\mu_*,\sigma)}{\partial \sigma}=0$ if and only if $\sigma=\sigma_0$,
which implies that $\sigma_*=\sigma_0$.
Thus we have established (<ref>), which completes the proof of L1std.

§ DERIVATION OF L2P

Here we prove L2P.
To start, we establish some notation and conditional-mean properties:

\begin{equation}
\begin{array}{r@{~}l@{\qquad}r@{~}l}
\hvec{x}\mmse &\defn \Exgy\{\vec{x}|\vec{y}\} \\
\vec{e}\mmse &\defn \vec{x}-\hvec{x}\mmse, & % fxn of x but not zP
\vec{0} &= \Exgy\{\vec{e}\mmse|\vec{y}\} \\
\hvec{x}_i(\vec{\theta}) &\defn G_{\vec{\theta}}(\vec{z}_i,\vec{y}), &
\ovec{x}(\vec{\theta}) &\defn \Ezigy\{\hvec{x}_i(\vec{\theta})|\vec{y}\} \\ % deterministic
\hvec{x}\avgP(\vec{\theta}) &\defn \textstyle \frac{1}{P} \sum_{i=1}^P \hvec{x}_i(\vec{\theta}), & % fxn of zP but not x
\ovec{x}(\vec{\theta}) &= \EzPgy\{\hvec{x}\avgP(\vec{\theta})|\vec{y}\} \\ % deterministic
\vec{d}_i(\vec{\theta})) &\defn \hvec{x}_i(\vec{\theta})-\ovec{x}(\vec{\theta}), &
\vec{0} &= \Ezigy\{\vec{d}_i(\vec{\theta})|\vec{y}\}~\forall\vec{\theta} \\ % fxn of zP but not x
\vec{d}\avgP(\vec{\theta}) &\defn \textstyle \frac{1}{P} \sum_{i=1}^P \vec{d}_i(\vec{\theta}), &
\vec{0} &= \EzPgy\{\vec{d}\avgP(\vec{\theta})|\vec{y}\}~\forall\vec{\theta} \\ % fxn of zP but not x
\end{array}
\label{eq:quantities}
\end{equation}

Our first step is to write (<ref>) as

\begin{align}
\LtwoP(\vec{\theta})
&= \Ey\big\{ \ExzPgy\{\|\vec{x}-\hvec{x}\avgP(\vec{\theta})\|_2^2 | \vec{y}\} \big\}
\label{eq:LtwoPerr1} .
\end{align}

Leveraging the fact that $\hvec{x}\mmse$ and $\ovec{x}(\vec{\theta})$ are deterministic given $\vec{y}$, we write the inner term in (<ref>) as

\begin{align}
\lefteqn{ \ExzPgy\{\|\vec{x}-\hvec{x}\avgP(\vec{\theta})\|_2^2 | \vec{y}\} }\nonumber\\
&= \ExzPgy\{\|\hvec{x}\mmse+\vec{e}\mmse-\ovec{x}(\vec{\theta})-\vec{d}\avgP(\vec{\theta})\|_2^2 | \vec{y}\} \\
&= \ExzPgy\{\|\hvec{x}\mmse-\ovec{x}(\vec{\theta})\|_2^2 | \vec{y}\}
\nonumber\\&\quad
+ 2\real\ExzPgy\{(\hvec{x}\mmse-\ovec{x}(\vec{\theta}))\herm(\vec{e}\mmse-\vec{d}\avgP(\vec{\theta}))|\vec{y}\}
\nonumber\\&\quad
+ \ExzPgy\{\|\vec{e}\mmse-\vec{d}\avgP(\vec{\theta})\|_2^2 | \vec{y}\} \\
&= \|\hvec{x}\mmse-\ovec{x}(\vec{\theta})\|_2^2
+ 2\real\big[(\hvec{x}\mmse-\ovec{x}(\vec{\theta}))\herm
\underbrace{\ExzPgy\{(\vec{e}\mmse-\vec{d}\avgP(\vec{\theta}))|\vec{y}\}}_{\displaystyle =\vec{0}}\big]
\nonumber\\&\quad
+ \ExzPgy\{\|\vec{e}\mmse-\vec{d}\avgP(\vec{\theta})\|_2^2 | \vec{y}\} \label{eq:Exdiff2a}\\
&= %\|\hvec{x}\mmse-\ovec{x}(\vec{\theta})\|_2^2
\|\hvec{x}\mmse-\Ezigy\{\hvec{x}_i(\vec{\theta})|\vec{y}\}\|_2^2
+ \ExzPgy\{\|\vec{e}\mmse-\vec{d}\avgP(\vec{\theta})\|_2^2 | \vec{y}\}
\label{eq:Exdiff2}.
\end{align}

where in (<ref>) we used the fact that $\vec{d}\avgP$ and $\vec{e}\mmse$ are both zero-mean when conditioned on $\vec{y}$.
We now leverage the fact that $\{\vec{z}_i\}$ are independent of $\vec{x}$ and $\vec{y}$ to write

\begin{align}
\lefteqn{ \ExzPgy\{\|\vec{e}\mmse-\vec{d}\avgP(\vec{\theta})\|_2^2 | \vec{y}\} }\nonumber\\
&= \ExzPgy\{\|\vec{e}\mmse\|_2^2 | \vec{y}\}
+ 2\real\ExzPgy\{\vec{e}\mmse\herm\vec{d}\avgP(\vec{\theta}) | \vec{y}\}
+ \ExzPgy\{\|\vec{d}\avgP(\vec{\theta})\|_2^2 | \vec{y}\} \\
&= \Exgy\{\|\vec{e}\mmse\|_2^2 | \vec{y}\}
+ 2\real\big[{\underbrace{\Exgy\{\vec{e}\mmse | \vec{y}\}}_{\displaystyle =\vec{0}}}\herm
\underbrace{\EzPgy\{\vec{d}\avgP(\vec{\theta}) | \vec{y}\}}_{\displaystyle =\vec{0}}\big] + \EzPgy\{\|\vec{d}\avgP(\vec{\theta})\|_2^2 | \vec{y}\}
\label{eq:Eediff2}.
\end{align}

Finally, we can leverage the fact that $\{\vec{z}_i\}$ are i.i.d. to write

\begin{align}
\EzPgy\{\|\vec{d}\avgP(\vec{\theta})\|_2^2 | \vec{y}\}
&= \textstyle \EzPgy\{\|\frac{1}{P}\sum_{i=1}^P\vec{d}_i(\vec{\theta})\|_2^2 | \vec{y}\} \\
&= \textstyle \frac{1}{P^2} \sum_{i=1}^P \Ezigy\{\|\vec{d}_i(\vec{\theta})\|_2^2 | \vec{y}\} \\
&= \tfrac{1}{P} \Ezigy\{\|\vec{d}_i(\vec{\theta})\|_2^2 | \vec{y}\} \text{~for any $i$} \\
&= \tfrac{1}{P} \Ezigy\{\tr[\vec{d}_i(\vec{\theta})\vec{d}_i(\vec{\theta})\herm] | \vec{y}\} \\
&= \tfrac{1}{P} \tr\big[ \Ezigy\{\vec{d}_i(\vec{\theta})\vec{d}_i(\vec{\theta})\herm\} | \vec{y}\} \big] \\
&= \tfrac{1}{P} \tr\big[ \Covzigy\{\hvec{x}_i(\vec{\theta}) | \vec{y}\} \big]
\label{eq:EeP2} .
\end{align}

Combining (<ref>), (<ref>), (<ref>), and (<ref>), we get the bias-variance decomposition

\begin{align}
\LtwoP(\vec{\theta})
&= \textstyle \Ey\Big\{
\|\hvec{x}\mmse-\Ezigy\{\hvec{x}_i(\vec{\theta})|\vec{y}\}\|_2^2
%+ \tfrac{1}{P} \Ezigy\{\|\vec{d}_i(\vec{\theta})\|_2^2 | \vec{y}\}
+ \tfrac{1}{P} \tr\big[ \Covzigy\{\hvec{x}_i(\vec{\theta}) | \vec{y}\} \big]
+ \Exgy\{\|\vec{e}\mmse\|_2^2 | \vec{y}\}
\Big\}
\label{eq:LtwoPerr} .
\end{align}

We now see that if $\vec{\theta}$ has complete control over the $\vec{y}$-conditional mean and covariance of $\hvec{x}_i(\vec{\theta})$, then minimizing (<ref>) over $\vec{\theta}$ will cause

\begin{align}
\Ezigy\{\hvec{x}_i(\vec{\theta})|\vec{y}\} &= \hvec{x}\mmse \\
\Covzigy\{\hvec{x}_i(\vec{\theta}) | \vec{y}\} &= \vec{0} ,
\end{align}

which proves L2P.

§ DERIVATION OF (<REF>)

To show that the expression for $\LvarP$ in (<ref>) holds, we first rewrite (<ref>) as

\begin{align}
\LvarP(\vec{\theta})
&= \textstyle \frac{1}{P-1}\sum_{i=1}^P \Ey\{ \EzPgy\{\|\hvec{x}_i(\vec{\theta})-\hvec{x}\avgP(\vec{\theta})\|_2^2 | \vec{y}\}
\label{eq:LvarPB}
\end{align}

where the definitions from (<ref>) imply

\begin{align}
\lefteqn{
\EzPgy\{\|\hvec{x}_i(\vec{\theta})-\hvec{x}\avgP(\vec{\theta})\|_2^2 | \vec{y}\}
&= \EzPgy\{\|\ovec{x}(\vec{\theta}) + \vec{d}_i(\vec{\theta}) - \vec{d}\avgP(\vec{\theta}) - \ovec{x}(\vec{\theta})\|_2^2 | \vec{y}\} \\
&= \textstyle \EzPgy\{\|\vec{d}_i(\vec{\theta}) - \frac{1}{P}\sum_{j=1}^P\vec{d}_j(\vec{\theta})\|_2^2 | \vec{y}\} \\
&= \textstyle \EzPgy\{\|(1-\frac{1}{P})\vec{d}_i(\vec{\theta}) - \frac{1}{P}\sum_{j\neq i}\vec{d}_j(\vec{\theta})\|_2^2 | \vec{y}\} \\
&= \textstyle (1-\frac{1}{P})^2 \Ezigy\{\|\vec{d}_i(\vec{\theta})\|_2^2 | \vec{y}\}
+ \frac{P-1}{P^2}\Ezigy\{\|\vec{d}_i(\vec{\theta})\|_2^2 | \vec{y}\}
\label{eq:ExdiffB}\\
&= \textstyle \frac{P-1}{P} \Ezigy\{\|\vec{d}_i(\vec{\theta})\|_2^2 | \vec{y}\} \text{~for any $i$}
\label{eq:ExdiffC}.
\end{align}

For (<ref>), we leveraged the zero-mean and i.i.d. nature of $\{\vec{d}_i(\vec{\theta})\}$ conditioned on $\vec{y}$.
By plugging (<ref>) into (<ref>), we get

\begin{align}
\LvarP(\vec{\theta})
&= \textstyle \frac{1}{P}\sum_{i=1}^P \Ey\{ \Ezigy\{\|\vec{d}_i(\vec{\theta})\|_2^2 | \vec{y}\} \}\\
&= \textstyle \Ey\{ \Ezigy\{\|\vec{d}_i(\vec{\theta})\|_2^2 | \vec{y}\} \}
\text{~for any $i$}
\label{eq:LvarPerrB} \\
&= \Ey\{ \tr [ \Covzigy\{\hvec{x}_i(\vec{\theta}) | \vec{y}\} ] \}
\label{eq:LvarPerrC} ,
\end{align}

where (<ref>) follows because $\{\vec{d}_i(\vec{\theta})\}$ are i.i.d. when conditioned on $\vec{y}$
and (<ref>) follows from manipulations similar to those used for (<ref>).

§ PROOF OF RATIO

Here we prove ratio.
Recall from (<ref>) that
$\hvec{x}\mmse \defn \E\{\vec{x}|\vec{y}\}$
and $\vec{e}\mmse \defn \vec{x} - \hvec{x}\mmse$.
To reduce clutter, we will abbreviate $\vec{e}\mmse$ by $\vec{e}$ in this appendix.
Also, for true-posterior samples $\hvec{x}_i\sim \pxgy(\cdot|\vec{y})$, we define

\begin{align}
\hvec{e}_i &\defn \hvec{x}_i-\hvec{x}\mmse
\label{eq:ei}.
\end{align}

Then using $\hvec{x}\avgP\defn\frac{1}{P}\sum_{i=1}^P\hvec{x}_i$ and from $\mc{E}_P$ from (<ref>), we have

\begin{align}
\mc{E}_P
&= \E\{ \|\hvec{x}\avgP-\vec{x}\|^2 |\vec{y}\} \\
&= \textstyle \E\{ \|(\frac{1}{P}\sum_{i=1}^P\hvec{x}_i)-\vec{x}\|^2 |\vec{y}\}
\\
&= \textstyle \E\{ \|\frac{1}{P}\sum_{i=1}^P(\hvec{x}_i-\vec{x})\|^2 |\vec{y}\}
\\
%&= \textstyle \frac{1}{P^2}\E\{ \|\sum_{i=1}^P(\hvec{x}_i-\vec{x})\|^2 |\vec{y}\} \\
&= \textstyle \frac{1}{P^2}\E\{ \|\sum_{i=1}^P(\hvec{x}_i-\hvec{x}\mmse+\hvec{x}\mmse-\vec{x})\|^2 |\vec{y}\} \\
&= \textstyle \frac{1}{P^2}\E\{ \|\sum_{i=1}^P(\hvec{e}_i-\vec{e})\|^2 |\vec{y}\} \\
&= \textstyle \frac{1}{P^2}\E\{ \sum_{i=1}^P(\hvec{e}_i-\vec{e})\herm
\sum_{j=1}^P(\hvec{e}_j-\vec{e})|\vec{y}\} \\
&= \textstyle \frac{1}{P^2}\sum_{i=1}^P \sum_{j=1}^P \E\{(\hvec{e}_i-\vec{e})\herm
(\hvec{e}_j-\vec{e})|\vec{y}\} \\
&= \textstyle \frac{1}{P^2}\sum_{i=1}^P \E\{(\hvec{e}_i-\vec{e})\herm
+ \frac{1}{P^2}\sum_{i=1}^P \sum_{j\neq i} \E\{(\hvec{e}_i-\vec{e})\herm
(\hvec{e}_j-\vec{e})|\vec{y}\} \\
&= \textstyle \frac{1}{P^2}\sum_{i=1}^P \big[
\E\{\|\hvec{e}_i\|^2|\vec{y}\}
+\E\{\|\vec{e}\|^2|\vec{y}\} \big]
\nonumber\\&\quad\textstyle
+ \frac{1}{P^2}\sum_{i=1}^P \sum_{j\neq i} \real\big[
\E\{\hvec{e}_i\herm\hvec{e}_j|\vec{y}\}
- \E\{\hvec{e}_i\herm\vec{e}|\vec{y}\}
- \E\{\vec{e}\herm\hvec{e}_j|\vec{y}\}
+ \E\{\|\vec{e}\|^2|\vec{y}\} \big] \\
&= \textstyle \frac{1}{P^2}\sum_{i=1}^P
\E\{\|\hvec{e}_i\|^2|\vec{y}\}
+ \frac{P(P-1)}{P^2} \E\{\|\vec{e}\|^2|\vec{y}\}
\label{eq:EP1},
\end{align}

where certain terms vanished because the i.i.d. and zero-mean properties of $\{\vec{e},\hvec{e}_1,\dots,\hvec{e}_P\}$ imply

\begin{align}
\E\{\hvec{e}_i\herm\hvec{e}_j|\vec{y}\}
&=\E\{\hvec{e}_i|\vec{y}\}\herm\E\{\hvec{e}_j|\vec{y}\} = 0\\
\E\{\hvec{e}_i\herm\vec{e}|\vec{y}\}
&=\E\{\hvec{e}_i|\vec{y}\}\herm\E\{\vec{e}|\vec{y}\} = 0\\
\E\{\vec{e}\herm\hvec{e}_j|\vec{y}\}
&=\E\{\vec{e}|\vec{y}\}\herm\E\{\hvec{e}_j|\vec{y}\} = 0 .
\end{align}

Finally, note that
$\E\{\|\vec{e}\|^2|\vec{y}\}=\mc{E}\mmse$ from (<ref>).
Furthermore, because $\{\vec{x},\hvec{x}_1,\dots,\hvec{x}_P\}$ are independent samples of $\pxgy(\cdot|\vec{y})$ under the assumptions of ratio, we have $\E\{\|\vec{e}\|^2|\vec{y}\}=\E\{\|\hvec{e}_i\|^2|\vec{y}\}$ and so (<ref>) becomes

\begin{align}
\mc{E}_P
&= \frac{1}{P^2}\sum_{i=1}^P \mc{E}\mmse
+\frac{1}{P} \mc{E}\mmse
+ \frac{P(P-1)}{P^2} \mc{E}\mmse
= \frac{P+1}{P} \mc{E}\mmse .
\end{align}

This result holds for any $P\geq 1$, which implies the ratio

\begin{align}
\frac{\mc{E}_1}{\mc{E}_P} = \frac{2P}{P+1} .
\end{align}

§ CFID IMPLEMENTATION DETAILS

With the Gaussian approximation described in cfid, where $\pxugyu$ and $\pxuhgyu$ are approximated by $\mc{N}(\muxugyu,\Sigxugyu)$ and $\mc{N}(\muxhugyu,\Sigxuhgyu)$, respectively, the CWD in (<ref>) reduces to

\begin{align}
\CFID &\defn \Ey\big\{ \|\muxugyu-\muxhugyu\|_2^2  + \tr\big[ \Sigxugyu + \Sigxuhgyu - 2 \big( \Sigxugyu^{1/2} \Sigxuhgyu \Sigxugyu^{1/2} \big)^{1/2} \big]\big\}
\label{eq:CFID} .
\end{align}

The values in (<ref>) are computed using

\begin{align}
\muxugyu
&= \muxu + \Sigxuyu\Sigyuyu^{-1}(\uvec{y}-\muyu)
\label{eq:muxugyu}\\
\Sigxugyu
&= \Sigxuxu - \Sigxuyu\Sigyuyu^{-1}\Sigxuyu\tran
\label{eq:sigxxgy}\\
\muxhugyu
&= \muxhu + \Sigxhuyu\Sigyuyu^{-1}(\uvec{y}-\muyu)\\
\Sigxuhgyu
&= \Sigxhuxhu - \Sigxhuyu\Sigyuyu^{-1}\Sigxhuyu\tran
\label{eq:sigxuxugy}.
\end{align}

Plugging (<ref>)-(<ref>) into (<ref>), the CFID can be written as <cit.>:

\begin{align}
\CFID
&= \|\muxu-\muxhu\|_2^2
+ \tr\Big[(\Sigxuyu - \Sigxhuyu)\Sigyuyu^{-1}(\Sigxuyu - \Sigxhuyu)\tran\Big]
\nonumber\\&\quad
+ \tr\Big[\Sigxugyu + \Sigxuhgyu - 2 \big( \Sigxugyu^{1/2} \Sigxuhgyu \Sigxugyu^{1/2} \big)^{1/2}\Big] ,
\label{eq:CFID_2}
\end{align}

where $\Sigyuyu^{-1}$ is typically implemented using a pseudo-inverse.

We now detail how the means and covariances in (<ref>) are computed.
We start with a dataset $\{(\vec{x}_t, \vec{y}_t)\}_{t=1}^n$ of truth/measurement pairs.
For each $\vec{y}_t$, we generate a set of $P$ posterior samples $\{\hvec{x}_{ti}\}_{i=1}^{P}$.
We merge these samples with $P$ repetitions of $\vec{x}_t$ and $\vec{y}_t$ to obtain $\{(\vec{x}_{ti}, \vec{y}_{ti}, \hvec{x}_{ti})\}_{i=1}^{P}$ for $t=1\dots n$.
These terms are processed by a feature-generating network to yield the feature embeddings $\{(\uvec{x}_{ti}, \uvec{y}_{ti}, \huvec{x}_{ti})\}_{i=1}^{P}$, which are then packed into matrices $\uvec{X}$, $\uvec{Y}$, and $\huvec{X}$ with $P n$ rows.
We used the VGG-16 feature-generating network <cit.> for our MRI experiments, since <cit.> found that it gave results that correlated much better with radiologists' perceptions, while we used the standard Inception-v3 network <cit.> for our inpainting experiments.
The embeddings are then used to compute the sample-mean values

\begin{align} \textstyle
\muxu \defn \frac{1}{P n} \vec{1}\tran \uvec{X},
\quad
\muyu \defn \frac{1}{P n} \vec{1}\tran \uvec{Y},
\quad
\muxhu \defn \frac{1}{P n} \vec{1}\tran \huvec{X}
\label{eq:meancfid}.
\end{align}

We then subtract the sample mean from each row of $\uvec{X}$, $\uvec{Y}$, and $\huvec{X}$ to give the zero-mean embedding matrices $\uvec{X}\zm\defn\uvec{X}-\vec{1}\muxu\tran$, $\uvec{Y}\zm\defn\uvec{Y}-\vec{1}\muyu\tran$, and $\huvec{X}\zm\defn\huvec{X}-\vec{1}\muxhu\tran$, which are then used to compute the sample covariance matrices

\begin{align}
\Sigxuxu
&\defn \textstyle \frac{1}{P n} \uvec{X}\zm\tran\uvec{X}\zm, \quad
\Sigyuyu
\defn \frac{1}{P n} \uvec{Y}\zm\tran\uvec{Y}\zm, \quad
\Sigxhuxhu
\defn \frac{1}{P n} \huvec{X}\zm\tran\huvec{X}\zm \\
\Sigxuyu
&\defn \textstyle \frac{1}{P n} \uvec{X}\zm\tran\uvec{Y}\zm , \quad
\Sigxhuyu
\defn \frac{1}{P n} \huvec{X}\zm\tran\uvec{Y}\zm .
\end{align}

We plug the sample statistics from (<ref>)-(<ref>) into (<ref>)-(<ref>), which yields the statistics needed to compute the CFID in (<ref>).
In <cit.>, the authors use $P=1$ in all of their experiments.
To be consistent with how we evaluated the other metrics, we use $P=32$ unless otherwise noted.

§ MR IMAGING DETAILS

We now give details on magnetic resonance (MR) image recovery.
Suppose that the goal is to recover the $N$-pixel MR image $\vec{i}\in\Complex^N$ from the multicoil measurements $\{\vec{k}_c\}_{c=1}^C$, where <cit.>

\begin{equation}
\vec{k}_c = \vec{MFS}_c\vec{i} + \vec{n}_c
\label{eq:yc}.
\end{equation}

In (<ref>), $C$ refers to the number of coils, $\vec{k}_c\in\Complex^M$ are the measurements from the $c$th coil, $\vec{M} \in \Real^{M\times N}$ is a sub-sampling operator containing rows from $\vec{I}_N$—the $N\times N$ identity matrix, $\vec{F} \in \Complex^{N\times N}$ is the unitary 2D discrete Fourier transform, $\vec{S}_c \in \Complex^{N\times N}$ is a diagonal matrix containing the sensitivity map of the $c$th coil,  and $\vec{n}_c \in \Complex^{M}$ is noise.
From (<ref>), it can be seen that the MR measurements are collected in the spatial Fourier domain, otherwise known as the “k-space.”
The sensitivity maps $\{\vec{S}_c\}$ are estimated from $\{\vec{k}_c\}$ using ESPIRiT <cit.> (in our case via SigPy <cit.>), which yields maps with the property $\sum_{c=1}^C\vec{S}_c\herm\vec{S}_c=\vec{I}_N$.
The ratio $R\triangleq \frac{N}{M}$ is known as the acceleration rate.

There are different ways that one could apply the generative posterior sampling framework to multicoil MR image recovery.
One is to configure the generator to produce posterior samples $\hvec{i}$ of the complex image $\vec{i}$. Another is to configure the generator to produce posterior samples $\hvec{x}$ of the stack $\vec{x}\defn [\vec{x}_1\tran,\dots,\vec{x}_C\tran]\tran$ of “coil images” $\vec{x}_c\defn\vec{S}_c\vec{i}$ and later coil-combining them to yield a complex image estimate $\hvec{i}\defn [\vec{S}\herm_1,\dots,\vec{S}\herm_C]\hvec{x}$.
We take the latter approach.
Furthermore, rather than feeding our generator with k-space measurements $\vec{k}_c$, we choose to feed it with aliased coil images $\vec{y}_c\defn \vec{F}\herm\vec{M}\tran\vec{k}_c$.
Writing (<ref>) in terms of the coil images, we obtain

\begin{equation}
\vec{y}_c = \vec{F}\herm\vec{M}\tran\vec{MFx}_c + \vec{w}_c
\label{eq:yc2},
\end{equation}

where $\vec{w}_c\defn \vec{F}\herm\vec{M}\tran\vec{n}_c$.
Then we can stack $\{\vec{y}_c\}$ and $\{\vec{w}_c\}$ column-wise into vectors $\vec{y}$ and $\vec{w}$, and set $\vec{A} = \vec{I}_C \otimes \vec{F}\herm\vec{M}\tran\vec{MF}\in\Complex^{NC\times NC}$, to obtain the formulation $\vec{y} = \vec{Ax}+\vec{w}$ described in intro.

To train our generator, we assume to have access to paired training examples $\{(\vec{x}_t,\vec{y}_t)\}$, where $\vec{x}_t$ is a stack of coil images and $\vec{y}_t$ is the corresponding stack of k-space coil measurements.
The fastMRI multicoil dataset <cit.> provides $\{(\vec{x}_t,\vec{k}_t)\}$, from which we can easily obtain $\{(\vec{x}_t,\vec{y}_t)\}$.

§ DATA-CONSISTENCY

In this section we describe a data-consistency procedure that can be optionally used when our cGAN is used to solve a linear inverse problem, i.e., to recover $\vec{x}$ from $\vec{y}$ under a model of the form

\begin{equation}
\vec{y}
= \vec{Ax}+\vec{w}
\label{eq:y} ,
\end{equation}

where $\vec{A}$ is a known linear operator and $\vec{w}$ is unknown noise.
The motivation is that, in some applications, such as medical imaging or inpainting, the end user may feel comfortable knowing that the generated samples $\{\hvec{x}_i\}$ are consistent with the measurements $\vec{y}$ in that

\begin{align}
\vec{y}=\vec{A}\hvec{x}_i
\label{eq:DC}.
\end{align}

When (<ref>) holds, $\vec{A}^+\vec{y}=\vec{A}^+\vec{A}\hvec{x}_i$ must also hold, where $(\cdot)^+$ denotes the pseudo-inverse.
The quantity $\vec{A}^+\vec{A}$ can be recognized as the orthogonal projection matrix associated with the row space of $\vec{A}$.
So, (<ref>) requires the component of $\hvec{x}_i$ in the row space of $\vec{A}$ to equal $\vec{A}^+\vec{y}$, while placing no constraints on the component of $\hvec{x}_i$ in the nullspace of $\vec{A}$.
This suggests the following data-consistency procedure:

\begin{align}
\hvec{x}_i
&= (\vec{I}-\vec{A}^+\vec{A})\hvec{x}_i\raw + \vec{A}^+\vec{y}
\label{eq:DC2}.
\end{align}

where $\hvec{x}_i\raw$ is the raw generator output.
We note that a version of this idea for point estimation was proposed in <cit.>.

The data-consistency procedure (<ref>) ensures that any generative method will generate only the component of $\vec{x}$ that lies in the nullspace of $\vec{A}$.
Consequently, (<ref>) is admissible only when $\vec{A}$ has a non-trivial nullspace.
Also, because no attempt is made to remove the noise $\vec{w}$ in $\vec{y}$, this approach is recommended for low-noise applications.
For high-noise applications, an extension based on the dual-decomposition approach <cit.> would be more appropriate, but we leave this to future work.

When applying (<ref>) to the MRI formulation in MRI, we note that $\vec{A}=\vec{I}_C \otimes \vec{F}\herm\vec{M}\tran\vec{MF}$ is an orthogonal projection matrix, and so $\vec{I}-\vec{A}^+\vec{A} = \vec{I}-\vec{A} = \vec{I}\otimes \vec{F}\herm(\vec{I}-\vec{M}\tran\vec{M})\vec{F}$.

§ IMPLEMENTATION DETAILS

§.§ MRI

§.§.§ cGAN training

At each training iteration, our cGAN's generator takes in $n\batch$ measurement samples $\vec{y}_t$ and $P\train$ code vectors for every $\vec{y}_t$, and performs an optimization step on the loss

\begin{equation}
\Ell_{\mathsf{G}}(\vec{\theta})
\defn \betaadv \Ladv(\vec{\theta},\vec{\phi}) + \LonePt(\vec{\theta}) - \betastd\LstdPt(\vec{\theta})
\label{eq:genloss} ,
\end{equation}

where by default we use $\betaadv = 1\text{e-5}$, $n\batch=36$, $P\train=2$, and update $\betastd$ via (<ref>) using $P\val=8$.
Then, using the $P\train n\batch$ generator outputs, our cGAN's discriminator performs an optimization step on the loss

\begin{equation}
\Ell_{\mathsf{D}}(\vec{\phi}) = -\Ladv(\vec{\theta},\vec{\phi}) + \alpha_1 \Lgp(\vec{\phi}) + \alpha_2 \Ldrift(\vec{\phi})
\label{eq:dloss} ,
\end{equation}

with gradient penalty $\Ell_\text{gp}$ from <cit.>.
As per <cit.>, $\Ell_\text{drift}$ is a drift penalty, $\alpha_1 = 10$, $\alpha_2 = 0.001$, and one discriminator update was used per generator update.
The models were trained for 100 epochs using the Adam optimizer <cit.> with a learning rate of 1e-3, $\beta_1=0$, and $\beta_2=0.99$, as in <cit.>.
Running PyTorch on a server with 4 Tesla A100 GPUs, each with 82 GB of memory, the training of an MRI cGAN took approximately 1 day.

Adler and Öktem's cGAN <cit.> uses generator loss $\betaadv\Ladler(\vec{\theta},\vec{\phi})$, where $\Ladler(\vec{\theta},\vec{\phi})$ was described in (<ref>), and discriminator loss $-\Ladler(\vec{\theta},\vec{\phi})+ \alpha_1 \Lgp(\vec{\phi}) + \alpha_2 \Ldrift(\vec{\phi})$
with the values of $\alpha_1=$10, $\alpha_2=$0.001, and $\betaadv=$1, as in the original paper.

Ohayon et al.'s cGAN <cit.> uses generator loss $\betaadv\Ladv(\vec{\theta},\vec{\phi}) + \LtwoPt(\vec{\theta})$, where $\LtwoPt(\vec{\theta})$ was described in (<ref>), and discriminator loss $-\Ladv(\vec{\theta},\vec{\phi})+ \alpha_1 \Lgp(\vec{\phi}) + \alpha_2 \Ldrift(\vec{\phi})$
with the values $\alpha_1=$10, $\alpha_2=$0.001, and $\betaadv=$1e-5. We modify $\betaadv$ to re-balance the loss due to an increased magnitude of our discriminator's outputs.

All three cGANs used the same generator and discriminator architectures (detailed below), except that Adler and Öktem's discriminator used extra input channels to facilitate the 3-input loss $\Ladler(\vec{\theta},\vec{\phi})$ from (<ref>).

§.§.§ cGAN generator architecture

For our MRI experiments, we take inspiration from the UNet architecture from <cit.>, using it as the basis for the cGAN generators.
The primary input, $\vec{y}$ is concatenated with the code vector $\vec{z}$ and fed through the UNet.
The network consists of 4 pooling layers with 128 initial channels.
However, instead of pooling, we opt to use convolutions with kernels of size $3\times3$, “same” padding, and a stride of 2 when downsampling.
Likewise, we upsample using transpose convolutions, again with kernels of size $3\times3$, “same” padding, and a stride of 2.
All other convolutions utilize kernels of size $3\times3$, “same” padding, and a stride of 1.

Within each encoder and decoder layer we include a residual block, the architecture of which can be found in <cit.>.
We use instance-norm for all normalization layers and parametric ReLUs as our activation functions, in which the network learns the optimal “negative slope.”
Finally, we include 5 residual blocks at the base of the UNet, in between the encoder and decoder. This is done in an effort to artificially increase the depth of the network and is inspired by <cit.>.
Our generator has $86\,734\,334$ trainable parameters.

§.§.§ cGAN discriminator architecture

Our discriminator is a standard CNN with 6 layers and 1 fully-connected layer.
In the first 3 layers, we use convolutions with kernels of size $3\times3$, “same” padding. We reduce spatial resolution with average pooling, using $2\times2$ kernels with a stride of 2.
We use batch-norm as our normalization layer and leaky ReLUs with a “negative-slope” of 0.2 as our activation functions. The network outputs an estimated Wasserstein score for the whole image.

§.§.§ UNet baseline

For the UNet baseline, we use the same architecture as the cGAN generator described above, but we train it to minimize the weighted combination of supervised $\ell_1$ loss and (negative) SSIM proposed in <cit.>.
In particular, the UNet is trained to minimize the following loss:

\begin{equation}
\mathcal{L}_{1,\mathsf{SSIM}}(\vec{\theta}) = (1-\alpha)\Exy\{\|\vec{x} - \hvec{x}\|_1\} - \alpha \Exy\{\text{SSIM}(\vec{x}, \hvec{x})\},
\label{eq:unetloss}
\end{equation}

where $\vec{x}$ denotes a ground truth multicoil MR image and $\hvec{x}$ denotes a UNet reconstruction. Following <cit.>, we choose $\alpha=0.84$ to balance the two terms in the loss.

§.§.§ Langevin approach

For Jalal et al.'s MRI approach <cit.>, we do not modify the original implementation from <cit.> other than replacing the default sampling pattern with the GRO undersampling mask.
We generated 32 samples for 72 different test images using a batch-size of 4, which took roughly 6 days.
These samples were generated on a server with 4 NVIDIA V100 GPUs, each with 32 GB of memory.
We used 4 samples per batch (and recorded the time to generate 4 samples in mriresults) because the code from <cit.> is written to generate one sample per GPU.

§.§ Inpainting

§.§.§ Our cGAN

For our generator and discriminator, we use the CoModGAN networks from <cit.>.
Unlike CoModGAN, however, we train our cGAN with $\LonestdPt$ regularization and we do not use MBSD at the discriminator.
We use the same general training and testing procedure described in MRIexperiments, but with $\betaadv=$ 5e-3, $n\batch=100$, and 110 epochs of cGAN training.
Running PyTorch on a server with 4 Tesla A100 GPUs, each with 82 GB of memory, the training takes approximately 2 days.

§.§.§ CoModGAN

We use the PyTorch implementation of CoModGAN from <cit.> and train the model

to inpaint a $128\times128$ centered square on $256\times256$ CelebA-HQ images.
The total training time on a server with 4 NVIDIA A100 GPUs, each with 82 GB of memory, is roughly 2 days.

§.§.§ Score-based SDE

For the inpainting experiment in inpainting, we compare against Song et al.'s more recent SDE technique <cit.>, for which we use the publicly available pretrained weights, the suggested settings for the $256\times256$ CelebA-HQ dataset, and the code from the official PyTorch implementation <cit.>.
We generate 32 samples for all $1\,000$ images in our test set, using a batch-size of 20 and generating 32 samples for each batch element concurrently.
The total generation time on a server with 4 NVIDIA A100 GPUs, each with 82 GB of memory, is roughly 9 days.

§ ADDITIONAL EXPERIMENTAL RESULTS

§.§ R=4 MRI Reconstruction

We tested our model on $R=4$ reconstruction of fastMRI $384 \times 384$ T2 brain images under the same conditions outlined in MRIexperiments.
The results, shown in mriresultsr4, are consistent with those from our $R=8$ experiment, however note that for this problem Ohayon et al.'s cGAN has even less variation than the other techniques.

Average results for $R\!=\!4$ MRI.

and APSD used $72$ test samples and $P\!=\!32$.
(C)FID$^2$ used $2\,376$ test samples and $P\!=\!8$.
CFID$^3$ used all $14\,576$ samples and $P\!=\!1$

Model     CFID$^1\!\!\downarrow$     CFID$^2\!\!\downarrow$     CFID$^3\!\!\downarrow$     FID$^1\!\!\downarrow$     FID$^2\!\!\downarrow$

APSD     Time (4)$\downarrow$

Langevin (Jalal et al. <cit.>)     5.29     -     -     6.12     -

5.9e-6     14 min

cGAN (Adler & Öktem <cit.>)     6.39     4.27     3.82     5.25     2.86

3.9e-6     217 ms

cGAN (Ohayon et al. <cit.>)     4.06     3.27     2.95     6.45     3.61

7.2e-8     217 ms

cGAN (Ours)     3.10     1.54     1.29     3.75     1.22

3.8e-6     217 ms

§.§ CFID decomposition into mean and covariance components

In this section we investigate the small-sample bias effects of CFID, which have been previously noted in <cit.>.
To do this, we write the CFID from (<ref>) as a sum of two terms: a term that quantifies the conditional-mean error and a term that quantifies the conditional-covariance error:

\begin{align}
\CFID
&= \CFID\mean + \CFID\cova \\
\CFID\mean
&\defn \Ey\{ \|\muxugyu-\muxhugyu\|_2^2 \} \\
\CFID\cova
&\defn  \tr\big[ \Sigxugyu + \Sigxuhgyu
- 2 \big( \Sigxugyu^{1/2} \Sigxuhgyu \Sigxugyu^{1/2} \big)^{1/2} \big]
\label{eq:CFIDcov}.
\end{align}

To verify that (<ref>) quantifies the error in $\Sigxuhgyu$, notice that (<ref>) equals zero when $\Sigxuhgyu=\Sigxugyu$ and is otherwise positive (by Cauchy Schwarz).

The mean and covariance components of CFID, along with the total CFID, for the MRI and inpainting experiments.
For the MRI experiment, CFID$^1$ used 72 test samples and $P=32$,
CFID$^2$ used 2 376 test samples and $P=8$, and
CFID$^3$ used all 14 576 samples and $P=1$.
For the inpainting experiment, CFID$^1$ used $1\,000$ test images and $P=32$,
CFID$^2$ used $3\,000$ test and validation images and $P=8$, and
CFID$^3$ used all 30 000 images and $P=1$.

CFID$^1\mean\!\downarrow$     CFID$^1\cova\!\downarrow$     CFID$^1\!\downarrow$
CFID$^2\mean\!\downarrow$     CFID$^2\cova\!\downarrow$     CFID$^2\!\downarrow$
CFID$^3\mean\!\downarrow$     CFID$^3\cova\!\downarrow$     CFID$^3\!\downarrow$

9c$R=4$ MRI
Langevin (Jalal <cit.>)        1.89     3.40     5.29     -     -     -     -     -     -

cGAN (Adler <cit.>)             3.12     3.27     6.39     2.79     1.48     4.27     2.71     1.10     3.82

cGAN (Ohayon <cit.>)     1.94     2.12     4.06     2.27     1.00     3.27     2.29     0.66     2.95

cGAN (Ours)     0.98     2.12     3.10     0.86     0.68     1.54     0.86     0.43     1.29

9c$R=8$ MRI
Langevin (Jalal <cit.>)        1.89     3.40     5.29     -     -     -     -     -     -

cGAN (Adler <cit.>)             5.00     5.10     10.10     4.16     2.14     6.30     4.09     1.63     5.72

cGAN (Ohayon <cit.>)     2.73     3.31     6.04     3.07     1.52     4.59     3.30     0.97     4.27

cGAN (Ours)     1.55     3.32     4.87     1.24     0.99     2.23     1.17     0.62     1.79

Score SDE (Song <cit.>)     0.97     38.69     39.66     -        -         -     0.90     4.21     5.11

CoModGAN (Zhao <cit.>)        0.42     41.21         41.63         0.35     25.39     25.74     0.32     4.98     5.29

cGAN (Ours)     0.32     39.41     39.73     0.25     22.32     22.58     0.24     4.45     4.69

In cfiderrors, we report $\CFID\mean$ and $\CFID\cova$ for the MRI and inpainting experiments, in addition to the total CFID (also shown Tables <ref>, <ref>, and <ref>).
As before, we computed CFID on three test sets for each experiment,

which contained 72, 1000, and 30 000 samples respectively for MRI, and
1000, 3000, and 30 000 samples respectively for inpainting.
Due to the slow sample-generation time of the Langevin/score-based methods <cit.>, we did not have the computational resources to evaluate them on all datasets, and that's why certain table entries are blank.

For both MRI experiments, cfiderrors shows our method outperforming the competing methods in both the mean and covariance components of CFID (and thus the total CFID) for all sample sizes.
And in the inpainting experiment, cfiderrors shows our method outperforming CoModGAN in both the mean and covariance components (and thus the total CFID) for all sample sizes.

For the inpainting experiment, cfiderrors shows our method outperforming the score-based approach in total CFID on the 3000- and 30 000-sample test sets but not on the 1000-sample test set.
However, we now argue that the 1000-sample inpainting experiment is heavily affected by small-sample bias, and therefore untrustworthy.
Looking at the mean component of CFID (i.e., $\CFID\mean^1$, $\CFID\mean^2$, and $\CFID\mean^3$) across the inpainting experiments, we see that the values are relatively small and stable with sample size.
But looking at the covariance component of CFID (i.e., $\CFID\cova^1$, $\CFID\cova^2$, and $\CFID\cova^3$) across the inpainting experiments, we see that the values are large and heavily dependent on sample size.
For the 1000-sample inpainting experiment, the total CFID is dominated by the covariance component and thus strongly affected by small-sample bias.
Consequently, for the 1000-sample inpainting experiment, the total CFID is not trustworthy.

§ ADDITIONAL RECONSTRUCTION PLOTS

§.§ R=8 MRI Reconstruction

§.§ Inpainting