[Paper Review] Geometry-Aware Generative Autoencoders for Warped Riemannian Metric Learning and Generative Modeling on Data Manifolds

Xingzhi Sun, Danqi Liao+8 more|arXiv (Cornell University)|Oct 16, 2024

Morphological variations and asymmetry4 citations

TL;DR

This paper introduces Geometry-Aware Generative Autoencoders (GAGA), a framework that integrates manifold learning with generative modeling by learning a warped Riemannian metric on the data space using both on-manifold and off-manifold points. GAGA enables geometry-aware data generation, geodesic interpolation, and population transport, achieving a 30% improvement over state-of-the-art methods in single-cell trajectory inference on real biological datasets.

ABSTRACT

Rapid growth of high-dimensional datasets in fields such as single-cell RNA sequencing and spatial genomics has led to unprecedented opportunities for scientific discovery, but it also presents unique computational and statistical challenges. Traditional methods struggle with geometry-aware data generation, interpolation along meaningful trajectories, and transporting populations via feasible paths. To address these issues, we introduce Geometry-Aware Generative Autoencoder (GAGA), a novel framework that combines extensible manifold learning with generative modeling. GAGA constructs a neural network embedding space that respects the intrinsic geometries discovered by manifold learning and learns a novel warped Riemannian metric on the data space. This warped metric is derived from both the points on the data manifold and negative samples off the manifold, allowing it to characterize a meaningful geometry across the entire latent space. Using this metric, GAGA can uniformly sample points on the manifold, generate points along geodesics, and interpolate between populations across the learned manifold using geodesic-guided flows. GAGA shows competitive performance in simulated and real-world datasets, including a 30% improvement over the state-of-the-art methods in single-cell population-level trajectory inference.

Motivation & Objective

• To address the challenge of generating data that respects complex, non-linear data geometries in high-dimensional biological datasets.
• To enable meaningful interpolation along geodesics for modeling continuous biological processes like development.
• To facilitate population transport across experimental conditions or time points using geodesic-guided flows.
• To overcome data imbalance and sparsity in single-cell and spatial omics data through volume-aware sampling.
• To unify manifold learning, metric learning, and generative modeling into a single, geometry-aware framework.

Proposed method

• GAGA trains a neural network autoencoder to embed data into a latent space that preserves intrinsic manifold distances.
• It learns a warped pullback Riemannian metric by contrasting on-manifold points with negative samples off the manifold, penalizing paths that exit data density.
• The warped metric is derived from the differential of the encoder map, inducing a Riemannian structure on the original data space.
• The framework uses this metric to perform uniform sampling on the manifold, generate points along geodesics, and transport populations via geodesic flows.
• It employs a distance-matching loss during training to align latent space distances with geodesic distances on the data manifold.
• The method generalizes to unseen data points and preserves the geometric structure of underlying dimensionality reduction techniques like PHATE and HeatGeo.

Experimental results

Research questions

• RQ1 Can a generative model be designed to respect the intrinsic geometry of high-dimensional data manifolds, even in sparse or noisy regions?
• RQ2 How can a Riemannian metric be learned that guides geodesic paths to remain within data-dense regions while avoiding spurious paths?
• RQ3 To what extent can a unified framework enable both data generation and trajectory inference on complex biological manifolds?
• RQ4 Can the integration of off-manifold negative samples improve the robustness and generalization of metric learning in data manifold modeling?
• RQ5 How does the warped Riemannian metric compare to local or density-regularized metrics in preserving geodesic accuracy on synthetic and real manifolds?

Key findings

• GAGA achieved a 30% improvement in population-level trajectory inference accuracy over state-of-the-art methods on single-cell RNA-seq datasets.
• On noisy single-cell data, GAGA outperformed standard autoencoders in preserving manifold structure, especially as noise increased (measured by DEMaP score).
• Volume-guided generation with GAGA produced point distributions whose densities correlated strongly with ground-truth volume elements, correcting for data imbalance observed in raw data.
• Geodesic paths learned by GAGA closely matched ground-truth geodesics on toy manifolds like the saddle and paraboloid, outperforming local metrics and density-regularized baselines.
• GAGA’s latent embeddings preserved both local clusters and global branching structures in real scRNA-seq data, such as embryoid body differentiation over 27 days.
• The framework demonstrated robustness in left-one-timepoint-out prediction tasks, accurately reconstructing missing timepoint populations using geodesic interpolation.