Constructing VAE Latent Spaces with Prescribed Topology

Abstract: Variational autoencoders (VAEs) learn low-dimensional latent representations of high-dimensional data. When the data lies on a manifold with non-Euclidean topology, the standard Gaussian prior introduces a topological mismatch that degrades reconstruction quality and prevents faithful representation. We present a constructive mathematical framework that resolves this mismatch for all manifolds that admit a product covering space. These are manifolds expressible as products of elementary factors (circles, intervals, or lines) or as quotients of such products by a finite symmetry group. The class includes cylinders, tori, Möbius strips, Klein bottles, and real projective spaces. Factorized distributions over the elementary factors yield product topologies with closed-form, decoupled KL divergences, so that each latent factor can be shaped independently while keeping training tractable. We catalogue reparametrizable encoder-prior pairs for periodic, bounded, and unbounded supports, and provide coordinate transformations that allow standard neural networks to output non-Euclidean parameters with smooth gradients. For quotient manifolds, the decoder receives group-invariant features of the covering-space coordinates, so that identified points produce identical outputs. Anchor constraints fix the coordinate system relative to the data or create soft topological holes. Experiments on synthetic manifolds and real-image datasets (rotated and cyclically shifted MNIST) confirm that a topology-matched prior aligns KL regularization with the data manifold. The resulting topology-aware models outperform the Gaussian baseline at all practically relevant regularization strengths. The code is available at https://github.com/JvHulst/VAE-Topology.

• The paper introduces a novel framework that enables the design of VAEs with latent spaces matching arbitrary prescribed topologies.
• It employs a four-step pipeline that pairs reparameterizable distributions with tailored encoder–prior selections and group-invariant decoding.
• Empirical tests on synthetic manifolds and modified MNIST datasets show improved topological fidelity and reconstruction accuracy over standard Gaussian VAEs.

Constructing VAE Latent Spaces with Prescribed Topology

Introduction

The paper "Constructing VAE Latent Spaces with Prescribed Topology" (2606.07058) presents a systematic and mathematically rigorous approach for designing variational autoencoders (VAEs) with latent spaces that match arbitrary prescribed topologies, specifically for any manifold admitting a product covering space structure. The methodology addresses the deficiencies of standard Gaussian VAEs, which induce a Euclidean topology incompatible with data exhibiting periodic, bounded, or non-orientable structure. The core contribution is a constructive framework enabling practitioners to instantiate VAEs whose latent variables possess precisely the topology required by the problem domain, while retaining tractable evidence lower bound (ELBO) training via decoupled KL divergences and reparameterizable sampling.

Mathematical Framework and Design Pipeline

The authors define a four-step pipeline for constructing topology-aware VAEs:

1. Product or Quotient Decomposition: The target manifold $\mathcal{M}$ is expressed as a product of elementary manifolds or as a quotient thereof by a finite symmetry group $G$, covering cases such as cylinders ($S^1 \times [0,1]$), tori ($T^n$), Möbius strips, and Klein bottles.

   Figure 1: Pipeline for designing a topology-aware VAE, illustrated with the example of a Möbius strip.

2. Encoder–Prior Distribution Selection: For each elementary factor, encoder and prior distributions are paired from a catalog of reparameterizable distributions with matched supports. For instance, Wrapped Normal distributions are assigned to $S^1$ factors, Kumaraswamy for bounded intervals, and Gaussian for $\mathbb{R}$.
3. Invariant Feature Map for Decoding: For quotient manifolds, a group-invariant feature map $\rho$ translates latent variables to Euclidean features for the decoder, enforcing topological identifications within the decoding process (e.g., for the Möbius strip, points related by $(h,\theta)\sim(1-h,\theta+\pi)$ yield identical decoder features).
4. Anchor Constraints: Anchor losses can fix coordinate systems in the latent representation (crucial for identifiability and ablation of coordinate ambiguities), or (in a repulsive form) induce soft topological holes not representable by product/quotient structure alone.

VAE Topology Support and Reparameterizable Distributions

The framework’s flexibility and tractability are rooted in decoupled, factorized priors and encoders. The KL divergence $\text{KL}(q_\phi(z|y) \| p(z))$ decomposes cleanly across dimensions or factors, which in turn supports mixtures of distribution families (e.g., Gaussian, Kumaraswamy, Wrapped Normal on $S^1$, von Mises-Fisher on $G$0). This allows latent spaces with rich topological structure to be realized while retaining closed-form or efficiently estimable KL terms.

Figure 2: Standard VAE architecture with Gaussian encoder and prior.

Figure 3: Enhanced VAE architecture for topology shaping, with per-factor encoder distributions and support-matched priors, regularized via KL and equipped for group-invariant decoding.

Central to practical success is the smooth and differentiable parameterizations: for periodic variables, angular encoders use normalization-based approaches over $G$1 outputs to sidestep gradient discontinuities.

Topology Shaping Results: Synthetic and Real Data

Synthetic Manifold Experiments

The authors validate the framework on three canonical manifolds with precisely known topology and semantics: the cylinder $G$2, the torus $G$3, and the Möbius strip $G$4. They demonstrate that topology-aware VAEs (i) faithfully recover the latent coordinates and identification, (ii) vastly reduce the mismatch between prior and posterior, (iii) achieve lower geodesic stress (measured via normalized Kruskal stress-1 between true and latent space geodesic distances), and (iv) exhibit improved prior consistency across a Pareto sweep in the $G$5-VAE regime.

Figure 4: Learned latent spaces for synthetic experiments, showing alignment between ground-truth factors and the VAE latent representations, with substantial improvement when using anchoring.

Figure 5: KL sweep reveals the topology-aware VAE dominates the standard Gaussian baseline for all non-negligible regularization regimes.

Real-World Applications: Rotational and Translational Invariances

The methodology is further substantiated on rotated MNIST (with ground-truth $G$6 latent) and shifted MNIST datasets (with $G$7 periodicity in the pixel shift dimensions):

• Rotated MNIST: The latent $G$8 is realized natively on $G$9, with anchoring yielding interpretable, well-aligned latent variables, improved prior consistency, and lower reconstruction errors. The Gaussian VAE baseline exhibits periodic boundary failures.
• Shifted MNIST: The Mixed-Torus VAE, with two $S^1 \times [0,1]$0-wrapped factors, accurately captures the double periodicity of translation, unlike the Gaussian baseline, which lacks periodic generalization.

  Figure 6: Rotated MNIST experiment; topology-aware VAEs preserve the circular structure in latent space, supporting robust and interpretable interpolation across rotations.

  

  Figure 7: Shifted MNIST experiment; the topology-aware VAE’s latent space accurately reflects the two-dimensional torus structure of digit shifts.

Role and Effects of Anchoring

Anchoring is empirically critical for coordinate frame alignment, resolving coordinate ambiguity on symmetric manifolds, preventing latent collapse (as shown in the Möbius strip example), and enhancing interpretability without constraining the model's generative capability. Notably, even when prior support is matched, failure to anchor may result in coordinate drift or collapse under symmetry.

Theoretical Implications and Limitations

The methodology theoretically supports any manifold expressible as a product of elementary supports (circle, interval, line, hypersphere) and finite quotients thereof. Notably, the procedure does not require computation of Jacobian determinants, Monte Carlo Kullback-Leibler estimation, or expensive group-averaged architectures. However, it does not learn topology from data (it must be specified a priori), and truly arbitrary manifolds not covered by these compositions are outside its scope. Disconnected manifolds would require further extension to mixture priors and importance-weighted bounds.

Implications and Future Directions

Practically, this framework enables the principled use of structured priors in downstream tasks where boundary and identification properties matter: e.g., periodic variables in physical systems identification, nontrivial configuration spaces in robotics, and representations required for equivariant or invariant learning. The framework cleanly integrates with normalizing flow-based VAEs, and extension to continuous symmetries or unsupervised topology discovery represents clear directions for future research. Additionally, the approach provides new tools for studying disentanglement and interpretability as geometric/structural properties of learned latent spaces.

Conclusion

This work formalizes and solutionizes the problem of latent space topology mismatch in VAEs, presenting a tractable, systematic framework grounded in product topologies, support-matched reparameterizable distributions, quotient constructions via invariant features, and coordinate anchoring. Rigorous synthetic and real-data empirical evidence establishes superior topological fidelity, regularization alignment, generative consistency, and coordinate interpretability relative to standard Gaussian baselines. This positions the framework as a canonical starting point for any application where latent space topology matters and tractable training must be preserved.