Learning the geometry of data at scale
High-dimensional datasets often concentrate near low-dimensional structures. Estimating the geometry of these structures from samples typically relies on graphs and kernels, which scale poorly with dataset size and dimension.
Riemannian metric matching
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, through diffusion geometry, gives 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. This allows sample-wise training and amortised inference at constant cost, without explicitly constructing a kernel.
Results
Metric matching rivals or improves the accuracy of k-nearest-neighbour diffusion geometry estimators, while enabling amortised inference that is up to 400 times faster. It also supports graph-free geometric analysis of high-dimensional images, where nearest-neighbour methods break down.
The paper is also on arXiv. See also carré du champ flow matching, which uses data geometry to improve generalisation in generative models.