Riemannian Metric Matching for Scalable Geometric Modeling of Distributions

2026-06-12

AI Summary

This work addresses the challenge of estimating geometric structures in high-dimensional data, a task traditionally reliant on graph- or kernel-based methods that struggle to scale to large, high-dimensional settings. The authors propose a novel approach grounded in Riemannian metric matching, wherein a neural network learns the intrinsic Riemannian geometry of the data. By leveraging the carré du champ operator from diffusion geometry—expressed as a conditional expectation under data perturbations—the method avoids explicit kernel matrix construction, enabling sample-wise training and constant-time inference. Empirical results demonstrate that the proposed technique matches or exceeds the accuracy of k-nearest-neighbor diffusion estimators while achieving up to a 400-fold speedup in inference. Notably, it enables graph-free geometric analysis on high-dimensional images for the first time.

Abstract

High-dimensional datasets often concentrate near low-dimensional structures, but estimating their geometry from samples typically relies on graphs and kernels that scale poorly with dataset size and dimension. We propose Riemannian metric matching: a denoising probabilistic framework for learning the Riemannian geometry of data using neural networks. Specifically, we learn the carré du champ operator, which, using diffusion geometry, gives us access to the Riemannian geometry toolkit for downstream machine learning and statistical tasks. Our key observation is that the carré du champ operator can be formulated as a conditional expectation over random perturbations of the data, which can be exploited for sample-wise training and constant cost, amortized inference without explicit kernel construction. Empirically, metric matching rivals or improves the accuracy of $k$-NN-based diffusion geometry estimators, while enabling amortized inference that is up to $400\times$ faster, and supports graph-free geometric analysis on high-dimensional images where nearest neighbors break down.

Problem

• Riemannian geometry
• high-dimensional data
• scalability
• geometric modeling
• diffusion geometry

Innovation

• Riemannian metric matching
• carré du champ operator
• diffusion geometry
• amortized inference
• graph-free geometric analysis