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 authors propose a new type of generative model. The new scheme is based on a kernel transfer operator that leads to a cheap method for distribution matching. The authors rely on rigorous theory on RKHS and propose a framework for transferring a prior distribution linearly (in RKHS) to the data distribution. The authors demonstrate that the proposed approach leads to improved approximations of observed distributions. Specifically, the new approach can lead to the generation of new images form a given distribution and requires less training time compared to existing baselines. The paper is mostly well written, and the method relies on solid justifications. The authors demonstrate the usefulness of the method and compare it to several existing methods.