Can Kernel Transfer Operators Help Flow based Generative Models?

1 INTRODUCTION . A flow-based generative model refers to a deep generative model composed using a set of invertible transformations . While GANs and VAEs remain the two dominant generative models in the community , flow based formulations have continually evolved and now offer competitive performance in applications including audio/speech synthesis Kim et al . ( 2019 ; 2020 ) , text to speech Miao et al . ( 2020 ) , photo-realistic image generation Kingma & Dhariwal ( 2018 ) , and learning cross-domain mappings Mahajan et al . ( 2020 ) . An important property of such models is the explicit use of a tractable likelihood function , which enables leveraging maximum likelihood principles during training as well as efficient/exact density estimation and sampling . The formulation is invertible by design but this involves higher memory requirements . For example , permitting the bijective mapping to be expressive enough involves increases in the memory footprint Lee et al . ( 2020 ) ; Kim et al . ( 2019 ) , an issue that is a focus of several recent results Jacobsen et al . ( 2018 ) ; Chen et al . ( 2016 ) . Moreover , in these models we need to calculate the inverse and backpropagate through all invertible transformations during training . Calculating the inverse incurs a multiplicative increase in cost , usually as a function of the feature dimension , relative to the calculation of the likelihood , an issue addressed to some extent in Dinh et al . ( 2017 ) ; Kingma & Dhariwal ( 2018 ) . At a high level , a flow-based generative model bijectively pushes the data density from a source to a target , i.e. , from a known simple distribution to an unknown ( may be intractable ) data distribution . During training , we seek to learn this bijective mapping by maximizing the likelihood of the mapped training samples . In the generation step , we need the inverse of this mapping ( given such an inverse exists ) to map from a sample drawn from the known distribution back to the input ( data ) space . When the Jacobian of the transformation mapping can be efficiently computed or estimated ( e.g. , having a lower triangular form ) , directly optimizing the likelihood of the training samples is possible . However , in training flow-based generative either we must restrict the expressiveness at each layer or fall back on more numerically heavy solutions , see ( Chen et al. , 2018 ) . Next , we discuss how several existing results may provide a simplification strategy . 1.1 RELATED WORKS AND RATIONALE . Our starting point is the existing literature on Koopman and Perron-Frobenius operators Song et al . ( 2009 ) ; Fukumizu et al . ( 2013 ) ; Klus et al . ( 2020 ) , that offers an arguably easier , optionally linear , procedure that can be used to analyze non-linear dynamics of measurements that evolve temporally . For instance , as described in Arbabi ( 2018 ) ; Lusch et al . ( 2018 ) , if we view the data/measurements as evaluations of functions of the state ( of a dynamical system ) – where the functions are also called observables – then the entire set of such functions forms a linear vector space . Transfer operators on this space describe a linear evolution of the dynamics , i.e. , finite-dimensional nonlinear dynamics are replaced by infinite-dimensional linear dynamics Brunton et al . ( 2017 ) , perfectly evolving one set of measurements to another over time if the space can be well characterized . Of course , this is not practically beneficial because constructing such infinite-dimensional spaces could be intractable . Nonetheless , results in optimal control demonstrate that the idea can still be effective in specific cases , using approximations with either spectral analysis of large but finite number of functions Williams et al . ( 2015 ) or via a search for potential eigenfunctions of the operators using neural networks Li et al . ( 2017 ) ; Lusch et al . ( 2018 ) . Within the last year , several results describe its potential benefits of such operators in machine learning problems as well Li et al . ( 2020 ) ; Azencot et al . ( 2020 ) . If we consider the transformations that flow-based generative models learn as a non-linear dynamics , also used in ( Chen et al. , 2018 ) , a data-driven approximation strategy one can consider is to map the given data ( or distribution ) into infinite dimension space of functions through the kernel trick , which may allow the use of well known results based on kernel methods , including old and new results on powerful neural kernels Neal ( 1994 ) ; Jacot et al . ( 2018 ) ; Arora et al . ( 2019 ) . Utilizing these results , a mean embedding on the corresponding Reproducing Kernel Hilbert Space ( RKHS ) would correspond to the distribution in the input space ( the distribution from which input samples are drawn ) . Therefore , the problem of identifying a nonlinear mapping ( or dynamics ) in the input space ( going from an intractable distribution to a known distribution or vice-versa ) reduces to estimating a linear mapping operator between two empirical mean kernel embeddings where recent results on kernel transfer operators Klus et al . ( 2020 ) could be relevant or applicable . However , due to the high variability of the data , estimation of the distribution directly in the input space , as we will see shortly , can be difficult . But , if the input space is low-dimensional or otherwise structured , this problem could be mitigated . Fortunately , for many image datasets , one can identify a low-dimensional latent space such that , in theory , the above pipeline could be instantiated , enabling us to learn a transfer operator . Conceptually , it is not difficult to see how the foregoing idea could potentially help ( or simplify ) flow-based generative models . In principle , using a transfer operator , one could push-forward the input data distribution to a target distribution of our choice , if both have already been mapped to a sufficiently high dimensional space . If – additionally – the operator could also be inverted , this strategy may , at least , be viable . Of course , several key components are missing . We need to ( a ) assess if setting up a suitable infinite dimensional space is possible , ( b ) identify if we can estimate the transfer operator and then finally , ( c ) check if the procedure works at all . In the following sections of the paper , we will verify these key components and show that using only a linear operator yield surprisingly competitive results on several image generation tasks . 2 PRELIMINARIES . Auto-encoders . Images often lie close to an ( unknown ) lower dimensional manifoldM⊂ Rm such that dim ( M ) m , and operating with densities in a lower dimensional setting is often much easier . VAEs leverage this explicitly via auto-encoders . If the parameterized empirical density is pτ ( x ) , in VAEs , we write it as ∫ pτ ( x|z ) p ( z ) dz where z is the low dimensional representation with a suitable prior . We then use a decoder distribution qτ ′ ( z|x ) and an encoder distribution pτ ( x|z ) and train the parameters τ and τ ′ . In practice , pτ ( · ) and qτ ′ ( · ) are assumed to be Gaussian , but it is known that jointly fitting the manifold as well as regressing to a particular prior distribution can be challenging in VAEs Kingma et al . ( 2016 ) ; Dai & Wipf ( 2019 ) . Now , consider the following approach , where we do not impose a distributional assumption on z . If a well-regularized auto-encoder is able to capture information about the data generating density Alain & Bengio ( 2014 ) , we can think of z as the input measurements ( likely , meaningful representations from the input data ) which is subsequently mapped to an infinite dimensional space ( of observables of these measurements ) . Now , if we could push-forward the embedded RKHS distribution to the RKHS mapping of a simpler distribution , similar to flow-based generation , one could easily sample from the simple ( e.g. , standard normal ) distribution and transform it via the learned mapping to samples on the latent space . In summary , instead of explicitly searching for the eigenfunctions , we propose to Step 1 : embed the density from an auto-encoder into a RKHS , Step 2 : learn a kernel transfer operator in RKHS in one step . In the remainder of this paper , we will provide the details to operationalize this idea and show that this simple approach , in fact , performs surprisingly well with highly favorable computational properties . We now introduce the definition of reproducing kernel Hilbert space ( RKHS ) and the kernel embedding of probability distributions which are the building blocks of this paper . Definition 1 ( RKHS Aronszajn ( 1950 ) ) . Given a set X and H to be a set of functions φ : X → R , H is called a reproducing kernel Hilbert space ( RKHS ) with corresponding product 〈. , .〉H if there exists a function k : X × X → R ( called a reproducing kernel ) such that ( i ) ∀X ∈ X , φ ∈ H , φ ( X ) = 〈φ , k ( X , . ) 〉H ( ii ) H = span ( { k ( X , . ) , X ∈ X } ) , where . is the completion . The kernel mean embedding can be used to embed a probability measure in a RKHS . Definition 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 µH ∈ H , is defined as µH =∫ k ( X , . ) p ( X ) , and the mean embedding operator E : L1 ( X ) → H is defined as µH = Ep . For a characteristic kernel , the mapping from an input space distribution to its corresponding kernel mean embedding is one-to-one . Thus , two distributions in the input space are identical if and only if their kernel mean embedding matches exactly . This property enables Maximum Mean Discrepancy ( MMD ) for distribution matching . For a finite number of samples { Xi } Ni=1 , drawn from the probability measure µ , an empirical estimation of µH is µ̂H = 1N ∑N i=1 k ( Xi , . ) . Definition 3 ( Flow-based generative model Dinh et al . ( 2014 ) ) . A flow-based generative model explicitly learns the data distribution by trying to bijectively map it to a tractable density via invertible transformations . Formally , given a random variable z following a tractable density , i.e. , z ∼ pθ ( z ) , a flow-based model learns an invertible mapping fφ such that the data sample x = fφ ( z ) and the corresponding data distribution , pθ̃ ( x ) = f∗pθ ( z ) , where , f∗ = | dz dx | is the push-forward operator .

Summary: This paper starts with an autoencoder trained on vision data. Autoencoders aren't generative models, strictly, so in order to make it so, they leverage a simple linear transformation on over RKHSs. The kernel they use is the NTK. In order to generate, the construct a reduced sample version of the Kernel using Nystrom's approximation, then use a geodesic interpolation algorithm to inverse map the transferred sample from the prior in the RKHS space back to the latent space for use in the decoder.