Technical Summary: Riemannian Generative Decoder

Problem Statement

Real-world data often possesses intrinsic non-Euclidean structures (e.g., evolutionary trees, social networks, periodic signals) that are poorly represented by standard Euclidean latent spaces ($\mathbb{R}^d$). Forcing such data into Euclidean configurations distorts their geometric relationships, leading to poor visualizations and missed signals.

Existing approaches to Riemannian representation learning, primarily Variational Autoencoders (VAEs), attempt to embed data onto manifolds by enforcing specific priors (e.g., von Mises–Fisher on spheres or Riemannian normals on hyperbolic spaces). However, these methods face significant challenges:

1. Numerical Brittleness: They require optimizing complex objectives involving intractable normalizing constants and Monte Carlo estimates for densities on arbitrary manifolds.
2. Scalability: These density approximations often limit scalability to general manifolds or high-dimensional latent spaces.
3. Complexity: Enforcing manifold constraints typically relies on an encoder to estimate densities, adding architectural and computational overhead.

Methodology

The authors propose the Riemannian Generative Decoder (RGD), a unifying framework that discards the encoder entirely. Instead of amortizing inference, the model treats latent codes as free parameters to be optimized directly on a chosen Riemannian manifold.

Core Formulation

The learning problem is formulated as a Maximum A Posteriori (MAP) estimation. Given a dataset $X = \{x_i\}_{i=1}^N$ and a differentiable decoder $f_\theta: \mathcal{M} \to \mathcal{X}$, the model jointly optimizes the decoder parameters $\theta$ and the latent codes $Z = \{z_i\}_{i=1}^N$ by minimizing the negative log-posterior:
$$\mathcal{L}(\theta, Z) = \sum_{i=1}^N \left( -\log p_\theta(x_i|z_i) - \log p(z_i) \right) - \log p(\theta)$$

• Likelihood: Modeled via reconstruction loss (e.g., squared error for Gaussian noise).
• Prior: For compact manifolds, a uniform prior is used; for non-compact manifolds, wrapped distributions or Riemannian normals are employed.
• Optimization: Latent codes are updated using Riemannian gradient descent (specifically Riemannian Adam), ensuring iterates remain on the manifold $\mathcal{M}$ via retraction maps. Decoder parameters are updated via standard Euclidean Adam.

Geometric Regularization

To align the decoder's smoothness with the manifold's geometry without requiring explicit density estimation, the authors introduce a geometry-aware regularization via input noise.

• During training, each latent $z$ is perturbed by Gaussian noise $\epsilon \sim \mathcal{N}(0, \sigma^2 G^{-1}(z))$, where $G(z)$ is the Riemannian metric tensor.
• This noise is injected via the exponential map (approximated as identity for small noise).
• Theoretical Derivation: Based on Bishop (1995) and An (1996), the expected loss under this noise perturbation yields an additive regularizer:
  $$\mathbb{E}_\epsilon[\mathcal{L}(z')] \approx \mathcal{L}(z) + \sigma^2 \text{Tr}\left( J(z)^\top G^{-1}(z) J(z) \right)$$
  where $J(z)$ is the Jacobian of the decoder. This term penalizes large output gradients weighted by the inverse metric, effectively encouraging the decoder to respect the local curvature of the manifold.

Flexibility

The framework supports any Riemannian manifold implemented in the geoopt library, including spheres, Poincaré balls, Lorentz models, and ProductManifolds (allowing heterogeneous latent geometries where different dimensions of the latent space reside on different manifolds).

Key Contributions

1. Unifying Framework: A method for representation learning on any Riemannian manifold by combining Riemannian optimization with an encoder-less generative model, eliminating the need for variational density approximations.
2. Scalable Geometric Regularization: A noise-based regularization technique that promotes coherence between the decoder function and the manifold metric, avoiding computationally expensive second-order curvature penalties.
3. Empirical Validation: Demonstration of the method on diverse datasets, showing improved alignment with intrinsic geometry and stability in high dimensions where variational baselines fail.

Results and Evaluation

The method was validated on three case studies:

1. Synthetic Branching Diffusion Process:

   • Task: Recovering tree-structured data.
   • Findings: Euclidean and UMAP projections failed to reveal the tree topology. Regularized hyperbolic embeddings (Poincaré disk) successfully recovered the hierarchy.
   • Ablation: Increasing the noise scale $\sigma$ improved the correlation between latent geodesic distances and data-space distances up to an optimal point ($\sigma \approx 0.9$), after which noise overwhelmed the signal.

2. Human Mitochondrial DNA (hmtDNA):

   • Task: Inferring migration patterns and lineage from genetic sequences.
   • Findings: Hyperbolic models (Lorentz/Poincaré) recovered haplogroup hierarchies and geographical correlations significantly better than Euclidean or UMAP baselines.
   • Downstream Utility: In classification tasks (predicting geographical regions and haplogroups), the hyperbolic RGD outperformed Euclidean and P-VAE baselines, particularly in high-dimensional settings.

3. Cell Cycle Stages (scRNA-seq):

   • Task: Modeling the cyclic progression of cell division.
   • Findings: Spherical ($S^2$) and toroidal ($S^1 \times S^1$) embeddings captured the cyclic nature of the data. While Euclidean spaces offered slightly better reconstruction fidelity (due to higher degrees of freedom), Riemannian models showed superior correlation with the biological phase geometry.
   • Scalability: The RGD demonstrated stable scaling to high latent dimensions (up to 500D), whereas variational baselines (P-VAE, $\Delta$VAE) became numerically unstable or infeasible.

Significance and Claims

The paper claims that the Riemannian Generative Decoder offers a simpler and more stable alternative to existing manifold learning methods. By removing the encoder, the method avoids the "numerically brittle" objective functions associated with density estimation on curved spaces.

• Interpretability: The learned latent spaces are explicitly aligned with the prescribed geometry, facilitating hypothesis-driven exploration in scientific domains.
• Generality: The approach is not limited to specific manifolds (like spheres or hyperbolic spaces) but extends to products of heterogeneous manifolds.
• Practicality: The method requires only a decoder and is compatible with existing neural network architectures.

The authors acknowledge limitations, noting that the per-sample parameterization (storing each latent explicitly) leads to linear memory growth with dataset size, which may be prohibitive for millions of points. They suggest hybrid schemes (e.g., amortized inference) as a future direction. The work is positioned as a tool for exploratory analysis and hypothesis testing in scientific fields, particularly where data topology is known or hypothesized.