Explorations in Homeomorphic Variational Auto-Encoding
Luca Falorsi, Pim de Haan, Tim R. Davidson, Nicola De Cao, Maurice Weiler, Patrick Forré, Taco S. Cohen
Introduction
Many complex probability distributions can be represented more compactly by introducing latent variables. Intuitively, the idea is that there is some simple underlying latent structure, which is mapped to the observation space by a potentially complex nonlinear function. It will come as no surprise then, that most research effort has aimed at using maximally simple priors for the latent variables (e.g. Gaussians), combined with flexible likelihood functions (e.g. based on neural networks).
However, it is not hard to see (Fig. 1.1) that if the data is concentrated near a low-dimensional manifold with non-trivial topology, there is no continuous and invertible mapping to a blob-like manifold (the region where prior mass is concentrated). We believe that for purposes of representation learning, the embedding map (encoder) should be homeomorphic (i.e. continuous and invertible, with continuous inverse), which means that although dimensionality reduction and geometrical simplification (flattening) may be possible, the topological structure should be preserved.
Once could encode such a manifold in a higher dimensional flat space with a regular variational auto-encoder (VAE, Kingma & Welling (2013); Rezende et al. (2014)), rather than learning a homeomorphism. This has two disadvantages. The prior on the flat space will put density outside of the embedding and traversals along the extra dimensions that are normal to the manifold will either leave the decoding invariant, or move out of the data manifold. This is because at each point there will be many more degrees of freedom than the dimensionality of the manifold.
In this paper we investigate this idea for the special case of Lie groups, which are symmetry groups that are simultaneously differentiable manifolds. Lie groups include rotations, translations, scaling, and other geometric transformations, which play an important role in many application domains such as robotics and computer vision. More specifically, we show how to constructOur implementation is available at https://github.com/pimdh/lie-vae. a VAE with latent variables that live on a Lie group, which is done by generalizing the reparameterization trick.
We will describe an approach for reparameterizing densities on , the group of 3D rotations, which can be extended to general compact and connected Lie group VAEs in a straightforward manner. The primary technical difficulty in the construction of this theory is to show that the pushforward measure induced by our reparameterization has a density that is absolutely continuous w.r.t. the Haar measure. Moreover, we show how to construct the encoder such that it can learn a homeomorphic map from the data manifold to the mean parameter of the posterior. Finally, we propose a decoder that uses the group action to further encourage the latent space to respect the group structure.
We perform experiments on two types of synthetic data: embedded into a high dimensional space through its group representation, and images of 3D rotations of a single colored cube. We find that a theoretically sound architecture is capable of continuously mapping the data manifold to the latent space. On the other hand, models that do not respect topological structure, and in particular those with a standard Gaussian latent space, show discontinuities when trajectories in the latent space are visualized. To better study this phenomenon, we introduce a way to measure the continuity of the embedding based on the concept of Lipschitz continuity. We empirically demonstrate that only a manifold-valued latent variable with the required topological structure is capable of fully solving the difficult task of the more complicated experiment.
Our main contributions in this work are threefold:
A reparameterization trick for distributions on the group of rotations in three dimensions.
An encoder for the mean parameter that learns a homeomorphism between the manifold embedded in the data and itself.
A decoder that uses the group action to respect the group structure.
Preliminary Concepts
In this section we will first cover a number of preliminary concepts that will be used in the rest of the paper.
The VAE is a latent variable model, in which denotes a set of observed variables, stochastic latent variables, and a parameterized model of the joint distribution called the generative model. Given a dataset , we typically wish to maximize the average marginal log-likelihood , w.r.t. the parameters. However when the model is parameterized by neural networks, the marginalization of this expression is generally intractable. One solution to overcome this issue is applying variational inference in order to maximize the Evidence Lower Bound (ELBO) for each observation:
where the approximate posterior belongs to the variational family . To make inference scalable an inference network is introduced that outputs a probability distribution for each data point , leading to the final objective
with representing the parameters of and . The ELBO can be efficiently approximated for continuous latent variable by Monte Carlo estimates using the reparameterization trick of (Kingma & Welling, 2013; Rezende et al., 2014).
2 Lie Groups and Lie Algebras
A group is a set equipped with a product that follows the four group axioms: the product is closed and associative, there exists an identity element, and every group element has an inverse. This is closely linked to symmetry transformations that leave some property invariant. For example, composing two symmetry transformations should still maintain the invariance. A Lie group has additional structure, as its set is also a smooth manifold. This means that we can, at least in local regions, describe group elements continuously with parameters. The number of parameters equals the dimension of the group. We can see (connected) Lie groups as continuous symmetries where we can continuously traverse between group elementsWe refer the interested reader to (Hall, 2003)..
Lie Algebra
The Lie algebra , of a dimensional Lie group is its tangent space at the identity, which is a vector space of dimensions. We can see the algebra elements as infinitesimal generators, from which all other elements in the group can be created. For matrix Lie groups we can represent vectors in the tangent space as matrices .
Exponential Map
The structure of the algebra creates a map from an element of the algebra to a vector field on the group manifold. This gives rise to the exponential map which maps an algebra element to the group element at unit length from the identity along the flow of the vector field. The zero vector is thus mapped to the identity. For compact connected Lie groups, such as , the exponential map is surjective.
3 The group SO(3)SO3\operatorname{SO}(3)
The special orthogonal Lie group of three dimensional rotations is defined as:
The elements of Lie algebra of group , are represented by the 3 dimensional vector space of the skew-symmetric matrices. We choose a basis for the algebra:
Since is a compact and connected Lie group this map is surjective, however it is not injective.
Reparameterizing SO(3)SO3\operatorname{SO}(3)
In this section we will explain our reparameterization trick by analogy to the classic version described in (Kingma & Welling, 2013; Rezende et al., 2014). An overview of the different steps and their relation to the classical case are given in Figure 2.1.
As further derived in Appendix B this implies the pushforward measure on to be absolutely continuous w.r.t. to the Haar measure where the density
is defined almost everywhere. Here and
Further, is defined as a principal branch and maps back the group element to the unique Lie algebra element next to the origin. Notice that even if the density is singular at , it still integrates to 1. After rotating by left multiplying with another element , we obtain the final sample:
where the second step is valid because of the left invariance of the Haar measure.
Encoder and Decoder networks
Having defined the reparameterizable density , we need to design encoder networks which map elements from the input space to the reparameterization parameters , and decoder networks which map group elements to the output prediction.
We consider the data as lying in a lower dimensional manifold , embedded in the input space . In our particular problem the manifold is assumed to be generated by , acting on a canonical object and a subsequent projection into ambient space (e.g. pixel space) which, for simplicity we assume to be injective. This means that we can make the simplifying assumption that can be recovered from its image in , i.e. that the map is a homeomorphism. The encoder is now meant to learn the inverse map, i.e. to learn a map from to , which when restricted to is a homeomorphism and thus preserves the topological structure of .
2 Group Action Decoder
Our decoder must be capable to map a group element and optionally additional latent structure back to the original high dimensional input space. When the factor of variation in the input is the pose of an object, and we learn a latent variable , we desire that a transformation applied to a latent object representation results in a corresponding transformation of the pose of the decoded object. The task of the decoder is thus to learn a three dimensional representation of the object, to rotate it according to the latent variable and finally to project it back to the two dimensional frame of the input image. A naive approach could be to simply provide the 9 elements of the rotation matrix to a neural network. However, although it may learn to reconstruct the input well, it provides no guarantee that the latent space accurately reflects the pose variations of the object. Therefore, we like to make the method more explicit.
Alternatively, we propose a method based on the representation theory of groups (Hall, 2003). Rather than learning the function , we directly learn its (band-limited) Fourier modes, which form a simple vector space. It can be shown (Chirikjian & Kyatkin, 2000) that rotations of a signal on the sphere correspond to a linear transformation of the Fourier modes. The transformed Fourier modes are subsequently fed through an image generative network, and the linear transformation is the Wigner-D-matrix, which is a function of the element. Technically, the Wigner-D-matrices form representations of the group. This means that as mapping to the linear transformations is a homomorphism, it preserves the group structure: , for and a Wigner-D-matrix. This method encourages the latent space to represent the actual pose of the input, while only requiring the construction of the matrices and performing a linear transformation. We refer to Figure 4.1 and Appendix F for details.
An additional advantage is that this decoder allows for disentangling of content and pose, as it is forced to encode the pose in a meaningful way. The Fourier modes are in that case also generated by the encoder. We leave this for future work.
Related Work
As VAEs utilize VI to recover some distribution on a latent manifold responsible for generating the observed data, the majority of extensions is focused on increasing the flexibility of the prior and approximate posterior. Although the majority of approaches make use of a normal Gaussian prior, recently there has been a surge to provide additional options to offset some of this distribution’s perceived limitations. Tomczak & Welling (2017) propose to directly tie the prior to the approximate posterior and learn it as a mixture over approximate posteriors. Nalisnick & Smyth (2017) introduce a non-parametric prior applying a truncated stick-breaking method. Research to support discrete latent variables was done in Jang et al. (2017); Maddison et al. (2017), while in Naesseth et al. (2017); Figurnov et al. (2018) recently novel techniques were introduced to reparameterize a suite of continuous distributions. In (Davidson et al., 2018), the reparameterization technique of Naesseth et al. (2017) is extended to explore the properties of the hyperspherical von Mises-Fisher distribution to better capture intrinsically hyperspherical data. This is done in the context of avoiding manifold mismatches, and as such is closely related to the motivation of this work.
The predominant procedure to generate a more complex approximate posterior is through normalizing flows (Rezende & Mohamed, 2015), in which a class of invertible transformations is applied sequentially to a reparameterizable density. This general idea has later been extended Kingma et al. (2016); Berg et al. (2018), to improve flexibility even further. As this framework does not hold any specific distributional requirements on the prior besides being reparameterizable, it would be interesting to investigate possible applications to in future work.
The problem of defining distributions on homogeneous spaces, including Lie groups, was investigated in (Chirikjian & Kyatkin, 2000; Chirikjian, 2010; Chirikjian & Kyatkin, 2016; Chirikjian, 2012). Cohen & Welling (2015) devised harmonic exponential families which are a powerful family of distributions defined on homogeneous spaces. These works did not concentrate on making the distributions reparameterizable.
Rendering complex scenes from multiple poses has been explored in (Eslami et al., 2018). However, this work assumes access to ground truth poses and does not do unsupervised pose learning as in the presented framework.
The idea of incorporating prior knowledge on mathematical structures into machine learning models has proven fruitful in many works. Cohen et al. (2018a) adapt convolutional networks to operate on spherical and valued data. Equivariant networks, investigated in Cohen & Welling (2016, 2017); Worrall et al. (2017); Weiler et al. (2018); Cohen et al. (2018b) reduce the complexity of a learning task by taking a quotient over group orbits which explain a subset of dimensions of the data manifold.
Experiments
We perform two experiments to investigate the importance of using a homeomorphic parameterization of the VAE in recovering the original underlying manifold. In both experiments we explore three main axes of comparison: (1) manifold topology, (2) decoder architecture, and (3) specifically for the models we compare different mean parameterizations as discussed in section 4.1. For each model we compute a tight bound on the negative log likelihood (NLL) through importance sampling following Burda et al. (2016).
For manifold topology we examine VAEs with the Gaussian parameterization (-VAE), the hyperspherical parameterization of Davidson et al. (2018) (-VAE), and the latent variable discussed above. The two decoder variants are a simple MLP versus the group action decoder described in section 4.2. Lastly we explore mean parameterizations through unit Quaternions (q), the Lie algebra (alg), (s2s1), and (s2s2). These parameterizations are chosen to be either valid () or invalid (q, alg, ) for the purpose of investigating the soundness of our theoretical considerations and to compare their behaviour. Details and derivations on the properties of these different parameterizations can be found in Appendix D.
In order to verify the ability of the models to correctly learn to encode from the embedded manifold to the manifold itself, we learn various variational and non-variational auto-encoders on this data set. The encoder is a 3 layer MLP, and for the decoder we use the group action decoder of Section 4.2. The same representation is used as in the data generation, but we learn . In addition to the models, we use a 3 dimensional normal, which we map to using the ZYZ-Euler angles, and a von Mises-Fisher, which we map to by identifying as the unit quaternions.
The quantitative results are shown in Table 1. We observe that the choice for the mean parametrization significantly impacts the ability of the model to correctly learn the manifold. The method strongly outperforms the competing methods in the non-variational Auto-Encoder achieving near-perfect reconstructions. Additionally, the metric indicating the continuity of the encoder, which we define in Appendix E, shows it is the only model that does not have discontinuities in the latent space. These results are in line with our theoretical prediction outlined in Section 4.1 and Appendix D.
The qualitative results in Figure 4.2 and Figures A.1, A.2 in Appendix A tell a similar story. These plots are created by taking a subgroup of and making a submanifold in the data space using the same process with which the data was generated. This embedded trajectory is then encoded and reconstructed. The trajectory is divided in four equally sized partitions, each shown in a different color. We clearly see that only the method is able to learn a continuous latent space.
2 Sphere-Cube
For this experiment we learn auto-encoders on renderings of a cube. The cube is made highly asymmetrical through the colors of the faces and the colored spheres at the vertices. This should make it easier for the encoder to detect the orientation. This sphere-cube is then rotated by applying uniformly sampled group elements from , to create a training set of 1M images. Ideally the model learns to correctly represent these encodings in the latent space.
The encoder consists of 5 convolutional layers, followed by one of the mean encoders and reparameterization methods. The decoder uses either the group action or a 3 layer MLP, both followed by a 5 deconvolutional layers. In order to balance reconstruction and the KL divergence in a controlled manner, we follow Burgess et al. (2018) and replace the negative KL term in the original VAE loss with a squared difference of the computed KL value and a target value. We found that a target value of 7 early in training to 15 at the end of the training gave good results. This allows the model to first organize the space and later become more certain of its predictions. We found that two additional regularizing loss terms were needed to correctly learn the latent space. Details can be found in Appendix G.
Quantitative results comparing the best performing parameterization to -VAEs of diff dimensionality are shown in Table 2. Although the higher dimensional -VAEs are able to achieve competitive metrics compared to the best model, they only learn to embed in a high dimensional space in an unstructured fashion. As can be seen in in 6.1, the latent space with mean parameterization learns a nearly perfect encoding, while the 10 dimensional Normal learns disconnected patches of the data manifold.Animated interpolations can be found at https://sites.google.com/view/lie-vae.
It can be seen in Table 3 that the results from the Toy experiment extend to this more complicated task. We observe that only the continuous encoding, , achieves low log likelihood and reconstruction losses compared to the other mean parameterizations.
Lastly, we observe that the group action decoder yields significantly higher performance than the MLP decoder. This is in line with the hypotheses that using the group action encourages structure in the latent space.
Discussion & Conclusion
In this paper we explored the use of manifold-valued latent variables, by proposing an extension of the reparameterization trick to compact connected Lie groups. We worked out the implementation details for the specific case of , and highlighted the various subtleties that must be taken into account to ensure a successful parameterization of the VAE. Through a series of experiments, we showed the importance of matching the topology of the latent data manifold with that of the latent variables to induce a continuous, well-behaved latent space. Additionally we demonstrated the improvement in learned latent space structure by using a group action decoder, and the need for care in choosing an embedding space for the posterior distribution’s mean parameter.
We believe that the use of and other well-known manifold-valued latent variables could present an interesting addition to tackling problems in such fields as model based RL and computer vision. Moving forward we thus aim to extend this theory to other Lie groups such as . A limitation of the current work, and reparameterizing distributions on specific manifolds in general, is that it relies on the assumption of a priori knowledge about the observed data’s latent structure. Hence in future work our ambition is to find a general theory to learn arbitrary manifolds not known in advance.
Acknowledgements
The authors would like to thank Rianne van den Berg, Jakub Tomczak, and Yvan Scher for their suggestions and support in improving this paper.
References
Appendix A Additional Figures
Appendix B Pushforward Measure SO(3)SO3\operatorname{SO}(3)
The proof then tells us how to compute the Radon-Nikodym derivative of the pushforward with respect to the Haar measure. In fact:
Defining we then have:
To then have an expression explicitly dependent on consider that
Where we have defined . Moreover we then have:
Appendix C Entropy computation
(Where we dropped dependency on the parameters for simplicity) Then using, Equation (24):
In the way we defined we obtain samples from it in the following way:
Substituting it in in the previous expression we get:
Notice that this expression depends only on the samples from in the lie algebra
Assuming the density decays quickly enough to zero, the above infinite summation can be truncated. This is always the case for exponentially decaying distributions, like the Normal. The truncated summation can then can be computed using the logsumexp trick:
Appendix D Mean parameterization
As discussed above, some requirements exist on for the encoder to correctly represent the data manifold.
We discuss the existence of a map such that it is a right inverse of (), which is necessary for the correct encoder to exist.
It’s inverses are the branches of the log map. However, a path in that is a full rotation around a fixed axis is continuous in but discontinuous in the algebra, when mapped with the log map. Thus the log map is not continuous.
For , to be continuous, its image must be closed (as it is a compact subset of a Hausdorff space). Thus so must the set . However, as for , is a hemisphere (times a point) that does not contain its entire boundary, thus it is not closed and is not continuous.
Notice that there exists a continuous and injective map . It simply consists of taking the first two rows of the matrix representation of the element (The third row is the vector product between the first two, so it can always be recovered). Moreover we have that
Appendix E Continuity Metric
Consider a map where , are metric spaces with metrics and respectively. In order to compute the proposed continuity metric we take a continuous path , defined as pairwise close points, and compute the relative distances
From this we further compute the quantities
By comparing these two values, we want to discover whether there is at least one outlier in the set of . Such outliers corresponds to a transition with a big jump, signalling a discontinuity point. We define a path to be discontinuous if .
In order to capture stochastic effects we repeat the above procedure with several paths. The final score is the fraction of discontinuous paths. In the practical implementation we chose paths, using and (th percentile).
Appendix F Group Action
where are the Fourier components and is the Wigner-D-matrix. We use identity , where are the first two Euler angles, to write the spherical harmonics that are the basis functions of the Fourier modes as Wigner-D-matrices. Then for a rotation , using the homomorphism property:
where corresponds to rotating a point on the sphere.
We see that our method of using representations in the decoder corresponds to having the content latent code represent the Fourier coefficients of a virtual signal on the sphere.
Appendix G Regularizers
Even when an appropriate mean parametrization is selected and proper behaviour of the decoder is encouraged by the group action decoder, the network can still learn a discontinuous latent space. To encourage it to learn the data manifold correctly, we employ two additional loss terms that act as regularizers. An ablative analysis of the effectiveness of these regularizers is shown in Table 4.
This equivariance is shown in Figure G.1. The relationship is exact if the object is centered, acts on all pixels and if the camera is orthographic (located infinitely far away from the subject). If the object is off center, the pixel rotation can be performed around a learned center point. If the images have a rotation-invariant background, a learned mask can be applied. If the camera is not orthographic, the equivariance relationship is not exact, but approximate. The decoder is regularized by enforcing Equation (42) through a mean squared error loss on the pixels for uniformly sampled and . We choose to correspond to rotation around the -axis.
This regularizer helps align all rotations in each orbit, but does not help in correctly aligning the orbits among each other. Thus we reduce the problem from aligning to aligning , since the cosets of after the orbit are identified, are homeomorphic to the sphere.
G.2 Continuity regularizer
If the learner is provided with pairs images that are nearby with respect to the manifold metric, the encoder can be regularized by penalizing differences in the encodings of the two inputs. This is done by penalizing the mean squared error of the Frobenius norms of the two encoded rotation matrices, which is a proper metric on the manifold.
This simplifies the problem from unsupervised learning on i.i.d. samples to learning a VAE on two frame samples from random trajectories of data lying on the manifold.