MaGNET: Uniform Sampling from Deep Generative Network Manifolds Without Retraining

Deep Generative Networks ( DGNs ) are extensively employed in Generative Adversarial Networks ( GANs ) , Variational Autoencoders ( VAEs ) , and their variants to approximate the data manifold and distribution . However , training samples are often distributed in a non-uniform fashion on the manifold , due to costs or convenience of collection . For example , the CelebA dataset contains a large fraction of smiling faces . These inconsistencies will be reproduced when sampling from the trained DGN , which is not always preferred , e.g. , for fairness or data augmentation . In response , we develop MaGNET , a novel and theoretically motivated latent space sampler for any pre-trained DGN , that produces samples uniformly distributed on the learned manifold . We perform a range of experiments on various datasets and DGNs , e.g. , for the state-of-the-art StyleGAN2 trained on FFHQ dataset , uniform sampling via MaGNET increases distribution precision and recall by 4.1 % & 3.0 % and decreases gender bias by 41.2 % , without requiring labels or retraining . As uniform distribution does not imply uniform semantic distribution , we also explore separately how semantic attributes of generated samples vary under MaGNET sampling . 1 INTRODUCTION . Deep Generative Networks ( DGNs ) are Deep Networks ( DNs ) trained to learn latent representations of datasets ; such frameworks include Generative Adversarial Networks ( GANs ) ( Goodfellow et al. , 2014 ) , Variational Autoencoders ( VAEs ) ( Kingma & Welling , 2013 ) , flow-based models such as NICE ( Dinh et al. , 2014 ) , and their variants ( Dziugaite et al. , 2015 ; Zhao et al. , 2016 ; Durugkar et al. , 2016 ; Arjovsky et al. , 2017 ; Mao et al. , 2017 ; Yang et al. , 2019 ; Fabius & van Amersfoort , 2014 ; van den Oord et al. , 2017 ; Higgins et al. , 2017 ; Tomczak & Welling , 2017 ; Davidson et al. , 2018 ; Dinh et al. , 2016 ; Grathwohl et al. , 2018 ; Kingma & Dhariwal , 2018 ) . A common assumption that we will carry through our study is that the datasets of interest are not uniformly distributed in their ambient space , but rather are concentrated on , or around , manifolds of lower intrinsic dimension , e.g. , the manifold of natural images ( Peyré , 2009 ) . Different DGN training methods have been developed and refined to obtain models that approximate as closely as possible the training set distribution . This becomes an Achilles heel when the training set , regardless of its size , is not representative of the true data distribution , i.e. , when the training samples have been curated based on cost or availability that result in implicit/explicit biases . In such scenarios , while the training samples will lie on the true data manifold , the density distribution of the training set will be different from the natural distribution of the data . Deploying a DGN trained with a biased data distribution can be catastrophic , in particular , when employed for tasks such as data augmentation ( Sandfort et al. , 2019 ) , controlled data generation for exploration/interpretation ( Thirumuruganathan et al. , 2020 ) , or estimation of statistical quantities of the data geometry , such as the Lipschitz constant of the data manifold ( Gulrajani et al. , 2017 ; Scaman & Virmaux , 2018 ) . Biased data generation from DGNs due to skewed training distributions also raises serious concerns in terms of fair machine learning ( Hwang et al. , 2020 ; Tan et al. , 2020 ) . While ensuring semantic uniformity in samples is an extremely challenging task , we take one step in the more reachable goal of controlling the DGN sampling distribution to be uniform in terms of the sample distribution on the data manifold . To that end , we propose MaGNET ( for Maximum entropy Generative NETwork ) , a simple and efficient modification to any DGN that adapts its latent space distribution to provably produce samples uniformly spread on the DGN learned manifold . Importantly , MaGNET can be employed on any pre-trained and differentiable DGN regardless of its training setting , reducing the requirement of fine-tuning or retraining of the DGN . This is crucial as many models , such as BigGAN ( Brock et al. , 2019 ) and StyleGAN ( Karras et al. , 2020 ) , have significant computational and energy requirements for training . A plug-and-play method is thus greatly preferred to ease deployment in any already built/trained deep learning pipeline . Previously , there has been rigorous work on DGNs aimed at improving the training stability of models , deriving theoretical approximation results , understanding the role of the DGN architectures , and numerical approximations to speed-up training and deployment of trained models ( Mao et al. , 2017 ; Chen et al. , 2018 ; Arjovsky & Bottou ; Miyato et al. , 2018 ; Xu & Durrett , 2018 ; Liu et al. , 2017 ; Zhang et al. , 2017 ; Biau et al. , 2018 ; Li et al. , 2017 ; Kodali et al. , 2017 ; Roy et al. , 2018 ; Andrés-Terré & Lió , 2019 ; Chen et al. , 2018 ; Balestriero et al. , 2020a ; Tomczak & Welling , 2016 ; Berg et al. , 2018 ) . Existing methods ( Metz et al. , 2016 ; Tanaka , 2019 ; Che et al. , 2020 ) also try to tackle mode dropping by improving approximation of the data distribution , but this can potentially increase the bias learned implicitly by the DGNs . We are the first to consider the task of providing uniform sampling on the DGN underlying manifold , which has far-reaching consequences , ranging from producing DGNs stable to data curation and capable of handling inconsistencies such as repeated samples in the training set . We provide a first-of-its-kind provable uniform sampling on the data manifold that can be used to speed up estimation of various geometric quantities , such as estimation of the Lipschitz constant . MaGNET applies to any ( pretrained ) DGN architecture ( GAN , VAE , NF , etc . ) using continuous piecewise affine ( CPA ) nonlinearities , such as the ( leaky ) ReLU ; smooth nonlinearities can be dealt with via a first-order Taylor approximation argument . Our main contributions are as follows : [ C1 ] We characterize the transformation incurred by a density distribution when composed with a CPA mapping ( Sec . 3.1 ) and derive the analytical sampling strategy that enables one to obtain a uniform distribution on a manifold that is continuous and piecewise affine ( Sec 3.2 ) . [ C2 ] We observe that current DGNs produce continuous piecewise affine manifolds , and we demonstrate how to leverage [ C1 ] to produce uniform sampling on the manifold of any DGN ( Sec . 3.2 ) . [ C3 ] We conduct several carefully controlled experiments that validate the importance of uniform sampling and showcase the performance of MaGNET on pretrained models such as BigGAN ( Brock et al. , 2019 ) , StyleGAN2 ( Karras et al. , 2020 ) , progGAN ( Karras et al. , 2017 ) , and NVAE ( Vahdat & Kautz , 2020 ) , e.g. , we show that MaGNET can be used to increase distribution precision and recall of StyleGAN2 by 4.1 % & 3.0 % and decrease gender bias by 41.2 % , without requiring labels or retraining ( Sec . 4.2 and Sec . 4.3 ) . Reproducible code for the various experiments and figures will be provided upon completion of the review process . Computation and software details are provided in Appendix H , with the proofs of our results in Appendix I . Discussion on the settings in which MaGNET is desirable and on possible limitations is provided in ( Sec . 5 ) . 2 BACKGROUND . Continuous Piecewise Affine ( CPA ) Mappings . A rich class of functions emerge from piecewise polynomials : spline operators . In short , given a partition Ω of a domain RS , a spline of order k is a mapping defined by a polynomial of order k on each region ω ∈ Ω with continuity constraints on the entire domain for the derivatives of order 0 , . . . , k− 1 . As we will focus entirely on affine splines ( k = 1 ) we only define this case for concreteness . An affine spline S produces its output via S ( z ) = ∑ ω∈Ω ( Aωz + bω ) 1 { z∈ω } , ( 1 ) with input z and Aω , bω the per-region slope and offset parameters respectively , with the key constraint that the entire mapping is continuous on the domain S ∈ C0 ( RS ) . Spline operators and especially affine spline operators have been extensively used in function approximation theory ( Cheney & Light , 2009 ) , optimal control ( Egerstedt & Martin , 2009 ) , statistics ( Fantuzzi et al. , 2002 ) and related fields . Deep Generative Networks . A deep generative network ( DGN ) is a ( nonlinear ) operator GΘ with parameters Θ mapping a latent input z ∈ RS to an observation x ∈ RD by composing L intermediate layer mappings . The only assumption we require for our study is that the nonlinearities present in the DGN are continuous piecewise affine as is the case with ( leaky- ) ReLU , absolute value , max-pooling . For smooth nonlinearities , our results hold from a first-order Taylor approximation argument . Precise definitions of DGN operators can be found in Goodfellow et al . ( 2016 ) . We will omit Θ from the GΘ operator for conciseness unless needed . It is also common to refer to z as the latent representation , and x as the generated/observed data , e.g. , a time-series or image . One property of DGNs that employ nonlinearities such as ( leaky- ) ReLU , max-pooling , and the likes , is that the entire input-output mapping becomes a CPA spline with z as both the argument and region parameter . 3 CONTINUOUS PIECEWISE AFFINE MAPPING OF A PROBABILITY DENSITY . In this section , we study the properties of a probability density that is transformed by a CPA mapping . Our goal is to derive the produced density and characterize its properties , such as how the per-region affine mappings in Eq . 1 impact the density concentration . We present some key results that serve as the backbone of our core result in the next section : how to sample uniformly from the manifold generated by DGNs . 3.1 DENSITY ON THE GENERATED MANIFOLD . Consider an affine spline operator S ( Eq . 1 ) going from a space of dimension S to a space of dimension D with D ≥ S. The image of this mapping is a CPA manifold of dimension at most S , the exact dimension is determined by the rank of the per-region slope matrices . Formally , the image of S is given by Im ( S ) , { S ( z ) : z ∈ RS } = ⋃ ω∈Ω Aff ( ω ; Aω , bω ) ( 2 ) with Aff ( ω ; Aω , bω ) = { Aωz+bω : z ∈ ω } the affine transformation of region ω by the per-region parameters Aωbω . From Eq . 2 we observe that the generated manifold i.e. , surface is made of regions that are the affine transformations of the latent space partition regions ω ∈ Ω based on the coordinate change induced by Aω and the shift induced by bω . We visualize this in Fig . 2 for a toy spline operator with 2- dimensional latent space and 3-dimensional ambient/output space . In the remainder of our study we will denote for conciseness S ( ω ) , Aff ( ω ; Aω , bω ) . Figure 2 : Visual depiction of Eq . 2 with a toy affine spline mapping S : R2 7→ R3 . Left : latent space partition Ω made of different regions shown with different colors and with boundaries shown in black . Right : affine spline image Im ( S ) which is a continuous piecewise affine surface composed of the latent space regions affinely transformed by the per-region affine mappings ( Eq . 1 ) . The per-region colors are the same on the left and on the right . When the input space is equipped with a density distribution , then this density is transformed by the mapping S and “ lives ” on the surface of the CPA manifold generated by S. Given a distribution pz over the latent space , we can explicitly compute the output distribution after the application of S , which leads to an intuitive result exploiting the CPA property of the generator . For this result , we require that operator S be bijective between its domain and its image . That is , each slope matrix Aω , ∀ω ∈ Ω should be full rank , and there should not be any folding of the generated CPA surface that intersects with itself , i.e. , S ( ω ) ∩ S ( ω′ ) 6= { } ⇐⇒ ω = ω′ . We now derive the key result of this section that characterizes the density distribution lying on the manifold . Lemma 1 . The volume of a region ω ∈ Ω denoted by µ ( ω ) is related to the volume affinely transformed region S ( ω ) by , µ ( S ( ω ) ) µ ( ω ) = √ det ( ATωAω ) , ( 3 ) where µ ( S ( ω ) ) is using the measure on the S-dimensional affine subspace spanned by the CPA mapping . ( Proof in Appendix I.1 . ) Theorem 1 . The probability density pS ( x ) generated by S for latent space distribution pz is given by , pS ( x ) = ∑ ω∈Ω pz ( ( ATωAω ) −1 ATω ( x− bω ) ) √ det ( ATωAω ) 1 { x∈S ( ω ) } . ( 4 ) ( Proof in Appendix I.2 . ) In words , the distribution obtained in the output space naturally corresponds to a piecewise affine transformation of the original latent space distribution , weighted by the change in volume of the perregion mappings from ( Eq . 3 ) . For Gaussian and Uniform pz , we use the above results to obtain the analytical form of the density covering the output manifold , we have provided proof and differential entropy derivations in Appendix B .

Summary: The authors propose a uniform sampling technique for deep generative networks (DGNs) inspired by the probabilistic change of variables formula. The technique works with any already trained DGN and does not involve any further training. (Though it does require back propagation w.r.t. the input $x$.) In essence, the algorithm works by drawing many samples N >> K from the DGN, then sampling from these $N$ samples with probability inversely related to their pushforward density (as computed by the change-of-variables formula).