pcaGAN: Improving Posterior-Sampling cGANs via Principal Component Regularization

TL;DR

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.

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.

Figures (16)

• Figure 1: Gaussian experiment. Wasserstein-2 distance versus (a) lazy update period $M$ for pcaGAN with $d=100=K$, (b) estimated eigen-components $K$ for pcaGAN with $d=100$ and $M=100$, and (c) problem dimension $d$ for all methods under test with $K=d$ and $M=100$.
• Figure 2: For (a) pcaGAN and (b) NPPC, this figure shows the true image $\boldsymbol{x}$, noisy measurements $\boldsymbol{y}$, the conditional mean $\widehat{\boldsymbol{\mu}_{\mathsf{x}|\mathsf{y}}}$, principal eigenvectors $\{\Hat{\boldsymbol{v}}_k\}$, and two perturbations of $\widehat{\boldsymbol{\mu}_{\mathsf{x}|\mathsf{y}}}$.
• Figure 3: Example MRI recoveries at $R=8$. Arrows highlight meaningful variations.
• Figure 4: Example of inpainting a randomly generated mask on a $256\!\times\!256$ FFHQ face image.
• Figure C.1: Example MRI recoveries at $R=4$. Arrows highlight meaningful variations.