Distribution Matching in Deep Generative Models with Kernel Transfer Operators

1 Introduction . Generative modeling , in its unconditional form , refers to the problem of estimating the data generating distribution : given i.i.d . samples X with an unknown distribution PX , a generative model seeks to find a parametric distribution that closely resembles PX . In modern deep generative models , we often approach this problem via a latent variable – i.e. , we assume that there is some variable Z ∈ Z associated with the observed data X ∈ X that follows a known distribution PZ ( also referred to as the prior in generative models ) . Thus , we can learn a mapping f : Z → X such that the distribution after transformation , denoted by Pf ( Z ) , aligns well with the data generating distribution PX . Therefore , sampling from PX becomes convenient since PZ can be efficiently sampled . Frequently , f is parameterized by deep neural networks and optimized with stochastic gradient descent ( SGD ) . Existing generative modeling methods variously optimize the transformation f , most commonly modeling it as a Maximum Likelihood Estimation ( MLE ) or distribution matching problem . For instance , given data X = { x1 , . . . , xn } , a variational autoencoder ( VAE ) ( Kingma & Welling , 2013 ) first constructs Z through the approximate posterior qZ|X and maximizes a lower bound of likelihood pf ( Z ) ( X ) . Generative adversarial networks ( GANs ) ( Goodfellow et al. , 2014 ) relies on a simultaneously learned discriminator such that samples of Pf ( Z ) are indistinguishable from X . Results in ( Arjovsky et al. , 2017 ; Li et al. , 2017 ) suggest that GANs minimize the distributional discrepancies between Pf ( Z ) and PX . Flow-based generative models optimize pf ( Z ) ( X ) explicitly through the change of variable rule and efficiently calculating the Jacobian determinant of the inverse mapping f−1 . In all examples above , the architecture or objective notwithstanding , the common goal is to find a suitable function f that reduces the difference between Pf ( Z ) and PX . Thus , a key component in many deep generative models is to learn a forward operator as defined below . Definition 1.1 ( Forward operator ) . A forward operator f ? ∈ C : Z → X is defined to be a mapping associated with some latent variable Z ∼ PZ such that f ? = arg minf∈C d ( Pf ( Z ) , PX ) for some function class C and a distance measure d ( · , · ) . Motivation : The specifics of the forward operator may differ from case to case . But its properties and how it is estimated numerically greatly influences the empirical performance of the model . For instance , mode collapse issues in GANs are well known and solutions continue to emerge ( Srivastava et al. , 2017 ) . To learn the forward operator , VAEs use an approximate posterior qZ|X that may sometimes fail to align with the prior ( Kingma et al. , 2016 ; Dai & Wipf , 2019 ) . Flow-based generative models enable direct access to the posterior likelihood , yet in order to tractably evaluate the Jacobian of the transformation during training , one must either restrict the expressiveness at each layer ( Dinh et al. , 2017 ; Kingma & Dhariwal , 2018 ) or use more involved solutions ( Chen et al. , 2018 ) . Of course , solutions to mitigate these weaknesses ( Ho et al. , 2019 ) remains an active area of research . The starting point of our work is to evaluate the extent to which we can radically simplify the forward operator in deep generative models . Consider some desirable properties of a hypothetical forward operator ( in Def . ( 1.1 ) ) : ( a ) Upon convergence , the learned operator f ? minimizes the distance between PX and Pf ( Z ) over all possible operators of a certain class . ( b ) The training directly learns the mapping from the prior distribution PZ , rather than a variational approximation . ( c ) The forward operator f ? can be efficiently learned and sample generation is also efficient . It would appear that these criteria violate the “ no free lunch rule ” , and some compromise must be involved . Our goal is to investigate this trade-off : which design choices can make this approach work ? Specifically , a well studied construct in dynamical systems , namely the Perron-Frobenius operator ( Lemmens & Nussbaum , 2012 ) , suggests an alternative linear route to model the forward operator . Here , we show that if we are willing to give up on a few features in existing models – this may be acceptable depending on the downstream use case – then , the forward operator in generative models can be efficiently approximated as the estimation of a closed-form linear operator in the reproducing kernel Hilbert space ( RKHS ) . With simple adjustments of existing results , we identify a novel way to replace the expensive training for generative tasks with a simple principled kernel approach . Contributions . Our results are largely based on results in kernel methods and dynamical systems , but we demonstrate their relevance in generative modeling and complement recent ideas that emphasize links between deep generative models and dynamical systems . Our contributions are ( a ) We propose a non-parametric method for transferring a known prior density linearly in RKHS to an unknown data density – equivalent to learning a nonlinear forward operator in the input space . When compared to its functionally-analogous module used in other deep generative methods , our method avoids multiple expensive training steps yielding significant efficiency gains ; ( b ) We evaluate this idea in multiple scenarios and show competitive generation performance and efficiency benefits with pre-trained autoencoders on popular image datasets including MNIST , CIFAR-10 , CelebA and FFHQ ; ( c ) As a special use case , we demonstrate the advantages over other methods in limited data settings . 2 Preliminaries We briefly introduce reproducing kernel Hilbert space ( RKHS ) and kernel embedding of probability distributions , concepts we will use frequently . Definition 2.1 ( RKHS ( Aronszajn , 1950 ) ) . For a set X , let H be a set of functions g : X → R. Then , H is a reproducing kernel Hilbert space ( RKHS ) with a product 〈· , ·〉H if there exists a function k : X × X → R ( called a reproducing kernel ) such that ( i ) ∀x ∈ X , g ∈ H , g ( x ) = 〈g , k ( x , · ) 〉H ; ( ii ) H = cl ( span ( { k ( x , · ) , x ∈ X } ) ) , where cl ( · ) is the set closure . The function φ ( x ) = k ( x , · ) : X → H is referred to as the feature mapping of the induced RKHS H. A useful identity derived from feature mappings is the kernel mean embedding : it defines a mapping from a probablity measure in X to an element in the RKHS . Definition 2.2 ( Kernel Mean Embedding ( Smola et al. , 2007 ) ) . Given a probability measure p on X with an associated RKHS H equipped with a reproducing kernel k such that supx∈X k ( x , x ) < ∞ , the kernel mean embedding of p in RKHS H , denoted by µp ∈ H , is defined as µp = Ep [ φ ( x ) ] = ∫ k ( x , · ) p ( x ) dx , and the mean embedding operator E : L1 ( X ) → H is defined as µp = Ep . Remark 1 . For characteristic kernels , the operator E is injective . Thus , two distributions ( p , q ) in X are identical iff Ep = Eq . This property allows using of Maximum Mean Discrepancy ( MMD ) for distribution matching ( Gretton et al. , 2012 ; Li et al. , 2017 ) and is common , see ( Muandet et al. , 2017 ; Zhou et al. , 2018 ) . For a finite number of samples { xi } ni=1 drawn from the probability measure p , an unbiased empirical estimate of µH is µ̂H = 1 n ∑n i=1 k ( xi , · ) such that limn→∞ 1 n ∑n i=1 k ( xi , · ) = µH . Next , we review the covariance/cross covariance operators , two widely-used identities in kernel methods ( Fukumizu et al. , 2013 ; Song et al. , 2013 ) and building blocks of our approach . Definition 2.3 ( Covariance/Cross-covariance Operator ) . Let X , Z be random variables defined on X × Z with joint distribution PX , Z and marginal distributions PX , PZ . Let ( l , φ , H ) and ( k , ψ , G ) be two sets of ( a ) bounded kernel , ( b ) their corresponding feature map , and ( c ) their induced RKHS , respectively . The ( uncentered ) covariance operator CZZ : H → H and cross-covariance operator CXZ : H → G are defined as CZZ , Ez∼PZ [ φ ( z ) ⊗ φ ( z ) ] CXZ , E ( x , z ) ∼PX , Z [ ψ ( x ) ⊗ φ ( z ) ] ( 1 ) where ⊗ is the outer product operator . 3 Simplifying the estimation of the forward operator . Forward operator as a dynamical system : The dynamical system view of generative models has been described by others ( Chen et al. , 2018 ; Grathwohl et al. , 2019 ; Behrmann et al. , 2019 ) . These strategies model the evolution of latent variables in a residual neural network in terms of its dynamics over continuous or discrete time t , and consider the output function f as the evaluation function at a predetermined boundary condition t = t1 . Specifically , given an input ( i.e. , initial condition ) z ( t0 ) , f is defined as f ( z ( t0 ) ) = z ( t0 ) + ∫ t1 t0 ∆t ( z ( t ) ) dt ( 2 ) where ∆t is a time-dependent neural network function and z ( t ) is the intermediate solution at t. This view of generative models is not limited to specific methods or model archetypes , but generally useful , for example , by viewing the outputs of each hidden layer as evaluations in discrete-time dynamics . After applying f on a random variable Z ∈ Z , the marginal density of the output over any subspace Λ ⊆ X can be expressed as∫ Λ pf ( Z ) ( x ) dx = ∫ z∈f−1 ( Λ ) pZ ( z ) dz ( 3 ) If there exists some neural network instance ∆ ? t such that the corresponding output function f ? satisfies PX = Pf ? ( Z ) , by Def . 1.1 , f ? is a forward operator . Let X be a set of i.i.d . samples drawn from PX . In typical generative learning , either maximizing the likelihood 1|X| ∑ x∈X pf ( Z ) ( x ) or minimizing the distributional divergence d ( Pf ( Z ) , PX ) requires evaluating and differentiating through f or f−1 many times . Towards a one-step estimation of forward operator : Since f and f−1 in ( 3 ) will be highly nonlinear in practice , evaluating and computing the gradients can be expensive . Nevertheless , the dynamical systems literature suggests a linear extension of f∗ , namely the Perron-Frobenius operator or transfer operator , that conveniently transfers pZ to pX . Definition 3.1 ( Perron-Frobenius operator ( Mayer , 1980 ) ) . Given a dynamical system f : X → X , the Perron-Frobenius ( PF ) operator P : L1 ( X ) → L1 ( X ) is an infinitedimensional linear operator defined as ∫ Λ ( PpZ ) ( x ) dx = ∫ z∈f−1 ( Λ ) pZ ( z ) dz for all Λ ⊆ X . Although in Def . 3.1 , the PF operator P is defined for self-maps , it is trivial to extend P to mappings f : Z → X by restricting the RHS integral ∫ z∈f−1 ( Λ ) pZ ( z ) dz to Z . It can be seen that , for the forward operator f∗ , the corresponding PF operator P satisfies pX = PpZ . ( 4 ) If P can be efficiently computed , transferring the tractable density pZ to the target density pX can be accomplished simply by applying P. However , since P is an infinite-dimensional operator on L1 ( X ) , it is impractical to instantiate it explicitly and exactly . Nonetheless , there exist several methods for estimating the Perron-Frobenius operator , including Ulam ’ s method ( Ulam , 1960 ) and the Extended Dynamical Mode Decomposition ( EDMD ) ( Williams et al. , 2015a ) . Both strategies project P onto a finite number of hand-crafted basis functions – this may suffice in many settings but may fall short in modeling highly complex dynamics . Kernel-embedded form of PF operator : A natural extension of PF operator is to represent P by an infinite set of functions ( Klus et al. , 2020 ) , e.g. , projecting it onto the bases of an RKHS via the kernel trick . There , for a characteristic kernel l , the kernel mean embedding uniquely identifies an element µX = ElpX ∈ G for any pX ∈ L1 ( X ) . Thus , to approximate P , we may alternatively solve for the dynamics from pZ to pX in their embedded form . Using Tab . 1 notations , we have the following linear operator that defines the dynamics between two embedded densities . Definition 3.2 ( Kernel-embedded Perron-Frobenius operator ( Klus et al. , 2020 ) ) . Given pZ ∈ L1 ( X ) and pX ∈ L1 ( X ) . Denote k as the input kernel and l as the output kernel . Let µX = ElpX and µZ = EkpZ be their corresponding mean kernel embeddings . The kernel-embedded Perron-Frobenius ( kPF ) operator , denoted by PE : H → G , is defined as PE = CXZC−1ZZ ( 5 ) Proposition 3.1 ( Song et al . ( 2013 ) ) . With the above definition , PE satisfies µX = PEµZ ( 6 ) under the conditions : ( i ) CZZ is injective ( ii ) µt ∈ range ( CZZ ) ( iii ) E [ g ( X ) |Z = · ] ∈ H for any g ∈ G. The last two assumptions can sometimes be difficult to satisfy for certain RKHS ( see Theorem 2 of Fukumizu et al . ( 2013 ) ) . In such cases , a relaxed solution can be constructed by replacing C−1ZZ by a regularized inverse ( CZZ + λI ) −1 or a Moore-Penrose pseudoinverse C † ZZ . The following proposition shows commutativity between the ( kernel-embedded ) PF operator and the mean embedding operator , showing its equivalence to P when l is characteristic . Proposition 3.2 ( ( Klus et al. , 2020 ) ) . With the above notations , El ◦ P = PE ◦ Ek . Transferring embedded densities with the kPF operator : The kPF operator is a powerful tool that allows transferring embedded densities in RKHS . The main steps are : ( 1 ) Use mean embedding operator El on pZ . Let us denote it by µZ . ( 2 ) Transfer µZ using kPF operator PE to get the mean embedded pX , given by µX . Of course , in practice with finite data , { xi } i∈ [ n ] ∼ PX and { zi } i∈ [ n ] ∼ PX , PE must be estimated empirically ( see Klus et al . ( 2020 ) for an error analysis ) . P̂E = ĈXZ ( ĈZZ ) −1 ≈ Ψ ( ΦTΦ + λnI ) −1ΦT ≈ Ψ ( ΦTΦ ) †ΦT where Φ = [ k ( z1 , · ) , · · · , k ( zn , · ) ] , Ψ = [ l ( x1 , · ) , · · · , l ( xn , · ) ] are simply the corresponding feature matrices for samples of PX and PZ , and λ is a small penalty term . Learning kPF for unconditional generative modeling : Some generative modeling methods such as VAEs and flow-based formulations explicitly model the latent variable Z as conditionally dependent on the data variable X . This allows deriving/optimizing the likelihood pf ( Z ) ( X ) . This is desirable but may not be essential in all applications . To learn a kPF , however , X and Z can be independent RVs . While it may not be immediately obvious why we could assume this independence , we can observe the following property for the empirical kPF operator , assuming that the empirical covariance operator ĈZZ is non-singular : P̂E µ̂Z = ĈXZ Ĉ−1ZZ µ̂Z = ΨΦ > ︸ ︷︷ ︸ ĈXZ ( ΦΦ > ︸ ︷︷ ︸ ĈZZ ) −1Φ1n = Ψ ( Φ > Φ ) −1Φ > Φ1n = Ψ1n = µ̂X ( 7 ) Suppose that { xi } i∈ [ n ] and { zj } i∈ [ n ] are independently sampled from the marginals PX and PZ . It is easy to verify that ( 7 ) holds for any pairing { ( xi , zj ) } ( i , j ) ∈ [ n ] × [ n ] . However , instantiating the RVs in this way rules out the use of kPF for certain downstream tasks such as controlled generation or mode detection , since Z does not contain information regarding X . Nevertheless , if sampling is our only goal , then this instantiation of kPF will suffice . Mapping Z to G : Now , since PE is a deterministic linear operator , we can easily set up a scheme to map samples of Z to elements of G where the expectation of the mapped samples equals µX Define φ ( z ) = k ( z , · ) and ψ ( x ) = l ( x , · ) as feature maps of kernels k and l. We can rewrite µX as µX = PEEkpZ = PEEZ [ φ ( Z ) ] = EZ [ PE ( φ ( Z ) ) ] = EZ [ ψ ( ψ−1 ( PEφ ( Z ) ) ) ] ( 8 ) Here ψ−1 is the inverse or the preimage map of ψ . Such an inverse , in general , may not exist ( Kwok & Tsang , 2004 ; Honeine & Richard , 2011 ) . We will discuss a procedure to approximate ψ−1 in §4.1 . In what follows , we will temporarily assume that an exact preimage map exists and is tractable to compute . Define Ψ∗ = P̂Eφ ( Z ) as the transferred sample in G using the empirical embedded PF operator P̂E . Then the next result shows that asymptotically the transferred samples converge in distribution to the target distribution . Proposition 3.3 . As n → ∞ , ψ−1 ( Ψ∗ ) D→ PX . That is , the preimage of the transferred sample approximately conforms to PX under previous assumptions when n is large . Proof . Since P̂E asymp.→ P , the proof immediately follows from ( 8 ) .

Summary: The paper proposes to leverage the Perron-Frobenius operator to simplify the mapping from Z -> X in latent variable generative models. Specifically, this “forward operator” can be approximated by a closed-form linear operator in the RKHS. They evaluate their method on both synthetic settings and high-dimensional image datasets such as MNIST, CIFAR-10, CelebA, and FFHQ.