Riemannian Metric Matching for Scalable Geometric Modeling of Distributions

Jacob Bamberger, Adam Gosztolai, Pierre Vandergheynst, Michael M. Bronstein, Iolo Jones
Proceedings of the 43rd International Conference on Machine Learning, PMLR 306:6108-6127, 2026.

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.

Cite this Paper

BibTeX
@InProceedings{pmlr-v306-bamberger26a, title = {{R}iemannian Metric Matching for Scalable Geometric Modeling of Distributions}, author = {Bamberger, Jacob and Gosztolai, Adam and Vandergheynst, Pierre and Bronstein, Michael M. and Jones, Iolo}, booktitle = {Proceedings of the 43rd International Conference on Machine Learning}, pages = {6108--6127}, year = {2026}, editor = {Zhang, Tong and Dudik, Miroslav and Jaggi, Martin and Agarwal, Alekh and Li, Sharon and Schuurmans, Dale and Zhu, Jerry and Berkenkamp, Felix and Dong, Hanze and Bietti, Alberto}, volume = {306}, series = {Proceedings of Machine Learning Research}, month = {06--11 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v306/main/assets/bamberger26a/bamberger26a.pdf}, url = {https://proceedings.mlr.press/v306/bamberger26a.html}, 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.} }
Endnote
%0 Conference Paper %T Riemannian Metric Matching for Scalable Geometric Modeling of Distributions %A Jacob Bamberger %A Adam Gosztolai %A Pierre Vandergheynst %A Michael M. Bronstein %A Iolo Jones %B Proceedings of the 43rd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2026 %E Tong Zhang %E Miroslav Dudik %E Martin Jaggi %E Alekh Agarwal %E Sharon Li %E Dale Schuurmans %E Jerry Zhu %E Felix Berkenkamp %E Hanze Dong %E Alberto Bietti %F pmlr-v306-bamberger26a %I PMLR %P 6108--6127 %U https://proceedings.mlr.press/v306/bamberger26a.html %V 306 %X 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.
APA
Bamberger, J., Gosztolai, A., Vandergheynst, P., Bronstein, M.M. & Jones, I.. (2026). Riemannian Metric Matching for Scalable Geometric Modeling of Distributions. Proceedings of the 43rd International Conference on Machine Learning, in Proceedings of Machine Learning Research 306:6108-6127 Available from https://proceedings.mlr.press/v306/bamberger26a.html.

Related Material