Reparameterization through Coverings and Topological Weight Priors

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.

• The paper presents a measure-theoretic formulation that generalizes the reparameterization trick to arbitrary latent topologies by leveraging measurable, sheet-wise, measure-preserving covering maps.
• It introduces KleinVAE, a novel variational autoencoder that utilizes the Klein bottle latent space, demonstrating improved topological fidelity in image reconstructions and alignment with persistent homology metrics.
• The approach extends reparameterization methods beyond Lie groups, offering a practical framework for incorporating topologically-informed latent spaces and weight priors in generative models and Bayesian neural networks.

Reparameterization through Coverings: Generalizing Latent Topologies in Variational Inference

Introduction and Motivation

This paper advances the theoretical and algorithmic foundations of variational inference by generalizing the reparameterization trick (RT) to latent spaces with arbitrary topological structure via measurable covering maps. Previous work extended RT to specific nontrivial latent topologies, e.g., spheres and Lie groups, but remained constrained by requirements such as group structure and smooth global reparameterizations. This work addresses these limitations by providing a measure-theoretic formulation based on the measurability and measure-preserving properties of covering maps, enabling reparameterization on spaces such as the Klein bottle, which lacks both Lie group structure and global orientability.

The manifold hypothesis—that natural data often exhibits low-dimensional manifold structure—is central to the motivation of this approach. The use of topologically-informed latent spaces is particularly relevant in vision, as exemplified by the appearance of Klein bottle topology in natural image patches and the design of topological convolutions in CNNs. The proposed method supports using such nontrivial latent spaces directly, supporting richer topological inductive biases in generative models, and facilitating the construction of topological weight priors in Bayesian learning.

Mathematical Framework: Reparameterization via Coverings

The core contribution is a general result: for any measurable, sheet-wise (non-increasing) measure-preserving covering map, KL divergence between pushforwards of probability measures is non-increasing relative to the original measures. This forms the backbone for analytic tractability of the ELBO in variational autoencoders (VAEs) with nontrivial latent topologies.

Let $f: X \to Y$ be a covering map between measurable spaces, and $p$, $q$ be probability densities on $X$ with reference measure $m$. Then, for pushforwards $p^*, q^*$ under $f$, the following holds under mild conditions:

$\mathrm{KL}(q^*\|p^*) \leq \mathrm{KL}(q\|p)$

This inequality allows substituting the intractable KL term on the nontrivial topology (e.g., Klein bottle) with a tractable computation in its universal covering (e.g., $\mathbb{R}^2$), as the covering map is typically a sheet-wise projection.

The sampling procedure for generative modeling with latent variable $z$ in a base manifold $p$0 proceeds by (1) sampling in the cover $p$1, where conventional reparameterization is feasible (e.g., Gaussian in $p$2), and (2) mapping the sample to the base via the covering map. Gradients through this operation are handled naturally, as the map is either globally differentiable except on a null set, or is trivial for practical purposes due to the measure-zero locus of discontinuity.

Klein Bottle Latent Spaces and Gabor-Klein Filters

The Klein bottle $p$3 provides a minimal, non-orientable, non-Lie example of nontrivial topology relevant to image analysis. The construction realizes $p$4 as a quotient of the square $p$5 under specific identifications of boundary points; its universal cover is $p$6, enabling the proposed pushforward sampling and variational inference machinery.

This structure is directly motivated by reflection symmetry properties in the parameter space of Gabor filters. Specifically, a family of Gabor-Klein filters parameterized by angles $p$7 satisfies a Klein bottle-structuring identification, as shown by the persistence homology of sampled filter coefficients.

Figure 2: Discretized Gabor-Klein filters illustrate the Klein bottle structure; persistence diagrams confirm nontrivial first homology consistent with $p$8.

Persistent homology diagrams, computed over various coefficient fields ($p$9, $q$0), show the presence and disappearance of torsion—a hallmark of the Klein bottle topology—in the filter distribution. This points to direct applicability for topological data analysis (TDA), both as a tool and as a conceptual guide for model design.

Kleinian Variational Autoencoders (KleinVAE) and Empirical Study

The paper introduces KleinVAE, a VAE with latent space $q$1, constructed by pushing forward reparameterized Gaussian VPs from $q$2 to the Klein bottle via the explicit covering map. The empirical study validates the method on artificially generated image data—circle images with Klein bottle topology in the location of their center.

Figure 1: Samples from the Klein-Circles dataset and their reconstructions by KleinVAE demonstrate the preservation of topology-induced structure in the generated images.

Persistent homology of the decoded images confirms that the learned latent representations adhere to the nontrivial topology of $q$3.

Figure 3: Persistence diagrams of images decoded by KleinVAE from uniformly sampled Klein-Circle latent positions, matching the theoretical homology of the Klein bottle.

Ablation studies compare KleinVAE to VAEs with Euclidean and toroidal latent spaces using the bottleneck distance between persistence diagrams of original and reconstructed images. The topological-fidelity metrics (e.g., persistent homology bottleneck distances) were not universally superior for KleinVAE; however, reconstruction performance and topological match to the ground truth data manifold were better aligned for topologically-aware models. The ELBO and latent code variance during training reveal that the topological covering does not adversely affect optimization, as expected from the measure-theoretic analysis.

Figure 4: Learning curves for ELBO during training for models with different latent topologies.

Figure 5: Bottleneck distances $q$4 over $q$5 and $q$6 exhibit the fidelity of topological structure preservation in the reconstructions.

Theoretical Implications and Broader Impact

The principal theoretical claim—KL non-increase under measurable, sheet-wise measure-preserving coverings—removes restrictive smoothness and group-structure assumptions, substantially broadening the class of topologies for which reparameterized variational inference remains analytically controlled. This is significant not only for VAEs but also for flows, Bayesian neural network priors, and any generative model where a latent manifold structure is desired.

Practically, this enables using latent spaces that match data-derived inductive biases, as identified by TDA or symmetry considerations, with immediate relevance for vision tasks and robust Bayesian learning. The algorithmic machinery readily extends to learnable topological weight priors, supporting the recent trend of incorporating domain-geometric priors into deep networks, as in Gabor-Klein filtered convolutional architectures.

Finally, the framework motivates new directions in topological deep learning: (1) latent geometry matching via generative modeling, (2) Bayesian priors reflecting topological data analysis outputs, and (3) model classes designed for manifolds beyond Lie groups, such as moduli and non-orientable spaces seen in chemistry and image analysis.

Conclusion

The paper makes a substantial advancement by generalizing the RT to arbitrary manifold topologies admissible by measurable coverings, removing requirements for global diffeomorphism or group structure. This framework is rigorously instantiated in the context of the Klein bottle, and validated by constructing and analyzing KleinVAE. Theoretical and empirical results collectively underscore the utility of topological priors for generative models and Bayesian neural networks. The approach provides a practical pathway to encode domain-relevant manifold structure, opening avenues for further research in topological machine learning, generative modeling, and informed initialization of deep learning architectures.

Reference: "Reparameterization through Coverings and Topological Weight Priors" (2604.23804).