Score-based Pullback Riemannian Geometry: Extracting the Data Manifold Geometry using Anisotropic Flows

TL;DR

This paper introduces a scalable score-based pullback Riemannian geometry framework that leverages anisotropic flows to accurately extract and visualize the intrinsic geometry of data manifolds, improving interpretability and efficiency.

Contribution

It proposes a novel, scalable method combining pullback Riemannian geometry and generative models to extract data manifold geometry with closed-form geodesics and intrinsic dimension estimation.

Findings

Produces high-quality geodesics through data support

Accurately estimates data manifold intrinsic dimension

Provides a global chart of the data manifold

Abstract

Data-driven Riemannian geometry has emerged as a powerful tool for interpretable representation learning, offering improved efficiency in downstream tasks. Moving forward, it is crucial to balance cheap manifold mappings with efficient training algorithms. In this work, we integrate concepts from pullback Riemannian geometry and generative models to propose a framework for data-driven Riemannian geometry that is scalable in both geometry and learning: score-based pullback Riemannian geometry. Focusing on unimodal distributions as a first step, we propose a score-based Riemannian structure with closed-form geodesics that pass through the data probability density. With this structure, we construct a Riemannian autoencoder (RAE) with error bounds for discovering the correct data manifold dimension. This framework can naturally be used with anisotropic normalizing flows by adopting isometry regularization during training. Through numerical experiments on diverse datasets, including image data, we demonstrate that the proposed framework produces high-quality geodesics passing through the data support, reliably estimates the intrinsic dimension of the data manifold, and provides a global chart of the manifold. To the best of our knowledge, this is the first scalable framework for extracting the complete geometry of the data manifold.