SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Matthew C. Bendel Dept. ECE
The Ohio State University Columbus, OH 43210 bendel.8@osu.edu
Rizwan Ahmad Dept. BME
The Ohio State University Columbus, OH 43210 ahmad.46@osu.edu
Philip Schniter Dept. ECE
The Ohio State University Columbus, OH 43201 Schniter.1@osu.edu
In image recovery problems, one seeks to infer 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 seek to do more than just produce one good image estimate. Rather, we aim to rapidly and accurately sample from the posterior distribution. To do this, 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. The code for our model can be found here: https://github.com/matt-bendel/rcGAN.
We consider image recovery, where one observes measurements  $\pmb{y} = \mathcal{M}(\pmb{x})$  of the true image  $\pmb{x}$  that may be masked, distorted, and/or corrupted  $\pmb{x}$  with noise, and the goal is to infer  $\pmb{x}$  from  $\pmb{y}$ . This includes linear inverse problems arising in, e.g., deblurring, super-resolution, inpainting, colorization, computed tomography (CT), and magnetic resonance imaging (MRI), where  $\pmb{y} = \pmb{A}\pmb{x} + \pmb{w}$  with known linear operator  $\pmb{A}$  and noise  $\pmb{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  $\pmb{x}$  from  $\pmb{y}$ .
Image recovery is often posed as finding the single "best" recovery  $\widehat{\pmb{x}}$ , which is known as a point estimate of  $\pmb{x}$  [1]. But point estimation is problematic due to the perception-distortion tradeoff [2], which establishes a fundamental tradeoff between distortion (defined as some distance between  $\widehat{\pmb{x}}$  and  $\pmb{x}$ ) and perceptual quality (defined as some distance between  $\widehat{\pmb{x}}$  and the set of clean images). For example, the minimum mean-squared error (MMSE) recovery  $\widehat{\pmb{x}}_{\mathrm{mmse}}$  is optimal in terms of  $\ell_2$  distortion, but can be unrealistically smooth. Although one could instead compute an approximation of the maximum a posteriori (MAP) estimate [3] or minimize some combination of perceptual and distortion losses, it's unclear which combination would be most appropriate.
Another major limitation with point estimation is that it's unclear how certain one can be about the recovered  $\widehat{\pmb{x}}$ . For example, with deep-learning-based recovery, it's possible to hallucinate a nice-looking  $\widehat{\pmb{x}}$ , but is it correct? Quantifying the uncertainty in  $\widehat{\pmb{x}}$  is especially important in medical applications such as MRI, where a diagnosis must be made based on the measurements  $\pmb{y}$ . Rather than simply reporting our best guess of whether a pathology is present or absent based on  $\widehat{\pmb{x}}$ , we might want to report the probability that the pathology is present (given all available data).
Yet another problem with point estimation is that the estimated  $\widehat{\pmb{x}}$  could pose issues with fairness [4]. For example, say we are inpainting a face within an image. With a racially heterogeneous prior distribution,  $\widehat{\pmb{x}}_{\mathrm{mmse}}$  (being the posterior mean) will be biased towards the most predominant race. The same will be true of many other point estimates  $\widehat{\pmb{x}}$ .
To address the aforementioned limitations of point estimation, we focus on generating samples from the posterior distribution  $p_{\mathbf{x}|\mathbf{y}}(\cdot |\mathbf{y})$ , which represents the complete state-of-knowledge about  $\mathbf{x}$  given the measurements  $\mathbf{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 [4]). Furthermore, the posterior directly facilitates uncertainty quantification via, e.g., pixel-wise standard deviations or pathology detection probabilities (see Appendix A). 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  $\widehat{\pmb{x}}_{\mathrm{mmse}}$  with large  $P$  and sharp-but-more-distorted  $\widehat{\pmb{x}}$  with small  $P$ . Additionally, posterior sampling unlocks other important capabilities such as adaptive acquisition [5] and counterfactual diagnosis [6].
Concretely, given a training dataset of image/measurement pairs  $\{(\pmb{x}_t,\pmb{y}_t)\}_{t=1}^T$ , our goal is to learn a generating function  $G_{\theta}$  that, for a new  $\pmb{y}$ , maps random code vectors  $\pmb{z} \sim \mathcal{N}(\pmb{0},\pmb{I})$  to posterior samples  $\widehat{\pmb{x}} = G_{\theta}(\pmb{z},\pmb{y}) \sim p_{\times|\pmb{y}}(\cdot|\pmb{y})$ . There exist several well-known approaches to this task, with recent literature focusing on conditional generative adversarial networks (cGANs) [7, 8, 9, 10], conditional variational autoencoders (cVAEs) [11, 12, 13], conditional normalizing flows (cNFs) [14, 15, 16], and score/diffusion/Langevin-based generative models [17, 18, 19, 20, 21]. 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 choose to 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 [9] and score/Langevin-based approaches [19, 20] are readily available. Third, the linear operator “ $A$ ” manifests very differently in these two applications, $^{1}$  which illustrates the versatility of our approach. To quantify performance, we focus on conditional Fréchet inception distance (CFID) [22], 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 [8, 23, 9] and the score/diffusion/Langevin approaches from [19] and [20] in all tested metrics, while generating samples  $\sim 10^{4}$  times faster than [19, 20].
2 Problem formulation and background
We build on the Wasserstein cGAN framework from [8]. The goal is to design a generator network  $G_{\theta}:\mathcal{Z}\times \mathcal{Y}\to \mathcal{X}$  such that, for fixed  $y$ , the random variable  $\widehat{\pmb{x}} = G_{\pmb{\theta}}(\pmb {z},\pmb {y})$  induced by  $z\sim p_z$  has a distribution that best matches the posterior  $p_{\mathrm{x|y}}(\cdot |\pmb {y})$  in Wasserstein-1 distance. Here,  $\mathcal{X}$ ,  $\mathcal{V}$ , and  $\mathcal{Z}$  denote the sets of  $x,y$ , and  $z$ , respectively, and  $z$  is drawn independently of  $y$ .
The Wasserstein-1 distance can be expressed as
$$
W _ {1} \left(p _ {\mathrm {x} | \mathrm {y}} (\cdot , \boldsymbol {y}), p _ {\widehat {\mathrm {x}} | \mathrm {y}} (\cdot , \boldsymbol {y})\right) = \sup  _ {D \in L _ {1}} \mathrm {E} _ {\mathrm {x} | \mathrm {y}} \left\{D (\boldsymbol {x}, \boldsymbol {y}) \right\} - \mathrm {E} _ {\widehat {\mathrm {x}} | \mathrm {y}} \left\{D (\widehat {\boldsymbol {x}}, \boldsymbol {y}) \right\}, \tag {1}
$$
where  $L_{1}$  denotes functions that are 1-Lipschitz with respect to their first argument and  $D:\mathcal{X}\times \mathcal{Y}\to$ $\mathbb{R}$  is a "critic" or "discriminator" that tries to distinguish between true  $\pmb{x}$  and generated  $\widehat{\pmb{x}}$  given  $\pmb{y}$
<sup>1</sup>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.
Since we want the method to work for typical values of  $\mathbf{y}$ , we define a loss by taking an expectation of (1) over  $\mathbf{y} \sim p_{\mathbf{y}}$ . Since the expectation commutes with the supremum in (1), we have [8]
$$
\begin{array}{l} \mathrm {E} _ {\mathcal {Y}} \left\{W _ {1} \left(p _ {\mathrm {x} | \mathrm {y}} (\cdot , \boldsymbol {y}), p _ {\widehat {\mathrm {x}} | \mathrm {y}} (\cdot , \boldsymbol {y})\right) \right\} = \sup  _ {D \in L _ {1}} \mathrm {E} _ {\mathrm {x}, \mathrm {y}} \left\{D (\boldsymbol {x}, \boldsymbol {y}) \right\} - \mathrm {E} _ {\widehat {\mathrm {x}}, \mathrm {y}} \left\{D (\widehat {\boldsymbol {x}}, \boldsymbol {y}) \right\} (2) \\ = \sup  _ {D \in L _ {1}} \mathrm {E} _ {\mathbf {x}, \mathbf {z}, \mathbf {y}} \left\{D (\mathbf {x}, \mathbf {y}) - D \left(G _ {\boldsymbol {\theta}} (\mathbf {z}, \mathbf {y}), \mathbf {y}\right) \right\}. (3) \\ \end{array}
$$
In practice,  $D$  is implemented by a neural network  $D_{\phi}$  with parameters  $\phi$ , and  $(\theta, \phi)$  are trained by alternately minimizing
$$
\mathcal {L} _ {\mathrm {a d v}} (\boldsymbol {\theta}, \phi) \triangleq \mathrm {E} _ {\mathrm {x}, \mathrm {z}, \mathrm {y}} \left\{D _ {\phi} (\boldsymbol {x}, \boldsymbol {y}) - D _ {\phi} \left(G _ {\boldsymbol {\theta}} (\boldsymbol {z}, \boldsymbol {y}), \boldsymbol {y}\right) \right\} \tag {4}
$$
with respect to  $\pmb{\theta}$  and minimizing  $-\mathcal{L}_{\mathrm{adv}}(\pmb {\theta},\phi) + \mathcal{L}_{\mathrm{gp}}(\phi)$  with respect to  $\phi$  , where  $\mathcal{L}_{\mathrm{gp}}(\phi)$  is a gradient penalty that is used to encourage  $D_{\phi}\in L_1$  [24]. Furthermore, the expectation over  $\pmb{x}$  and  $\pmb{y}$  in (4) is replaced in practice by a sample average over the training examples  $\{(x_t,y_t)\}_{t = 1}^T$
One of the main challenges with the cGAN framework in image recovery problems is that, for each measurement example  $\mathbf{y}_t$ , there is only a single image example  $\mathbf{x}_t$ . Thus, with the previously described training methodology, there is no incentive for the generator to produce diverse samples  $G(z, \mathbf{y})|_{z \sim p_z}$  for a fixed  $\mathbf{y}$ . This can lead to the generator learning to ignore the code vector  $z$ , which causes a form of "mode collapse."
Although issues with stability and mode collapse are also present in unconditional GANs (uGANs) or discretely conditioned cGANs [25], the causes are fundamentally different than in continuously conditioned cGANs like ours. With continuously conditioned cGANs, there is only one example of a valid  $\boldsymbol{x}_t$  for each given  $\boldsymbol{y}_t$ , whereas with uGANs there are many  $\boldsymbol{x}_t$  and with discretely conditioned cGANs there are many  $\boldsymbol{x}_t$  for each conditioning class. As a result, most strategies that are used to combat mode-collapse in uGANs [26, 27, 28] are not well suited to cGANs. For example, mini-batch discrimination strategies like MBSD [29], where the discriminator aims to distinguish a mini-batch of true samples  $\{\boldsymbol{x}_t\}$  from a mini-batch of generated samples  $\{\widehat{\boldsymbol{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 & Oktem [8] proposed to use a three-input discriminator  $D_{\phi}^{\mathrm{adler}}: \mathcal{X} \times \mathcal{X} \times \mathcal{Y} \rightarrow \mathbb{R}$  and replace  $\mathcal{L}_{\mathrm{adv}}$  from (4) with the loss
$$
\begin{array}{l} \mathcal {L} _ {\mathrm {a d v}} ^ {\mathrm {a d l e r}} (\boldsymbol {\theta}, \phi) \triangleq \mathrm {E} _ {\mathrm {x}, \mathrm {z} _ {1}, \mathrm {z} _ {2}, \mathrm {y}} \left\{\frac {1}{2} D _ {\phi} ^ {\mathrm {a d l e r}} (\boldsymbol {x}, G _ {\boldsymbol {\theta}} (\boldsymbol {z} _ {1}, \boldsymbol {y}), \boldsymbol {y}) + \frac {1}{2} D _ {\phi} ^ {\mathrm {a d l e r}} (G _ {\boldsymbol {\theta}} (\boldsymbol {z} _ {2}, \boldsymbol {y}), \boldsymbol {x}, \boldsymbol {y}) \right. \\ \left. - D _ {\phi} ^ {\operatorname {a d l e r}} \left(G _ {\boldsymbol {\theta}} \left(\boldsymbol {z} _ {1}, \boldsymbol {y}\right), G _ {\boldsymbol {\theta}} \left(\boldsymbol {z} _ {2}, \boldsymbol {y}\right), \boldsymbol {y}\right) \right\}, \tag {5} \\ \end{array}
$$
which rewards variation between the first and second inputs to  $D_{\phi}^{\mathrm{adler}}$ . They then proved that minimizing  $\mathcal{L}_{\mathrm{adv}}^{\mathrm{adler}}$  in place of  $\mathcal{L}_{\mathrm{adv}}$  does not compromise the Wasserstein cGAN objective, i.e.,  $\arg \min_{\theta} \mathcal{L}_{\mathrm{adv}}^{\mathrm{adler}}(\theta, \phi) = \arg \min_{\theta} \mathcal{L}_{\mathrm{adv}}(\theta, \phi)$ . As we show in Section 4, this approach does prevent complete mode collapse, but it leaves much room for improvement.
3 Proposed method
3.1 Proposed regularization: supervised-  $\ell_1$  plus SD reward
We now propose a novel cGAN regularization framework. To train the generator, we propose to solve
$$
\arg \min  _ {\boldsymbol {\theta}} \left\{\beta_ {\text {a d v}} \mathcal {L} _ {\text {a d v}} (\boldsymbol {\theta}, \phi) + \mathcal {L} _ {1, \mathrm {S D}, P _ {\text {t r a i n}}} (\boldsymbol {\theta}, \beta_ {\mathrm {S D}}) \right\} \tag {6}
$$
with appropriately chosen  $\beta_{\mathrm{adv}},\beta_{\mathrm{SD}} > 0$  and  $P_{\mathrm{train}}\geq 2$  , where the regularizer
$$
\mathcal {L} _ {1, \mathrm {S D}, P _ {\text {t r a i n}}} (\boldsymbol {\theta}, \beta_ {\mathrm {S D}}) \triangleq \mathcal {L} _ {1, P _ {\text {t r a i n}}} (\boldsymbol {\theta}) - \beta_ {\mathrm {S D}} \mathcal {L} _ {\mathrm {S D}, P _ {\text {t r a i n}}} (\boldsymbol {\theta}) \tag {7}
$$
is constructed from the  $P_{\mathrm{train}}$ -sample supervised- $\ell_1$  loss and standard-deviation (SD) reward terms
$$
\mathcal {L} _ {1, P _ {\text {t r a i n}}} (\boldsymbol {\theta}) \triangleq \mathrm {E} _ {\mathrm {x}, \mathrm {z} _ {1}, \dots , \mathrm {z} _ {\mathrm {P}}, \mathrm {y}} \left\{\| \boldsymbol {x} - \widehat {\boldsymbol {x}} _ {(P _ {\text {t r a i n}})} \| _ {1} \right\} \tag {8}
$$
$$
\mathcal {L} _ {\mathrm {S D}, P _ {\text {t r a i n}}} (\boldsymbol {\theta}) \triangleq \sqrt {\frac {\pi}{2 P _ {\text {t r a i n}} \left(P _ {\text {t r a i n}} - 1\right)}} \sum_ {i = 1} ^ {P _ {\text {t r a i n}}} \operatorname {E} _ {\mathbf {z} _ {1}, \dots , \mathbf {z} _ {\mathrm {P}}, \mathbf {y}} \left\{\left\| \widehat {\boldsymbol {x}} _ {i} - \widehat {\boldsymbol {x}} _ {(P _ {\text {t r a i n}})} \right\| _ {1} \right\}, \tag {9}
$$
and where  $\{\widehat{\pmb{x}}_i\}$  denote the generated samples and  $\widehat{\pmb{x}}_{(P)}$  their  $P$ -sample average:
$$
\widehat {\boldsymbol {x}} _ {i} \triangleq G _ {\boldsymbol {\theta}} \left(\boldsymbol {z} _ {i}, \boldsymbol {y}\right), \quad \widehat {\boldsymbol {x}} _ {(P)} \triangleq \frac {1}{P} \sum_ {i = 1} ^ {P} \widehat {\boldsymbol {x}} _ {i}. \tag {10}
$$
The use of supervised- $\ell_1$  loss and SD reward in (7) is not heuristic. As shown in Proposition 3.1, it encourages the samples  $\{\widehat{\pmb{x}}_i\}$  to match the true posterior in both mean and covariance.
Figure 1: The contours show the regularizer value versus  $\pmb{\theta} = [\mu, \sigma]^\top$  for four different regularizers: (a) supervised- $\ell_1$  plus SD reward with  $\beta_{\mathrm{SD}} = \beta_{\mathrm{SD}}^{\mathcal{N}}$  at  $P_{\mathrm{train}} = 2$ , (b) supervised- $\ell_1$  plus SD reward with  $\beta_{\mathrm{SD}} = \beta_{\mathrm{SD}}^{\mathcal{N}}$  at  $P_{\mathrm{train}} = 8$ , (c) supervised- $\ell_2$  at  $P_{\mathrm{train}} = 8$ , and (d) supervised- $\ell_2$  plus variance reward at  $P_{\mathrm{train}} = 8$ . The red star shows the true posterior parameters  $[\mu_0, \sigma_0]^\top$ .
Proposition 3.1. Suppose  $P_{\mathrm{train}} \geq 2$  and  $\pmb{\theta}$  has complete control over the  $\pmb{y}$ -conditional mean and covariance of  $\hat{\pmb{x}}_i$ . Then the parameters  $\pmb{\theta}_{*} = \arg \min_{\pmb{\theta}} \mathcal{L}_{1,\mathrm{SD},P_{\mathrm{train}}}(\pmb{\theta},\beta_{\mathrm{SD}}^{\mathcal{N}})$  with
$$
\beta_ {\mathsf {S D}} ^ {\mathcal {N}} \triangleq \sqrt {\frac {2}{\pi P _ {\text {t r a i n}} \left(P _ {\text {t r a i n}} + 1\right)}} \tag {11}
$$
yield generated statistics
$$
\mathrm {E} _ {\mathbf {z} _ {| \mathcal {Y}}} \left\{\widehat {\boldsymbol {x}} _ {i} \left(\boldsymbol {\theta} _ {*}\right) | \boldsymbol {y} \right\} = \mathrm {E} _ {\mathbf {x} | \mathcal {Y}} \left\{\boldsymbol {x} | \boldsymbol {y} \right\} = \widehat {\boldsymbol {x}} _ {\text {m m s e}} \tag {12a}
$$
$$
\operatorname {C o v} _ {\mathbf {z} _ {| \mathbf {y}}} \{\hat {\boldsymbol {x}} _ {i} \left(\boldsymbol {\theta} _ {*}\right) | \boldsymbol {y} \} = \operatorname {C o v} _ {\mathbf {x} | \mathbf {y}} \left\{\boldsymbol {x} | \boldsymbol {y} \right\} \tag {12b}
$$
when the elements of  $\widehat{\pmb{x}}_i$  and  $\pmb{x}$  are independent Gaussian conditioned on  $\pmb{y}$ . Thus, minimizing  $\mathcal{L}_{1,\mathrm{SD},P_{\mathrm{train}}}$  encourages the  $\pmb{y}$ -conditional mean and covariance of  $\widehat{\pmb{x}}_i$  to match those of the true  $\pmb{x}$ .
See Appendix B for a proof. In imaging applications,  $\widehat{\pmb{x}}_i$  and  $\pmb{x}$  may not be independent Gaussian conditioned on  $\pmb{y}$ , and so the value of  $\beta_{\mathrm{SD}}$  in (11) may not be appropriate. Thus we propose a method to automatically tune  $\beta_{\mathrm{SD}}$  in Section 3.2.
Figure 1 shows a toy example with parameters  $\theta = [\mu, \sigma]^{\top}$ , generator  $G_{\theta}(z,y) = \mu + \sigma z$ , and  $z \sim \mathcal{N}(0,1)$ , giving generated posterior  $p_{\widehat{x} |y}(x|y) = \mathcal{N}(x; \mu, \sigma^2)$ . Assuming the true  $p_{x|y}(x|y) = \mathcal{N}(x; \mu_0, \sigma_0^2)$ , Figs. 1(a)-(b) show that, by minimizing the proposed  $\mathcal{L}_{1,\mathrm{SD},P_{\mathrm{train}}}(\theta, \beta_{\mathrm{SD}}^{\mathcal{N}})$  regularization over  $\theta = [\mu, \sigma]^{\top}$  for any  $P_{\mathrm{train}} \geq 2$ , we recover the true  $\theta_0 = [\mu_0, \sigma_0]^{\top}$ . They also show that the cost function steepens as  $P_{\mathrm{train}}$  decreases, with agrees with our empirical finding that  $P_{\mathrm{train}} = 2$  tends to work best in practice.
We note that regularizing a cGAN with supervised- $\ell_1$  loss alone is not new; see, e.g., [7]. 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 [30]. But regularizing a cGAN using supervised- $\ell_1$  loss alone can push the generator towards mode collapse, for reasons described below. For example, in [7],  $\ell_1$ -induced mode collapse led the authors to use dropout to induce generator variation, instead of random  $z_i$ .
Why not supervised- $\ell_2$  regularization? One may wonder: Why regularize using supervised- $\ell_1$  loss plus an SD reward in (7) 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 [7]. In this case, to train the generator, one would aim to solve  $\arg \min_{\theta} \{\mathcal{L}_{\mathrm{adv}}(\pmb{\theta}, \phi) + \lambda \mathcal{L}_2(\pmb{\theta})\}$  with some  $\lambda > 0$  and
$$
\mathcal {L} _ {2} (\boldsymbol {\theta}) \triangleq \mathrm {E} _ {\mathrm {x}, \mathrm {y}} \left\{\| \boldsymbol {x} - \mathrm {E} _ {\mathrm {z}} \left\{G _ {\boldsymbol {\theta}} (\boldsymbol {z}, \boldsymbol {y}) \right\} \| _ {2} ^ {2} \right\}. \tag {13}
$$
Ohayon et al. [23] revived this idea for the explicit purpose of fighting mode collapse. In practice, however, the  $\mathrm{E}_z$  term in (13) must be implemented by a finite-sample average, giving
$$
\mathcal {L} _ {2, P _ {\text {t r a i n}}} (\boldsymbol {\theta}) \triangleq \mathrm {E} _ {\mathrm {x}, \mathrm {z} _ {1}, \dots , \mathrm {z} _ {\mathrm {P}}, \mathrm {y}} \left\{\left\| \boldsymbol {x} - \frac {1}{P _ {\text {t r a i n}}} \sum_ {i = 1} ^ {P _ {\text {t r a i n}}} G _ {\boldsymbol {\theta}} \left(\boldsymbol {z} _ {i}, \boldsymbol {y}\right) \right\| _ {2} ^ {2} \right\} \tag {14}
$$
for some  $P_{\mathrm{train}} \geq 2$ . For example, Ohayon's implementation [31] used  $P_{\mathrm{train}} = 8$ . As we show in Proposition 3.2,  $\mathcal{L}_{2,P_{\mathrm{train}}}$  induces mode collapse rather than prevents it.
Proposition 3.2. Say  $P_{\mathrm{train}}$  is finite and  $\pmb{\theta}$  has complete control over the  $y$ -conditional mean and covariance of  $\hat{\pmb{x}}_i$ . Then the parameters  $\pmb{\theta}_{*} = \arg \min_{\pmb{\theta}}\mathcal{L}_{2,P_{\mathrm{train}}}(\pmb{\theta})$  yield generated statistics
$$
\mathrm {E} _ {\mathbf {z} _ {| \mathcal {Y}}} \left\{\widehat {\boldsymbol {x}} _ {i} \left(\boldsymbol {\theta} _ {*}\right) | \boldsymbol {y} \right\} = \mathrm {E} _ {\mathbf {x} | \mathcal {Y}} \left\{\boldsymbol {x} | \boldsymbol {y} \right\} = \widehat {\boldsymbol {x}} _ {\text {m m s e}} \tag {15a}
$$
$$
\operatorname {C o v} _ {\mathbf {z} _ {i} | \mathbf {y}} \left\{\widehat {\boldsymbol {x}} _ {i} \left(\boldsymbol {\theta} _ {*}\right) | \boldsymbol {y} \right\} = \mathbf {0}. \tag {15b}
$$
Thus, minimizing  $\mathcal{L}_{2,P_{\mathrm{train}}}$  encourages mode collapse.
Figure 2: Example PSNR of  $\widehat{\alpha}_{(P)}$  versus  $P$ , the number of averaged outputs, for several training  $\beta_{\mathrm{SD}}$  and MRI recovery at  $R = 4$ . Also shown is the theoretical behavior for true-posterior samples.
Number of averaged outputs,  $P$ , on a log scale
The proof (see Appendix C) follows from the bias-variance decomposition of (14), i.e.,
$$
\begin{array}{l} \mathcal {L} _ {2, P _ {\text {t r a i n}}} (\boldsymbol {\theta}) \\ = \mathrm {E} _ {\mathrm {y}} \left\{\| \widehat {\boldsymbol {x}} _ {\text {m m s e}} - \mathrm {E} _ {\mathrm {z} _ {| |} \mathrm {y}} \left\{\widehat {\boldsymbol {x}} _ {i} (\boldsymbol {\theta}) | \boldsymbol {y} \right\} \| _ {2} ^ {2} + \frac {1}{P _ {\text {t r a i n}}} \operatorname {t r} \left[ \operatorname {C o v} _ {\mathrm {z} _ {| |} \mathrm {y}} \left\{\widehat {\boldsymbol {x}} _ {i} (\boldsymbol {\theta}) | \boldsymbol {y} \right\} \right] + \mathrm {E} _ {\mathrm {x} _ {| |} \mathrm {y}} \left\{\| \boldsymbol {e} _ {\text {m m s e}} \| _ {2} ^ {2} | \boldsymbol {y} \right\} \right\}, \tag {16} \\ \end{array}
$$
where  $e_{\mathrm{mmse}} \triangleq x - x_{\mathrm{mmse}}$  is the MMSE error. Figure 1(c) shows that  $\mathcal{L}_{2,P_{\mathrm{train}}}$  regularization causes mode collapse in the toy example, and Section 4.2 shows that it causes mode collapse in MRI.
Why not supervised  $\ell_{2}$  plus a variance reward? To mitigate  $\mathcal{L}_{2,P_{train}}$  's incentive for mode-collapse, the second term in (16) could be canceled using a variance reward, giving