Invertible generative models for  inverse problems: mitigating representation error and dataset bias

1 INTRODUCTION . Generative deep neural networks have shown remarkable performance as natural signal priors in imaging inverse problems , such as denoising , inpainting , compressed sensing , blind deconvolution , and phase retrieval . These generative models can be trained from datasets consisting of images of particular natural signal classes , such as faces , fingerprints , MRIs , and more ( Karras et al. , 2017 ; Minaee and Abdolrashidi , 2018 ; Shin et al. , 2018 ; Chen et al. , 2018 ) . Some such models , including variational autoencoders ( VAEs ) and generative adversarial networks ( GANs ) , learn an explicit low-dimensional manifold that approximates a natural signal class ( Goodfellow et al. , 2014 ; Kingma and Welling , 2013 ; Rezende et al. , 2014 ) . We will refer to such models as GAN priors . With an explicit parameterization of the natural signal manifold by a low dimensional latent representation , these generative models allow for direct optimization over a natural signal class . Consequently , they can obtain significant performance improvements over non-learning based methods . For example , GAN priors have been shown to outperform sparsity priors at compressed sensing with 5-10x fewer measurements . Additionally , GAN priors have led to theory for signal recovery in the linear compressive sensing and nonlinear phase retrieval problems ( Bora et al. , 2017 ; Hand and Voroninski , 2017 ; Hand et al. , 2018 ) , and they have also shown promising results for the nonlinear blind image deblurring problem ( Asim et al. , 2018 ) . A significant drawback of GAN priors for solving inverse problems is that they can have representation error or bias due to architecture and training . This can happen for many reasons , including because the generator only approximates the natural signal manifold , because the natural signal manifold is of higher dimensionality than modeled , because of mode collapse , or because of bias in the training dataset itself . As many aspects of generator architecture and training lack clear principles , representation error of GANs may continue to be a challenge even after substantial hand crafting and engineering . Additionally , learning-based methods are particularly vulnerable to the biases of their training data , and training data , no matter how carefully collected , will always contain degrees of bias . As an example , the CelebA dataset ( Liu et al. , 2015 ) is biased toward people who are young , who do not have facial hair or glasses , and who have a light skin tone . As we will see , a GAN prior trained on this dataset learns these biases and exhibits image recovery failures because of them . In contrast , invertible neural networks can be trained as generators with zero representation error . These networks are invertible ( one-to-one and onto ) by architectural design ( Dinh et al. , 2016 ; Gomez et al. , 2017 ; Jacobsen et al. , 2018 ; Kingma and Dhariwal , 2018 ) . Consequently , they are capable of recovering any image , including those significantly out-of-distribution relative to a biased training set ; see Figure 1 . We call the domain of an invertible generator the latent space , and we call the range of the generator the signal space . These must have equal dimensionality . Flow-based invertible generative models are composed of a sequence of learned invertible transformations . Their strengths include : their architecture allows exact and efficient latent-variable inference , direct loglikelihood evaluation , and efficient image synthesis ; they have the potential for significant memory savings in gradient computations ; and they can be trained by directly optimizing the likelihood of training images . This paper emphasizes an additional strength : because they lack representation error , invertible models can mitigate dataset bias and improve performance on inverse problems with out-of-distribution data . In this paper , we study generative invertible neural network priors for imaging inverse problems . We will specifically use the Glow architecture , though our framework could be used with other architectures . A Glow-based model is composed of a sequence of invertible affine coupling layers , 1x1 convolutional layers , and normalization layers . Glow models have been successfully trained to generate high resolution photorealistic images of human faces ( Kingma and Dhariwal , 2018 ) . We present a method for using pretrained generative invertible neural networks as priors for imaging inverse problems . The invertible generator , once trained , can be used for a wide variety of inverse problems , with no specific knowledge of those problems used during the training process . Our method is an empirical risk formulation based on the following proxy : we penalize the likelihood of an image ’ s latent representation instead of the image ’ s likelihood itself . While this may be couterintuitive , it admits optimization problems that are easier to solve empirically . In the case of compressive sensing , our formulation succeeds even without direct penalization of this proxy likelihood , with regularization occuring through initialization of a gradient descent in latent space . We train a generative invertible model using the CelebA dataset . With this fixed model as a signal prior , we study its performance at denoising , compressive sensing , and inpainting . For denoising , it can outperform BM3D ( Dabov et al. , 2007 ) . For compressive sensing on test images , it can obtain higher quality reconstructions than Lasso across almost all subsampling ratios , and at similar reconstruction errors can succeed with 10-20x fewer measurements than Lasso . It provides an improvement of about 2x fewer linear measurements when compared to Bora et al . ( 2017 ) . Despite being trained on the CelebA dataset , our generative invertible prior can give higher quality reconstructions than Lasso on out-of-distribution images of faces , and , to a lesser extent , unrelated natural images . Our invertible prior outperforms a pretrained DCGAN ( Radford et al. , 2015 ) at face inpainting and exhibits qualitatively reasonable results on out-of-distribution human faces . We provide additional experiments in the appendix , including for training on other datasets . 2 METHOD AND MOTIVATION . We assume that we have access to a pretrained generative invertible neural network G : Rn → Rn . We write x = G ( z ) and z = G−1 ( x ) , where x ∈ Rn is an image that corresponds to the latent representation z ∈ Rn . We will consider a G that has the Glow architecture introduced in Kingma and Dhariwal ( 2018 ) . It can be trained by direct optimization of the likelihood of a collection of training images of a natural signal class , under a standard Gaussian distribution over the latent space . We consider recovering an image x from possibly-noisy linear measurements given by A ∈ Rm×n , y = Ax+ η , where η ∈ Rm models noise . Given a pretrained invertible generator G , we have access to likelihood estimates for all images x ∈ Rn . Hence , it is natural to attempt to solve the above inverse problem by a maximum likelihood formulation given by min x∈Rn ‖Ax− y‖2 − γ log pG ( x ) , ( 1 ) where pG is the likelihood function over x induced by G , and γ is a hyperparameter . We have found this formulation to be empirically challenging to optimize ; hence we study the following proxy : min z∈Rn ‖AG ( z ) − y‖2 + γ‖z‖ . ( 2 ) Unless otherwise stated , we initialize ( 2 ) at z0 = 0 . The motivation for formulation ( 2 ) is as follows . As a proxy for the likelihood of an image x ∈ Rn , we will use the likelihood of its latent representation z = G−1 ( x ) . Because the invertible network G was trained to map a standard normal in Rn to a distribution over images , the log-likelihood of a point z is proportional to ‖z‖2 . Instead of penalizing ‖z‖2 , we alternatively penalize the unsquared ‖z‖ . In Appendix B , we show comparable performance for both the squared and unsquared formulations . In principle , our formulation has an inherent flaw : some high-likelihood latent representations z correspond to low-likelihood images x . Mathematically , this comes from the Jacobian term that relates the likelihood in z to the likelihood in x upon application of the map G. For multimodel distributions , such images must exist , which we will illustrate in the discussion . This proxy formulation relies on the fact that the set of such images has low probability and that they are inconsistent with enough provided measurements . Surprisingly , despite this potential weakness , we will observe image reconstructions that are superior to BM3D and GAN-based methods at denoising , and superior to GAN-based and Lasso-based methods at compressive sensing . In the case of compressive sensing and inpainting , we take γ = 0 in formulation ( 2 ) . The motivation for such a formulation initialized at z0 = 0 is as follows . There is a manifold of images that are consistent with the provided measurements . We want to find the image x of highest likelihood on this manifold . Our proxy turns the likelihood maximization task over an affine space in x into the geometric task of finding the point on a manifold in z-space that is closest to the origin with respect to the Euclidean norm . In order to approximate that point , we run a gradient descent in z down the data misfit term starting at z0 = 0 . In the case of GAN priors for G : Rk → Rn , we will use the formulation from Bora et al . ( 2017 ) , which is the formulation above in the case where the optimization is performed over Rk , γ = 0 , and initialization is selected randomly . All the experiments that follow will be for an invertible model we trained on the CelebA dataset of celebrity faces , as in Kingma and Dhariwal ( 2018 ) . Similar results for models trained on birds and flowers ( Wah et al. , 2011 ; Nilsback and Zisserman , 2008 ) can be found in the appendix . Due to computational considerations , we run experiments on 64 × 64 color images with the pixel values scaled between [ 0 , 1 ] . The train and test sets contain a total of 27,000 and 3,000 images , respectively . We trained a Glow architecture ( Kingma and Dhariwal , 2018 ) ; see Appendix A for details . Once trained , the Glow prior is fixed for use in each of the inverse problems below . We also trained a DCGAN for the same dataset . We solve ( 2 ) using LBFGS , which was found to outperform Adam ( Kingma and Ba , 2014 ) . DCGAN results are reported for an average of 3 runs because we observed some variance due to random initialization . 3 APPLICATIONS . 3.1 DENOISING . We consider the denoising problem with A = I and η ∼ N ( 0 , σ2I ) , for images x in the CelebA test dataset . We evaluate the performance of a Glow prior , a DCGAN prior , and BM3D for two different noise levels . Figure 2 shows the recovered PSNR values as a function of γ for denoising by the Glow and DCGAN priors , along with the PSNR by BM3D . The figure shows that the performance of the regularized Glow prior increases with γ , and then decreases . If γ is too low , then the network fits to the noise in the image . If γ is too high , then data fit is not enforced strongly enough . The left panel reveals that an appropriately regularized Glow prior can outperform BM3D by almost 2 dB . The experiments also reveal that appropriately regularized Glow priors outperform the DCGAN prior , which suffers from representation error and is not aided by the regularization . The right panel confirms that with smaller noise levels , less regularization is needed for optimal performance . A visual comparison of the recoveries at the noise level σ = 0.1 using Glow , DCGAN priors , and BM3D can be seen in Figure 3 . Note that the recoveries with Glow are sharper than BM3D . See Appendix B for more quantitative and qualitative results .

Summary: This paper proposes to employ the likelihood of the latent representation of images as the optimization target in the Glow (Kingma and Dhariwal, 2018) framework. The authors argue that to optimize the ''proxy for image likelihood'' has two advantages: First, the landscapes of the surface are more smooth; Second, a latent sample point in the regions that have a low likelihood is able to generate desired outcomes. In the experimental analysis, the authors compare their proposed method with several baselines and show prior performance.