pcaGAN: Improving Posterior-Sampling cGANs via Principal Component Regularization

Authors: matthew c bendel
dept  ece
ahmad  rizwan
  dept
  bme
philip  schniter
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 17:54:10

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.

Main paper: 
 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] [2] [3] [4] [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] [9] [10] [11] [12] [13] [14] [15] [16] [17] [18] [19] [20] [21] [22] [23] [24] [25] [26] [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 (cVAEs) [12] [13] [14] , conditional normalizing flows (cNFs) [8] [9] [10] [11] , conditional generative adversarial networks (cGANs) [15] [16] [17] [18] [19] , and Langevin/score/diffusion-based generative models [20] [21] [22] [23] [24] [25] [26] [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] [26] [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.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. 
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 (1)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 formEQUATIONusing 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 minimizingEQUATIONover θ 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 Öktem proposed a two-sample discriminator that encourages diverse generator outputs without compromising the Wasserstein objective (2) . In [19] To do this, [19] replaced L adv (θ, ϕ) with the regularized adversarial lossEQUATIONwhich involves the P rc -sample supervised-ℓ 1 loss and standard-deviation (SD) reward terms L 1,P (θ) ≜ E x,z1,...,z P ,y ∥x -x (P ) ∥ 1 (7)EQUATIONwhere typically P rc = 2. Here, { x i } are the generated samples and x (P ) is their P -sample average:x i ≜ G θ (z i , y) for i = 1, . . . , P rc and x (P ) ≜ 1EQUATIONThe 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,...,z P ,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] .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. 
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:EQUATIONEQUATIONEQUATIONHere, {( 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 z i ∼ N (0, I) for i = 1, . . . , P rc 5:EQUATIONx i ← G θ (z i , y b ) for i = 1, . . . , P rc 6: L(θ) ← L(θ) -β adv Prc i=1 D ϕ ( x i , y b ) 7: x (Prc) = 1 Prc Prc i=1 x i 8: L(θ) ← L(θ) + ∥x b -x (Prc) ∥ 1 -β SD Prc i=1 E z1,...,z P ,y ∥ x i -x (Prc) ∥ 1 9: 10:if e ≥ E evec and s mod M = 0 then 11: z j ∼ N (0, I) for j = 1, . . . , P pca 12:x j ← G θ (z j , y b ) for j = 1, . . . , P pca 13: if e ≥ E eval and s mod M = 0 then 20:µ ← StopGrad( 1 Ppca Ppca j=1 x j ) 14: U S V ⊤ ← SVD([ x 1 -µ, . . . , x Ppca -µ] ⊤ ) 15: v k ← [ V ] :,k for k = 1, . . . , K 16: L(θ) ← L(θ) -β pca K k=1 [ v ⊤ k (x b -µ)]λ k ← [ S] 2 kk for k = 1, . . . , K 21: X ← [x b -µ, x 1 -µ, x 2 -µ, . . . , x Ppca -µ] ⊤ 22: L(θ) ← L(θ) + β pca K k=1 1 -1 λ k StopGrad( 1 Ppca+1 ∥ v ⊤ k X∥ 2 ) 2 23:end if 24: end for 25:26: θ ← Adam(θ, ∇L(θ))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 } Prc 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 } Ppca j=1 are generated in line 12 (where nominally P pca = 10K), 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 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.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Answer: [NA] 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. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
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.Here our goal is to recover 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 tox ∼ N (µ x , Σ x ) ∈ R d from y = M x + w ∈ R d ,EQUATION)For the cGANs, we compute µ x|y and Σ x|y empirically from 10d 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.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 (a) pcaGAN (K = 10) 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}.v 1 v 2 v 3 v 4 v 5 y α=-3 α=-2 α=0 α=2 α=3 x µ x|y µ x|y +α v4 µ x|y +α v1 (b) NPPC v 1 v 2 v 3 v 4 v 5 y α=-3 α=-2 α=0 α=2 α=3 x µ x|y µ x|y +α v4 µ x|y +α v1Competitors. 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 kth column equals the kth 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. 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 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.R = 4 R = 8 Model CFID 1 ↓ CFID 2 ↓ CFID 3 ↓ FID↓ APSD Time (4)↓ CFID 1 ↓ CFID 2 ↓ CFID 3 ↓ FID↓ APSD Time (4)↓ E2E-VarNet (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 & Öktem'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 ≜ ( 1P 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.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.  
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] [16] [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 = 10K 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] [61] [62] . Third, the application to MRI is preliminary; additional tuning and validation is needed before it can be considered for clinical practice. 
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.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. 
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. 
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). 
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). 
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). 
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. • 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. 
 
 Appendix: A Synthetic Gaussian priors [Q, R] = QRdecomp([u (d) 1 u (d) 2 . . . u (d) d ]) 13: v (d) k = [Q] :k for k = 1, . . . ,We quantify posterior-approximation accuracy using conditional Fréchet inception distance (CFID) [38] , which approximates the conditional Wasserstein distance 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] .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, 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.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.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 Öktem's cGAN.cGAN architecture. All four cGANs used the same generator and discriminator architectures described in [19] , except that Adler and Öktem's discriminator used extra input channels to facilitate the 3-input loss. [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.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. NeurIPS Paper Checklist38th Conference on Neural Information Processing Systems (NeurIPS 2024). 


Decision: In this paper, the authors present a method for using a conditional GAN to solve inverse problems by posterior sampling.  Whereas prior work enforces accuracy in the posterior mean and trace-covariance, this work further enforces accuracy in a user-specified number of principal components of the posterior covariance matrix.  The proposed work yields more accurate posterior samples than baselines with only modest increases in computational cost during training (relative to the most similar method to the proposed work).  Based on reviewer comments, there is some room for improvement in terms of thorough ablations and explanation of prior literature.  Based on the author-reviewer discussions, these concerns have been resolved enough to warrant acceptance.  The authors should implement the promised changes in the camera-ready.

Reviews: review: paper_summary: paper_summary: This 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).
strength_weakness: strength_weakness: - 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.
- 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.
questions: questions: - 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).
limitations: limitations: The 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.
paper_summary: paper_summary: The 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.
The 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.
strength_weakness: strength_weakness: Major Issues:
* 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. 
Furthermore, 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.  

* 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.

*  Unlike work [21], the paper present no proposition or theoretical guarantees regarding the estimated statistics, limiting the novelty of the paper. 

* The description of work [21] is imprecise. Specifically, the usage of the notations $P$ and $P_{rc}$ is incorrect, leading to confusion.

* 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.

* 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.

Minor:
* 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.
questions: questions: * Please address the weaknesses above.
* I suggest conducting an ablation study to assess the impact of key priors in the method, such as the different terms and $P_{rc}$.
* I suggest expanding the experiments to include comparisons with more recent posterior samplers in standard image inverse problems, like super-resolution and inpainting.
* Please provide plots illustrating the variation in the generated statistics over the training epochs.
* Do the additional regularization terms theoretically change the optimal solution, or does it remain the same?
limitations: limitations: See weaknesses.
paper_summary: paper_summary: The 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.
strength_weakness: strength_weakness: 1. 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.
2. 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).
3. 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.
4. Some minor errors:
    - In line 28, “There the goal" should be revised.
    - In lines 33-34, there is a grammar mistake.
questions: questions: 1. 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.

References:

[1] Song, Jiaming, Chenlin Meng, and Stefano Ermon. "Denoising diffusion implicit models." (2020).
limitations: limitations: The authors explained the limitations of their work.
paper_summary: paper_summary: To 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.
strength_weakness: strength_weakness: Even 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.

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

2) 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.

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

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

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

6) 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?

7) 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.


[1] Reasons for the Superiority of Stochastic Estimators over Deterministic Ones: Robustness, Consistency and Perceptual Quality, Guy Ohayon et al., ICML 2023

[2] From Posterior Sampling to Meaningful Diversity in Image Restoration, Noa Cohen et al, ICLR 2024.

[3] High Perceptual Quality Image Denoising with a Posterior Sampling CGAN, Guy Ohayon et al., ICCV workshops 2021.

[4] High-Perceptual Quality JPEG Decoding via Posterior Sampling, Sean Man et al, CVPR workshops 2023.
questions: questions: Please address each of the weaknesses.
limitations: limitations: The 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.

Comments: title: Post-Rebuttal
comment: We thank the reviewer for their valuable feedback and continued discussion.

- **Reviewer**: I thank the author for addressing my concerns by evaluating several diffusion algorithms against their method. It appears that the DPS algorithm yields results nearly identical to the proposed approach, however, DPS heavily relies on a specific scale factor. Additionally, while the DPS algorithm is useful, it is not currently the SOTA for posterior sampling. Although I recognize the novel and valuable contribution of this idea for GANs, I am not convinced that it outperforms the performance of current diffusion models. Therefore, I remain my score as it is.

    **Response**: Yes, our method and DPS had comparable performance after the latter was finely tuned for the FFHQ inpainting application. While we agree with the reviewer that diffusion models are powerful and hold significant promise for solving inverse problems, we emphasize that sampling speed is critical for many applications. Even with accelerated sampling methods, such as those based on DDIM, diffusion methods are still orders of magnitude slower than cGANs. For example, in most clinical setups, MR images are expected to be available immediately after the data acquisition. This level of latency cannot be achieved with current diffusion models. Therefore, posterior sampling methods that can compete with existing diffusion models while offering orders of magnitude faster image generation hold great practical value.
Thank you for delving deeper into this. Indeed, it seems like there is not a direct mapping between the inverse problem structure and the posterior covariance, as the (unknown) prior plays a heavy role in determining the structure of the posterior. After carefully considering the author's comments and the other reviewer's discussions, I believe that my score still accurately reflects my view of the paper. Hence I will keep my score.
I sincerely thank the authors for the thorough discussion.

I have given this a lot of thought. I reviewed the paper again and took the following into account:

1) Our discussion (including the revised introduction).
2) The additional experiments and results.
3) The reviews made by other reviewers, as well as the authors' responses.

While I think the results are appealing, I believe the paper is not 100% ready for publication due to insufficient soundness and presentation (not only in the introduction). Therefore, I am:

1) Decreasing my scores for soundness and presentation.
2) Increasing my score for contribution.
3) Increasing my overall score to borderline reject (*given that the authors will incorporate the new introduction section*).

Despite my opinion about the paper's soundness and presentation, I would like to emphasize that the paper's contribution is, in my opinion, worthy (given that all the results are 100% correct, which I have not confirmed). However, unfortunately, I can't confidently recommend acceptance for the current state of the paper.
I thank the authors for their rebuttal and clarifications, which have been very helpful in addressing most of my concerns.  However, to further strengthen the paper, I still encourage an ablation study on the $P_{rc}$  parameter, or at least a more detailed discussion of its applicability to the proposed method, as the analysis in [21] may not directly transfer.

While I acknowledge the novelty of pcaGAN in aiming for better posterior sampling, I believe theoretical guarantees would make the contribution even more impactful, similar to rcGAN. Additionally, an ablation study on the different regularization terms is recommended, as some original rcGAN terms might be redundant with the new ones and could potentially be removed to facilitate the training process.
Currently, I believe my score accurately reflects the merits of the work.
I thank the author for addressing my concerns by evaluating several diffusion algorithms against their method. It appears that the DPS algorithm yields results nearly identical to the proposed approach, however, DPS heavily relies on a specific scale factor. Additionally, while the DPS algorithm is useful, it is not currently the SOTA for posterior sampling. Although I recognize the novel and valuable contribution of this idea for GANs, I am not convinced that it outperforms the performance of current diffusion models. Therefore, I remain my score as it is.
Thanks, this is a great question. In general, the structure of the posterior covariance matrix will be governed by the interaction between the likelihood (defined by $\\boldsymbol{y}=\\mathcal{M}(\\boldsymbol{x})$ if we think of $\\mathcal{M}(\\cdot)$ as stochastic) and the prior on $\\boldsymbol{x}$.  One can construct toy problems where the posterior covariance has an intuitive structure, but generally speaking the likelihood, prior, and their interaction will be complicated and non-intuitive.

As a toy example, consider a noiseless inpainting problem, where the first half of the entries in $\\boldsymbol{x}$ are directly observed.  Then the posterior distribution over those first-half entries would be a Dirac delta, while the posterior distribution over the last-half entries would be the *prior* distribution of the last-half entries conditioned on the first-half entries.
For some priors, knowing the first-half entries tells you everything about the last-half entries, while for other priors it tells you nothing, and for most cases it gives you partial knowledge.
So, little can be said without a good understanding the prior and, in particular, its interaction with the likelihood.
    
As far as multi-modality is concerned, it does not directly impact the covariance structure, since covariance is based only on second-order statistics.  Again, consider a noiseless inpainting problem.  If the prior is a high-variance white Gaussian, then the posterior will be a high-variance white Gaussian in the unobserved dimensions.  If the prior is instead a Gaussian-mixture with many randomly located component means, the posterior will also be a Gaussian mixture and thus will generally be multi-modal in the unobserved dimensions.  But the posterior *covariance* will just measure the uncertainty *energy* along different dimensions, which might be very similar in the uni-modal and multi-modal cases.
[This is a continuation of our previous comment.]

Regarding hallucinations, we do find that, relative to posterior-mean recoveries, posterior samples tend to hallucinate fine details that are not present in the true image, and we show evidence of that in Figures 3, C.1, E.4, E.5, E.6, and E.7. While a posterior sample as a whole may not be "wrong," it can fabricate fine details. Note that individual posterior samples are farther from the truth than the posterior mean recovery (in terms of average L2 distance) since they have 3dB worse MSE.

In addition to revising our Introduction, we plan to add the following after L79 in Section 2:

> We note that the $\\mathcal{L}\_{1,P}(\\theta)$ regularization proposed in [21] is closely related to the regularization $\\mathcal{L}\_{2,P}(\\theta) \\triangleq E\_{x,z1,...,zP,y}\\{\\|\\boldsymbol{x}-\\hat{\\boldsymbol{x}}\_{(P)}\\|_2^2\\}$ proposed earlier by Ohayon et al. in [Ohayon'21].
> A detailed discussion of the advantages and disadvantages of various cGAN regularizations can be found in [21].
We thank the reviewer for the valuable feedback and continued discussion. Due to length constraints, we'll need to split our answer in two. Below is our proposed revision to the Introduction.

> &nbsp;&nbsp;&nbsp;&nbsp;In image recovery, the goal is to recover the true image $\\boldsymbol{x}$ from noisy/distorted/incomplete measurements $\\boldsymbol{y} = \\mathcal{M}(\\boldsymbol{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 $\\boldsymbol{x}$ from $\\boldsymbol{y}$.
>
> &nbsp;&nbsp;&nbsp;&nbsp;In much of the literature, image recovery is posed as point estimation, where the goal is to return a single best estimate $\\hat{\\boldsymbol{x}}$. But there are shortcomings to this approach. For example, it’s not clear how to define “best,” since L2- or L1-minimizing $\\hat{\\boldsymbol{x}}$ are often regarded as too blurry, while efforts to make $\\hat{\\boldsymbol{x}}$ perceptually pleasing can sacrifice agreement with the true image $\\boldsymbol{x}$ and cause hallucinations [1–5], as we show later. Also, many applications demand not only a recovery $\\hat{\\boldsymbol{x}}$ but also some quantification of the uncertainty in that recovery [7,8].
>
> &nbsp;&nbsp;&nbsp;&nbsp;As an alternative to point estimation, there is posterior-sampling-based image recovery [10–28], which aims to generate $P$ samples $\\{\\hat{\\boldsymbol{x}}\_i\\}\_{i=1}^P$ from the posterior distribution $p\_{x|y}(\\cdot|\\boldsymbol{y})$. Posterior sampling facilitates numerous strategies to quantify the uncertainty in estimating $\\boldsymbol{x}$, or any function of $\\boldsymbol{x}$, from $\\boldsymbol{y}$ [7,8]. 
> It also can help with visualizing uncertainty and increasing robustness to adversarial attacks [Ohayon,ICML'23].
That said, the design of accurate and computationally-efficient posterior samplers remains an open problem.
>
> &nbsp;&nbsp;&nbsp;&nbsp;Given a training dataset of image/measurement pairs $\{(\\boldsymbol{x}\_t, \\boldsymbol{y}\_t
)\}_{t=1}^T$, our goal is to build a generator $G\_{\\boldsymbol{\\theta}}$ that, for a given $\\boldsymbol{y}$, maps random code vectors $\\boldsymbol{z}\\sim\mathcal{N}(\\boldsymbol{0}, \\boldsymbol{I})$ to samples of the posterior, i.e., $\\hat{\\boldsymbol{x}} = G\_{\\boldsymbol{\\theta}}(\\boldsymbol{z}, \\boldsymbol{y}) \\sim p\_{x|y}(·|\\boldsymbol{y})$. 
> There are many ways to approximate the ideal generator. 
> The recent literature has focused on conditional variational autoencoders (cVAEs) [14–16], conditional normalizing flows (cNFs) [10–13], conditional generative adversarial networks (cGANs) [17–21,Ohayon'21,Man'24], and diffusion posterior samplers [22–28,Wang'23]. Although diffusion samplers 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.
>
> &nbsp;&nbsp;&nbsp;&nbsp;With the goal of fast and accurate posterior sampling, recent progress in cGAN training has been made through regularization.
> For example, Ohayon et al. [Ohayon'21] proposed to enforce correctness in the generated $\\boldsymbol{y}$-conditional mean using a novel form of L2 regularization.
> Later, Bendel et al. [21] proposed to enforce correctness in the generated $\\boldsymbol{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. [Man'24] showed that L2 regularization and a variance reward are effective when training cGANs for JPEG decoding.
> More details are given in Section 2.
> 
> &nbsp;&nbsp;&nbsp;&nbsp;In a separate line of work, Nehme et al. [29] trained a Neural Posterior Principal Components (NPPC) network to directly estimate the eigenvectors and eigenvalues of the $\\boldsymbol{y}$-conditional covariance matrix, which allows powerful insights into the nature of uncertainty in an inverse problem.
>
> &nbsp;&nbsp;&nbsp;&nbsp;**Our contributions.** 
> Inspired by regularized cGANs and NPPC, we propose a novel “pcaGAN” that encourages correctness in the $K$ principal components of the $\\boldsymbol{y}$-conditional covariance matrix, as well as the $\\boldsymbol{y}$-conditional mean and trace-covariance, when sampling from the posterior. 
> We demonstrate that pcaGAN outperforms existing cGANs [18, 19, 21, Ohayon'21] and diffusion methods [24, 27, 28, Wang'23] on posterior sampling tasks like denoising, large-scale inpainting, and accelerated MRI recovery, while sampling orders-of-magnitude faster than [24, 27, 28, Wang'23]. 
> We also demonstrate that pcaGAN recovers principal components more accurately than NPPC using approximately the same runtime.
Thank you to the authors for their detailed rebuttal and for responding to my comments. I appreciate that it is still a non-trivial task to extend rcGAN to fit to the principal components and singular values and there are certainly empirical benefits to the proposed approach. I was hoping to ask further about the choice of $K$.

It is mentioned that the choice of $K$ behaves similarly to PCA on the posterior covariance, which makes sense. But I think I was asking more specifically about if there is insight from the form of the inverse problem $y = \mathcal{M}(x)$ that can give insight into the conditioning of the posterior covariance. For example, do ill-conditioned linear forward models induce posterior covariance matrices that have slower decaying singular values? Similarly, if the inverse problem is non-linear, inducing a multi-modal posterior, is it necessarily the case that we should expect the posterior covariance to also have slower decaying singular values (as compared to a less, multi-modal posterior)? I ask these questions from the perspective of giving the reader some insight into this choice of $K$.
1. That sounds better and more accurate to me. There are still some things I would personally rephrase, which could be misleading to some readers. For example: "while perceptually pleasing $\hat{x}$ may suffer from hallucinations." The notion "hallucinations" carries a heavy load in several fields, so this could mislead readers about the meaning of producing perceptually pleasing $\hat{x}$. For example, researchers from the world of LLM would think that image reconstruction models that produce perceptually pleasing outputs often produce very wrong outputs, but any sample from the posterior distribution is actually not "wrong". 

2. Note that you also need to revise L43 onwards, where you start talking about the posterior mean regularization.

3. In the rebuttal, you wrote: "We don't think that rcGAN is "based on" [Ohayon'21] or [Man'24], though. The rcGAN paper [21] cites the 2017 Isola cGAN paper [17] for posterior-mean regularization." But now you are saying that the Isola cGAN paper did not propose posterior-mean regularization? While your work may have good results, I think it is equally important to write a proper (and correct) introduction section where this work is properly positioned compared to previous/related literature.

Could you please try to include here the entire new introduction so I could make sure it all makes sense to me?
We thank the reviewer for the continued discussion and opportunity to improve our submission.

1. We appreciate the reviewer's invitation to further our discussion around the PD tradeoff and its connection to point estimates and posterior sampling. Our main point is that an estimator that returns $P\\gg 1$ posterior samples lets us trade perception with distortion in a way that a point estimator (which returns a single sample, such as a single posterior sample) cannot. Regardless of whether the reviewer agrees with us on the specifics of this point, we would like to emphasize that **this discussion is immaterial to the main ideas of our paper**, which revolve around enforcing cGAN consistency with the principal components of the true posterior covariance matrix. To avoid directly tying point estimation to PD or the motivation for posterior sampling to PD, we have revised our Introduction; please see our response to your next point.

2. We are removing the discussion of the PD tradeoff from our Introduction (not our entire paper) to definitively address the first and second critiques in your original review.  In particular, we propose to revise the second and third paragraphs of our Introduction as follows.

    > In much of the literature, image recovery is posed as point estimation, where the goal is to return a single best estimate $\\hat{\\boldsymbol{x}}$. But there are shortcomings to this approach. For example, it’s not clear how to define “best,” since L2- or L1-minimizing $\\hat{\\boldsymbol{x}}$ are often regarded as too blurry, while perceptually pleasing $\\hat{\\boldsymbol{x}}$ may suffer from hallucinations [1–5]. Also, many applications demand not only a recovery $\\hat{\\boldsymbol{x}}$ but also some quantification of the uncertainty in that recovery [7,8].
    >
    > As an alternative to point estimation, posterior-sampling-based image recovery [10–28] has been proposed. There the goal is to generate $P$ samples $\\{\\hat{\\boldsymbol{x}}\_i\\}\_{i=1}^P$ from the posterior distribution $p\_{x|y}(·|\\boldsymbol{y})$. Posterior sampling facilitates numerous strategies to quantify the uncertainty in estimating $\\boldsymbol{x}$, or any function of $\\boldsymbol{x}$, from $\\boldsymbol{y}$ [7,8]. 
    It also can help visualize uncertainty, and it can increase robustness to adversarial attacks [Ohayon,ICML'23].
    That said, the design of accurate and computationally-efficient posterior samplers remains an open problem.

    We would be interested to hear your feedback.

3. We should clarify what we wrote in our rebuttal.  The rcGAN paper [21] defines the "supervised-L1 loss" as $E\\{\\|\\boldsymbol{x}-\hat{\\boldsymbol{x}}\_{(P)}\\|\_1\\}$ with $\\boldsymbol{x}$ the true (training) image and $\hat{\\boldsymbol{x}}\_{(P)}$ the $P$-sample output average. On page 4, [21] says "regularizing a cGAN with supervised-L1 loss alone is not new; see, e.g., [Isola]".  But looking closely at [Isola], it seems that only the $P=1$ case is described there (i.e., no averaging), and so perhaps [21] is the first to propose this "supervised-L1 loss".  The paper [Ohayon'21] considers the squared-L2 version $E\\{\\|\\boldsymbol{x}-\hat{\\boldsymbol{x}}\_{(P)}\\|\_2^2\\}$, for which only the $P=1$ case is described in [Isola]. In any case, if accepted, our final paper will cite [Ohayon'21] for what appears to be the first cGAN regularization involving  $\hat{\\boldsymbol{x}}\_{(P)}$ with $P>1$.
I'd like to thank the authors for their thorough response.

1. I disagree regarding your explanation about the perception-distortion tradeoff. First of all, the paper does not only consider *single* draw estimators. In fact, the definition of distortion used in the paper is a mean taken over all Y's, and for each Y the mean is taken over all the random draws from the posterior and all the random draws from the estimator:
$$\mathbb{E}[\Delta(X,\hat{X})]=\int\int\int \Delta(x,\hat{x})p_{X|Y}(x|y)p_{\hat{X}|Y}(\hat{x}|y)p_{Y}(y)dxd\hat{x}dy$$
Namely, when talking about stochastic estimators (e.g., the posterior sampler), distortion (as defined in this paper) actually considers all possible draws from the estimator, for each Y. So, I still don't see why the perception-distortion tradeoff tells us anything "bad" about point estimators. Moreover, the perception-distortion tradeoff proves that the perception-distortion function,
$$P(D)=\min_{p_{\hat{X}|Y}}d(p_{X},p_{\hat{X}})\quad s.t.\quad \mathbb{E}[\Delta(X,\hat{X})]\leq D$$
is monotonically non-increasing in $D$, but it was never proven that this general function is *strictly* decreasing. For example, if we take the trivial distortion $\forall x,\hat{x}:\Delta(x,\hat{x})=0$, then $P(D)$ is constant in $D$ (it always equals zero). Thus, perception and distortion (as defined in the perception-distortion tradeoff paper) are not always at odds with each other. Lastly, I am sorry, but I don't understand what you mean by "with large $P$, there is essentially no PD tradeoff."

2. I don't see why would you remove the discussion about the perception-distortion tradeoff from your paper. It is central to motivating special types of regularizations when training with GANs/cGANS and not just a "regular" MSE loss. So actually, in my opinion, without discussing the perception-distortion tradeoff, the motivation of such work (e.g., pcaGAN) would be unclear.

3. Could you please point out to me where in the cGAN paper by Isola et al. "Image-to-Image Translation with Conditional Adversarial Networks" I can find the posterior mean regularization?

Thank you!