Atlas-based Manifold Representations for Interpretable Riemannian Machine Learning

TL;DR

This work tackles learning on latent nonlinear manifolds by introducing Atlas, a differentiable atlas data structure that maintains multiple local charts and differentiable transitions to enable direct Riemannian optimization on learned manifolds. The Atlas-Learn heuristic constructs an atlas from point-cloud data via local quadratic tangent approximations, MVEE-based chart membership, and a patching strategy that tolerates imperfect overlaps. Empirically, the approach speeds online Fréchet-mean estimation on Grassmann manifolds, preserves nontrivial geometry on the Klein bottle better than standard dimensionality-reduction methods, enables interpretable Riemannian Classification of image patches, and improves RNA velocity integration by constraining dynamics to the learned manifold. The results highlight improvements in efficiency, geometry preservation, and interpretability, while pointing to future work on varifolds, adaptive charting, and rigorous error analysis for atlas-based Riemannian optimization.

Abstract

Despite the popularity of the manifold hypothesis, current manifold-learning methods do not support machine learning directly on the latent $d$-dimensional data manifold, as they primarily aim to perform dimensionality reduction into $\mathbb{R}^D$, losing key manifold features when the embedding dimension $D$ approaches $d$. On the other hand, methods that directly learn the latent manifold as a differentiable atlas have been relatively underexplored. In this paper, we aim to give a proof of concept of the effectiveness and potential of atlas-based methods. To this end, we implement a generic data structure to maintain a differentiable atlas that enables Riemannian optimization over the manifold. We complement this with an unsupervised heuristic that learns a differentiable atlas from point cloud data. We experimentally demonstrate that this approach has advantages in terms of efficiency and accuracy in selected settings. Moreover, in a supervised classification task over the Klein bottle and in RNA velocity analysis of hematopoietic data, we showcase the improved interpretability and robustness of our approach.

Figures (11)

• Figure 1: Chart domains for a differentiable atlas of the Klein bottleA. Charts depicted in polar coordinates show points belonging (blue) or not (red) to the chart, and indexed by the $(\theta, \phi)$ coordinates of the chart center $k(\theta,\phi)$ (Sec. \ref{['sec:comparison']}). B. The same coordinate charts shown in the Karcher representation.
• Figure 2: Quasi-Euclidean updates in the Atlas representation of the Grassmannian $\mathbf{Gr}_{n,k}$ converge faster to the population Fréchet mean than other first-order update schemes. Logarithms and retractions are executed by each method one observation at a time for 1,000 iterations, each observation being a $\mathbb{R}^{n, k}$ Stiefel matrix. Stiefel matrices $X \in \mathbb{R}^{n \times k}$ drawn from the distribution $\mathbf{GPD}(\mathcal{X},p)$ with fixed Fréchet mean $\mathcal{X}$ and $p\in\{2,3\}$ (Sec. \ref{['sec:geodesic_power_distribution']}). $\mathbf{Gr}_{n,k}$-distance to $\mathcal{X}$ from each method's iterate is plotted against cumulative runtime.
• Figure 3: The Atlas data structure preserves geodesic distances better than common dimensionality-reduction techniques.Isomap, PCA and $t$-SNE were computed on a mesh of 1 million points. Atlas-Learn was asked to produce 64 charts, given only a 10,000-points mesh. For 100 randomly sampled pairs of points on the Klein bottle, each scatter plot shows each pair (dot) according to a precise estimate of their true geodesic distance (x axis, "real dist.") versus the embedding distance (y axis, "dist. in representation") for each embedding (panel title) (Supp. Materials, Sec. \ref{['app:klein']}). Metric distortion given in top left corner. Dashed line, fit with $y$-intercept zero.
• Figure 4: The RPB algorithm effectively learns a discriminating boundary on the manifold of high-contrast image patches. A 64-chart Atlas data structure was generated by Atlas-Klein from a 10,000-point mesh of the Klein bottle manifold for high-contrast image patches. The RPB implementation was then applied to learn a discriminating boundary (colored by discretized ODE iterations of the algorithm) between the convex (red) and concave (black) patches, as illustrated in Karcher (left) and polar (right) coordinates.
• Figure 5: Integrating the RNA velocity vector field over the Atlas representation rather than the ambient space more closely matches observations in hematopoietic data. For 309 distinct initial conditions, corresponding to labeled HSCs, a vector field over human hematopoietic cells is integrated, either over the ambient space (blue) or over a 5-dimensional Atlas representation with 30 charts (red). At every iteration, distances from each of the 309 evolutions to the closest datum in 30-PC space is computed; these are shown as violin plots. Dashed line indicates average distance between a datum and nearest Atlas point.

Theorems & Definitions (30)

• Definition 1
• Lemma 1
• proof
• Theorem 1
• proof
• Lemma 2
• proof
• Claim 1
• proof
• Claim 2