Title: pcaGAN: Improving Posterior-Sampling cGANs via Principal Component Regularization
Venue: NeurIPS 2024
Accepted: yes
Keywords: Image recovery

inverse problems

MRI

posterior sampling

GAN
Abstract: In ill-posed imaging inverse problems, there can exist many hypotheses that fit both the observed measurements and prior knowledge of the true image. Rather than returning just one hypothesis of that image, posterior samplers aim to explore the full solution space by generating many probable hypotheses, which can later be used to quantify uncertainty or construct recoveries that appropriately navigate the perception/distortion trade-off. In this work, we propose a fast and accurate posterior-sampling conditional generative adversarial network (cGAN) that, through a novel form of regularization, aims for correctness in the posterior mean as well as the trace and K principal components of the posterior covariance matrix. Numerical experiments demonstrate that our method outperforms competitors in a wide range of ill-posed imaging inverse problems.

pcaGAN: Improving Posterior-Sampling cGANs via Principal Component Regularization

Abstract

In ill-posed imaging inverse problems, there can exist many hypotheses that fit both the observed measurements and prior knowledge of the true image. Rather than returning just one hypothesis of that image, posterior samplers aim to explore the full solution space by generating many probable hypotheses, which can later be used to quantify uncertainty or construct recoveries that appropriately navigate the perception/distortion trade-off. In this work, we propose a fast and accurate posterior-sampling conditional generative adversarial network (cGAN) that, through a novel form of regularization, aims for correctness in the posterior mean as well as the trace and K principal components of the posterior covariance matrix. Numerical experiments demonstrate that our method outperforms contemporary cGANs and diffusion models in imaging inverse problems like denoising, large-scale inpainting, and accelerated MRI recovery. The code for our model can be found here: https://github.com/matt-bendel/pcaGAN .

1 Introduction

In image recovery, the goal is to recover the true image x from noisy/distorted/incomplete measurements y = M ( x ) . This arises in linear inverse problems such as denoising, deblurring, inpainting, and magnetic resonance imaging (MRI), as well as in non-linear inverse problems like phase-retrieval and image-to-image translation. For all such problems, it is impossible to perfectly recover x from y .

In much of the literature, image recovery is posed as point estimation, where the goal is to return a single best estimate ̂ x . However, there are several shortcomings of this approach. First, it's not clear how to define 'best,' since L2- or L1-minimizing ̂ x are often regarded as too blurry, while efforts to make ̂ x perceptually pleasing can sacrifice agreement with the true image x and cause hallucinations [1-5], as we show later. Also, many applications demand not only a recovery ̂ x but also some quantification of uncertainty in that recovery [6, 7].

As an alternative to point estimation, posterior-sampling -based image recovery [8-27] aims to generate P ≥ 1 samples { ̂ x i } P i =1 from the posterior distribution p x | y ( ·| y ) . Posterior sampling facilitates numerous strategies to quantify the uncertainty in estimating x , or any function of x , from y [6, 7]. It also can help with visualizing uncertainty and increasing robustness to adversarial attacks [28]. That said, the design of accurate and computationally-efficient posterior samplers remains an open problem.

Given a training dataset of image/measurement pairs { ( x t , y t ) } T t =1 , our goal is to build a generator G θ that, for a given y , maps random code vectors z ∼ N ( 0 , I ) to samples of the posterior, i.e., ̂ x = G θ ( z , y ) ∼ p x | y ( · , y ) . There are many ways to approximate the ideal generator. The recent literature has focused on conditional variational autoencoders (cV AEs) [12-14], conditional

normalizing flows (cNFs) [8-11], conditional generative adversarial networks (cGANs) [15-19], and Langevin/score/diffusion-based generative models [20-27]. Although diffusion models have garnered a great share of recent attention, they tend to generate samples several orders-of-magnitude slower than their cNF, cVAE, and cGAN counterparts.

With the goal of fast and accurate posterior sampling, recent progress in cGAN training has been made through regularization. For example, Ohayon et al. [29] proposed to enforce correctness in the generated y -conditional mean using a novel form of L2 regularization. Later, Bendel et al. [19] proposed to enforce correctness in the generated y -conditional mean and trace-covariance using L1 regularization plus a correctly weighted standard-deviation (SD) reward, and gave evidence that their 'rcGAN' competes with contemporary diffusion samplers on MRI and inpainting tasks. Soon after, Man et al. [30] showed that L2 regularization and a variance reward are effective when training cGANs for JPEG decoding. More details are given in Section 2.

In a separate line of work, Nehme et al. [31] trained a Neural Posterior Principal Components (NPPC) network to directly estimate the eigenvectors and eigenvalues of the y -conditional covariance matrix, which allows powerful insights into the nature of uncertainty in an inverse problem.

Our contributions. Inspired by regularized cGANs and NPPC, we propose a novel 'pcaGAN' that encourages correctness in the K principal components of the y -conditional covariance matrix, as well as the y -conditional mean and trace covariance, when sampling from the posterior. We demonstrate that our pcaGAN outperforms existing cGANs [16, 17, 19, 29] and diffusion models [22, 25-27] on posterior sampling tasks like denoising, large-scale inpainting, and accelerated MRI recovery, while sampling orders-of-magnitude faster than those diffusion models. We also demonstrate that pcaGAN recovers the principal components more accurately than NPPC using approximately the same runtime.

2 Background

Our approach builds on the rcGAN regularization framework from [19], which itself builds on the cGAN framework from [16], both of which we now summarize. Let X , Y , and Z denote the domains of x , z , and y , respectively. The goal is to design a generator G θ : Z × Y → X where, for a given y , the random ̂ x = G θ ( z , y ) induced by the code vector z ∼ p z (with z independent of y ) has a distribution that is close to the true posterior p x | y ( · , y ) in the Wasserstein-1 distance, given by

Here, D : X × Y → R is a 'critic' or 'discriminator' that tries to distinguish between the true x and generated ̂ x given y , and L 1 denotes the set of 1-Lipschitz functions. A loss is constructed by averaging (1) over y ∼ p y , which takes the form

using the fact that the expectation commutes with the supremum [16]. D is then implemented as a neural network D ϕ with parameters ϕ . Finally, ( θ , ϕ ) are learned by minimizing

over θ and minimizing -L adv ( θ , ϕ ) + L gp ( ϕ ) over ϕ , where L gp ( ϕ ) is a gradient penalty that encourages D ϕ ∈ L 1 [32]. In practice, the expectation in (3) is approximated by a sample average over the training data { ( x t , y t ) } .

In the typical case that the training data includes only a single image x t for each measurement vector y t , minimizing L adv ( θ , ϕ ) alone does not encourage the generator to produce diverse samples. Rather, it leads to a form of mode collapse where the code z is ignored. In [16], Adler and ¨ Oktem proposed a two-sample discriminator that encourages diverse generator outputs without compromising the Wasserstein objective (2). In [19], Bendel et al. instead proposed to regularize L adv ( θ , ϕ ) in a way that encourages correct posterior means and trace-covariances, i.e.,

To do this, [19] replaced L adv ( θ , ϕ ) with the regularized adversarial loss

which involves the P rc -sample supervisedℓ 1 loss and standard-deviation (SD) reward terms

where typically P rc = 2 . Here, { ̂ x i } are the generated samples and ̂ x ( P ) is their P -sample average:

The reward weight β SD in (6) is automatically adjusted to accomplish (5) during training [19]. We note that the L 1 ,P ( θ ) regularization proposed in [19] is closely related to the regularization L 2 ,P ( θ ) ≜ E x , z1 ,..., zP , y {∥ x -̂ x ( P ) ∥ 2 2 } proposed earlier by Ohayon et al. in [29]. A detailed discussion of the advantages and disadvantages of various cGAN regularizations can be found in [19].

3 Proposed method

Whereas rcGAN aimed for correctness in the posterior mean and posterior trace-covariance statistics, our proposed pcaGAN also aims for correctness in the K principal components of the posterior covariance matrix Σ ̂ x | y , where K is user-selectable. To do this, pcaGAN adds two additional regularization terms to the rcGAN objective:

Here, { ( ̂ v k , ̂ λ k ) } K k =1 denote the principal eigenvectors and eigenvalues of the θ -dependent generated covariance matrix Σ ̂ x | y and { ( v k , λ k ) } K k =1 denote the principal eigenvectors and eigenvalues of the true covariance matrix Σ x | y . Because (11) is the classical PCA objective [33], minimizing L evec ( θ ) over θ will drive the generated principal eigenvector ̂ v k towards the true principal eigenvector v k for each k = 1 , . . . , K . Likewise, minimizing L eval ( θ ) over θ will drive the generated principal eigenvalue ̂ λ k towards the true principal eigenvalue λ k for each k = 1 , . . . , K . Based on our experiments, putting ̂ λ k in the denominator works better than the numerator and the squared error in (12) works better than an absolute value.

In practice, the expectations in (11)-(12) are replaced by sample averages over the training data. In the typical case that the training data includes only a single image x t for each measurement vector y t , the quantities µ x | y and { λ k } in (11)-(12) are unknown and non-trivial to estimate for each y t . Hence, when training the pcaGAN, we approximate them with learned quantities. This must be done carefully, however. For example, if µ x | y in (11) was simply replaced by the θ -dependent quantity µ ̂ x | y , then minimizing L evec ( θ ) over θ would encourage µ ̂ x | y to become overly large in order to drive L evec ( θ ) to a large negative value.

Algorithm 1 details our proposed approach to training the pcaGAN. In particular, it describes the steps used to perform a single update of the generator parameters θ based on the training batch { ( x b , y b ) } B b =1 . Before diving into the details, we offer a brief summary of Algorithm 1. For the initial epochs, the rcGAN objective L rcGAN alone is optimized, which allows the generated posterior mean µ ̂ x | y to converge to the vicinity of µ x | y . Starting at E evec epochs, the L evec ( θ ) regularization from (11) is added, but with µ x | y approximated as StopGrad ( µ ̂ x | y ) . The use of StopGrad forces L evec ( θ ) to be minimized by manipulating the eigenvectors { ̂ v k } K k =1 and not the generated posterior mean µ ̂ x | y . These eigenvectors are computed using an SVD of centered approximate-posterior samples. To reduce the computational burden imposed by this SVD, a 'lazy regularization' [34] approach is adopted, which computes L evec ( θ ) only once every M training steps. Training proceeds in this manner until the eigenvectors { ̂ v k } converge. Starting at E eval epochs, the L eval ( θ ) regularization from (12) is added, but with λ k approximated as

Algorithm 1 pcaGAN generator-training iteration

Require: number of estimated eigen-components K , number of samples used for eigenvector and eigenvalue regularization P pca , number of samples used for rcGAN regularization P rc , epoch at which to involve eigenvector regularization E evec , epoch at which to involve eigenvalue regularization E eval , lazy update period M , adversarial loss weight β adv , std regularization weight β SD , eigenvector and eigenvalue regularization weight β pca , training batch { ( x b , y b ) } B b =1 , current model parameters θ , current training epoch e , the current training step s

where StopGrad is used so that the optimization focuses on { ̂ λ k } . The rational behind (13) is that, when ̂ v k = v k and µ ̂ x | y = µ x | y , the terms [ ̂ v ⊤ k ( x b -µ ̂ x | y )] 2 and [ ̂ v ⊤ k ( ̂ x j -µ ̂ x | y )] 2 ∀ j all equal λ k in y b -conditional expectation. This expectation is approximated using a (1 + P pca ) -term sample average in (13) via the squared norm. The eigenvalues { ̂ λ k } in (12) are computed using the previously described SVD and the regularization schedule is again M -lazy.

We now provide additional details on Algorithm 1. After the loss is initialized in line 1, the following steps are executed for each measurement vector y b in the batch. First, approximate posterior samples { ̂ x i } P rc i =1 are generated in line 5, where P rc = 2 as per the suggestion in [19]. Using these samples, the adversarial component of the loss is added in line 6 and the rcGAN regularization is added in line 8. Starting at epoch E evec , lines 11-16 are executed whenever the training iteration is a multiple of M . Nominally, E evec is set where the validation PSNR of ̂ µ (an empirical approximation of µ ̂ x | y ) stabilizes and M = 100 . Within those lines, samples { ̂ x j } P pca j =1 are generated in line 12 (where nominally P pca = 10 K ), their sample mean is computed in line 13, and the SVD of the centered samples is computed in line 14. The top K right singular vectors are then extracted in line 15 in order to construct the L evec ( θ ) regularization, which is added to the overall generator loss L ( θ ) in line 16. Starting at epoch E eval , where nominally E eval = E evec +25 , lines 20-22 are executed whenever the training iteration is a multiple of M . In line 20, the top K eigenvalues { ̂ λ k } are constructed from the previously computed singular values and, in line 22, the L eval ( θ ) regularization is constructed

(a)

W

2

(

p

x

|

y

, p

̂

x

|

y

)

vs.

M

(b)

W

2

(

p

x

|

y

, p

̂

x

|

y

)

vs.

K

(c)

W

2

(

p

x

|

y

, p

̂

x

|

y

)

/d

vs.

d

0.5

0.5

0.4

0.4

100

rcGAN

NPPC

0.3

0.3

10-1

0.2

0.2

10-2

0.1

0.1

10-3

0.0

0

25

5075100125150175200

0.0

101

102

20

40

60

80

100

and added to the overall training loss. The construction of L eval ( θ ) was previously described around (13). Finally, once the losses for all batch elements have been incorporated, the gradient ∇L ( θ ) is computed using back-propagation and a gradient-descent step of θ is performed using the Adam optimizer [35] in line 26.

4 Numerical experiments

We now present experiments with Gaussian data, MNIST denoising, MRI, and FFHQ face inpainting. Additional implementation and training details for each experiment are provided in Appendix D.

4.1 Recovering synthetic Gaussian data

Here our goal is to recover x ∼ N ( µ x , Σ x ) ∈ R d from y = Mx + w ∈ R d , where M masks x at even indices and noise w ∼ N ( 0 , σ 2 I ) is independent of x with σ 2 = 0 . 001 . Since x and y are jointly Gaussian, the posterior is Gaussian posterior with µ x | y = µ x + Σ xy Σ -1 y ( y -µ y ) and Σ x | y = Σ x -Σ xy Σ -1 y Σ yx , where µ x , µ y , Σ x , Σ y are marginal and Σ xy , Σ yx are joint statistics.

We generate random µ x ∼ N ( 0 , I ) and Σ x with half-normal eigenvalues λ k (see additional details in App. A), and consider a sequence of problem sizes d = 10 , 20 , 30 , . . . , 100 . For each d , we generate 70 000 training, 20 000 validation, and 10 000 test samples. The generator and discriminator are simple multilayer perceptrons (see App. D.1) trained for 100 epochs with K = d , E evec = 10 , β adv = 10 -5 , and β pca = 10 -2 .

Competitors. We compare the proposed pcaGAN to rcGAN [19] and NPPC [31]. rcGAN uses the same generator and discriminator architectures as pcaGAN and is trained according to (6) with β adv = 10 -5 and P rc = 2 . For NPPC, we use the authors' implementation [36] with K = d and some minor modifications to work with vector data. To evaluate performance, we use the Wasserstein-2 (W2) distance between p x | y and ̂ p x | y , which in the Gaussian case reduces to

For the cGANs, we compute ̂ µ x | y and ̂ Σ x | y empirically from 10 d samples, while for NPPC we use the conditional mean, eigenvalues, and eigenvectors returned by the approach.

Results. Figure 1a examines the impact of the lazy update period M on pcaGAN's W2 distance at d = 100 with K = d . Based on this figure, to balance performance with training overhead, we set M = 100 for all future experiments. Figure 1b examines the impact of K on W2 distance for the pcaGAN with d = 100 . It shows that using K < d causes a relatively mild increase in W2 distance, as expected due to the half-normal distribution on the true eigenvalues λ k . Figure 1c shows that the proposed pcaGAN outperforms rcGAN and NPPC in W2 distance for all problem sizes d .

4.2 MNIST denoising

Now our goal is to recover an MNIST digit x ∈ [0 , 1] 28 × 28 from noisy measurements y = x + w with w ∼ N ( 0 , I ) . We randomly split the MNIST training fold into 50 000 training and 10 000

K

= 10

)

̂

v

1

̂

v

2

̂

v

3

̂

v

4

̂

v

5

y

α

=

-

3

α

=

-

2

α

=0

α

=2

α

=3

x

̂

µ

x

|

y

̂

µ

x

|

y

+

α

̂

v

4

̂

µ

x

|

y

+

α

̂

v

1

(b) NPPC

̂

v

1

̂

v

2

̂

v

3

̂

v

4

̂

v

5

y

α

=

-

3

α

=

-

2

α

=0

α

=2

α

=3

x

̂

µ

x

|

y

̂

µ

x

|

y

+

α

̂

v

4

̂

µ

x

|

y

+

α

̂

v

1

validation images, and we use the entire MNIST fold set for testing. For pcaGAN and rcGAN, we use a UNet [37] generator and the encoder portion of the same UNet followed by one dense layer as the discriminator. pcaGAN was trained for 125 epochs with E evec = 25 , β adv = 10 -5 , β pca = 10 -1 , and K ∈ { 5 , 10 } .

Competitors. We again compare the proposed pcaGAN to rcGAN and NPPC. For rcGAN we used the same generator and discriminator architectures as pcaGAN and trained according to (6) with β adv = 10 -5 and P rc = 2 . For NPPC, we used the authors' MNIST implementation from [36].

Following the NPPC paper [31], we evaluate performance using root MSE (rMSE) E x , y {∥ x -̂ µ x | y ∥ 2 } and residual error magnitude (REM 5 ) E x , y {∥ ∥ ( I -̂ V 5 ̂ V 5 ⊤ ) e ∥ ∥ 2 } , where e = x -̂ µ x | y and ̂ V 5 is an 28 2 × 5 matrix whose k th column equals the k th principal eigenvector ̂ v k . For the cGANs, we use ̂ µ x | y = ̂ x ( P ) and compute { ̂ v k } from the SVD of a matrix of centered samples { ̂ x i } P i =1 , both with P = 100 . For NPPC, we use the conditional means and eigenvectors returned by the approach. For performance evaluation, we also consider Conditional Frechet Inception Distance (CFID) [38] with InceptionV3 features. CFID is analogous to Frechet Inception Distance (FID) [39] but applies to conditional distributions (see App. B for more details).

Results. Table 1 shows rMSE, REM 5 , CFID, and the reconstruction time for a batch of 128 images on the test fold. (NPPC does not generate image samples and so CFID does not apply.) The table shows that the proposed pcaGAN wins in all metrics, except for rMSE where NPPC wins. This is not surprising because NPPC computes ̂ µ x | y using a dedicated network trained to minimize MSE loss. NPPC also generates its eigenvectors slightly quicker than pcaGAN generates samples. Table 1 also shows that pcaGAN performance improves as K increases from 5 to 10, despite the fact that REM 5 uses only the top 5 eigenvectors. Figure E.1 shows examples of the 5 principal eigenvectors and posterior mean learned by pcaGAN and NPPC. The eigenvectors of pcaGAN are more structured and less noisy than those of NPPC. Figure E.1 also shows ̂ µ x | y + α v k for α ∈ [ -3 , 3] and k ∈ { 1 , 4 } . Additional figures can be found in App. E.1.

4.3 Accelerated MRI

We now consider accelerated MRI, where the goal is to recover a complex-valued multicoil image x from masked frequency-domain (i.e., 'k-space') measurements y . To build the image data { x t } , we follow the approach in the rcGAN paper [19], which uses the first 8 slices of all fastMRI [40] T2

(ours)

rcGAN

(Bendel)

pscGAN

(Ohayon)

cGAN

(Adler)

(Jalal)

E2E-VarNet

Truth

Truth

Sample

Sample

Average

(P=32)

brain volumes with at least 8 coils, crops to 384 × 384 pixels, and compresses to 8 virtual coils [41]. This yields 12 200 training, 2 376 testing, and 784 validation images. To create each y t , we transform x t to the k-space, subsample using the Cartesian GRO mask [42] at accelerations R = 4 and R = 8 , and transform the zero-filled k-space measurements back to the image domain.

We train pcaGAN for 100 epochs with K = 1 , E evec = 25 , β adv = 10 -5 , and β pca = 10 -2 and select the final model using validation CFID computed with VGG-16 features.

Competitors. We compare the proposed pcaGAN to rcGAN [19], pscGAN [29], Adler & ¨ Oktem's cGAN [16], the Langevin approach [22], and the E2E-VarNet [43]. All cGANs use the generator and discriminator architectures from [19] and enforce data-consistency [44]. For rcGAN and the Langevin approach, we did not modify the authors' implementations from [45] and [46] except to use the GRO sampling mask. For E2E-VarNet, we use the GRO mask, hyperparameters, and training procedure from [22].

Following [19], we convert the multicoil outputs ̂ x i to complex-valued images using SENSE-based coil combining [47] with ESPIRiT-estimated [48] coil sensitivity maps, and compute performance on magnitude images. All feature-based metrics (CFID, FID, LPIPS, DISTS) were computed with AlexNet features to show that pcaGAN does not overfit to the VGG-16 features used for validation. It was shown in [49] that image-quality metrics computed using ImageNet-trained feature generators like AlexNet and VGG-16 perform comparably to metrics computed using MRI-trained feature generators in terms of correlation with radiologists' scores.

Results. Table 2 shows CFID, FID, APSD ≜ ( 1 P ∑ P i =1 1 N ∥ ̂ x ( P ) -̂ x i ∥ 2 ) 1 / 2 , and 4-sample generation time for the methods under test. Due to its slow sample-generation time, we evaluate the CFID, FID, and APSD of the Langevin technique [22] using the 72-image test from [19]. But due to the bias of CFID at small sample sizes [38], we evaluate the other methods using all 2 376 test images (CFID 2 ) and again using all 14 576 training and test images (CFID 3 ). Table 2 shows that pcaGAN yields better CFID and FID than the competitors. All cGANs generated samples 3-4 orders-of-magnitude faster than the Langevin approach [22].

Table 3 shows PSNR, SSIM, LPIPS [50], and DISTS [51] for the P -sample average ̂ x ( P ) at P ∈ { 1 , 2 , 4 , 8 , 16 , 32 } and R = 8 . (See App. C for R = 4 .) It has been shown that DISTS correlates particularly well with radiologist scores [52]. The E2E-VarNet achieves the best PSNR, but the proposed cGAN achieves the best LPIPS and DISTS when P = 2 and the best SSIM when P = 8 . This P -dependence is related to the perception-distortion trade-off [53] and consistent with that reported in [19].

Figure 3 shows zoomed versions of two recoveries ̂ x i and the sample average ̂ x ( P ) with P = 32 . Appendices E.2 and E.3 show additional plots of ̂ x ( P ) that visually demonstrate the perceptiondistortion trade-off.

4.4 Large-scale inpainting

Our final goal is to inpaint a face image with a large randomly generated masked region. For this task, we use 256 × 256 FFHQ face images [54] and the mask generation procedure from [17]. We randomly split the FFHQ training fold into 45 000 training and 5 000 validation images, and we use the remaining 20 000 images for testing.

For pcaGAN, we use CoModGAN's [17] generator and discriminator architecture and train for 100 epochs using K = 2, E evec = 25, β adv = 5 × 10 -3 , and β pca = 10 -3 .

Competitors. We compare with CoModGAN [17], pscGAN [29], rcGAN [19], and state-of-theart diffusion methods DDRM (20 NFEs) [25], DDNM (100 NFEs) [27], and DPS (1000 NFEs) [26]. CoModGAN, pscGAN, and rcGAN differ from pcaGAN only in generator regularization and CoModGAN's use of discriminator MBSD [55]. For rcGAN, DDNM, DDRM, and DPS, we use the authors' implementations from [45], [56], [57], and [58] with mask generation from [17]. FID and CFID were evaluated on our 20 000 image test set with P =1 .

Results. Table 4 shows test CFID, FID, LPIPS, and 40-sample generation time. The table shows that the proposed pcaGAN wins in CFID, FID and LPIPS, and that the four cGANs generate samples 3-4 orders-of-magnitude faster than DPS. Figure 4 shows five generated samples for each method under test, along with the true and masked image. pcaGAN shows better subjective quality than the competitors, as well as good diversity. Additional figures can be found in App. E.4.

Original

SgMaskedN

SgCoModGANN

SgDDNM N

SgDDRM N

5 Discussion

When training a cGAN, the overall goal is that the samples { ̂ x i } generated from a particular y accurately represent the true posterior p x | y ( ·| y ) . Achieving this goal is challenging when training from paired data { ( x t , y t ) } , because such datasets provide only one example of x for each given y . Early methods like [15-17] focused on providing some variation among { ̂ x i } , but did not aim for the correct variation. The rcGAN from [19] focused on providing the correct amount of variation by enforcing tr( Σ ̂ x | y ) = tr( Σ x | y ) , and the proposed pcaGAN goes farther by encouraging Σ ̂ x | y and Σ x | y to agree along K principal directions. Our experiments demonstrate that pcaGAN yields a notable improvement over rcGAN and outperforms contemporary diffusion approaches like DPS [26].

PCA principles have also been used in unconditional GANs, where the goal is to train a generator G θ that turns codes z ∼ N ( 0 , I ) into outputs ̂ x = G θ ( z ) that match the true marginal distribution p x from which the training samples { x t } are drawn. For example, the eigenGAN from [59] aims to train in such a way that semantic attributes are learned (without supervision) and can be independently controlled by manipulating individual entries of z . But their goal is clearly different from ours.

Limitations. We acknowledge several limitations of our work. First, generating P pca = 10 K samples during training can impose a burden on memory when x is high dimensional. In the multicoil MRI experiment, x ∈ R d for d = 2 . 4e6 , which limited us to K = 1 at batch size 2. Second, although our focus is on designing a fast and accurate posterior sampler, more work is needed on how to best use the generated samples across different applications. Using them to compute rigorous uncertainty

intervals seems like a promising direction [60-62]. Third, the application to MRI is preliminary; additional tuning and validation is needed before it can be considered for clinical practice.

6 Conclusion

In this work, we proposed pcaGAN, a novel image-recovery cGAN that enforces correctness in the K principal components of the conditional covariance matrix Σ ̂ x | y , as well as in the conditional mean µ ̂ x | y and trace-covariance tr( Σ ̂ x | y ) . Experiments with synthetic Gaussian data showed pcaGAN outperforming both rcGAN [19] and NPPC [31] in Wasserstein-2 distance across a range of problem sizes. Experiments on MNIST denoising, accelerated multicoil MRI, and large-scale image inpainting showed pcaGAN outperforming several other cGANs and diffusion models in CFID, FID, PSNR, SSIM, LPIPS, and DISTS metrics. Furthermore, pcaGAN generates samples 3-4 orders-of-magnitude faster than the tested diffusion models. The proposed pcaGAN thus provides fast and accurate posterior sampling for image recovery problems, which enables uncertainty quantification, fairness in recovery, and easy navigation of the perception/distortion trade-off.

Acknowledgments and Disclosure of Funding

The authors are funded in part by the National Institutes of Health under grant R01-EB029957.

References

C. Belthangady and L. A. Royer, 'Applications, promises, and pitfalls of deep learning for fluorescence image reconstruction,' Nature Methods , vol. 16, no. 12, pp. 1215-1225, 2019. 1

D. P. Hoffman, I. Slavitt, and C. A. Fitzpatrick, 'The promise and peril of deep learning in microscopy,' Nature Methods , vol. 18, no. 2, pp. 131-132, 2021. 1

M. J. Muckley, B. Riemenschneider, A. Radmanesh, S. Kim, G. Jeong, J. Ko, Y. Jun, H. Shin, D. Hwang, M. Mostapha, et al. , 'Results of the 2020 fastMRI challenge for machine learning MR image reconstruction,' IEEE Trans. Med. Imag. , vol. 40, no. 9, pp. 2306-2317, 2021. 1

S. Bhadra, V. A. Kelkar, F. J. Brooks, and M. A. Anastasio, 'On hallucinations in tomographic image reconstruction,' IEEE Trans. Med. Imag. , vol. 40, no. 11, pp. 3249-3260, 2021. 1

N. M. Gottschling, V. Antun, A. C. Hansen, and B. Adcock, 'The troublesome kernel-On hallucinations, no free lunches and the accuracy-stability trade-off in inverse problems,' arXiv:2001.01258 , 2023. 1

M. Abdar, F. Pourpanah, S. Hussain, D. Rezazadegan, L. Liu, M. Ghavamzadeh, P. Fieguth, X. Cao, A. Khosravi, U. R. Acharya, et al. , 'A review of uncertainty quantification in deep learning: Techniques, applications and challenges,' Info. Fusion , vol. 76, pp. 243-297, 2021. 1

B. Lambert, F. Forbes, S. Doyle, H. Dehaene, and M. Dojat, 'Trustworthy clinical AI solutions: A unified review of uncertainty quantification in deep learning models for medical image analysis,' Artificial Intelligence in Medicine , p. 102830, 2024. 1

L. Ardizzone, C. L¨ uth, J. Kruse, C. Rother, and U. K¨ othe, 'Guided image generation with conditional invertible neural networks,' arXiv:1907.02392 , 2019. 1, 2

C. Winkler, D. Worrall, E. Hoogeboom, and M. Welling, 'Learning likelihoods with conditional normalizing flows,' arXiv:1912.00042 , 2019. 1, 2

H. Sun and K. L. Bouman, 'Deep probabilistic imaging: Uncertainty quantification and multi-modal solution characterization for computational imaging,' in Proc. AAAI Conf. Artificial Intell. , vol. 35, pp. 2628-2637, 2021. 1, 2

J. Wen, R. Ahmad, and P. Schniter, 'A conditional normalizing flow for accelerated multi-coil MR imaging,' in Proc. Intl. Conf. Mach. Learn. , 2023. 1, 2

V. Edupuganti, M. Mardani, S. Vasanawala, and J. Pauly, 'Uncertainty quantification in deep MRI reconstruction,' IEEE Trans. Med. Imag. , vol. 40, pp. 239-250, Jan. 2021. 1

F. Tonolini, J. Radford, A. Turpin, D. Faccio, and R. Murray-Smith, 'Variational inference for computational imaging inverse problems,' J. Mach. Learn. Res. , vol. 21, no. 179, pp. 1-46, 2020. 1

K. Sohn, H. Lee, and X. Yan, 'Learning structured output representation using deep conditional generative models,' in Proc. Neural Info. Process. Syst. Conf. , 2015. 1

P. Isola, J.-Y . Zhu, T. Zhou, and A. A. Efros, 'Image-to-image translation with conditional adversarial networks,' in Proc. IEEE Conf. Comp. Vision Pattern Recog. , pp. 1125-1134, 2017. 1, 2, 9

J. Adler and O. ¨ Oktem, 'Deep Bayesian inversion,' arXiv:1811.05910 , 2018. 1, 2, 7, 8, 9, 14

S. Zhao, J. Cui, Y. Sheng, Y. Dong, X. Liang, E. I.-C. Chang, and Y. Xu, 'Large scale image completion via co-modulated generative adversarial networks,' in Proc. Intl. Conf. Learn. Rep. , 2021. 1, 2, 8, 9, 16

H. Zhao, H. Li, S. Maurer-Stroh, and L. Cheng, 'Synthesizing retinal and neuronal images with generative adversarial nets,' Med. Image Analysis , vol. 49, 07 2018. 1, 2

M. Bendel, R. Ahmad, and P. Schniter, 'A regularized conditional GAN for posterior sampling in inverse problems,' in Proc. Neural Info. Process. Syst. Conf. , 2023. 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 13, 14, 15, 16

M. Welling and Y. W. Teh, 'Bayesian learning via stochastic gradient Langevin dynamics,' in Proc. Intl. Conf. Mach. Learn. , pp. 681-688, 2011. 1, 2

Y. Song and S. Ermon, 'Improved techniques for training score-based generative models,' in Proc. Neural Info. Process. Syst. Conf. , 2020. 1, 2

A. Jalal, M. Arvinte, G. Daras, E. Price, A. Dimakis, and J. Tamir, 'Robust compressed sensing MRI with deep generative priors,' in Proc. Neural Info. Process. Syst. Conf. , 2021. 1, 2, 7, 8, 14, 15

Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole, 'Score-based generative modeling through stochastic differential equations,' in Proc. Intl. Conf. Learn. Rep. , 2021. 1, 2

Y. Song, L. Shen, L. Xing, and S. Ermon, 'Solving inverse problems in medical imaging with score-based generative models,' in Proc. Intl. Conf. Learn. Rep. , 2022. 1, 2

B. Kawar, M. Elad, S. Ermon, and J. Song, 'Denoising diffusion restoration models,' in Proc. Neural Info. Process. Syst. Conf. , 2022. 1, 2, 8

H. Chung, J. Kim, M. T. McCann, M. L. Klasky, and J. C. Ye, 'Diffusion posterior sampling for general noisy inverse problems,' in Proc. Intl. Conf. Learn. Rep. , 2023. 1, 2, 8, 9, 16

Y. Wang, J. Yu, and J. Zhang, 'Zero-shot image restoration using denoising diffusion null-space model,' in Proc. Intl. Conf. Learn. Rep. , 2023. 1, 2, 8

G. Ohayon, T. J. Adrai, M. Elad, and T. Michaeli, 'Reasons for the superiority of stochastic estimators over deterministic ones: Robustness, consistency and perceptual quality,' in Proc. Intl. Conf. Mach. Learn. , pp. 26474-26494, 2023. 1

G. Ohayon, T. Adrai, G. Vaksman, M. Elad, and P. Milanfar, 'High perceptual quality image denoising with a posterior sampling CGAN,' in Proc. IEEE Intl. Conf. Comput. Vis. Workshops , vol. 10, pp. 1805-1813, 2021. 2, 3, 7, 8, 14

S. Man, G. Ohayon, T. Adrai, and M. Elad, 'High-perceptual quality JPEG decoding via posterior sampling,' Proc. IEEE Conf. Comp. Vision Pattern Recog. Workshop , 2023. 2

E. Nehme, O. Yair, and T. Michaeli, 'Uncertainty quantification via neural posterior principal components,' in Proc. Neural Info. Process. Syst. Conf. , 2023. 2, 5, 6, 10, 15

I. Gulrajani, F. Ahmed, M. Arjovsky, V. Dumoulin, and A. Courville, 'Improved training of Wasserstein GANs,' in Proc. Neural Info. Process. Syst. Conf. , p. 5769-5779, 2017. 2

I. T. Jolliffe, Principal Component Analysis , vol. 2. New York: Springer Verlag, 2002. 3

T. Karras, S. Laine, M. Aittala, J. Hellsten, J. Lehtinen, and T. Aila, 'Analyzing and improving the image quality of StyleGAN,' in Proc. IEEE Conf. Comp. Vision Pattern Recog. , pp. 8110-8119, 2020. 3, 16

D. P. Kingma and J. Ba, 'Adam: A method for stochastic optimization,' in Proc. Intl. Conf. Learn. Rep. , 2015. 5, 15

E. Nehme, O. Yair, and T. Michaeli, 'Uncertainty quantification via neural posterior principal components.' Downloaded from https://github.com/EliasNehme/NPPC , Jan. 2023. 5, 6, 15

O. Ronneberger, P. Fischer, and T. Brox, 'U-Net: Convolutional networks for biomedical image segmentation,' in Proc. Intl. Conf. Med. Image Comput. Comput. Assist. Intervent. , pp. 234-241, 2015. 6

M. Soloveitchik, T. Diskin, E. Morin, and A. Wiesel, 'Conditional Frechet inception distance,' arXiv:2103.11521 , 2021. 6, 7, 13

M. Heusel, H. Ramsauer, T. Unterthiner, B. Nessler, and S. Hochreiter, 'GANs trained by a two time-scale update rule converge to a local Nash equilibrium,' in Proc. Neural Info. Process. Syst. Conf. , vol. 30, 2017. 6, 13

J. Zbontar, F. Knoll, A. Sriram, M. J. Muckley, M. Bruno, A. Defazio, M. Parente, K. J. Geras, J. Katsnelson, H. Chandarana, Z. Zhang, M. Drozdzal, A. Romero, M. Rabbat, P. Vincent, J. Pinkerton, D. Wang, N. Yakubova, E. Owens, C. L. Zitnick, M. P. Recht, D. K. Sodickson, and Y. W. Lui, 'fastMRI: An open dataset and benchmarks for accelerated MRI,' arXiv:1811.08839 , 2018. 6

T. Zhang, J. M. Pauly, S. S. Vasanawala, and M. Lustig, 'Coil compression for accelerated imaging with Cartesian sampling,' Magn. Reson. Med. , vol. 69, no. 2, pp. 571-582, 2013. 7

M. Joshi, A. Pruitt, C. Chen, Y. Liu, and R. Ahmad, 'Technical report (v1.0)-pseudo-random cartesian sampling for dynamic MRI,' arXiv:2206.03630 , 2022. 7

A. Sriram, J. Zbontar, T. Murrell, A. Defazio, C. L. Zitnick, N. Yakubova, F. Knoll, and P. Johnson, 'End-to-end variational networks for accelerated MRI reconstruction,' in Proc. Intl. Conf. Med. Image Comput. Comput. Assist. Intervent. , pp. 64-73, 2020. 7, 8, 14, 15

C. K. Sønderby, J. Caballero, L. Theis, W. Shi, and F. Husz´ ar, 'Amortised MAP inference for image super-resolution,' in Proc. Intl. Conf. Learn. Rep. , 2017. 7

M. Bendel, R. Ahmad, and P. Schniter, 'A regularized conditional GAN for posterior sampling in inverse problems.' Downloaded from https://github.com/matt-bendel/rcGAN , May 2023. 7, 8, 13

A. Jalal, M. Arvinte, G. Daras, E. Price, A. Dimakis, and J. Tamir, 'csgm-mri-langevin.' https:// github.com/utcsilab/csgm-mri-langevin , 2021. Accessed: 2021-12-05. 7, 15

K. P. Pruessmann, M. Weiger, M. B. Scheidegger, and P. Boesiger, 'SENSE: Sensitivity encoding for fast MRI,' Magn. Reson. Med. , vol. 42, no. 5, pp. 952-962, 1999. 7

M. Uecker, P. Lai, M. J. Murphy, P. Virtue, M. Elad, J. M. Pauly, S. S. Vasanawala, and M. Lustig, 'ESPIRiT-an eigenvalue approach to autocalibrating parallel MRI: Where SENSE meets GRAPPA,' Magn. Reson. Med. , vol. 71, no. 3, pp. 990-1001, 2014. 7

P. M. Adamson, A. D. Desai, J. Dominic, C. Bluethgen, J. P. Wood, A. B. Syed, R. D. Boutin, K. J. Stevens, S. Vasanawala, J. M. Pauly, A. S. Chaudhari, and B. Gunel, 'Using deep feature distances for evaluating MRimage reconstruction quality,' in Proc. Neural Info. Process. Syst. Workshop , 2023. 7

R. Zhang, P. Isola, A. A. Efros, E. Shechtman, and O. Wang, 'The unreasonable effectiveness of deep features as a perceptual metric,' in Proc. IEEE Conf. Comp. Vision Pattern Recog. , pp. 586-595, 2018. 8, 13

K. Ding, K. Ma, S. Wang, and E. P. Simoncelli, 'Image quality assessment: Unifying structure and texture similarity,' IEEE Trans. Pattern Anal. Mach. Intell. , vol. 44, no. 5, pp. 2567-2581, 2020. 8, 13

S. Kastryulin, J. Zakirov, N. Pezzotti, and D. V . Dylov, 'Image quality assessment for magnetic resonance imaging,' IEEE Access , vol. 11, pp. 14154-14168, 2023. 8

Y. Blau and T. Michaeli, 'The perception-distortion tradeoff,' in Proc. IEEE Conf. Comp. Vision Pattern Recog. , pp. 6228-6237, 2018. 8

T. Karras, S. Laine, and T. Aila, 'A style-based generator architecture for generative adversarial networks,' in Proc. IEEE Conf. Comp. Vision Pattern Recog. , pp. 4396-4405, 2019. 8

T. Karras, T. Aila, S. Laine, and J. Lehtinen, 'Progressive growing of GANs for improved quality, stability, and variation,' in Proc. Intl. Conf. Learn. Rep. , 2018. 8, 14, 15

Y. Wang, 'DDNM.' Downloaded from https://github.com/wyhuai/DDNM, Sept. 2023. 8, 16

B. Kawar, M. Elad, S. Ermon, and J. Song, 'Denoising diffusion restoration models.' Downloaded from https://github.com/bahjat-kawar/ddrm , May 2022. 8, 16

H. Chung, J. Kim, M. T. McCann, M. L. Klasky, and J. C. Ye, 'diffusion-posterior-sampling.' https: //github.com/DPS2022/diffusion-posterior-sampling , Mar. 2023. 8, 16

Z. He, M. Kan, and S. Shan, 'EigenGAN: Layer-wise eigen-learning for GANs,' in Proc. IEEE Intl. Conf. Comput. Vis. , pp. 14408-14417, 2021. 9

A. N. Angelopoulos, A. P. Kohli, S. Bates, M. I. Jordan, J. Malik, T. Alshaabi, S. Upadhyayula, and Y. Romano, 'Image-to-image regression with distribution-free uncertainty quantification and applications in imaging,' in Proc. Intl. Conf. Mach. Learn. , 2022. 10

J. Teneggi, M. Tivnan, J. W. Stayman, and J. Sulam, 'How to trust your diffusion model: A convex optimization approach to conformal risk control,' 2023. 10

D. Narnhofer, A. Habring, M. Holler, and T. Pock, 'Posterior-variance-based error quantification for inverse problems in imaging,' SIAM J. Imag. Sci. , vol. 17, no. 1, pp. 301-333, 2024. 10

Y. Zeng, 'co-mod-gan-pytorch.' Downloaded from https://github.com/zengxianyu/ co-mod-gan-pytorch , Sept. 2022. 16

Supplementary Materials

A Synthetic Gaussian priors

Algorithm 2 captures how we construct the Gaussian priors used in Sec. 4.1.

Algorithm 2 Gaussian prior generation

We begin with dimension d = 100 , generating random mean µ (100) x ∈ R 100 and eigenvalues { λ (100) k } 100 k =1 . To construct the eigenvectors { v (100) k } 100 k =1 , we perform a QR decomposition on a 100 × 100 matrix with i.i.d. N (0 , 1) entries and set v (100) k as the k th column of Q . For each remaining d ∈ { 90 , 80 , . . . , 10 } , we construct µ ( d ) x , { λ ( d ) k } , and u ( d ) k by truncating the previous quantities to ensure some level of continuity across d .

B Conditional Fr´ echet inception distance

We quantify posterior-approximation accuracy using conditional Fr´ echet inception distance (CFID) [38], which approximates the conditional Wasserstein distance

In (B.1), W 2 ( p a , p b ) denotes the Wasserstein-2 distance between distributions p a and p b , defined as

where Π( p a , p b ) ≜ { p a , b : p a = ∫ p a , b d b and p b = ∫ p a , b d a } denotes the set of joint distributions p a , b with marginal distributions p a and p b . Similarly to how FID [39] is computed for marginal distributions, CFID approximates (B.1) for conditional distributions. In particular, the random vectors x , ̂ x , and y are replaced by low-dimensional embeddings x , ̂ x , and y , generated by the convolutional layers of a deep network (e.g., InceptionV3 or VGG-16 or AlexNet), and the embedding distributions p x | y and p ̂ x | y are approximated by multivariate Gaussians. We compute CFID using the code from [45], which is described in [19].

C Additional MRI results for acceleration R = 4

Table 5 shows PSNR, SSIM, LPIPS [50], and DISTS [51] for the P -sample average ̂ x ( P ) at P ∈ { 1 , 2 , 4 , 8 , 16 , 32 } for R = 4 . In this case, the E2E-VarNet attains the best PSNR and SSIM, while pcaGAN performs best in LPIPS and DISTS when P = 2 and P = 1 , respectively. These results,

(ours)

rcGAN

(Bendel)

pscGAN

(Ohayon)

cGAN

(Adler)

(Jalal)

E2E-VarNet

Truth

Truth

Sample

Sample

Average

(P=32)

in conjunction with the R = 8 results discussed in Sec. 4.3, show that pcaGAN yields a notable improvement over rcGAN in all metrics.

Figure C.1 shows zoomed versions of two recoveries ̂ x i and the sample average ̂ x ( P ) with P = 32 .

D Implementation details

In each experiment, all cGANs were trained using the Adam optimizer with a learning rate of 10 -3 , β 1 = 0 , and β 2 = 0 . 99 as in [55].

D.1 Synthetic Gaussian recovery

cGAN training. We choose β adv = 10 -5 , n batch = 64 , P rc = 2 , and train for 100 epochs for both rcGAN and pcaGAN. Running PyTorch on a server with 4 Tesla A100 GPUs, each with 82 GB of memory, the cGAN training for d = 100 takes approximately 8 hours, with training time decreasing with smaller d . For pcaGAN, we choose K = d for each d in this experiment (unless otherwise noted) and β pca = 10 -2 .

cGAN architecture. We exploit the Gaussian nature of the problem, constructing G θ with two dense layers; one which takes in y as input and one which takes in z as input. The output of each layer is added, yielding ̂ x . Similarly, D ϕ is comprised of a single dense layer which takes in the concatenation of x / ̂ x and y and outputs a scalar. We use this architecture for both rcGAN and pcaGAN. Note that there is no listed license for rcGAN.

NPPC. For NPPC, we use the suggested hyperparameters from [31] and opt to train the MMSE network before training NPPC. We use the suggested architectures from their Gaussian denoising experiment and train for 100 epochs with n batch = 64 . We leverage the authors' implementation in [36], modifying it for this Gaussian problem. There is no listed license for NPPC.

D.2 MNIST denoising

The MNIST dataset is available under the GNU general public license, which we respect through our use.

cGAN training. We choose β adv = 10 -5 , n batch = 64 , P rc = 2 , and train for 125 epochs for both rcGAN and pcaGAN. Running PyTorch on a server with 4 Tesla A100 GPUs, each with 82 GB of memory, the cGAN training for d = 100 takes approximately 8 hours, with training time decreasing with smaller d . For pcaGAN, we train two models, one with K = 5 and one with K = 10 . In both cases, β pca = 10 -1 , E evec = 25 , and E eval = 50 .

cGAN architecture. For both cGANs, the generator is the standard UNet which takes in the concatenation of y and code z . The network consists of 3 pooling layers and 32 initial channels. The convolutions use a kernel of size 3 × 3 , instance normalization, and leaky ReLU activations with a negative slope of 0 . 1 . For the encoder portion of the UNet, we use max pooling with a kernel size of 2 × 2 to reduce spatial resolution by a factor of 2. For the decoder portion of the UNet, we use nearest-neighbor interpolation to increase spatial resolution by a factor of 2. The discriminator is simply the encoder portion of the UNet with an additional dense layer appended to map the UNet's latent space to a scalar output. Note that there is no listed license for rcGAN.

NPPC. For NPPC, we do not modify the authors' implementation in [36] in any way. We first train the MMSE reconstruction network and then train the NPPC network. There is no listed license for NPPC.

D.3 Accelerated MRI

For our MRI experiments, we use the fastMRI dataset which is available under the MIT license, which we respect through our use.

cGAN training. We choose β adv = 10 -5 , β pca = 10 -2 , n batch = 2 , P rc = 2 , K = 1 , E evec = 25 , and E eval = 50 for pcaGAN. For rcGAN, pscGAN , and Adler's cGAN, we use the hyperparameters and training procedure described in [19]. All models were trained for 100 epochs using the Adam optimizer [35] with a learning rate of 10 -3 , β 1 = 0 , and β 2 = 0 . 99 , as in [55]. Running PyTorch on a server with 4 Tesla A100 GPUs, each with 82 GB of memory, the training of each MRI cGAN took approximately 1 day. Note that there is no listed license for rcGAN, pscGAN, or Adler and ¨ Oktem's cGAN.

cGAN architecture. All four cGANs used the same generator and discriminator architectures described in [19], except that Adler and ¨ Oktem's discriminator used extra input channels to facilitate the 3-input loss.

E2E-VarNet. For the Sriram et al.'s E2E-VarNet [43], we use the same training procedure and hyperparameters outlined in [22] with modification to the sampling pattern as in [19]. As in [22], we use the SENSE-based coil-combined image as the ground truth instead of the RSS image. The E2E-VarNet is available under the MIT license, which we respect.

Langevin approach. For Jalal et al.'s MRI approach [22], we do not modify the authors' implementation from [46] other than replacing the default sampling pattern with the GRO sampling mask. We borrow both generated samples and results from [19]. The authors' code is available under the MIT license, which we respect.

D.4 FFHQ Inpainting

For our inpainting experiment, we use the FFHQ dataset which is available under the Creative Commons BY-NC-SA 4.0 license, which we respect through our use.

cGAN training. We choose β adv = 5 × 10 -3 , β pca = 10 -3 , n batch = 5 , P rc = 2 , K = 2 , E evec = 25 , and E eval = 50 for pcaGAN. Running PyTorch on a server with 4 Tesla A100 GPUs, each with 82 GB of memory, the training of our cGAN took approximately 1.5 days.

cGAN architecture. As in [19], we use the CoModGAN networks from [17] which extend the StyleGAN2 [34] network. The StyleGAN2 architecture is available under the NVIDIA Source Code License, which we respect.

rcGAN. We follow the training procedure outlined in [19], only modifying the inpainting mask to be random. The total training time on a server with 4 NVIDIA A100 GPUs, each with 82 GB of memory, is roughly 1 day. There is no listed license for rcGAN.

pscGAN. We use the same training procedure outlined in [19], modifying the inpainting masks to be random and using the L 2 ,P objective described briefly in Sec. 2 with P =8 . The total training time on a server with 4 NVIDIA A100 GPUs, each with 82 GB of memory, is roughly 1.5 days. There is no listed license for pscGAN.

CoModGAN. We use the PyTorch implementation of CoModGAN from [63] and train the model. The total training time on a server with 4 NVIDIA A100 GPUs, each with 82 GB of memory, is roughly 1 day. There is no listed license for CoModGAN, beyond the NVIDIA Source Code License.

D.4.1 Diffusion Methods

For all three diffusion methods, we use the pretrained weights from [26].

DPS. We use the suggested settings for the 256 × 256 FFHQ dataset and the code from the official PyTorch implementation [58]. We found the LPIPS-minimizing step-size ζ via grid search over a 1000 image validation set. We generate 1 sample for all 20 000 images in our test set, using a batch-size of 1 and 1000 NFEs. The total generation time on a server with 4 NVIDIA A100 GPUs, each with 82 GB of memory, is roughly 9 days. There is no listed license for DPS.

DDNM. We use the code from the official PyTorch implementation [56]. We generate 1 sample for all 20 000 images in our test set, using a batch-size of 1 and 100 NFEs. The total generation time on a server with 4 NVIDIA A100 GPUs, each with 82 GB of memory, is roughly 1.5 days. There is no listed license for DDNM.

DDRM. We use the code from the official PyTorch implementation [57]. We generate 1 sample for all 20 000 images in our test set, using a batch-size of 1 and 20 NFEs. The total generation time on a server with 4 NVIDIA A100 GPUs, each with 82 GB of memory, is roughly 5.5 hours. There is no listed license for DDRM.

E Additional reconstruction plots

E.1 MNIST denoising

K

= 10

)

̂

v

1

̂

v

2

̂

v

3

̂

v

4

̂

v

5

y

α

=

-

3

α

=

-

2

α

=0

α

=2

α

=3

x

̂

µ

x

|

y

̂

µ

x

|

y

+

α

̂

v

4

̂

µ

x

|

y

+

α

̂

v

1

(b) NPPC

̂

v

1

̂

v

2

̂

v

3

̂

v

4

̂

v

5

y

α

=

-

3

α

=

-

2

α

=0

α

=2

α

=3

x

̂

µ

x

|

y

̂

µ

x

|

y

+

α

̂

v

4

̂

µ

x

|

y

+

α

̂

v

1

南

41

h

h

h

h

h

h

h

h

h

h

4

4

K

= 10

)

̂

v

1

̂

v

2

̂

v

3

̂

v

4

̂

v

5

y

α

=

-

3

α

=

-

2

α

=0

α

=2

α

=3

x

̂

µ

x

|

y

̂

µ

x

|

y

+

α

̂

v

4

̂

µ

x

|

y

+

α

̂

v

1

(b) NPPC

̂

v

1

̂

v

2

̂

v

3

̂

v

4

̂

v

5

y

α

=

-

3

α

=

-

2

α

=0

α

=2

α

=3

x

̂

µ

x

|

y

̂

µ

x

|

y

+

α

̂

v

4

̂

µ

x

|

y

+

α

̂

v

1

7

7

7

7

7

7

7

7

7

7

7

K

= 10

)

̂

v

1

̂

v

2

̂

v

3

̂

v

4

̂

v

5

y

α

=

-

3

α

=

-

2

α

=0

α

=2

α

=3

x

̂

µ

x

|

y

̂

µ

x

|

y

+

α

̂

v

4

̂

µ

x

|

y

+

α

̂

v

1

(b) NPPC

̂

v

1

̂

v

2

̂

v

3

̂

v

4

̂

v

5

y

α

=

-

3

α

=

-

2

α

=0

α

=2

α

=3

x

̂

µ

x

|

y

̂

µ

x

|

y

+

α

̂

v

4

̂

µ

x

|

y

+

α

̂

v

1

13.

9.

9

口

9

9

9

9

3

子

a

9

a

a

a

3

9

a

a

a

E.2 MRI at acceleration R = 4

E2E-VarNet

(ours)

rcGAN

(Bendel)

pscGAN

(Ohayon)

cGAN

(Adler)

(Jalal)

Truth

Truth

Sample

Sample

Average

(P=2)

Average

(P=4)

Average

(P=32)

2e-05

0e+00

E2E-VarNet

(ours)

rcGAN

(Bendel)

pscGAN

(Ohayon)

cGAN

(Adler)

(Jalal)

Truth

Truth

Sample

Sample

Average

(P=2)

Average

(P=4)

Average

(P=32)

2e-05

0e+00

E.3 MRI at acceleration R = 8

E2E-VarNet

(ours)

rcGAN

(Bendel)

pscGAN

(Ohayon)

cGAN

(Adler)

(Jalal)

Truth

Truth

Sample

Sample

Average

(P=2)

Average

(P=4)

Average

(P=32)

5e-05

0e+00

E2E-VarNet

(ours)

rcGAN

(Bendel)

pscGAN

(Ohayon)

cGAN

(Adler)

(Jalal)

Truth

Truth

Sample

Sample

Average

(P=2)

Average

(P=4)

Average

(P=32)

5e-05

0e+00

E.4 Inpainting

Original

SgMaskedN

SgCoModGANN

SgDDNM N

SgDDRM N

Original

SgMaskedN

SgCoModGANN

SgDDNM N

SgDDRM N

Original

SgMaskedN

SgCoModGANN

SgDDNM N

SgDDRM N

Original

SgMaskedN

SgCoModGANN

SgDDNM N

SgDDRM N

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The contribution and results described in the abstract and introduction are detailed in Section 3 and Section 4, respectively.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: The limitations are discussed in Section 5.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate 'Limitations' section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [NA]

Justification: We do not provide any theoretical results in this work.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: Experimental settings are described in Section 4 with additional details given in Appendix D. Should this work be accepted, we will release our source code and link to it in the camera-ready submission.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: Upon acceptance, we will release source code to reproduce each experiment.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: Experimental settings are described in Section 4 with additional details given in Appendix D.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [No]

Justification: Error bars are omitted due to computational constraints.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer 'Yes' if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g., negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: Compute resources are described in Appendix D.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: This research abides by the NeurIPS Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: The broader impacts are discussed in the Introduction (our method facilitates uncertainty quantification, fairness, and perception-distortion trade-offs) and MRI-specific impacts are discussed in Section 5.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: This paper does not pose such a risk.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: Licenses are discussed in Appendix D and are properly respected. Credit is given for the various data/code/models we use throughout the paper when appropriate.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: Upon release, the source code will be well-documented with instructions for reproducing the experimental results and extending the code.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: This research did not involve crowdsourcing nor research with human subjects. Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: This research did not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nTo solve image inverse problems (denoising, super resolution, etc.), many works attempt to sample from the posterior distribution of ground truth images given a degraded measurement. One line of works propose to train a stochastic neural network as a CGAN (conditioned on the degraded image and an additional random seed as inputs), while regularizing the conditional mean and the conditional covariance matrix of the outputs. This paper proposes yet another, improved regularization term for such a CGAN approach. Specifically, the authors propose to correct (regularize) the K principal components of the conditional covariance matrix in additional to the conditional mean and the trace of the covariance.\n\nSTRENGTHS:\nThe paper is mostly written well. The proposed method is demonstrated on toy examples (such as Gaussian random signals and MNIST images), which I find appealing to gain intuition into the advantages of the proposed approach. The paper shows results on MRI reconstruction, which is not very common to report but an incredibly important inverse problem to solve.\n\nWEAKNESSES:\nEven though I personally like this type of work and the proposed approach, unfortunately I see several major weaknesses in this paper, starting even with wrong/inaccurate claims in the introduction & motivation. I'll list some of the weaknesses below.\n\n1) I disagree with most of the claims made in L17 onwards regarding the shortcomings of point estimators (introduction). The perception-distortion tradeoff theorem in [6] does not prove anything specific about point estimators (a claim made in L20-L21). In fact, this theorem proves that distortion can only be at odds with perceptual quality (but now always), regardless of whether the estimator is \"deterministic\" or not. Moreover, [6] defines perceptual quality as the statistical distance (e.g., KL) between $p_{\\hat{X}}$ (the distribution of the reconstructions) and $p_{X}$ (the distribution of the source images), and not as the distance between $\\hat{x}$ and the manifold of natural images (as claimed in L22). What is the definition of such a distance? Lastly, the perception-distortion tradeoff is not always a strict one, as hinted in L21. This tradeoff  depends on the distortion measure, as discussed by the authors of [6]. Some distortion measures (e.g., \"perceptual\" losses in the deep image features space) virtually do not contradict with perceptual quality (as demonstrated experimentally in [6]). I would be pleased to see a complete revision of this paragraph in which the errors are corrected, or otherwise a clear explanation for why am I wrong. This is the entry point of the paper, it should at least be accurate.\n\n2) I also do not agree with L27: \"to address these shortcomings, posterior-sampling-based image recovery has been proposed\". Posterior sampling as a solution for inverse problems has been proposed and suggested, e.g., with Langevin dynamics (though, it didn't work properly for high-dimensional data) much before we knew anything about the perception-distortion tradeoff or fairness issues in image reconstruction. As I understand it, the true motivation for the rise of posterior-sampling is the following: We are often dealing with ill-posed inverse problems, and performing posterior sampling allows us to conduct uncertainty quantification of the inverse task at hand, explore the space of solutions visually, and make more informed decisions in down-stream tasks. It has also been shown that stochastic estimators are more robust to adversarial attacks [1]. I don't see why the rise of posterior sampling is linked to the perception-distortion tradeoff and/or to fairness. Please provide evidence (citations and explanation) if that's the case.\n\n3) L30: \"posterior samples will display attributes that reflect minority and majority populations in the training data\". As far as I know, this is often incorrect. In fact, we may often require to produce an intractable amount of posterior samples to observe an outcome for a minority. Please see [2] as an example.\n\n4) L43: You are missing citations (and comparison) with highly related work that proposed to use posterior sampling using a CGAN with conditional mean and conditional covariance regularizations, prior to the work of Bendel et al. For example, [3] and [4]. In fact, the work of Bendel et al. is based on the idea proposed in [2] (to regularize the neural network by enforcing agreement of the average posterior sample with the posterior mean), and in [3] (to use a second-moment penalty of the covariance matrix as well), as you suggest in L75. As far as I understand, your algorithm is a different (and maybe improved) version of [4], but I don't know which one is better since a comparison with [3] and [4] is not reported in the paper. Please address the differences in your work so I could better understand why it's not appropriate to cite such papers and/or include them in your analysis.\n\n5) I would suggest explaining L76-L79 better, by adding appropriate definitions for all the terminology used. For instance, the definition of a $P_{rc}$-sample supervised-$l_{1}$ loss is missing and unclear. I would be happy to see a more self-contained background section. Make sure every notation you are using is clearly defined before using it.\n\n6) Some results seem wrong to me. For example, the results of DPS in Figure 4. I have been using DPS in many different image inverse problems, and never observed such inconsistencies between the outputs and the inputs (see, e.g., Figure 14 in the DPS paper appendices, which shows almost perfect consistency with the measurements for severe inpainting tasks). The FID results of DPS also seem odd to me. I find it very hard to believe that a CGAN algorithm would outperform DPS, given that the hyper-parameters of DPS are properly tuned. Can you please provide the code you used to execute these experiments so I could verify them? Maybe the hyperparemeters you used to run DPS are suboptimal?\n\n7) The contribution of this work is quite limited in my opinion. The proposed algorithm only slightly improves its predecessor (Badler et al.), and the experimental evaluation is limited. A comparison with more works is missing, such as [3] and [4], and other, more recent state-of-the-art stochastic image restoration methods that produce better results than DPS. What about DDNM, DDRM, etc.? Why are these not included in the analysis? (besides the fact that they don't \"promise\" posterior sampling). I would be happy if you could clarify the contribution of your work and its impact on the field, as it seems quite incremental to me.\n\n\n[1] Reasons for the Superiority of Stochastic Estimators over Deterministic Ones: Robustness, Consistency and Perceptual Quality, Guy Ohayon et al., ICML 2023\n\n[2] From Posterior Sampling to Meaningful Diversity in Image Restoration, Noa Cohen et al, ICLR 2024.\n\n[3] High Perceptual Quality Image Denoising with a Posterior Sampling CGAN, Guy Ohayon et al., ICCV workshops 2021.\n\n[4] High-Perceptual Quality JPEG Decoding via Posterior Sampling, Sean Man et al, CVPR workshops 2023.\n\nQUESTIONS:\nPlease address each of the weaknesses.\n\nLIMITATIONS:\nThe authors do provide a dedicated limitations section, but I observe additional limitations in this paper which are not addressed, such as incremental novelty of the proposed algorithm over previous work, limited experimental section (comparing only to one predecessor), and more, as explained in the weaknesses section.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-06-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors address inverse problems and aim to improve GAN-based posterior samplers by adopting regularization to correct the principal components of the covariance matrix of the approximated posterior distribution.\n\nSTRENGTHS:\n1. The paper is well written and easy to follow.\n2. The idea presented in the paper is novel and clear, especially given the tendency of GAN-based posterior samplers to experience mode collapse. This serves as a strong motivation for the proposed method in the paper.\n3. The authors conducted empirical study with a variety of datasets.\n\nWEAKNESSES:\n1. For each dataset in the empirical study, a different set of baseline methods was chosen, possibly due to the requirement of pretraining. However, for a fair comparison, this is necessary for a complete research.\n2. The authors claim that the proposed method outperforms existing diffusion models (line 52). However, the experiments conducted with diffusion models are not comparable due to the much smaller test size (line 210 in MRI experiments and line 235 in inpainting experiments).\n3. Figure 4 presents inpainting posterior samples, but for the diffusion model referred to as DPS, the masking was not applied, which results in incorrect and misleading visualization.\n4. Some minor errors:\n    - In line 28, “There the goal\" should be revised.\n    - In lines 33-34, there is a grammar mistake.\n\nQUESTIONS:\n1. While diffusion models are slower compared to GANs, they have demonstrated success over GAN-based models in enhancing posterior richness for conditional tasks. A fair comparison with diffusion models would be valuable, as it remains uncertain whether the proposed method outperforms diffusion models. For faster sampling approaches, the authors might consider DDIM[1], which accelerates the diffusion process.\n\nReferences:\n\n[1] Song, Jiaming, Chenlin Meng, and Stefano Ermon. \"Denoising diffusion implicit models.\" (2020).\n\nLIMITATIONS:\nThe authors explained the limitations of their work.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-06-19"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper focuses on learning posterior samplers for inverse problems, specifically aiming to accurately estimate the principal components of the posterior covariance matrix. To achieve this, the authors consider training a conditional generative adversarial network (cGAN), building upon previous work [21]. They introduce two new regularization terms: one for the eigenvalues and another for the eigenvectors. Recognizing the challenge of estimating these statistics from a single ground truth solution per input measurement, the authors propose gradually adding these terms during training, while utilizing intermediate results to estimate the statistics.\nThe authors present a series of experiments to demonstrate their method, ranging from toy examples to real-world MRI data. Comparisons are primarily made with rcGAN and NPPC.\n\nSTRENGTHS:\n* The paper is well-written, with clear motivation, method presentation, and background (mostly).\n* The detailed algorithm description, beyond pseudocode, significantly enhances understanding and reproducibility.\n* The authors nicely addressed the problem of estimating the desired statistics from a single ground truth vector.\n* The authors provide sufficient evidence to demonstrate their method (but additional evidence could further strengthen their claims).\n\nWEAKNESSES:\nMajor Issues:\n* The paper lacks an ablation study which is crucial. While the authors followed [21] and set $P_{rc}=2$, I still expect an ablation study on this parameter, especially since using 2 samples seems insufficient to estimate statistics properly. \nFurthermore, there is no ablation study regarding the different regularization terms and their relative weights. Given the newly regularization terms, it is possible that other terms, introduced in [21], are redundant, especially when $K=d$. Since accurately recovering the eigenvalues implies perfect recovery of the trace, for example.  \n\n* I find the comparison to competing methods limited. I suggest adding comparisons to additional recent posterior sampling techniques. For example, the authors of NPPC presented comparison to DDRM, DDNM, RePaint, and MAT.\n\n*  Unlike work [21], the paper present no proposition or theoretical guarantees regarding the estimated statistics, limiting the novelty of the paper. \n\n* The description of work [21] is imprecise. Specifically, the usage of the notations $P$ and $P_{rc}$ is incorrect, leading to confusion.\n\n* The authors state in lines 103-104 that they apply the regularization term (11) only after allowing the generated posterior mean to converge near the true mean. This raises the question of whether (11) is necessary or if its impact is marginal, as it only leads to small variations in the generated posterior mean. The same question applies to the generated eigenvalues. I suggest adding plots to illustrate the variation in generated statistics over training epochs, comparing them to their final values or their values before applying the new regularization term during training.\n\n* While the relevance of NPPC to this work is clear, the experimental comparison seems unfair as NPPC is not designed for generation tasks. Furthermore, the authors demonstrate a speedup against DPS, but neglect to discuss the computational load of their method compared to NPPC.\n\nMinor:\n* The authors in lines 24-26 correctly point out that bias can be an issue in single-sample recovery. However, this is also a potential concern for recovery methods that yield multiple solutions.\n\nQUESTIONS:\n* Please address the weaknesses above.\n* I suggest conducting an ablation study to assess the impact of key priors in the method, such as the different terms and $P_{rc}$.\n* I suggest expanding the experiments to include comparisons with more recent posterior samplers in standard image inverse problems, like super-resolution and inpainting.\n* Please provide plots illustrating the variation in the generated statistics over the training epochs.\n* Do the additional regularization terms theoretically change the optimal solution, or does it remain the same?\n\nLIMITATIONS:\nSee weaknesses.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work considers using posterior sampling to solve inverse problems via conditional GANs. Previous work has shown that cGANs are competitive in terms of reconstruction quality as compared to other generative models to sample from the posterior (e.g., diffusion models), while also having the advantage that, once trained, sampling is fast as it simply requires a forward pass. In this paper, the authors build on such approaches by adding regularization terms to encourage the cGAN to learn K principal components of the posterior covariance matrix, where K is user-defined quantity. Due to the computational burden of such an approach, a training scheme is devised where learning of the eigenvectors and eigenvalues is done after optimizing a regularized cGAN objective. The method is validated on a variety of linear inverse problems (denoising, accelerated MRI, inpainting) and datasets (mnist, fastmri, ffhq).\n\nSTRENGTHS:\n- The paper is well-written and easy to read. In particular, the idea of adding regularization to fit to principal components of the posterior distribution is a simple and natural idea. It is also one of importance, as being able to visualize uncertainty in reconstructions is crucial in understanding which regions of the reconstruction can be “trusted”\n- The method is validated on a nice range of problem types and datasets.\n- Using a conditional cGAN allows for fast sample generation, especially as compared to diffusion models. \n- The additional use of the pca regularization term appears to improve performance of rcGAN, while also allowing for visualization of directions of uncertainty.\n\nWEAKNESSES:\n- The proposed approach is a slight, potentially somewhat incremental, improvement to the prior rcGAN approach of [21]. In particular, rcGAN also aims to better model the posterior distribution by matching certain statistics of the posterior mean and covariance. This work takes this reasoning a step further by aiming to fit the principle components of the covariance. This has also been done (in a different way) in prior work [29], that the authors compare to. The bonus for this work, however, is that the authors can sample quickly, whereas in [29] sampling is not possible.\n- Coupled with the concern regarding novelty, the additional regularization terms are computationally expensive to optimize, especially for high-dimensional problems. This is evidenced by, for example, only being able to set K = 1 in the MRI experiments and K = 2 in the large-scale inpainting experiments, and is discussed by the authors in their limitations. Moreover, the training strategy becomes a bit more complicated in order to aid in efficiency, as one needs to first optimize the rcGAN objective and then integrate the PCA terms into the learning process.\n\nQUESTIONS:\n- As the parameter $K$ is fixed prior to training with knowledge of the inverse problem, it would be good to discuss how the choice of $K$ relates to the particular inverse problem itself. For example, how does the choice of $K$ relate to how challenging the inverse problem is? I presume that ill-posedness plays a role here (e.g., comparing a linear inverse problem with a non-linear inverse problem, such as phase retrieval).\n\nLIMITATIONS:\nThe authors discuss some of the above weaknesses I mentioned near the end of their paper, which I appreciate. They also discuss other issues regarding how to properly incorporate this approach for rigorous uncertainty quantification.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}