Riemannian AmbientFlow: Towards Simultaneous Manifold Learning and Generative Modeling from Corrupted Data

TL;DR

This work tackles the challenge of learning both the data distribution and its latent nonlinear geometry from corrupted measurements. It introduces Riemannian AmbientFlow, a framework that integrates AmbientFlow-style variational learning with data-driven pullback geometry to recover a smooth, bi-Lipschitz manifold via a Riemannian Autoencoder (RAE). The authors provide recoverability guarantees under geometric regularization and a range-restricted isometry condition, derive an interpretable decoder-based prior for inverse problems, and validate the approach on synthetic manifolds and MNIST. The results demonstrate that joint manifold learning and probabilistic modeling from corrupted data is feasible, yielding geometry-aware generative models that can improve inverse-problem reconstruction and interpretability in scientific imaging tasks. The work points to future extensions to more expressive diffeomorphisms, diverse corruption models, and geometry-guided applications in science and engineering.

Abstract

Modern generative modeling methods have demonstrated strong performance in learning complex data distributions from clean samples. In many scientific and imaging applications, however, clean samples are unavailable, and only noisy or linearly corrupted measurements can be observed. Moreover, latent structures, such as manifold geometries, present in the data are important to extract for further downstream scientific analysis. In this work, we introduce Riemannian AmbientFlow, a framework for simultaneously learning a probabilistic generative model and the underlying, nonlinear data manifold directly from corrupted observations. Building on the variational inference framework of AmbientFlow, our approach incorporates data-driven Riemannian geometry induced by normalizing flows, enabling the extraction of manifold structure through pullback metrics and Riemannian Autoencoders. We establish theoretical guarantees showing that, under appropriate geometric regularization and measurement conditions, the learned model recovers the underlying data distribution up to a controllable error and yields a smooth, bi-Lipschitz manifold parametrization. We further show that the resulting smooth decoder can serve as a principled generative prior for inverse problems with recovery guarantees. We empirically validate our approach on low-dimensional synthetic manifolds and on MNIST.

Figures (7)

• Figure 1: After 500 iterations of minimizing the Riemannian AmbientFlow problem \ref{['eq:full-rie-ambientflow-loss']} we have learned both a generative model and a smooth manifold that matches both the clean and corrupted data distribitions.
• Figure 2: After 250 iterations of minimizing the Riemannian AmbientFlow problem \ref{['eq:full-rie-ambientflow-loss']} we have learned a generative model that matches the corrupted data distribition, but is still finetuning the model that approximates clean data distribution. This is also noticeable in the quality of the learned manifolds -- both for the clean data manifold as the corrupted data manifold --, which are not quite as accurate as the ones from Figure \ref{['fig:sinusoid_views']}.
• Figure 3: After 500 iterations of minimizing the Riemannian AmbientFlow problem \ref{['eq:full-rie-ambientflow-loss']} with $\lambda = 0$ we have learned both a generative model and a smooth manifold that matches both the clean and corrupted data distribitions.
• Figure 4: After 500 iterations of minimizing the Riemannian AmbientFlow problem \ref{['eq:full-rie-ambientflow-loss']} with $\mu = 0$ we have learned a generative model that matches the corrupted data distribition, but does not match the clean data distribution. Instead the distribution is offset in the null space of the forward operator. This is also noticeable in the quality of the learned manifolds -- both for the clean data manifold as the corrupted data manifold --, which are do not even come close to the ones learned in Figure \ref{['fig:sinusoid_views']}.
• Figure 5: Top row: $10$ randomly selected corrupted training images. Bottom row: $10$ randomly generated samples from our learned prior $p_{\theta}$.

Theorems & Definitions (23)

• Proposition 1: Distributional vs expected projection error
• proof
• Lemma 1: Tangent space covariance matrix
• proof
• Lemma 2: Expected tangent space projection error
• proof
• Theorem 1: Expected projection error
• proof
• Definition 2: Restricted Isometry Property
• Theorem 3: Recoverability