🤖 AI Summary
Traditional variational autoencoders (VAEs) are constrained by latent spaces endowed with Euclidean or Lie group structures, limiting their ability to model data with nontrivial topologies such as the Klein bottle. This work proposes a “reparameterization via covering” framework that generalizes the reparameterization trick to arbitrary smooth manifolds using covering maps. By integrating measure-preserving transformations and pushforward density techniques, the method enables analytical computation of the KL divergence term in the evidence lower bound (ELBO). To the best of our knowledge, this is the first approach to achieve end-to-end VAE training with latent spaces possessing non–Lie-group topologies. The resulting KleinVAE effectively learns complex topological structures on synthetic data and demonstrates promising potential as a Bayesian prior over convolutional network weights.
📝 Abstract
We generalise the reparameterization trick applied in variational autoencoders (VAEs) letting these have latent spaces of non-trivial topology - i.e. that of base manifolds covered with other ones, on which some technique for RT is available. That is possible since covering maps are measurable - moreover, in case of particular measure preservation property holding for the covering, one can establish an inequality on KL-divergence between pushforward (PF) densities on the base latent manifold, making the KL-term of VAE's ELBO analytically tractable, despite the topological non-triviality of the supporting latent manifold. Our development follows a route close but somewhat alternative to reparameterization on Lie groups, the latest proposal for which is to reparameterize PFs of normal densities from the Lie algebra - "through" the exponential map, seen by us as sometimes a particular case of what we propose to call reparameterization through a covering. Covering maps need not be global diffeomorphisms (although Lie-exp maps, in general, need not either, but, to date only smooth ones were considered in this context, to the best of our knowledge), which makes many non-trivial topologies tamable to our proposed technique, that we detail on a particular such example. We demonstrate the working of our approach by constructing a VAE with the latent space of Klein bottle (not a Lie group) topology, which we call KleinVAE, successfully learning an appropriate artificial dataset. We discuss potential applicability of such topology-informed generative models as weight priors in Bayesian learning, particularly for convolutional vision models, where said manifold was peculiarly shown to have some relevance.