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: The authors build a generator that builds on top of the latent space of a “well-trained auto-encoder”. The generator consists of several steps: 1) sampling an latent element from the spherical latent space, 2) using a kernel Perron-Frobenius operator to embed the sampled latent element 3) selecting latent representations of real samples based on the indices of the highest values of the newly calculated embedding 4) calculating the geodesic interpolation of these latent representations 5) using the decoder of the autoencoder to generate new samples from the latent element resulting from geodesic interpolation. The authors describe the related work, the methods they propose and present experimental results on 4 datasets.