Follow the Energy, Find the Path: Riemannian Metrics from Energy-Based Models

TL;DR

This work addresses the challenge of shortest-path computation for high-dimensional data that lie on curved manifolds by deriving Riemannian metrics from pretrained Energy-Based Models (EBMs). It introduces two conformal metrics, $\mathbf{G}_{E_{\theta}}(\mathbf{x})$ and $\mathbf{G}_{1/p_{\theta}}(\mathbf{x})$, and uses a neural interpolant to approximate geodesics that stay close to the data manifold while respecting curvature. Across toy 2D mixtures, rotated character manifolds, and AFHQ latent spaces, the EBM-derived metrics consistently outperform baselines like $\mathbf{G}_{LAND}$ and $\mathbf{G}_{RBF}$ in terms of manifold alignment and geodesic fidelity, with the $\mathbf{G}_{E_{\theta}}$ variant often providing the strongest performance. By grounding geometry in a learned energy landscape, the approach enables scalable, data-aware geodesics that can improve generative modeling, trajectory planning, and cognitive-neuroscience interpretations of high-dimensional data geometry.

Abstract

What is the shortest path between two data points lying in a high-dimensional space? While the answer is trivial in Euclidean geometry, it becomes significantly more complex when the data lies on a curved manifold -- requiring a Riemannian metric to describe the space's local curvature. Estimating such a metric, however, remains a major challenge in high dimensions. In this work, we propose a method for deriving Riemannian metrics directly from pretrained Energy-Based Models (EBMs) -- a class of generative models that assign low energy to high-density regions. These metrics define spatially varying distances, enabling the computation of geodesics -- shortest paths that follow the data manifold's intrinsic geometry. We introduce two novel metrics derived from EBMs and show that they produce geodesics that remain closer to the data manifold and exhibit lower curvature distortion, as measured by alignment with ground-truth trajectories. We evaluate our approach on increasingly complex datasets: synthetic datasets with known data density, rotated character images with interpretable geometry, and high-resolution natural images embedded in a pretrained VAE latent space. Our results show that EBM-derived metrics consistently outperform established baselines, especially in high-dimensional settings. Our work is the first to derive Riemannian metrics from EBMs, enabling data-aware geodesics and unlocking scalable, geometry-driven learning for generative modeling and simulation.

Figures (27)

• Figure 1: Geodesics visualization for the URC dataset. Trajectories and samples are projected in the PCA space for visualization.
• Figure 2: Geodesics on UCG and WCG datasets.(a, c): Some geodesics obtained on UCG (a) and WCG (c), for $6$ different Riemannian metrics. The contour plots represent the energy landscape given by $-\log p_{\mathcal{M}}$. (b, d) Quantitative evaluation of geodesics on UCG (b) and WCG (d). We report (i) the accumulated probability along the geodesic (the higher the better) and ii) RMSE between each geodesic and its corresponding baseline (i.e., $\boldsymbol{G}_{E_{\mathcal{M}}}$ for $\textcolor{google_blue}{\mathbf{ G_{E_{\theta}}}}$, and $\boldsymbol{G}_{1/p_{\mathcal{M}}}$ for $\textcolor{google_green}{\mathbf{ G_{1/p_{\theta}}}}$, $\textcolor{google_orange}{\mathbf{ G_{LAND}}}$ and $\textcolor{google_red}{\mathbf{ G_{RBF}}}$). See Supp. \ref{['app:toy_2sig']} for the $2$-$\sigma$ error.
• Figure 3: Step size along geodesics in the WCG dataset. Log-based metrics ($\textcolor{google_blue}{\mathbf{ G_{E_{\theta}}}}$ and $\boldsymbol{G}_{E_{\mathcal{M}}}$) produce sharper variations, reflecting stronger sensitivity to density curvature.
• Figure 4: Geodesics on the URC dataset.(a) Geodesics computed with different Riemannian metrics, projected into pixel space for visualization. $\textcolor{google_red}{\mathbf{ G_{RBF}}}$ and $\textcolor{google_orange}{\mathbf{ G_{LAND}}}$ often deviate from the intended path, sometimes drifting toward other characters (e.g., the letter F). (b) Quantitative evaluation using two metrics: (i) $\mathcal{D}$-RMSE, which measures proximity to the dataset manifold, and (ii) $\gamma$-RMSE, which measures the deviation from an ideal smooth rotation. See Supp. \ref{['app:rot_2sig']} for the $2$-$\sigma$ error.
• Figure 5: Step size along geodesics in the WCG dataset. Log-based metric ($\textcolor{google_blue}{\mathbf{ G_{E_{\theta}}}}$) produces sharper variations, reflecting stronger sensitivity to density curvature.