Riemannian Energy Matching for Fixed-Density Transport

An energy-based model defines a target density, but efficient sampling also depends on how particles move through it. Energy Matching uses the same scalar function to define both the density and the transport field. Improving transport by changing the energy therefore also changes the distribution being sampled. We ask whether learning the geometry of transport can improve sampling while keeping the energy fixed. We introduce Riemannian Energy Matching (\REM), which learns a symmetric positive-definite mobility from optimal-transport velocities. This mobility redirects the frozen energy gradient, while a unit-determinant constraint rules out uniform time rescaling. We characterize the velocities this model can represent and derive an objective for fitting them along transport paths. A divergence correction ensures that the resulting diffusion preserves the original Boltzmann density in continuous time. On a curved multimodal target, full mobility captures motion that diagonal models miss. Removing the correction from the same fitted mobility substantially increases sampling error. A high-dimensional study shows that geometry can improve field fitting without changing the energy-defined density. For image generation, we use learned geometry during deterministic transport and return to Euclidean dynamics during stochastic refinement. This phase-gated strategy improves generation quality under matched solver budgets. \REM provides a framework for separating errors in density learning from limitations in transport geometry, and for testing whether a better fitted field leads to better samples.