Neural Sampling from Boltzmann Densities: Fisher-Rao Curves in the Wasserstein Geometry

Part of International Conference on Representation Learning 2025 (ICLR 2025) Conference

Bibtex Paper Supplemental

Authors

Jannis Chemseddine, Christian Wald, Richard Duong, Gabriele Steidl

Abstract

We deal with the task of sampling from an unnormalized Boltzmann density $\rho_D$by learning a Boltzmann curve given by energies $f_t$ starting in a simple density $\rho_Z$.First, we examine conditions under which Fisher-Rao flows are absolutely continuous in the Wasserstein geometry.Second, we address specific interpolations $f_t$ and the learning of the related density/velocity pairs $(\rho_t,v_t)$.It was numerically observed that the linear interpolation, which requires only a parametrization of the velocity field $v_t$,suffers from a "teleportation-of-mass" issue.Using tools from the Wasserstein geometry,we give an analytical example,where we can precisely measure the explosion of the velocity field.Inspired by Máté and Fleuret, who parametrize both $f_t$ and $v_t$, we propose aninterpolation which parametrizes only $f_t$ and fixes an appropriate $v_t$. This corresponds tothe Wasserstein gradient flow of the Kullback-Leibler divergence related to Langevin dynamics. We demonstrate by numerical examples that our model provides a well-behaved flow field which successfully solves the above sampling task.