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

We deal with the task of sampling from an unnormalized Boltzmann density ρD by learning a Boltzmann curve given by energies ft starting in a simple density ρZ. First, we examine conditions under which Fisher-Rao flows are absolutely continuous in the Wasserstein geometry. Second, we address specific interpolations ft and the learning of the related density/velocity pairs (ρt,vt). It was numerically observed that the linear interpolation, which requires only a parametrization of the velocity field vt, 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'at'e and Fleuret, who parametrize both ft and vt, we propose an interpolation which parametrizes only ft and fixes an appropriate vt. This corresponds to the 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.