Iso-Riemannian Optimization on Learned Data Manifolds

TL;DR

This work addresses optimization on learned data manifolds, where standard Riemannian methods struggle due to non-constant geodesic speed and loss of geodesic convexity under pullback mappings. It introduces iso-Riemannian optimization, powered by an iso-connection, and develops iso-Riemannian descent with new notions of iso-monotonicity and iso-Lipschitzness, yielding linear convergence guarantees. The framework enables iso-barycentres and iso-K-means clustering with theoretical and empirical support, and extends to iso-convex optimization with applications to inverse problems, including denoising and compressed-sensing-like tasks on MNIST and synthetic data. Overall, iso-geometry provides a principled, interpretable approach to optimize on learned data manifolds, mitigating distortions from learned mappings and enabling efficient, reliable downstream tasks.

Abstract

High-dimensional data that exhibit an intrinsic low-dimensional structure are ubiquitous in machine learning and data science. While various approaches allow for learning the corresponding data manifold from finite samples, performing downstream tasks such as optimization directly on these learned manifolds presents a significant challenge. This work introduces a principled framework for optimization on learned data manifolds using iso-Riemannian geometry. Our approach addresses key limitations of classical Riemannian optimization in this setting, specifically, that the Levi-Civita connection fails to yield constant-speed geodesics, and that geodesic convexity assumptions break down under the learned pullback constructions commonly used in practice. To overcome these challenges, we propose new notions of monotonicity and Lipschitz continuity tailored to the iso-Riemannian setting and propose iso-Riemannian descent algorithms for which we provide a detailed convergence analysis. We demonstrate the practical effectiveness of those algorithms on both synthetic and real datasets, including MNIST under a learned pullback structure. Our approach yields interpretable barycentres, improved clustering, and provably efficient solutions to inverse problems, even in high-dimensional settings. These results establish that optimization under iso-Riemannian geometry can overcome distortions inherent to learned manifold mappings.

Figures (11)

• Figure 1: If $p_{\theta^*} = p_{\text{data}}$ holds, any geodesic $\gamma^{\varphi_{\theta^*}}_{\mathbf{x}, \mathbf{y}}$ (orange) between data points $\mathbf{x},\mathbf{y}\in \mathbb{R}^2$ (blue) moves through regions with higher likelihood -- visualized by the level set curves of $p_{\text{data}}$ -- than the end points.
• Figure 2: Under the pullback structure $(\mathbb{R}, (\cdot, \cdot)^\varphi)$ with $\varphi(x) = \sinh(x +1)$ the convex function $f(x)= \frac{1}{2}x^2$ is not geodesically convex.
• Figure 3: Barycentre results on the synthetic river and spiral data sets under modeled pullback structures. Both the Riemannian and iso-barycentre (computed with \ref{['alg:iso-bary-line-search']}) yield more interpretable data representations -- i.e., lying within the data set -- compared to the Euclidean mean. Moreover, based on the observed iso-monotonicity and iso-Lipschitz ratios, we generally expect the parameter $\alpha$ to be much smaller than 1, while $L$ tends to be much larger.
• Figure 4: Barycentre results for MNIST under different geometries. Both the Riemannian and iso-barycentre (\ref{['alg:iso-bary-line-search']}) give very similar and interpretable data representations -- looking somewhat like an 8 -- whereas the Euclidean mean just yields a blurry mix. Notably, the iso-barycentre and the Riemannian barycentre are visually very similar.
• Figure 5: Clustering results on the synthetic river and spiral data sets under modeled pullback structures. Both the Riemannian and iso-barycentre-based methods yield more interpretable data representations compared to the Euclidean mean. However, only iso-K-means clustering (\ref{['alg:iso-riem-kmeans']}) also gives meaningful clusters. That is, in contrast to Riemannian K-means, which gives odd clusters due to the fact that the Riemannian distance does not encode proximity in an $\ell^2$-sense.

Theorems & Definitions (42)

• Theorem 1: convergence of (\ref{['eq:basic-first-order-scheme']})
• Definition 1: iso-monotone vector fields
• Definition 2: iso-Lipschitz vector fields
• Definition 3: locally iso-monotone vector fields
• Definition 4: locally iso-Lipschitz vector fields
• Lemma 1: iso-Taylor expansion of vector fields
• proof
• Theorem 2: implications of local iso-monotonicity and iso-Lipschitzness
• proof
• Proposition 1: Properties of iso-geodesic velocity fields