FLAT MANIFOLD VAES

1 INTRODUCTION . Latent-variable models ( LVMs ) are a viable tool in data analysis : a set of observations is explained by a simpler set of latent variables in conjunction with a map from the latent space to the space of observation . Methods from this family ( e.g. , principal component analysis ( Wold et al. , 1987 ) , non-negative matrix factorisation ( Lee & Seung , 2001 ) , generalised discriminant analysis ( Baudat & Anouar , 2000 ) , etc . ) are standard tools , serving either as feature extractors for subsequent data processing pipelines , density estimators or dimensionality reducers for visualisation . Despite the maturity of the field , research has far from halted . While kernel methods ( KernelPCA ( Schölkopf et al. , 1997 ) , KernelNMF ( Li & Ding , 2006 ) , etc . ) have been used to improve the applicability of LVMs to data inhibiting non-linear phenomena , neural formulations such as the variational autoencoder ( VAE ) ( Kingma & Welling , 2014 ; Rezende et al. , 2014 ) or the generative adversarial network ( GAN ) ( Goodfellow et al. , 2014 ) have become popular recently , especially due to their enormous success on modelling natural images . Here , a simple prior distribution ( such as a multivariate standard normal or a uniform distribution ) is mapped to the space of observations by means of a powerful deep neural network . A learning algorithm then finds weights for that neural network such that the data distribution is approximated well . In case of GANs , this is a minimax game , while the evidence lower bound is maximised in the case of VAEs . If the data distribution is relatively complicated and the prior is relatively simple , the map from the latent to the observable space , the decoder , has to be sufficiently complex . In fact , it has to mimic the inverse CDF in parts . This results in highly non-linear neural networks . As an example , the separation of two modes in the probability landscape has to be implemented by a flat CDF , which in turn requires an infinitely steep inverse CDF . Not only does this often pose difficulties for gradient-based learning . It also hinders the calculation of geodesics , the shortest paths from one point to another , as measured by the rate of change in the decoder along the path in latent space . Geodesics are often solved numerically , i.e . through a gradient-based optimisation of first or second order . In a cartography scenario , it is crucial to project the three-dimensional earth onto a two-dimensional map . When applying VAEs , the problem is summarised to map from a two-dimensional Euclidean latent space ( the map ) to the three-dimensional observation space ( the earth surface ) . If , by construction , the decoder exhibits regions of high curvature , the stage is especially bad for such methods . We aim to improve the state of geodesics in deep latent-variable models with respect to runtime and use the geodesic as a distance metric . We expect that short ( approximate ) geodesics under the learned model indicate similarity of data points in question . If we assume that a simple prior and a simple decoder are insufficient to represent complex data distributions and a complex decoder is detrimental to the calculation of geodesics , we need the prior to be sufficiently complex to allow the decoder to be sufficiently simple . Our solution is the use of a powerful hierarchical prior representation in the context of VAEs and the simple penalisation of the curvature of the decoder . We name this approach flat manifold variational autoencoder since the Riemannian manifold of the decoder is isometric to Euclidean space . We show empirically that the resulting model features geodesics which are approximated very well by the Euclidean distance . This effectively removes the need for numerical optimisation , reducing the calculation of an approximate geodesic to that of a simple Euclidean distance calculation in latent space . This is accompanied by a speedup of several orders of magnitude , rendering the method practical for applications in real-time scenarios . 2 VARIATIONAL AUTOENCODERS WITH RIEMANNIAN MANIFOLD REGULARISATION . 2.1 BACKGROUND ON VAES WITH HIERARCHICAL PRIORS . Latent-variable models ( LVMs ) are defined as p ( x ) = ∫ p ( x|z ) p ( z ) dz , ( 1 ) where z ∈ RNz represents latent variables and x ∈ RNx the observable data . The integral in Eq . ( 1 ) is usually intractable but it can be approximated by maximising the evidence lower bound ( ELBO ) ( Kingma & Welling , 2014 ; Rezende et al. , 2014 ) : EpD ( x ) [ log pθ ( x ) ] ≥ EpD ( x ) [ Eqφ ( z|x ) [ log pθ ( x|z ) ] −KL ( qφ ( z|x ) ‖ p ( z ) ) ] , ( 2 ) where pD ( x ) = 1N ∑N i=1 δ ( x− xi ) represents the empirical distribution of the data D. The distribution parameters of the approximate posterior qφ ( z|x ) and the likelihood pθ ( x|z ) are represented by neural networks . The prior p ( z ) is usually defined as the standard normal distribution . This model is commonly referred to the variational autoencoder ( VAE ) . However , a standard normal prior often leads to an over-regularisation of the approximate posterior , which results in a less informative learned latent representation of the data ( Tomczak & Welling , 2018 ; Klushyn et al. , 2019 ) . To enable the model to learn an informative latent representation , Klushyn et al . ( 2019 ) propose to use a flexible hierarchical prior pΘ ( z ) = ∫ pΘ ( z|ζ ) p ( ζ ) dζ , where p ( ζ ) is the standard normal distribution . Based on the insight that the optimal prior is the aggregated posterior ( Tomczak & Welling , 2018 ) , the above integral is approximated by an importance-weighted ( IW ) bound ( Burda et al. , 2015 ) using samples from qφ ( z|x ) . This leads to a model with two stochastic layers and the following upper bound on the Kullback-Leibler ( KL ) term : EpD ( x ) KL ( qφ ( z|x ) ‖ p ( z ) ) ≤ F ( φ , Θ , Φ ) ≡ EpD ( x ) Eqφ ( z|x ) [ log qφ ( z|x ) − Eζ1 : K∼qΦ ( ζ|z ) [ log 1 K K∑ i=1 pΘ ( z|ζi ) p ( ζi ) qΦ ( ζi|z ) ] ] , ( 3 ) where K is the number of importance samples . Since it has been shown that high ELBO values do not necessarily correlate with informative latent representations ( Alemi et al. , 2018 ; Higgins et al. , 2017 ) —which is also the case for hierarchical models ( Sønderby et al. , 2016 ) —different optimisation approaches have been introduced ( Bowman et al. , 2016 ; Sønderby et al. , 2016 ) . Klushyn et al . ( 2019 ) follow the line of argument in ( Rezende & Viola , 2018 ) and reformulate the resulting ELBO as the Lagrangian of a constrained optimisation problem : LVHP ( θ , φ , Θ , Φ ; λ ) ≡ F ( φ , Θ , Φ ) + λ ( EpD ( x ) Eqφ ( z|x ) [ Cθ ( x , z ) ] − κ2 ) , ( 4 ) with the optimisation objective F ( φ , Θ , Φ ) , the inequality constraint EpD ( x ) Eqφ ( z|x ) [ Cθ ( x , z ) ] ≤ κ2 , and the Lagrange multiplier λ. Cθ ( x , z ) is defined as the reconstruction-error-related term in − log pθ ( x|z ) . Thus , we obtain the following optimisation problem : min Θ , Φ min θ max λ min φ LVHP ( θ , φ , Θ , Φ ; λ ) s.t . λ ≥ 0 . ( 5 ) Building on that , the authors propose an optimisation algorithm—including a λ-update scheme— to achieve a tight lower bound on the log likelihood . This approach is referred to as variational hierarchical prior VHP-VAE . 2.2 BACKGROUND ON RIEMANNIAN GEOMETRY FOR VAES . A manifold M is a space which is differentiable and locally Euclidean . Given M , a Riemannian manifold is ( M , G ) , where G ∈ RNz×Nz is the Riemannian metric tensor . At each point z ∈M in the latent space , the corresponding metric tensor G defines an inner product in the tangent space z′ ∈ TzM : 〈z′ , z′〉z ≡ z′T G ( z ) z′ . ( 6 ) Chen et al . ( 2018a ) ; Arvanitidis et al . ( 2018 ) define the latent space of a VAE as a Riemannian manifold , which allows for computing the observation space distance based on distances in the latent space . Given a smooth trajectory γ : [ 0 , 1 ] → RNz in the Riemannian ( latent ) space and the corresponding Nx-dimensional Euclidean ( observation ) space , the Riemaninan distance of the curve can be written as L ( γ ) = ∫ 1 0 φ ( t ) dt , φ ( t ) ≡ √〈 γ′ ( t ) , γ′ ( t ) 〉 γ ( t ) = √ γ′ ( t ) TG ( γ ( t ) ) γ′ ( t ) , ( 7 ) where G depends on γ , φ ( t ) denotes the Riemannian velocity and γ′ ( t ) represents the time-derivative of the trajectory . In the VAE models , γ is transformed by a continuous function f ( γ ( t ) ) ( decoder ) to x , and the metric tensor is defined as G = JTJ , where J is the Jacobian of the decoder . The geodesic is obtained by minimising L ( γ ) . The magnification factor ( MF ( z ) ≡ √ detG ( z ) ) ( Bishop et al. , 1997 ) shows the sensitivity of the likelihood functions . When projecting from the Riemannian ( latent ) to the Euclidean ( observation ) space , the MF can be considered a scaling coefficient . 2.3 FLAT MANIFOLD VAE . In this work , we aim to compute geodesics directly in the latent space by measuring the Euclidean distance between encoded data points . For this purpose , the metric tensor , which describes our latent space , needs to be G ∝ 1—hence a Euclidean metric . This simplifies the computation of geodesics ( Eq . ( 7 ) ) to L ( γ ) = ∫ 1 0 φ ( t ) dt ∝ ‖z ( 1 ) − z ( 0 ) ‖ . ( 8 ) We refer to a manifold M with this property as flat manifold ( Lee , 2006 ) . As a consequence , our model must be capable of learning such flat manifold latent spaces , which typically requires complex latent representations of the data ( see experiments in Sec . 4 ) . Therefore , we propose the following approach : ( i ) to enable our model to learn complex latent representations , we introduce a flexible prior , which is learned by the model ( empirical Bayes ) ; and ( ii ) we penalise the curvature of the decoder such that G ∝ 1 . For this purpose , we extend the VHP-VAE introduced in Sec . 2.1 by a Jacobian-regularisation term . Defining the regularisation term as part of the constraint is in line with the constrained optimisation setting . The resulting objective function is L = F ( φ , Θ , Φ ) + λ ( EpD ( x ) Eqφ ( z|x ) [ Cθ ( x , z ) + β ∥∥J ( z ) TJ ( z ) − c21∥∥ ] − κ2 ) , ( 9 ) where β is a hyper-parameter determining the influence of the regularisation . c2 is defined to be the mean over the batch samples and diagonal elements of JTJ , which we view as a normalisation process . Additionally , we use a stochastic approximation ( first order Taylor expansion ) of the Jacobian ( Rifai et al. , 2011b ) to improve the computational efficiency : J ( z ) = lim σ→0 1 σ EpD ( x ) Eqφ ( z|x ) [ f ( z + ) − f ( z ) ] , ( 10 ) where ∼ N ( 0 , σ2I ) . This approximation method allows for a faster computation of the gradient and avoids the second-derivative problem of piece-wise linear layers ( Chen et al. , 2018a ) during optimisation . However , the regularisation term in Eq . ( 9 ) only effects the decoder function in regions where data is available . To overcome this issue , we propose to use mixup , a data-augmentation method ( Zhang et al. , 2018 ) , which was introduced in the context of supervised learning . We extend this method to the VAE framework ( unsupervised learning ) by applying it to encoded data in the latent space . Our aim is to augment data by interpolating between two encoded data points zi and zj : gaug ( zi , zj ) = ( 1− α ) zi + α zj , ( 11 ) with xi , xj ∼ pD ( x ) , zi ∼ qφ ( z|xi ) , zj ∼ qφ ( z|xj ) , and α ∼ U ( −α0 , 1 + α0 ) . In contrast to ( Zhang et al. , 2018 ) , where α ∈ [ 0 , 1 ] limits the data augmentation to only convex combinations , we define α0 > 0 to take into account the outer edge of the data manifold . We obtain the objective function of our flat manifold VAE ( FMVAE ) by combining mixup ( Eq . ( 11 ) ) with Eq . ( 9 ) : LVHP-FMVAE = LVHP + λβ Exi , xj∼pD ( x ) Ezi∼qφ ( z|xi ) , zj∼qφ ( z|xj ) [ ∥∥G ( gaug ( zi , zj ) ) − c21∥∥ ] . ( 12 ) By using augmented data when minimising ‖G − c21‖ , we regularise G to be a scaled identity matrix for the entire latent space enclosed by our data manifold . Hence , our VAE learns a scaled Euclidean latent space , where c is the scale factor . Therefore , the function f ( z ) ( decoder ) is—up to the scale factor c—isometry/distance-preserving since Dx ( f ( zi ) , f ( zj ) ) ≈ cDz ( zi , zj ) , where D refers to the distance between two data points in the observation and latent space , respectively . The decoder of the proposed approach satisfies the Lipschitz continuity condition . Given the Lipschitz continuity conditionDx ( f ( zi ) , f ( zj ) ) ≤ aDz ( zi , zj ) , where a is the Lipschitz constant , we consider the decoder function , and hence the latent space as smooth if ∃ c ≤ a .

Summary: The paper is concerned with the geometry of latent spaces in VAEs. In particular, it is argued that since geodesics (shortest paths) in the Riemannian interpretation of latent spaces are expensive to compute, then it might be beneficial to regularize the decoder (generator) to be flat, such that geodesics are straight lines. One such regularization is proposed.