Title: LEARNING FLAT LATENT MANIFOLDS WITH VAES

Abstract: Latent-variable models represent observed data by mapping a prior distribution over some latent space to an observed space. Often, the prior distribution is specified by the user to be very simple, effectively shifting the burden of a learning algorithm to the estimation of a highly non-linear likelihood function. This poses a problem for the calculation of a popular distance function-the geodesic between data points in the latent space-as this is often solved iteratively via numerical methods. These are less effective if the problem at hand is not well captured by first or secondorder approximations. In this work, we propose less complex likelihood functions by allowing complex distributions and explicitly penalising the curvature of the decoder. This results in geodesics which are approximated well by the Euclidean distance in latent space, decreasing the runtime by a factor of 1,000 with little loss in accuracy. Additionally, we apply our method to a state-of-the-art tracking algorithm using real world image data, showing that our unsupervised method performs similar to supervised learning methods.

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

EQUATION

where z ∈ R Nz represents latent variables and x ∈ R Nx 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) :

EQUATION

where

p D (x) = 1 N N i=1 δ(x -x i )

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:

EQUATION

)

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:

EQUATION

with the optimisation objective F(φ, Θ, Φ), the inequality constraint E p D (x) E q φ (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:

EQUATION

Building on that, the authors propose an optimisation algorithm-including a λ-update schemeto 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 ∈ R Nz×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 ∈ T z M :

EQUATION

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] → R Nz in the Riemannian (latent) space and the corresponding N x -dimensional Euclidean (observation) space, the Riemaninan distance of the curve can be written as

EQUATION

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 = J T J, where J is the Jacobian of the decoder. The geodesic is obtained by minimising L(γ).

The magnification factor (MF(z) ≡ det G(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

EQUATION

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

EQUATION

)

where β is a hyper-parameter determining the influence of the regularisation. c 2 is defined to be the mean over the batch samples and diagonal elements of J T J, 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:

EQUATION

where ∼ N (0, σ 2 I). 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 z i and z j :

EQUATION

with x i , x j ∼ p D (x), z i ∼ q φ (z|x i ), z j ∼ q φ (z|x j ), 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):

EQUATION

)

By using augmented data when minimising Gc 2 1 , 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

D x (f (z i ), f (z j )) ≈ c D z (z i , z j )

, 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 condition

D x (f (z i ), f (z j )) ≤ a D z (z i , z j )

, where a is the Lipschitz constant, we consider the decoder function, and hence the latent space as smooth if ∃ c ≤ a.

3. RELATED WORK:

Latent space of VAEs. In general, the latent space of VAEs is considered to be Euclidean (e.g. Kingma et al., 2016; Higgins et al., 2017) , but they are not constrained to be Euclidean. This can be problematic if a precise metric is required in the latent space. Some recent works (Mathieu et al., 2019; Grattarola et al., 2018) adapted the latent space to be non-Euclidean to match the data structure.

We solve the problem from another perspective by enforcing the latent space to be Euclidean.

Jacobian and Hessian regularisation. In (Rifai et al., 2011a) , the authors propose to regularise the Jacobian and Hessian of the encoder. However, the encoders of VAEs are already regularised by the KL term. Furthermore, it is more difficult augment data in the observation space than in the latent space for data augmentation. Encoder regularisation enables the model to perform better in case of, e.g., object recognition from the latent space. By contrast, decoder regularisation enables the model to do tasks such as generating motions based on the latent space. In (Hadjeres et al., 2017) , the Jacobian of the decoder was regularised to be as small as possible/zero. By contrast, we regularise the Jacobian to be constand, and hence the Hessian to be zero leading to a correct metric in the latent space. Nie & Patel (2019) regularised the Jacobian with respect to the weights of both the encoder and decoder for GANs. In terms of supervised learning, Jakubovitz & Giryes (2018) regularised the Jacobian to improve the robustness for classification.

Metric learning. Various metric learning approaches for both deep supervised and unsupervised models were proposed. For instance, deep metric learning (Hoffer & Ailon, 2015) used a triplet network for supervised learning. Karaletsos et al. (2016) introduced an unsupervised metric learning method, where a VAE is combined with triplets. However, a human oracle is still required. By contrast, our approach is completely based on unsupervised learning, using the tangent space of the decoder as a distance metric. Our proposed method is similar to the metric learning methods such as Large Margin Nearest Neighbor (Weinberger & Saul, 2009) , which pulls target neighbours together and pushes imposters away, but our approach is an unsupervised method.

Constraints in latent space. Constraints on time (e.g. Wang et al., 2007; Chen et al., 2016; 2015) allow to obtain similar distance metrics in the latent space. However, our method can be used for general datasets without sequential data. Additionally, constraints on time cannot guarantee that the metric is correct in between of different time steps. By contrast, in case of sequential data, our method can be used to obtain a correct metric through data augmentation.

Data augmentation. In regions without training data, the latent space is trained arbitrarily. Our method is able to augment data in the latent space, so that we can smoothly interpolate between two points even in case the data in between is missing in the training dataset using mixup. Various follow-up studies of mixup were developed, such as (Verma et al., 2018; Beckham et al., 2019) . GANs which generates fake data in the latent space are a similar approach as our approach.

Geodesic. There have been some recent studies on geodesics for generative models using both stochastic methods (e.g. Tosi et al., 2014; Arvanitidis et al., 2018; Hauberg, 2018; Chen et al., 2018b) and deterministic approaches (e.g., Chen et al., 2018a; 2019) . The main difference is that the stochastic methods work for the regions without data, because the RBF layer generates high MF for those-however, it is less general. The uncertainty does not emerge from a principled way (such as in a Bayesian model) but is instead driven by certain assumptions. The deterministic method requires other strategies to guarantee that the geodesic is within the data manifold. In our proposed method, we regularise the latent space to have a similar metric, so that we do not need to consider low density regions. In previous work, methods were introduced for computing/finding the geodesic in the latent space. However, it is a novel approach to use the geodesic/Riemannian distance for influencing the latent representation. Tenenbaum et al. (2000) projected the latent space to a new latent space where the geodesic is equivalent to the Euclidean interpolation. However, the two separate processes-VAEs and projection-probably cannot allow the model to find the latent features autonomously.

4. EXPERIMENTS:

We test our method on artificial pendulum images, human motion, the MNIST and the MOT16 datasets. We measure the performance in terms of equidistances, interpolation smoothness and geodesics. Additionally, our method is applied to a real-world environment-a tracking system from the context of autonomous driving. Consequently, the tracking and re-identification capabilities are evaluated.

Riemannian metric tensor has many intrinsic properties of a manifold and measures local angles, length, surface area and volumes (Bronstein et al., 2017) . Therefore, the models are quantified using the Riemannian metric tensor-condition numbers and MFs. The condition number which shows the ratio of the most elongated to the least elongated direction is defined as k(G) = Smax(G) Smin(G) , where S is an eigenvalue of G. Since we cannot directly compare the MFs of different models, the MFs are normalised through dividing by their mean. Accordingly, we can measure how the MFs spread out from their mean. The model is more invariant with respect to the metric tensor if the condition number is smaller and the normalised MF is closer to one. Figure 2 : Boxplot of the condition number and the normalised MF for pendulum and MNIST datasets.

4.1. PENDULUM IMAGE DATASET:

The pendulum and MNIST have 10,000 and 1,000 samples generated in from the latent space, respectively.

The pendulum dataset (Klushyn et al., 2019; Chen et al., 2018a) consists of 16 × 16-pixel images generated by a pendulum simulator. We generated 15 • 10 3 images with joint angle in the ranges of [0, 360) degrees. Additionally, we added 0.05 Gaussian noise to each pixel to avoid overfitting.

Fig. 1 shows the equidistance plots for five different encoded data points. VHP-FMVAE smoothens the MF, while VHP-VAE has large area of high MF in the middle. Without regularisation, the contour of the equidistances are significant different from high MF areas to low MF areas. Fig. 2a shows that the condition number of the VHP-FMVAE is smaller than that of the VAE-VHP in terms of the Riemannian tensor. Additionally, the normalised MF of the VHP-FMVAE is closer to one.

In order to avoid bias visualisation in Fig. 1 , 3 and 7, we reset the range of the MF for plotting while not changing the MF values. In Fig. 1a , Fig. 3a and 7 , the upper range is set to be max(MF2(grid_area))•mean(MF1(data)) mean(MF2(data))

. MF 1 and MF 2 refer to the MF of VHP-FMVAE with/without Jacobian normalisation/mixup and VAE-VHP, respectively. MF(data) and MF(grid_area) are the MF of the training data and the MF of the grid area, respectively.

4.2. HUMAN MOTION:

To evaluate our approach, CMU human motion dataset (http://mocap.cs.cmu.edu) is used. Walking (subject 35), jogging (subject 35), balancing (subject 49), punching (subject 143) and kicking (subject 74) are selected for the experiment. After data pre-processing, the input data is a 50-dimensional vector of the joint angles. The dataset is unbalanced-walking dataset size is larger than that of jogging.

Table 1 : The length ratio of Euclidean interpolation to geodesic. The Riemannian distance and the distance in the latent space are computed. We randomly sample 100 pairs of points and interpolate between each pair. The mean and the standard deviation of the ratios are listed below. Equidistance. We randomly select a point from each class as the centre of the equidistance. As shown in Fig. 3 , the proposed method has more similar equidistance at different locations in the latent space, while in the model without regularisation, the equidistance contour is distorted in the high MF Figure 3 : Equidistance in the latent space of the human motion dataset. (a) Jogging is a large-range movement compared with walking, so that jogging is reasonably distributed on a larger area in the latent space than that of walking. (b) In contrast, without regularisation, walking is larger than the jogging in the latent space. For FMVAE, the Euclidean interpolatioins are much closer to the geodesics. area. The latent space of VHP-FMVAE reflects the true distribution of the data (see more details Fig. 3 ). Moreover, we compute the Riemannian tensors for 3,000 samples randomly generated from the latent space. In different locations of a latent space, the Riemannian tensors of VHP-FMVAE is more invariant than that of VAE-VHP (see Fig. 4 ). Smoothness. We randomly sample 100 pair points and linearly interpolate between each pair. The second derivative of each trajectory is defined as the smoothness factor. Fig. 5 illustrates that VHP-FMVAE significantly outperforms the latent space of VAE-VHP in terms of the smoothness. Fig. 6 shows five examples of the interpolated trajectories.

V H P -F M V A E V H P -V A E V H P -F M V A E w / o

Geodesic. We compare the proposed method with the graph-based geodesic approach (Chen et al., 2019) which approximates the geodesic using a graph in a generative model. The graph-based approach is much faster than previous geodesic search method such as (Chen et al., 2018a) . The graph of the baseline has 14,400 nodes which are sampled in the latent space using uniform distribution. Each node has 12 neighbours. The regularisation of singular value decomposition (SVD) (Chen et al., 2018a) is 0.001 for VAE-VHP while it is adapted to VHP-FMVAE based the mean value of the square of the singular.

Table 1 shows the ratios from Euclidean interpolations to geodesics. If the ratio of the distance is close to one, the Euclidean interpolation is able to approximate the geodesic. Table 1 demonstrates that the Euclidean interpolation of VHP-FMVAE is more close to geodesic, compared with VAE-VHP. Additionally, the proposed method is 1,000 times faster than the graph-based method in terms of searching for the geodesics. Fig. 3 depicts five examples of the geodesics.

Influence of the data augmentation and the Jacobian normalisation. Fig. 7a shows the influence of the data augmentation. The samples of the regularisation term are the same as L VHP . Fig. 7b illustrates the influence of the Jacobian normalisation. We removed the normalisation, and consequently the regularisation term is G(g aug (z i , z j )) . The c 2 1 term in the regularisation is necessary; otherwise it only has dissimilarity constraints, but cannot reduce the distance for points with high similarities. For instance, the walking is not squeezed in the latent space (see Fig. 4 ). By contrast, regularising the Jacobian to be constant elongates the distance in the latent space with high MF areas while squeezing the distance with low MF areas.

4.3. MNIST:

A fixed binarised version of the MNIST digit dataset (Larochelle & Murray, 2011) is used to evaluate our approach. The dataset consists of 50,000 training and 10,000 test images of handwritten digits (zero to nine) with 28 × 28 pixels in size.

The equidistances of VHP-FMVAE are more invariant and smoother than that of VAE-VHP (see Fig. 8 and Fig. 4 ). Similar as the human motion dataset, the geodesic of our method is more similar to Euclidean interpolation, compared with VAE-VHP, which indicates that the latent space of the VHP-FMVAE is able to approximate geodesic (see Table 1 ). [-106.21, 369.38 ] and [-365.64, 164 .08], respectively. For better visualisation, we crop out the less dense areas.

4.4. TRACKING:

We evaluate our approach on the MOT16 object-tracking database (Milan et al., 2016) , which is a large-scale person re-identification dataset, containing both static and dynamic scenes from diverse cameras.

Table 2 : Comparisons between different descriptors for the purposes of object tracking and reidentification (Ristani et al., 2016) . The bold and the red numbers denote the best results among all methods and among non-supervised methods, respectively. We compare with two baselines: SORT (Bewley et al., 2016) and DeepSORT (Wojke et al., 2017) . SORT is a simple online and realtime tracking method, which uses bounding box intersection-overunion (IOU) for associating detections between frames and Kalman filters for the track predictions. It relies on good two-dimensional bounding box detections from a separate detector, and suffers from ID switching when tracks overlap in the image. DeepSORT extends the original SORT algorithm to integrate appearance information based on a deep appearance descriptor, which helps with reidentification in the case of such overlaps or missed detections. The deep appearance descriptor is trained using a supervised cosine metric learning approach (Wojke & Bewley, 2018) . The candidate object locations of the pre-generated detections for both SORT, DeepSORT and our method are taken from (Yu et al., 2016) . Further details regarding the implementation can be found in App. A.3.

:



Figures and Tables:

Latent representation of VHP-FMVAE.

Latent representation of VHP-VAE.

Figure1: Equidistance in the latent space of the pendulum dataset. The black curves are points of equal distance to a center. The distance is computed using (7).

Latent representation of VAE-VHP.

of the condition number and the normalised MF of human motion dataset.

Figure 5: Smoothness of the human dataset. The mean and standard deviation are shown. The smaller the value is, the smoother the model is.

(a) Linear interpolation of VHP-FMVAE. (b) Linear interpolation of VHP-VAE.

Figure 6: Generated movements of the human motion dataset. The abrupt motions are marked by blue boxes.

Latent representation of VHP-FMVAE without mixup. In the area which has training data, the equidistance contour is smooth. However, in the area of data missing, e.g., between two movements, the MF is high and the equidistances are distorted. Latent representation of VHP-FMVAE without the Jacobian normalisation. Although it does not have extreme sharp equidistance contour, the equidistance is still scaled in various locations of the latent space. Additionally, the distribution of walking in the latent space is still larger than that of jogging.

Figure 7: Influence of the data augmentation and the Jacobian normalisation. The movements are coloured the same as Fig. 3.

Latent representation of VHP-VAE.

Figure 8: in the latent space of MNIST dataset. (b) The data ranges on z 1 and z 2 of the VHP-VAE are[-106.21, 369.38] and[-365.64, 164.08], respectively. For better visualisation, we crop out the less dense areas.

Peer review (iclr2020):
1. The idea of explicitly forcing the encoding space to be flat by putting constraint on metric tensor is simple but neat. 2. The use of Jacobi regularization in Eq. (9) is effective but the choice of using interpolation to extend this in the entire decoding space is kind of adhoc. Can authors please justify? 3. Not sure how authors put the Lipschitz continuity constraint on f. Please explain. 4. The title of “Flat manifold VAEs” is misleading as it potentially means VAEs for flat manifold 5. I wonder what will happen if you put an unfolding constraint in the encoding space like LLE, ISOMAP etc.. The loss function is data driven so this should give atleast similar behavior. 6. Overall I like the experimental setup, but the tracking experiment is kind of distracting. The authors may want to remove this experiment. 7. In Fig. 7, the authors have shown with and without Jacobi normalization which I am really not convinced with, need better explanation.