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

## Abstract

Abstract We deal with the task of sampling from an unnormalized Boltzmann density ρ D subscript 𝜌 𝐷 \rho_{D} italic_ρ start_POSTSUBSCRIPT italic_D end_POSTSUBSCRIPT by learning a Boltzmann curve given by energies f t subscript 𝑓 𝑡 f_{t} italic_f start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT starting in a simple density ρ Z subscript 𝜌 𝑍 \rho_{Z} italic_ρ start_POSTSUBSCRIPT italic_Z end_POSTSUBSCRIPT .
First, we examine conditions under which Fisher-Rao flows are absolutely continuous in the Wasserstein geometry.
Second, we address specific interpolations f t subscript 𝑓 𝑡 f_{t} italic_f start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT and the learning of the related density/velocity pairs ( ρ t , v t ) subscript 𝜌 𝑡 subscript 𝑣 𝑡 (\rho_{t},v_{t}) ( italic_ρ start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT , italic_v start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) .
It was numerically observed that the linear interpolation,
which requires only a parametrization of the velocity field v t subscript 𝑣 𝑡 v_{t} italic_v start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ,
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 subscript 𝑓 𝑡 f_{t} italic_f start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT and v t subscript 𝑣 𝑡 v_{t} italic_v start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT , we propose an
interpolation which parametrizes only f t subscript 𝑓 𝑡 f_{t} italic_f start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT and fixes an appropriate v t subscript 𝑣 𝑡 v_{t} italic_v start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT .
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.

## 1 Introduction

In this paper, we consider the problem of sampling from a Boltzmann density $\rho_{D}=e^{-f_{D}}/Z_{D}$ with unknown normalizing constant $Z_{D}$. Our approach is in the spirit of recent developments in generative modelling and aims to construct a curve $\rho_{t}=e^{-f_{t}}/Z_{t}$ interpolating between a simple density $\rho_{Z}$ and the target $\rho_{D}$. If such curve admits a velocity field $v_{t}$ and there exists a solution $\varphi$ of

$$ $\partial_{t}\varphi_{t}=v_{t}(\varphi_{t}),\quad\varphi_{0}=\mathrm{Id}\quad \text{where}\quad(\varphi_{t})_{\sharp}(\rho_{Z}\mathrm{d}x)=\rho_{t}\mathrm{d }x,$ (1) $$

then we can use $\varphi_{1}$ and $\rho_{Z}$ to sample from $\rho_{D}$. An important question is, whether for a given family of functions $f_{t}$, such velocity fields exists.
While, for a fixed time $t$, the existence of such velocity fields is well established, see e.g. Laugesen et al. (), global existence with integrability of $v:[0,1]\times\mathbb{R}^{d}\to\mathbb{R}$ in time and space is addressed in this paper. This is directly related to the question whether a large class of Fisher-Rao curves $\rho_{t}$ is absolutely continuous in the Wasserstein geometry, and in particular,
fulfills a continuity equation.
This can be reduced to finding solutions of a certain family of PDEs.
With the appropriate Hilbert spaces at hand, we show the existence of the solutions of these PDEs and the
existence of an integrable velocity field $v_{t}$ in time and space which admits moreover a minimality property
of $\|v_{t}\|_{L_{2}(\mathbb{R}^{d},\rho_{t})}$.

Determining $\varphi_{1}$ requires both finding a curve $f_{t}$ and the associated vector field $v_{t}$. One way is to choose a curve interpolating between $f_{Z}$ and $f_{D}$ first and then to learn the velocity field. For the linear interpolation
$f_{t}=(1-t)f_{Z}+tf_{D}$, using above considerations, we show that there indeed exists an integrable velocity field implying the absolute continuity of the curve in the Wasserstein space.
Yet, it was numerically shown that the corresponding velocity field is often badly behaved as noted in (Máté & Fleuret, , Figure 4 and 6).
We highlight by an analytic example the bad regularity of this flow and show in particular that its velocity field may have an exploding norm $\|v_{t}\|_{L_{2}(\mathbb{R}^{d},\rho_{t})}$.
Another approach, recently proposed by Máté & Fleuret (),
considers $f_{t}=(1-t)f_{Z}+tf_{D}+t(1-t)\psi_{t}$ with unknown $\psi$ and
learns $(\psi_{t},v_{t})$ simultaneously. Henceforth, we refer to this method as the ”learned interpolation”.
However, it is theoretically unclear whether the learned curve is well-behaved or the velocity field is optimal with respect to the above norm.
As an alternative approach, we propose instead to deal *directly* with a well-behaved and mathematically accessible curve on $[0,T]$ and consider the parameterization
$f_{t}=\frac{T-t}{T}f_{D}+t\psi_{t}$. Fixing the velocity field
as $v_{t}\coloneqq\nabla(f_{t}-f_{Z})$ we only have to learn $\psi_{t}$.
The resulting PDE is the well-known Fokker-Planck equation related to Langevin dynamics
which has a solution with Boltzmann densities $\rho_{t}$, if, e.g., $\rho_{Z}$ is a Gaussian. Then, the corresponding SDE is the Ornstein-Uhlenbeck process. Here, the backward ODE of (1) must be applied for sampling.
Finally, we learn the networks for the above three cases and
demonstrate by numerical examples the effectiveness of our approach.

#### Contributions.

- 1.
We prove for the first time that under mild conditions every reasonable curve of Boltzmann densities is an absolutely continuous Wasserstein curve, see
Theorems 1 and 2.
- 2.
We verify by an analytic example
that the velocity field corresponding to the linear interpolation
can admit an exploding norm $\|v_{t}\|_{L_{2}(\mathbb{R}^{d},\rho_{t})}$, leading
to a ”non-smooth” particle transport.
- 3.
We propose to learn a curve using the parameterization
$f_{t}=\frac{T-t}{T}f_{D}+t\psi_{t}$, which we call ”gradient flow interpolation” by the following reason. With an appropriate velocity field which can be computed directly from $f_{t}$, this resembles the
Wasserstein *gradient flow* of the Kullback-Leibler divergence
with fixed second argument $\rho_{Z}$. This curve can be also described as the Fokker-Planck equation of a Langevin SDE.
- 4.
We learn the networks for the linear, learned and gradient flow interpolations
and demonstrate the performance of these three methods on certain sampling problems.

#### Related Work.

Classical approaches for sampling from an unnormalized density are based on the Markov Chain Monte Carlo (MCMC) method, see Gilks et al. (), including variants as the Hamiltonian Monte Carlo (HMM) in Hoffman & Gelman () or the Metropolis Adjusted Langevin Algorithm (MALA) in Girolami & Calderhead (). Arising problems with the correct distribution of mass among different modes lead to the use of importance sampling Neal () and sequential Monte Carlo samplers Moral et al. ().
Newer approaches include the simulation of gradient flows of probability measures. In Chen et al. () Gaussian mixtures are used to approximate the Fisher-Rao gradient flow of the Kullback-Leibler divergence. This corresponds to a time rescaling of the linear interpolation.
In Nüsken (); Maurais & Marzouk () they approximate the latter curve by kernelizing the associated Poisson equation and simulating a corresponding interacting particle system. Closely related Stein variational gradient descent Liu & Wang () kernelizes the Wasserstein gradient flow of the reverse KL divergence. Recently in Chehab & Korba () they consider a different path which they term the ”dilation path”. The associated vector field is given as the gradient of the dilation of the target density. The convolution of the density of the dilation path with (scaled) Gaussian kernels corresponds to the density of an Ornstein-Uhlenbeck process.
Recently, (stochastic) normalizing flows have been widely used to sample from unnormalized densities Noé et al. (); Wu et al. (); Hagemann et al. ().
In Midgley et al. () this problem is tackled without using a curve of probability measures by augmenting normalizing flows, which are trained using an $\alpha$-divergence loss, with annealed importance sampling.
In Hertrich & Gruhlke () a functional $F$ with minimum $\rho_{D}$ is considered. They then learn a vector field by approximating a JKO scheme for the Wasserstein gradient flow of $F$ and combine it with an importance based rejection method. They focus on $F$ being the Kullback Leibner divergence with target distribution $\rho_{D}$. Using KL for sampling is a natural choice and it is shown in Chen et al. () that it is the only $f$-divergence such that the Wasserstein gradient flow does not depend on the normalization constant of $\rho_{D}$.
Note that parametrizing the vector field using the score $\nabla\log p_{t}$ is common practice for sampling via SDEs. In Zhang & Chen () for example where they formulate the sampling problem in terms of a stochastic control problem and explicitly include the score into the parametrization of the vector field. The stochastic control viewpoint is also used in (Berner et al., , Section 3) where the KL divergence between path measures associated to a controlled and an uncontrolled SDE is used to approximate the score. A different approach using the KL on path measures is formulated in Vargas et al. ().
Another way of approximating a flow ending in $\rho_{D}$ is to consider the variance exploding diffusion SDE and approximating the score $\nabla\log p_{t}$ which is done in Akhound-Sadegh et al. ().

The authors of Máté & Fleuret () present a loss based on the continuity equation and the special form $\frac{e^{-f_{t}}}{Z_{t}}$ of Boltzmann densities. This loss allows for using fixed interpolations $f_{t}$ and only learning $v_{t}$ as well as learning $f_{t}$ and $v_{t}$ simultaneously. When finishing this paper, we became aware of the recent preprint of Sun et al. () comparing a variety of PDE based sampling methods.
They included a similar approach as the one proposed here, using the same parametrization to simulate an Ornstein-Uhlenbeck SDE. However, coming from the Wasserstein perspective, we came to the conclusion that the associated gradient flow path is especially well-behaved. We simulate the corresponding *ODE* and highlight its benefits.
For the proof that many Boltzmann density curves are absolutely continuous Wasserstein curves, the main object of interest is

$$ $\displaystyle-\nabla\cdot\left(\rho_{t}\nabla s_{t}\right)=\left(\alpha_{t}- \overline{\alpha_{t}}\right)\rho_{t},\quad t\in[0,1]$ (2) $$

which relates Fisher-Rao curves and Wasserstein absolutely continuous curves.
This equation was studied for a fixed time $t$ in various contexts. In nonlinear filtering it is used to study weak solutions of nonlinear filters see e.g. (Laugesen et al., , Section 2.1). Their theory for fixed time steps $t$ can be generalized to the case of a time dependent equation as seen in Section A.1. The close relationship between the Poincare inequality and fixed time solutions of (2) is discussed in Dhara & Kałamajska ().

## 2 Wasserstein Flows of Boltzmann Densities

### 2.1 General Background

Let $\mathcal{P}(\mathbb{R}^{d})$ denote the space of probability measures on $\mathbb{R}^{d}$.
For $\mu\in\mathcal{P}(\mathbb{R}^{d})$, let $L^{p}(\mathbb{R}^{d},\mu)$, $p\in[1,\infty)$, denotes the Banach space of (equivalence classes of) real-valued functions on $\mathbb{R}^{d}$ with finite norm
$\|f\|_{L_{p}(\mathbb{R}^{d},\mu)}\coloneqq\left(\int_{\mathbb{R}^{d}}|f|^{p}\,
\mathrm{d}\mu\right)^{1/p}$,
and $L_{p}(\mathbb{R}^{d},\mathbb{R}^{d},\mu)$ be
the Banach space of Borel vector fields $v:\mathbb{R}^{d}\to\mathbb{R}^{d}$
with finite norm
$\|v\|_{L_{p}(\mathbb{R}^{d},\mathbb{R}^{d},\mu)}\coloneqq(\int_{\mathbb{R}^{d}
}|\|v\||^{p}\,\mathrm{d}\mu)^{1/p}$,
where $\|\cdot\|$ is the Euclidean norm on $\mathbb{R}^{d}$. For the Lebesgue measure, we skip the $\mu$ in the notation.
The space $\mathcal{P}_{2}(\mathbb{R}^{d})$ of probability measures having finite second moments
equipped with the Wasserstein(-2) metric

$$ $W_{2}(\mu_{0},\mu_{1})\coloneqq\min_{\pi\in\Gamma(\mu_{0},\mu_{1})}\Big{(}\int _{\mathbb{R}^{d}\times\mathbb{R}^{d}}\|x-y\|^{2}\,\mathrm{d}\pi(x,y)\Big{)}^{ \frac{1}{2}},$ (3) $$

where $\Gamma(\mu_{0},\mu_{1})$ denotes the
set of couplings or plans $\pi\in\mathcal{P}(\mathbb{R}^{d}\times\mathbb{R}^{d})$
having marginals $\mu_{i}\in\mathcal{P}_{2}(\mathbb{R}^{d})$, $i=0,1$,
is a complete metric space.
In complete metric spaces, we can consider absolutely continuous curves depending on the metric of the space.
In the above Wasserstein space,
a weakly continuous curve $\mu:[0,1]\to\mathcal{P}_{2}(\mathbb{R}^{d})$, $t\mapsto\mu_{t}$ is *absolutely continuous*, if there exists a Borel vector field
$v:[0,1]\times\mathbb{R}^{d}\to\mathbb{R}^{d}$
with
$\int_{0}^{1}\|v_{t}\|_{L_{2}(\mathbb{R}^{d},\mathbb{R}^{d},\mu_{t})}\,\mathrm{
d}t<\infty$
such that the *continuity equation*

$$ $\partial_{t}\mu_{t}+\nabla\cdot(\mu_{t}v_{t})=0$ (4) $$

is fulfilled in the sense of distributions
(^1^11$\int_{0}^{1}\int_{\mathbb{R}^{d}}\partial_{t}\phi+\langle\nabla_{x}\phi,v_{t} \rangle\,\mathrm{d}\mu_{t}\mathrm{d}t=0~{}\text{ for all }\phi\in C_{c}^{ \infty}\left((0,1)\times\mathbb{R}^{d}\right)$)
see, e.g. Ambrosio et al. ().
While for fixed $\mu_{t}$ there are many vector fields such that the continuity equation is fulfilled,
there exists a unique one such that $\|v_{t}\|_{L_{2}(\mathbb{R}^{d},\mu_{t})}$ becomes minimal for almost every $t\in[0,1]$.
In other words, if $\mu_{t}$ fulfills the continuity equation for some vector field,
then there exists a unique vector field that solves

$$ $\operatorname*{arg\,min}_{v_{t}\in L_{2}(\mathbb{R}^{d},\mathbb{R}^{d},\mu_{t} )}\int_{\mathbb{R}^{d}}\|v_{t}\|^{2}\,\mathrm{d}\mu_{t}\quad\text{s.t.}\quad \partial_{t}\mu_{t}+\nabla\cdot(\mu_{t}v_{t})=0.$ (5) $$

Moreover, this optimal vector field is determined
by the condition that $v_{t}\in T_{\mu_{t}}\mathcal{P}_{2}(\mathbb{R}^{d})$ for almost every $t\in[0,1]$, with the *regular tangent space*

$$ $\displaystyle T_{\mu}\mathcal{P}_{2}(\mathbb{R}^{d})$ $\displaystyle\coloneqq\overline{\left\{\nabla\phi:\phi\in C^{\infty}_{\mathrm{ c}}(\mathbb{R}^{d})\right\}}^{L_{2}(\mathbb{R}^{d},\mu)},$ (6) $$

see (Ambrosio et al., , § 8).
Finally, if $\mu_{t}$ is an absolutely continuous curve with Borel vector field $v_{t}$ such that for every compact Borel set $B\subset\mathbb{R}^{d}$ it holds
$\int_{0}^{1}\sup_{B}\|v_{t}\|+\text{Lip}(v_{t},B)\,\mathrm{d}t<\infty$,
then there exists a solution $\varphi:[0,1]\times\mathbb{R}^{d}\to\mathbb{R}^{d}$ of

$$ $\partial_{t}\varphi(t,x)=v_{t}(\varphi(t,x)),\quad\varphi(0,x)=x\\$ (7) $$

which fulfills $\mu_{t}=\varphi_{t,\sharp}\mu_{0}\coloneqq\mu_{0}\circ\varphi_{t}^{-1}$, see e.g. (Ambrosio et al., , Proposition 8.1.8).

### 2.2 Wasserstein meets Fisher-Rao Flows

In the following, we are interested in absolutely continuous probability measures
with respect to the Lebesgue measure, i.e., the space
$\mathcal{P}^{ac}(\mathbb{R}^{d})\coloneqq\{\mu=\rho\,\mathrm{d}x:\rho\in L_{1}
(\mathbb{R}^{d})$,
$\rho\geq 0$,
$\int_{\mathbb{R}^{d}}\rho\,\mathrm{d}x=1\}$
and we will address such measures just by their densities.
Moreover, we will only consider Boltzmann/Gibbs densities depending on functions
$f:[0,1]\times\mathbb{R}^{d}\to\mathbb{R}$ with $e^{-f_{t}}\in L_{1}(\mathbb{R}^{d})$ and $f(\cdot,x)\in C^{1}([0,1])$ for all $x\in\mathbb{R}^{d}$ and $\partial_{t}f(t,\cdot)\in L_{2}(\mathbb{R}^{d},\rho_{t})$
such that

$$ $\rho_{t}\coloneqq\frac{e^{-f_{t}}}{Z_{t}},\quad Z_{t}\coloneqq\int_{R^{d}}e^{- f_{t}}\,\mathrm{d}x.$ (8) $$

In this case, straightforward computation gives

$$ $\displaystyle\partial\mu_{t}=\partial_{t}\rho_{t}$ $\displaystyle=-\left(\partial_{t}f_{t}-\mathbb{E}_{\rho_{t}}[\partial_{t}f_{t} ]\right)\rho_{t},$ (9) $$

where it maybe useful to note that

$$ $\mathbb{E}_{\rho_{t}}[\partial_{t}f_{t}]=-\partial_{t}\log Z_{t}.$ $$

This differential equation is a Fisher-Rao flow equation.
More precisely, a weakly continuous curve $\rho:[0,1]\to\mathcal{P}^{ac}(\mathbb{R}^{d})$,
$t\mapsto\rho_{t}$ is a *Fisher-Rao curve*, if there exists a measurable map $\alpha:[0,1]\times\mathbb{R}^{d}\to\mathbb{R}$
such that $\alpha_{t}:=\alpha(t,\cdot)\in L_{2}(\mathbb{R}^{d},\rho_{t})$ and we have in a distibutional sense

$$ $\partial_{t}\rho_{t}=(\alpha_{t}-\bar{\alpha}_{t})\rho_{t},\quad\bar{\alpha}_{ t}\coloneqq\mathbb{E}_{\rho_{t}}[\alpha_{t}].$ (10) $$

Obviously, (9) delivers a Fisher-Rao curve with $\alpha_{t}\coloneqq-\partial_{t}f_{t}$.
For more information on Fisher-Rao flows see, e.g., Gallouët & Monsaingeon (); Wang & Li ().
We are interested in the
question under which conditions a Fisher-Rao curve determined by (9) is also an absolutely continuous curve in the Wasserstein space.
For later purposes, note that for Boltzmann densities (8), the second summand of the continuity equation (4) reads as

$$ $\nabla\cdot(\mu_{t}v_{t})=\nabla\cdot(\rho_{t}v_{t})=\left(-\langle\nabla f_{t },v_{t}\rangle+\nabla\cdot v_{t}\right)\rho_{t}.$ (11) $$

Nevertheless, using (9), for our Fisher-Rao curves the continuity equation becomes

$$ $\left(\partial_{t}f_{t}-\mathbb{E}_{\rho_{t}}[\partial_{t}f_{t}]\right)\rho_{t }=\nabla\cdot(\rho_{t},v_{t}),$ (12) $$

and in the ideal case that we can find a vector field $v_{t}=\nabla s_{t}$, compare (6), this
can be rewritten as the family of Poisson equations

$$ $\left(\partial_{t}f_{t}-\mathbb{E}_{\rho_{t}}[\partial_{t}f_{t}]\right)\rho_{t }=\nabla\cdot(\rho_{t}\nabla s_{t}),$ (13) $$

which needs to be solved for $s_{t}$.
For *fixed time* $t$, the solution of the Poisson equation has already been examined, e.g., in Laugesen et al. ()
and for more recent work, see also Reich (); Taghvaei & Mehta ().
However, the pointwise solution for each fixed $t$ does not ensure that $(\rho_{t},\nabla s_{t})$ gives rise to an absolutely
continuous Wasserstein curve, since by definition this requires additional properties of $\nabla_{x}s:[0,1]\times\mathbb{R}^{d}\to\mathbb{R}^{d}$, globally in $t$:
i) $\nabla_{x}s$ is Borel measurable on $[0,1]\times\mathbb{R}^{d}$, and ii)
$t\mapsto\|\nabla_{x}s(t,\cdot)\|_{L_{2}(\rho_{t},\mathbb{R}^{d})}^{2}\in L_{1}
([0,1])$.
If also, iii) $\nabla_{x}s(t,\cdot)\in T_{\rho_{t}}\mathcal{P}_{2}(\mathbb{R}^{d})$ for a.e. $t\in[0,1]$, then it is known that the velocity field $\nabla_{x}s$ is the *optimal one* in the sense of (5). To our best knowledge, the following two theorems, which are detailed in Appendix A.1, tackle a *rigorous* solution of this problem for the first time.

###### Theorem 1 .

Assume that $\boldsymbol{\rho}$ determined by
$\int_{[0,1]\times\mathbb{R}^{d}}f(t,x)\,\mathrm{d}\boldsymbol{\rho}=\int_{0}^{
1}\int_{\mathbb{R}^{d}}f(t,x)\,\mathrm{d}\rho_{t}\mathrm{d}t$
satisfies a so-called partial Poincare inequality, see Definition A.7, for some $K>0$, and that $\frac{1}{\boldsymbol{\rho}}$ is locally integrable.
Furthermore, suppose that $\alpha_{t}-\overline{\alpha_{t}}\in L_{2}([0,1]\times\mathbb{R}^{d},
\boldsymbol{\rho})$.
Then, there exists a unique weak solution $\boldsymbol{s}$ in a certain Hilbert space ${\mathcal{X}_{0}^{1}}$ of the problem

$$ $\displaystyle\int_{0}^{1}\int_{\mathbb{R}^{d}}\langle\nabla\boldsymbol{s}_{t}, \nabla\psi_{t}\rangle\,\mathrm{d}\rho_{t}\mathrm{d}t=\int_{0}^{1}\int_{\mathbb {R}^{d}}\psi_{t}\left(\alpha_{t}-\overline{\alpha_{t}}\right)\,\mathrm{d}\rho_ {t}\mathrm{d}t,\quad\text{for all }\psi\in{\mathcal{X}_{0}^{1}}.$ (14) $$

###### Theorem 2 .

Let $\rho_{0},\rho_{1}\in\mathcal{P}(\mathbb{R}^{d})$ and assume that $\rho_{t}$ is a Fisher-Rao curve defined by (10),
Let $\boldsymbol{\rho}$ and $\alpha$ satisfy the assumptions of Theorem 1.
Then, $\rho_{t}$ is a Wasserstein absolutely continuous curve with vector field $\nabla_{x}\boldsymbol{s}$, where $\boldsymbol{s}\in{\mathcal{X}_{0}^{1}}$ is the unique weak solution of (14). If $\rho_{t}$ is bounded from above, it holds that $\nabla_{x}\boldsymbol{s}_{t}\in T_{\rho_{t}}\mathcal{P}_{2}(\mathbb{R}^{d})$ for a.e. $t\in[0,1]$.
Furthermore, it holds that

$$ $\displaystyle\int_{0}^{1}\int_{\mathbb{R}^{d}}\|\nabla_{x}\boldsymbol{s}_{t}\| ^{2}\,\mathrm{d}\rho_{t}\mathrm{d}t\leq K\int_{0}^{1}\int_{\mathbb{R}^{d}} \left(\alpha_{t}-\overline{\alpha_{t}}\right)^{2}\,\mathrm{d}\rho_{t}\mathrm{d }t.$ (15) $$

Note that by the Benamou-Brenier formula, see (Ambrosio et al., , Equation 8.0.3), the left-hand side of (15) is an upper bound for the Wasserstein distance $W_{2}^{2}(\rho_{0},\rho_{1})$, while the double-integral on the right-hand side defines the Fisher-Rao *action* of $\rho_{t}$.

###### Theorem 1 .

###### Theorem 2 .

## 3 Neural Sampling from Boltzmann Densities

Our aim is to sample from a Boltzmann measure $\mu_{D}=\rho_{D}\,\mathrm{d}x$,
where we only have access to $f_{D}:\mathbb{R}^{d}\to\mathbb{R}$,
but not to the normalization constant $Z_{D}$,
by following a measure flow starting in a simple to sample measure $\mu_{Z}=\rho_{Z}\,\mathrm{d}x$ like the standard Gaussian one.
We apply the continuity equation (12) of Fisher-Rao curves, which can be rewritten by (11) and division
by $\rho_{t}>0$ as

$$ $\partial_{t}f_{t}-\underbrace{\mathbb{E}_{\rho_{t}}[\partial_{t}f_{t}]}_{C_{t} }=-\langle\nabla f_{t},v_{t}\rangle+\nabla\cdot v_{t}.$ (16) $$

Indeed every spatially constant function $C_{t}$ which fulfills the above equation must be of the form
$C_{t}=\mathbb{E}_{\rho_{t}}[\partial_{t}f_{t}]$, see Máté & Fleuret ().
This enables us to learn $C_{t}$, so that (16) does no longer depend on normalization constants.
We will deal with the following parameterizations of $f_{t}$ and $v_{t}$ and how to learn them:

- 1.
linear interpolation: $f_{t}\coloneqq(1-t)f_{Z}+tf_{D}$  $\hookrightarrow$ learn $v_{t}$,
- 2.
learned interpolation by Máté and Fleuret: $f_{t}\coloneqq(1-t)f_{Z}+tf_{D}+t(1-t)\psi_{t}$  $\hookrightarrow$ learn $(\psi_{t},v_{t})$,
- 3.
gradient flow interpolation: $f_{t}\coloneqq\frac{T-t}{T}f_{D}+t\psi_{t}$, $v_{t}\coloneqq\nabla(f_{t}-f_{Z})$  $\hookrightarrow$ learn $\psi_{t}$.

While by definition the first two settings interpolate between
$\rho_{Z}$ and $\rho_{D}$ and $\rho_{0}=\rho_{Z}$, $\rho_{1}=\rho_{D}$, we will see that the newly proposed third variant interpolates between $\rho_{D}$ and $\rho_{Z}$ by the choice of $v_{t}$, i.e., $\rho_{0}=\rho_{D}$ and
$\lim_{T\to\infty}\rho_{T}=\rho_{Z}$. In practice we reparametrize the latter curve into unit time, see Remark 3.

### 3.1 Linear & Learned Interpolation

#### Linear interpolation.

The simplest interpolation is the linear one
$f_{t}\coloneqq(1-t)f_{Z}+tf_{D}$, i.e., $\rho_{t}\propto\rho_{0}^{1-t}\rho_{1}^{t}$.
Under suitable conditions on the end points this density fulfills the assumptions of Theorem 2, see Corollary A.17 in the Appendix, and is therefore a Wasserstein curve. By (9), it also defines a Fisher-Rao curve, and the corresponding Fisher-Rao flow equation (10) is given by

$$ $\displaystyle\partial_{t}\rho_{t}$ $\displaystyle=\left(f_{0}-f_{1}-\mathbb{E}_{\rho_{t}}[f_{0}-f_{1}]\right)\rho_ {t}$ (17) $\displaystyle=-\Big{(}\log(\frac{\rho_{0}}{\rho_{1}})-\mathbb{E}_{\rho_{t}}[ \log(\frac{\rho_{0}}{\rho_{1}})]\Big{)}\rho_{t}.$ (18) $$

Note that, $\rho_{t}$ is even a Fisher-Rao *gradient flow* of the negative log likelihood, see also Maurais & Marzouk (), given by

$$ $\displaystyle\mathcal{E}_{\rm NLL}(\rho):=-\mathbb{E}_{\rho}\left[\log(\frac{ \rho_{1}}{\rho_{0}})\right].$ (19) $$

There are a variety of approaches for approximating the vector field corresponding to the linear interpolation.
In Maurais & Marzouk (); Nüsken () a kernel-based method for solving (13) was proposed. In Máté & Fleuret () they propose to solve (16) which in case of the linear interpolation becomes

$$ $\displaystyle f_{1}-f_{0}-C_{t}+\langle\nabla f_{t},v_{t}\rangle-\nabla\cdot v _{t}=0.$ (20) $$

This equation has many solutions $v_{t}$
for the same $\rho_{t}$, and in order to enforce uniqueness we could demand that $v_{t}=\nabla s_{t}$.

This approach leads to learning neural networks $(v^{\theta_{1}}_{t},C^{\theta_{2}}_{t})$ by minimizing the loss function

$$ $\displaystyle L(\theta)$ $\displaystyle\coloneqq\mathbb{E}_{t\in\mathcal{U}[0,1],x\in\mathcal{U}[a,b]^{d }}\left[\mathcal{E}(\theta,x,t)\right],$ (21) $\displaystyle\mathcal{E}(\theta,x,t)$ $\displaystyle\coloneqq|f_{1}-f_{0}-C_{t}^{\theta_{2}}+\langle\nabla f_{t},v_{t }^{\theta_{1}}\rangle-\nabla\cdot v_{t}^{\theta_{1}}|^{2},\quad\theta\coloneqq (\theta_{1},\theta_{2}).$ (22) $$

Figure: Figure 1: Evolution of the probability densities $\rho_{t}\propto\rho_{0}^{1-t}\rho_{1}^{t}$, where $\rho_{0}\propto e^{-|x|}$ and $\rho_{1}\propto e^{-2\min\{|x|,|x-m|\}}$ for $m=50$. For different values of $m$ see 4.
Refer to caption: extracted/5901341/imgs/teleportation/combined_density_row_m_50.png

*A teleportation issue with the linear interpolation.*
Máté & Fleuret () gave examples, where for an asymmetric constellation of target and prior measure, this path transports mass into distant modes only at very late times which may cause problems while learning the corresponding vector field. Mathematically, this corresponds to the fact that the vector field $v_{t}(x)$ develops a singularity in $x\in\mathbb{R}^{d}$ for late times $t$, see (Máté & Fleuret, , Figure 6). In this work, we demonstrate that this phenomenon is *more severe*: even when averaging over the space $x$, the norm $\|v_{t}\|_{L_{2}(\rho_{t})}$ of the optimal field (5) explodes for late times $t$. For an example in $\mathbb{R}^{1}$ involving a Laplacian $\rho_{0}\propto e^{-|x|}$ centered at 0 0 and a Laplacian-like distribution $\rho_{1}\propto e^{-2\min\{|x|,|x-m|\}}$ with a mode at 0 0 and a second mode at $m>0$, see Figure 1, we are able to give an explicit and analytic formula for the calculation of $\|v_{t}\|_{L_{2}(\rho_{t})}$ in the Appendix A.4. Using this formula (71), we demonstrate that $\|v_{t}\|_{L_{2}(\rho_{t})}$ explodes when $t$ is close to $1$, and that this explosion happens more severely and at later times, the further the second mode $m>0$ is away from 0 0, see Figure 2.

Figure: Figure 2: In the first three figures, the evolution of $\|v_{t}\|_{L_{2}(\rho_{t})}^{2}$ belonging to the path in Figure 1 is depicted for different values of the second mode $m>0$. For larger $m$, the norm $\|v_{t}\|_{L_{2}(\rho_{t})}^{2}$ approaches the limit $\|v_{1}\|_{L_{2}(\rho_{1})}^{2}$ at later times, hence with a steeper slope. The last figure shows the *log scale* of $\|v_{1}\|_{L_{2}(\rho_{1})}^{2}$ depending on $m$. It demonstrates that $\|v_{1}\|_{L_{2}(\rho_{1})}$ roughly grows *exponentially* with $m$.
Refer to caption: extracted/5901341/imgs/teleportation/velo_5.0.png

#### Learned Interpolation.

To cope with the above disadvantages of the linear interpolation, Máté & Fleuret () proposed to parameterize also the flow by $f_{t}\coloneqq(1-t)f_{0}+tf_{1}+t(1-t)\psi_{t}$ with a neural network $\psi_{t}^{\theta_{1}}$
and to minimize according to (16) the above loss with

$$ $\displaystyle\mathcal{E}(\theta,x,t):=|\partial_{t}f^{\theta_{1}}_{t}(x)-C^{ \theta_{3}}_{t}+\langle\nabla f^{\theta_{1}}_{t},v^{\theta_{2}}_{t}\rangle- \nabla\cdot v^{\theta_{2}}_{t}|^{2},\quad\theta\coloneqq(\theta_{1},\theta_{2} ,\theta_{3}).$ (23) $$

Since the approximation of high dimensional integrals is difficult, for both the linear and learned interpolation Máté & Fleuret () propose to use an iterative approach by solving the ODE $\partial_{t}\varphi_{t}^{\theta^{\prime}}=v^{\theta^{\prime}}_{t}(\varphi_{t}^
{\theta^{\prime}})$, $\varphi_{0}^{\theta^{\prime}}=\mathrm{Id}$ for $\theta^{\prime}$ from the previous minimization and to solve

$$ $\displaystyle L(\theta)\coloneqq\mathbb{E}_{t\in\mathcal{U}[0,1],z\sim\rho_{0} }[\mathcal{E}(\theta,\varphi_{t}^{\theta^{\prime}}(z),t)].$ (24) $$

Heuristically this penalizes $\mathcal{E}(\theta,x,t)$ stronger in regions where the intermediate solution has more mass. Finally, sampling can be done in both cases – linear and learned interpolation – assuming sufficient regularity of the learned $v_{t}^{\theta_{i}}$, by simulating the ODE (7). Note that in general there are infinitely many solutions of (20) and neither the parameterized curve $\rho_{t}$ nor the vector field $v_{t}$ is unique. Further $v_{t}^{\theta_{2}}$ obtained in this way may not be minimal in any sense. The proof of the validity of the interpolation with respect to Theorem 2
is future work.

### 3.2 Gradient Flow Interpolation

Recall that the linear interpolation of the energies $f_{0},f_{1}$ leads to a multiplicative interpolation of densities $\rho_{t}(x)\propto\rho_{Z}(x)^{1-t}\rho_{D}(x)^{t}$. Note that this does not take account of any *local* behaviour of $\rho_{Z},\rho_{D}$ in a neighbourhood of $x$, resulting in irregular behavior of the vector field as shown in Figure 1 and 2.
In contrast, widely used paths in generative modelling Song et al. (); Lipman et al. () originate from an interpolation of densities via (spatial) convolutions
$\rho_{t}\sim\rho_{Z}(\frac{\cdot}{a_{t}})\ast\rho_{D}(\frac{\cdot}{b_{t}})$ for time schedules $a_{t},b_{t}$, corresponding to an interpolation $X_{t}=a_{t}Z+b_{t}X_{0}$ of associated independent random variables $Z\sim\rho_{Z},X_{0}\sim\rho_{D}$ on the *spatial level*. We aim to learn such a curve and an associated vector field.

In this paper, we propose to use $f_{t}\coloneqq\frac{T-t}{T}f_{D}+t\psi_{t}$ and $v_{t}\coloneqq\nabla(f_{t}-f_{Z})$
and to parameterize $\psi_{t}$ by a neural network.
Then, equation (16) becomes

$$ $\displaystyle\partial_{t}f_{t}-\mathbb{E}_{\rho_{t}}[\partial_{t}{f}_{t}]+ \langle\nabla{f}_{t},\nabla({f}_{t}-f_{Z})\rangle-\Delta({f}_{t}-f_{Z})=0.$ (25) $$

Indeed, we will show in the paragraph below that this equation makes perfect sense, in particular, it has a solution for $\rho_{Z}\sim\mathcal{N}(0,1)$.
Then we can compute $(\psi_{t}^{\theta_{1}},C_{t}^{\theta_{2}})$ for $f_{t}^{\theta_{1}}\coloneqq\frac{T-t}{T}f_{D}+t\psi_{t}^{\theta_{1}}$ by minimizing the loss
$L$ in (21) with the function

$$ $\displaystyle\mathcal{E}(\theta,x,t)\coloneqq|\partial_{t}f_{t}^{\theta_{1}}-C _{t}^{\theta_{2}}+\langle\nabla f_{t}^{\theta_{1}},\nabla(f_{t}^{\theta_{1}}-f _{Z})\rangle-\Delta(f^{\theta_{1}}_{t}-f_{Z})|^{2},\quad\theta\coloneqq(\theta _{1},\theta_{2}).$ (26) $$

Finally, to sample from the target distribution,
we simulate the ODE (7) backwards on $[0,T]$
to obtain $\tilde{\varphi}:[0,T]\times\mathbb{R}^{d}\to\mathbb{R}^{d}$ satisfying

$$ $\partial_{t}{\tilde{\varphi}}(t,x)=-v_{t}({\tilde{\varphi}}(t,x)),\quad{\tilde {\varphi}}(T,x)=x$ (27) $$

which fulfills $\mu_{t}=\tilde{\varphi}_{T-t,\sharp}\mu_{Z}$.

#### Relation to Wasserstein gradient flows.

Equation (25) is well-known in connection with Wasserstein gradient flows.
A *Wasserstein gradient flow* is an absolutely continuous curve
whose vector field
$v_{t}=-\text{grad}_{\rho_{t}}F\in T_{\rho_{t}}(\mathcal{P}_{2}(\mathbb{R}^{d}))$
is given by the negative Wasserstein gradient (or more generally, *subdifferential*) of a function $F:\mathcal{P}_{2}(\mathbb{R}^{d})\to\mathbb{R}$.
Then, the continuity equation

$$ $\partial_{t}\rho_{t}-\nabla\cdot\left(\rho_{t}\text{grad}_{\rho_{t}}F\right)=0$ (28) $$

produces a flow towards a minimizer of $F$.
A prominent example, see, e.g. (Ambrosio et al., , Chapter 10.4), is
$F(\rho)\coloneqq\text{KL}(\rho|\rho_{Z})$ arising from the Kullback-Leibler divergence
which has the unique minimizer $\rho_{Z}$. Since it is easy to show, see (Ambrosio et al., , Lemma 10.4.1),
that $\text{grad}_{\rho_{t}}F=\nabla\log\frac{\rho_{t}}{\rho_{Z}}$ , we can be rewrite
(28) as

$$ $\displaystyle\partial_{t}\rho_{t}=-\nabla\cdot\left(\rho_{t}(\nabla\log\rho_{Z }-\nabla\log\rho_{t})\right),\quad\rho_{0}=\rho_{D},$ (29) $$

which is also known as the Fokker-Planck equation.
Inserting $\rho_{t}\coloneqq e^{-f_{t}}/Z_{t}$ results exactly in (25). The existence and uniqueness of a solution of (29) is shown in (Ambrosio & Savaré, , Theorem 6.6).
However, it is not clear if Boltzmann densities stay Boltzmann densities. But in the special case, when $\rho_{Z}$ is a Gaussian, this is the case as we will see next.
There is a close relation between (29) and SDEs, namely the Langevin SDE

$$ $\displaystyle\mathrm{d}X_{t}=\nabla\log\rho_{Z}(X_{t})\mathrm{d}t+\sqrt{2} \mathrm{d}B_{t}$ (30) $$

corresponds to the Fokker-Planck equation (29), see e.g., (Chewi et al., , Chapter 6.2), i.e., $X_{t}\sim\rho_{t}$. For $\rho_{Z}\propto\mathcal{N}(0,I_{d})$,
(30) becomes
$\mathrm{d}X_{t}=-X_{t}\mathrm{d}t+\sqrt{2}\mathrm{d}B_{t}$.
This is an Ornstein-Uhlenbeck process, which has the closed form solution

$$ $X_{t}=e^{-t}Z+e^{-t}X_{0},\quad Z\sim\mathcal{N}\left(0,(e^{2t}-1)I_{d}\right) ,\quad X_{0}\sim\rho_{D},$ $$

where $Z$ is independent from $X_{0}$, see (Herrmann & Massin, , Eq. (2.4)).
From the closed form solution it follows that
$\rho_{t}\sim\mathcal{N}\left(0,(1-e^{-2t})I_{d}\right)*\rho_{D}\left(e^{t}
\cdot\right)>0$, and thus $\rho_{t}\propto e^{-f_{t}}$ which implies that (25) has a solution.

###### Remark 3 .

The dynamics of both the linear and learned interpolations are on the unit interval. To obtain dynamics on the unit interval for the gradient flow interpolation we apply an approach from the score based diffusion literature Song et al. (). We use a linear time schedule $\beta(t):=\beta_{\text{min}}+t\left(\beta_{\text{max}}-\beta_{\text{min}}\right)$ and the associated SDE

$$ $\displaystyle\mathrm{d}X_{t}=-\frac{1}{2}\beta(t)X_{t}\mathrm{d}t+\sqrt{\beta( t)}\mathrm{d}B_{t}\quad X_{0}\sim\rho_{1}.$ (31) $$

The SDE (31) has a closed form solution given by

$$ $X_{t}=\sqrt{1-e^{-g(t)}}Z+e^{\frac{-g(t)}{2}}X_{0},$ (32) $$

where $g(t):=\int_{0}^{t}\beta(s)\,\mathrm{d}s$ and $Z\sim\mathcal{N}(0,I_{d})$.
As done in Song et al. () we choose $\beta_{\text{min}}=0.1$ and $\beta_{\text{max}}=20$. These choices ensure that the distribution of $X_{1}$ is approximately $\mathcal{N}(0,I_{d})$.
We then aim to learn the solution of the associated Fokker-Planck equation, this corresponds to setting $v_{t}\coloneqq\frac{\beta(t)}{2}\nabla(f_{t}-\frac{\|\cdot\|^{2}}{2})$.

###### Remark 3 .

## 4 Experiments

In this section we apply the different approaches to common sampling problems. We compare the performance of using the linear, learned or gradient flow interpolation.

For both the linear interpolation and the learned interpolation, we can choose to learn the vector field directly, or learn a potential and obtain the vector field as its gradient. In practice, we choose to parameterize it directly as in Máté & Fleuret (). Note that in practice for the gradient flow interpolation we apply Remark 3. Furthermore we improve performance by using learned scheduling function $g(t)$, satisfying $g(0)=1$ and $g(1)=0$, and parametrizing the corresponding $f_{t}$ as $f_{t}:=g(t)f_{D}+t\psi_{t}$ .

A major design choice is at which points to evaluate the pointwise error $\mathcal{E}(\theta,x,t)$. The approach taken in Máté & Fleuret () is to use samples along the trajectory induced by the current estimate of $v_{t}$.
For the gradient flow interpolation, we chose a different approach. We sample points uniformly from a large enough domain and interpolate them with samples from our latent distribution according to the formula given in (32), where we replace the samples from the data distribution with random uniform ones. For the linear and learned interpolation we tried similar approaches sampling the points from scaled uniform domains. However, we found in our experiments that the method employed in Máté & Fleuret () led to better performance.

In the following, we evaluate the methods using the effective sample size, negative log likelihood and the energy distance Székely & Rizzo (). See A.5 for further details on the implementation and evaluation.

#### Gaussian Mixture Model.

We use the experimental setup from Midgley et al. (). The target distribution consists of a mixture of 40 evenly weighted Gaussians in 2 dimensions. The means are distributed uniformly over $[-40,40]^{2}$.

Figure: Figure 3: Results for a Gaussian Mixture Model with 40 modes. For the linear and learned interpolation we showed the results for $\sigma=30$, for the gradient flow interpretation we used $\sigma=1$.
Refer to caption: extracted/5901341/imgs/GT.png

In Figure 3 we plot the countour lines of the target distribution as well as 1000 samples from the simulated flow and the ground truth, respectively.

**Table 1: Comparison of effective sample size (ESS), negative log likelihood (NLL), and the energy distance for different interpolations. Arrows indicate if ascending/descending values are better.**
|  | ESS $\uparrow$ | NLL $\downarrow$ | Energy Distance $\downarrow$ |
| --- | --- | --- | --- |
| linear ($\sigma=30$) | 0.183$\pm$ 0.05 | 8.657 $\pm$ 0.012 | 0.247$\pm$ 0.016 |
| learned ($\sigma=20$) | 0.605 $\pm$ 0.316 | 6.894$\pm$ 0.005 | 0.058$\pm$ 0.006 |
| learned ($\sigma=30$) | 0.922$\pm$ 0.002 | 6.913 $\pm$ 0.003 | 0.121$\pm$ 0.009 |
| gradient flow ($\sigma=1$) | 0.992 $\pm$ 5e-5 | 6.864 $\pm$ 0.005 | 0.0021 $\pm$ 0.001 |

For this experiment an important factor is the choice of latent distribution. We restricted ourselves to using a Gaussian latent distribution. However as outlined in Section 3.1 the linear interpolation faces difficulties when some modes are far away from the latent while others are close. In order to mitigate this problem, we experimented with using varying standard deviations $\sigma$; the impact of this choice is visualized in A.5. We found that using a mode covering Gaussian led to the best results. In contrast, the gradient flow interpolation works well when starting in a standard Gaussian which is what we do.

#### Many Well Distribution.

The target is a $8$-dimensional many well distribution and we use the implementation of the target energy from Midgley et al. (). More precisely we consider the $4$-fold product distribution of the two dimensional density given by

$$ $\displaystyle\log p(x_{1},x_{2})=-x_{1}^{4}+6x_{1}^{2}+\frac{1}{2}x_{1}-\frac{ 1}{2}x_{2}^{2}+\text{constant}$ (33) $$

as target distribution. The results in Table 2 show that both the learned and gradient flow interpolations significantly outperform the linear interpolation. The learned interpolation works especially well in this setting. Here the choice of latent distribution is not as relevant as starting in a standard Gaussian is sufficient for all methods, since the modes of the many well distribution are concentrated close to zero.

**Table 2: Comparison of effective sample size, negative log likelihood, and the energy distance for different interpolations.**
|  | ESS $\uparrow$ | NLL $\downarrow$ | Energy Distance $\downarrow$ |
| --- | --- | --- | --- |
| linear | 0.985$\pm$ 0.016 | 9.689 $\pm$ 0.01 | 0.278 $\pm$ 0.003 |
| learned | 0.999 $\pm$ 1e-5 | 6.876 $\pm$ 0.01 | 9e-5 $\pm$ 3e-5 |
| gradient flow | 0.991 $\pm$ 3e-4 | 6.989 $\pm$ 0.01 | 8e-4 $\pm$ 1e-4 |

## 5 Conclusions

#### Discussion.

We studied different paths in the probability space in order to sample from unnormalized densities and compared their behaviour and computibility. The linear interpolation of energies suffers from an exploding vector field, and learning the vector field *and* energy path simultaneously (*learned interpolation*) in practice solves this problem. However, there is no theoretical knowledge about the path which we want to approximate and in turn no conclusions about the regularity of the corresponding vector field can be drawn. In contrast, the proposed gradient flow approach allows for a thorough mathematical description of the velocity field and the path itself, opening the doors for a rigorous analysis via Wasserstein gradient flows and Fokker-Planck equations. We show numerical performance on sampling from Gaussian mixture models and a many well distribution. Furthermore we close a gap in the literature by showing that a large class of Fisher Rao curves and in particular curves of Boltzmann densities are in fact absolutely continuous curves in the Wasserstein geometry. Hence the existence of a corresponding vector field driving the evolution is implied, justifying sampling from the target density via solving an ODE.

#### Limitations and Outlook.

The methods proposed require knowledge of the domain of the target distribution. Furthermore solving PDEs in high dimensional settings remains challenging. These issues can be addressed by developing more efficient strategies to sample points at which to evaluate the loss $\mathcal{E}(\theta,x,t)$, such as combining uniform sampling and trajectory based sampling. From a theoretical point of view further investigation of the behaviour and regularity of the gradient flow path is needed, especially the behaviour of the norm $\|v_{t}\|_{L_{2}(\rho_{t})}$ would be of interest.

## Appendix A Appendix

### A.1 Wasserstein meets Fisher-Rao: Solution of a family of Poisson Equations

We will show that under mild assumptions on $f$, the Fisher-Rao curve $\rho_{t}$ from (9) is also a Wasserstein absolutely continuous curve with a tangent vector field $\nabla s_{t}$(^2^22Overall, let $\nabla:=\nabla_{x}$ denote differentiation in direction $x\in\mathbb{R}^{d}$.). In order to do that, we will describe $s_{t}$ as the unique solution of a Poisson equation in a suitable Hilbert space ${\mathcal{X}_{0}^{1}}$.
More precisely, we show that there exists a unique $s\in L_{2}([0,1]\times\mathbb{R}^{d};\boldsymbol{\rho})$ (with $\boldsymbol{\rho}$ given by Definition A.1) such that $\nabla s_{t}:=\nabla s(t,\cdot)$ fulfills the Poisson equation

$$ $\displaystyle\nabla\cdot(\rho_{t}\nabla s_{t})=\Big{(}\partial_{t}f-\mathbb{E} _{\rho_{t}}\left[\partial_{t}f\right]\Big{)}\rho_{t}$ (34) $$

in a distributional sense in ${\mathcal{X}_{0}^{1}}$.
For fixed time $t$, this has already been done, e.g., in Laugesen et al. (). But note that a pointwise-in-time treatment does not allow us to conclude Wasserstein absolute continuity of our curve $\rho_{t}$. Therefore, we generalize their methods to a global-in-time ansatz; in particular, under suitable assumptions on $f$, we find a unique weak solution $s$ of (34) such that

- 1.
$\nabla_{x}s$ is Borel measurable on $[0,1]\times\mathbb{R}^{d}$,
- 2.
$t\mapsto\|\nabla_{x}s(t,\cdot)\|_{L_{2}(\rho_{t})}^{2}\in L_{1}([0,1])$,  i.e., $\int_{0}^{1}\|\nabla_{x}s(t,\cdot)\|_{L_{2}(\rho_{t})}^{2}\,\mathrm{d}t<\infty$,
- 3.
$\nabla_{x}s(t,\cdot)\in T_{\rho_{t}}\mathcal{P}_{2}(\mathbb{R}^{d})$ for a.e. $t\in[0,1]$.

These conditions ensure that the Fisher-Rao curve (9) is absolutely continuous in the Wasserstein sense, and that
the velocity field $\nabla_{x}s$ is the optimal one fulfilling the continuity equation w.r.t. our curve $\rho_{t}$ as described in (5).

As a by-product, we obtain the estimate

$$ $\displaystyle\int_{0}^{1}\int_{\mathbb{R}^{d}}\|\nabla_{x}s(t,\cdot)\|^{2}\, \mathrm{d}\rho_{t}\mathrm{d}t\leq K\int_{0}^{1}\int_{\mathbb{R}^{d}}\left( \alpha_{t}-\overline{\alpha_{t}}\right)^{2}\,\mathrm{d}\rho_{t}\mathrm{d}t,$ (35) $$

linking (an upper bound of) the Wasserstein-2 distance $W_{2}^{2}(\rho_{0},\rho_{1})$ on the one side, and the Fisher-Rao action on the other.

Next we will define the appropriate spaces for our analysis.

### A.2 Prerequisites for solving the Poisson equation

###### Definition A.1 .

Let $t\mapsto\rho_{t}$ be a weakly continuous curve $[0,1]\to\mathcal{P}^{\rm ac}(\mathbb{R}^{d})$. Then, for all Borel measurable $B\subseteq\mathbb{R}^{d}$, the map $t\mapsto\rho_{t}(B)$ is measurable: Indeed, it is the pointwise limit of continuous, hence measurable maps
$t\mapsto\int_{\mathbb{R}^{d}}f_{k}\,\mathrm{d}\rho_{t}$ with appropriate $f_{k}$. Thus, $\boldsymbol{\rho}$ defined by the disintegration $\int_{[0,1]\times\mathbb{R}^{d}}f(t,x)\,\mathrm{d}\boldsymbol{\rho}=\int_{0}^{
1}\int_{\mathbb{R}^{d}}f(t,x)\,\mathrm{d}\rho_{t}\mathrm{d}t$, is a probability measure on $[0,1]\times\mathbb{R}^{d}$, see e.g., (Ambrosio et al., , Chapter 5.3).

###### Definition A.2 .

Let $\rho\,\mathrm{d}x\in\mathcal{P}^{\rm ac}(\mathbb{R}^{d})$ with density $\rho$. We consider the Sobolev space

$$ $\displaystyle H^{1}(\rho):=\Big{\{}f\in L_{2}(\mathbb{R}^{d},\rho):\text{ the weak derivative }\partial_{x_{i}}f\text{ (w.r.t. $\mathcal{L}$) exists and }$ (36) $\displaystyle\partial_{x_{i}}f\in L_{2}(\rho),~{}\forall i\in\{1,\ldots,d\} \Big{\}}$ (37) $$

with inner product $\langle f,g\rangle_{H^{1}(\rho)}:=\langle f,g\rangle_{L_{2}(\rho)}+\langle
\nabla f,\nabla g\rangle_{L_{2}(\rho,\mathbb{R}^{d})}$.
Furthermore, define the space of Sobolev functions with zero mean by

$$ $\displaystyle H_{0}^{1}(\rho):=\left\{f\in H^{1}(\rho):\int_{\mathbb{R}^{d}}f \,\mathrm{d}\rho=0\right\}.$ (38) $$

###### Lemma A.3 .

Assume that $\rho\,\mathrm{d}x\in\mathcal{P}^{\rm ac}(\mathbb{R}^{d})$ and $\frac{1}{\rho}\in L_{1,{\rm loc}}(\mathbb{R}^{d})$. Then, $(H^{1}(\rho),\langle\cdot,\cdot\rangle_{H^{1}(\rho)})$ is a Hilbert space and $H_{0}^{1}(\rho)$ is a closed subspace.

###### Proof.

The first claim is (Kufner & Opic, , Theorem 1.11).
For the second claim note that for $f_{n}\to f$ in $H^{1}(\rho)$ with $f_{n}\in H^{1}_{0}(\rho)$, we have that

$$ $\displaystyle\left|\int_{\mathbb{R}^{d}}f\,\mathrm{d}\rho\right|$ $\displaystyle=|\langle f,1\rangle_{L_{2}(\rho)}|\leq|\langle f-f_{n},1\rangle_ {L_{2}(\rho)}|+|\langle f_{n},1\rangle_{L_{2}(\rho)}|$ (39) $\displaystyle\leq\|f_{n}-f\|_{H^{1}(\rho)}+0\to 0.$ (40) $$

∎

###### Definition A.4 .

Consider a weakly continuous curve $[0,1]\to\mathcal{P}^{\rm ac}(\mathbb{R}^{d}),~{}t\mapsto\rho_{t}$, and the associated measure $\boldsymbol{\rho}\in\mathcal{P}^{\rm ac}([0,1]\times\mathbb{R}^{d})$ given by Definition A.1.
Define

$$ $\displaystyle H_{x}^{1}(\boldsymbol{\rho})\coloneqq\Big{\{}f\in L_{2}([0,1] \times\mathbb{R}^{d},\boldsymbol{\rho}):\text{ the weak derivative }\partial_{ x_{i}}f\text{ exists and }$ (41) $\displaystyle\partial_{x_{i}}f\in L_{2}(\boldsymbol{\rho}),~{}\forall i\in\{2, \ldots,d+1\}\Big{\}}$ (42) $$

with respect to the inner product $\langle f,g\rangle_{H_{x}^{1}(\boldsymbol{\rho})}:=\langle f,g\rangle_{L_{2}(
\boldsymbol{\rho})}+\langle\nabla_{x}f,\nabla_{x}g\rangle_{L_{2}(\boldsymbol{
\rho})}$, where $f=f(t,x)$, $g=g(t,x)$.
Now, consider the subspace

$$ $\displaystyle{\mathcal{X}_{0}^{1}}:=\Big{\{}s\in H_{x}^{1}(\boldsymbol{\rho}): s(t,\cdot)\in H_{0}^{1}(\rho_{t})\ \text{ for almost every}\ t\in[0,1]\Big{\}}.$ (43) $$

For $s\in{\mathcal{X}_{0}^{1}}$ we will also write $s_{t}$ for denoting the function $s(t,\cdot)$.

###### Lemma A.5 .

Assume that $\frac{1}{\boldsymbol{\rho}}\in L_{1,{\rm loc}}([0,1]\times\mathbb{R}^{d})$. Then, the space $H^{1}_{x}(\boldsymbol{\rho})$ is a Hilbert space, and ${\mathcal{X}_{0}^{1}}$ is a closed subspace of $H_{x}^{1}(\boldsymbol{\rho})$.

###### Proof.

The Hilbert space property of $H^{1}_{x}(\boldsymbol{\rho})$ again follows by (Kufner & Opic, , Theorem 1.11) and we only prove the closedness of ${\mathcal{X}_{0}^{1}}$.
Now, let $(s_{n})\subset{\mathcal{X}_{0}^{1}}$ be a sequence with $s_{n}\to s$ in $H^{1}_{x}(\boldsymbol{\rho})$. Then,

$$ $\displaystyle\int_{0}^{1}\int_{\mathbb{R}^{d}}|s_{n}(t,\cdot)-s(t,\cdot)|^{2}+ \|\nabla_{x}s_{n}(t,\cdot)-\nabla_{x}s(t,\cdot)\|^{2}\,\mathrm{d}\rho_{t} \mathrm{d}t\to 0.$ $$

This implies pointwise convergence for a.e. $t\in[0,1]$ (along a subsequence), i.e.,

$$ $\displaystyle\int_{\mathbb{R}^{d}}|s_{n}(t,\cdot)-s(t,\cdot)|^{2}+\|\nabla_{x} s_{n}(t,\cdot)-\nabla_{x}s(t,\cdot)\|^{2}\,\mathrm{d}\rho_{t}\to 0\quad\text{ for a.e. }t\in[0,1].$ $$

Thus, we have that $s_{n}(t,\cdot)\to s(t,\cdot)$ converges in $H^{1}(\rho_{t})$ for almost all $t\in[0,1]$. Since $H_{0}^{1}(\rho_{t})$ is closed, we have that the limit $s(t,\cdot)$ is in $H_{0}^{1}(\rho_{t})$ for a.e. $t$.
∎

###### Definition A.6 .

On ${\mathcal{X}_{0}^{1}}$ we define an bilinear form by setting

$$ $\langle f,g\rangle_{{\mathcal{X}_{0}^{1}}}=\langle\nabla_{x}f,\nabla_{x}g \rangle_{L_{2}(\boldsymbol{\rho})}.$ $$

###### Definition A.7 .

We say that a probability measure $\boldsymbol{\rho}$ fulfills the partial Poincare inequality (PPI) if there exists $K>0$ such that

$$ $\displaystyle\|g\|^{2}_{L_{2}(\boldsymbol{\rho})}\leq K\|\nabla_{x}g\|^{2}_{L_ {2}(\boldsymbol{\rho})}$ (PPI) $$

for all $g\in{\mathcal{X}_{0}^{1}}$. Note that usually it is formulated as $Var_{{\mu}}[g]\leq K\|\nabla_{x}g\|^{2}_{L_{2}({\mu})}$, but since we have $\mathbb{E}_{{\rho_{t}}}[g_{t}]=0$ for $g\in{\mathcal{X}_{0}^{1}}$ both notions coincide.

###### Lemma A.8 .

Assume that $\boldsymbol{\rho}$ fulfills the (PPI), and that $\frac{1}{\boldsymbol{\rho}}\in L_{1,{\rm loc}}([0,1]\times\mathbb{R}^{d})$. Then, $({\mathcal{X}_{0}^{1}},\langle\cdot,\cdot\rangle_{{\mathcal{X}_{0}^{1}}})$ is a Hilbert space.

###### Proof.

We have to show that $\langle\cdot,\cdot\rangle_{{\mathcal{X}_{0}^{1}}}$ is positive definite and ${\mathcal{X}_{0}^{1}}$ is complete. The first claim follows directly from the (PPI). The second follows from the fact that $\|\cdot\|_{H^{1}_{x}(\boldsymbol{\rho})}\leq\sqrt{K+1}\|\cdot\|_{{\mathcal{X}_
{0}^{1}}}\leq\sqrt{K+1}\|\cdot\|_{H^{1}_{x}(\boldsymbol{\rho})}$ and ${\mathcal{X}_{0}^{1}}$ is closed in $H^{1}_{x}(\boldsymbol{\rho})$.
∎

This assumption (PPI) is usually fulfilled in our setting as we discuss in the following.

###### Lemma A.9 .

(Bakry et al., , Corollary 1.6) Let $\mu=e^{-V(x)}\,\mathrm{d}x$ be a probability measure on $\mathbb{R}^{d}$ and assume that $V\in C^{2}(\mathbb{R}^{d})$ such that $V$ is lower bounded. Assume furthermore that there exists $\eta>0,\tilde{R}\geq 0$ such that for all $\|x\|>\tilde{R}$ we have

$$ $\langle x,\nabla V(x)\rangle\geq\eta\|x\|.$ $$

Then, there exists a constant $C=C\big{(}d,\eta,\|V\|_{L^{\infty}(B(R))},\|\nabla V\|_{L^{\infty}(B(R))}\big{
)}>0$, where $R:=\max\{\frac{2d}{\eta},\tilde{R}\}$, such that for all $\phi\in H^{1}_{0}(\mu)$, it holds

$$ $\displaystyle\int_{\mathbb{R}^{d}}\phi^{2}(x)\,\mathrm{d}\mu\leq C\int_{ \mathbb{R}^{d}}\|\nabla\phi\|^{2}\,\mathrm{d}\mu.$ (44) $$

The constant $C$ can be estimated as follows: Let $\beta:=\beta_{V,R}:=2\|V\|_{L^{\infty}(B(R))}$. There exist $C_{1,\eta,R},C_{2,\eta,R},C_{3,\eta,R}>0$ (not depending on $V$) and a universal constant $D>0$ such that

$$ $\displaystyle C\leq\frac{4d}{\eta^{2}}\Big{(}1+\left(C_{1,\eta,R}+C_{2,\eta,R} \|\nabla V\|_{L^{\infty}(B(R))}+C_{3,\eta,R}\right)D\,R^{2}e^{\beta}\Big{)}.$ (45) $$

###### Proof.

We first show the claim for all smooth $\phi\in C^{\infty}(\mathbb{R}^{d})\cap H^{1}(\mu)$ with zero mean.
Unfortunately in (Bakry et al., , Corollary 1.6), the constants are not explicitly computed which is why we track them here. They consider for $\gamma=\eta/2$ the function given by $W(x)=W_{\eta,R}(x):=e^{\gamma\|x\|}$ for $\|x\|>R$, and smooth continuation for $\|x\|\leq R$. They show that, if for $L:=\Delta-\langle\nabla V,\nabla\cdot\rangle$, it holds that

$$ $LW\leq-\theta W(x)+b\chi_{B(R)},$ $$

then, the inequality $Var_{\mu}[\phi]\leq C\int_{\mathbb{R}^{d}}\|\nabla\phi^{2}\|\,\mathrm{d}\mu$ is true with constant $C=\frac{1}{\theta}(1+b\kappa_{R})$, where the constant $\kappa_{R}$ can be estimated by $D\,R^{2}e^{\beta}$. We are left to upper bound the constants $\theta$ and $b$. Let $\|x\|\leq R$. Then, we have that

$$ $|LW(x)|\leq|\Delta W(x)|+|\langle\nabla V(x),\nabla W(x)|\leq\|\Delta W\|_{L^{ \infty}(B(R))}+\|\nabla V\|_{L^{\infty}(B(R))}\|\nabla W\|_{L^{\infty}(B(R))},$ (46) $$

and for $\|x\|>R$, by choice of $R$,

$$ $\displaystyle LW(x)$ $\displaystyle=\frac{\eta}{2}\left(\frac{d-1}{\|x\|}+\frac{\eta}{2}-\frac{ \langle x,\nabla V\rangle}{\|x\|}\right)W(x)\leq\frac{\eta}{2}\left(\frac{d-1} {R}+\frac{\eta}{2}-\eta\right)W(x)$ (47) $\displaystyle\leq\frac{\eta}{2}\left(\frac{\eta(d-1)}{2d}-\frac{\eta d}{2d} \right)W(x)=-\frac{\eta^{2}}{4d}W(x),$ (48) $$

and thus the claim is true for $\phi\in C^{\infty}(\mathbb{R}^{d})\cap H^{1}(\mu)$ with zero mean, using the constants $\theta=\frac{\eta^{2}}{4d}>0$ and $b=C_{1,\eta,R}+C_{2,\eta,R}\|\nabla V\|_{L^{\infty}(B(R))}+C_{3,\eta,R}$, where $C_{1,\eta,R}:=\|\Delta W_{\eta,R}\|_{L^{\infty}(B(R))}$, $C_{2,\eta,R}:=\|\nabla W_{\eta,R}\|_{L^{\infty}(B(R))}$ and
$C_{3,\eta,R}:=\theta\|W_{\eta,R}\|_{L^{\infty}(B(R))}$.
In order to extend this result to $\phi\in H^{1}_{0}(\mu)$, note that $C_{c}^{\infty}(\mathbb{R}^{d})$ is dense in $H^{1}(\mathbb{R}^{d})$, see (Adams & Fournier, , Cor. 3.23). Since by assumption, $V$ is lower bounded, i.e., $e^{-V}$ is upper bounded, $C_{c}^{\infty}(\mathbb{R}^{d})$ is also dense in $H^{1}(\mu)$ by Lebesgue’s dominated convergence theorem. Now take a sequence $(\phi_{n})\in C_{c}^{\infty}(\mathbb{R}^{d})$ converging to $\phi$ in $H^{1}(\mu)$. Then, also $\tilde{\phi_{n}}:=\phi_{n}-\int_{\mathbb{R}^{d}}\phi_{n}\,\mathrm{d}\mu$ converges to $\phi$ in $H^{1}(\mu)$, and $\tilde{\phi_{n}}\in C^{\infty}(\mathbb{R}^{d})\cap H^{1}(\mu)$ has zero mean. The first part now yields the claim.
∎

For the specific path $\rho_{t}$ given by the *linear* potential interpolation $V_{t}=(1-t)f_{0}+tf_{1}$, the disintegration measure $\boldsymbol{\rho}$ fulfills the partial Poincare inequality (PPI) given suitable assumptions on the end points $f_{0},f_{1}$, as the next proposition shows.

###### Proposition A.10 .

Let $f_{0},f_{1}\in C^{2}(\mathbb{R}^{d})$ be lower bounded such that there exists $\alpha,\eta>0$ and $\tilde{R}>0$ such that for all $\|x\|\geq\tilde{R}$ it holds the *drift condition*

$$ $\langle x,\nabla f_{i}(x)\rangle\geq\eta\|x\|,\quad i=0,1,$ (49) $$

and the *linear growth condition*

$$ $f_{i}(x)\geq\alpha\|x\|,\quad i=0,1.$ (50) $$

Let $V_{t}=(1-t)f_{0}+tf_{1}$ and $\rho_{t}=Z_{t}^{-1}e^{-V_{t}}$ for $Z_{t}=\int_{\mathbb{R}^{d}}e^{-V_{t}}\,\mathrm{d}x$, and $\boldsymbol{\rho}$ as in Definition A.1. Then, $\boldsymbol{\rho}$ satisfies the partial Poincare inequality (PPI), i.e.,
there exists $C>0$ such that for all $\phi\in{\mathcal{X}_{0}^{1}}$ it holds

$$ $\displaystyle\int_{[0,1]\times\mathbb{R}^{d}}\phi^{2}\,\mathrm{d}\boldsymbol{ \rho}\leq C\int_{[0,1]\times\mathbb{R}^{d}}\|\nabla_{x}\phi\|^{2}\,\mathrm{d} \boldsymbol{\rho}.$ (51) $$

###### Proof.

First note that by the lower boundedness and linear growth assumption (50) on $f_{0},f_{1}$, the normalization constant $Z_{t}$ is well-defined, and $\rho_{t}$ is indeed a probability measure for all $t\in[0,1]$.
Further note that $\boldsymbol{\rho}$ exists if $\rho_{t}$ is weakly continuous, see again Definition A.1. We now show the weak continuity of $\rho_{t}$:
Let $g$ be a bounded and continuous function on $\mathbb{R}^{d}$. By the linear growth (50) of $f_{0},f_{1}$ and dominated convergence, the map $t\mapsto Z_{t}=\int_{\mathbb{R}^{d}}e^{-V_{t}}\,\mathrm{d}x$ is continuous on $[0,1]$. Since it is also strictly positive on $[0,1]$, also its inverse $t\mapsto Z_{t}^{-1}$ is continuous on $[0,1]$. By the same token, the map $t\mapsto\int_{\mathbb{R}^{d}}ge^{-V_{t}}\,\mathrm{d}x$ is continuous, and altogether, $t\mapsto\rho_{t}$ is weakly continuous.
In order to prove the inequality, we conclude by (49) that $\langle x,\nabla V_{t}(x)\rangle\geq(1-t)\eta\|x\|+t\eta\|x\|\geq\eta\|x\|$ for all $\|x\|\geq\tilde{R}$. Let $R:=\max\{\frac{2d}{\eta},\tilde{R}\}$ and note that for $V(t,x):=V_{t}(x)+\ln(Z_{t})$, it is now easy to check that $V(t,\cdot)\in C^{2}(\mathbb{R}^{d})$, $\|\nabla_{x}V\|_{L^{\infty}([0,1]\times B(R))}<\infty$ and $\beta:=2\|V\|_{L^{\infty}([0,1]\times B(R))}<\infty$.
Then, by Proposition A.9 we know that

$$ $\int_{\mathbb{R}^{d}}\phi(t,\cdot)^{2}\,\mathrm{d}\rho_{t}\leq C\int_{\mathbb{ R}^{d}}\|\nabla_{x}\phi(t,\cdot)\|^{2}\,\mathrm{d}\rho_{t}$ $$

for all $\phi\in{\mathcal{X}_{0}^{1}}$ and all $t\in[0,1]$, with a constant $C$ which does *not* depend on $t$, hence the claim.
∎

###### Remark A.11 .

(i) Gaussian distributions given by $f_{i}:=\frac{1}{2\sigma^{2}}\|x-m\|^{2}$ satisfy the assumptions of Proposition A.10.
(ii)
In Proposition A.10, the regularity assumptions on the end points, i.e., $f_{0},f_{1}\in C^{2}(\mathbb{R}^{d})$ can be relaxed. It is enough to assume that there exist sequences of functions $(f_{0,n}),(f_{1,n})\subset C^{2}(\mathbb{R}^{d})$, which converge pointwise to $f_{0}$ and $f_{1}$, respectively, such that the assumptions of Proposition A.10 are satisfied uniformly in $n$ and such that $\|\nabla_{x}V_{n}\|_{L^{\infty}},\|V_{n}\|_{L^{\infty}}$ are bounded in $n$.
Hereby, also the Fisher-Rao curve given by linear interpolation of $f_{0}:=\|x\|$ and $f_{1}:=\min\{\|x\|,\|x-m\|\}$ with $m\in\mathbb{R}^{d}$, i.e., the interpolation between two Laplacian (type) distributions, satisfies (PPI).

###### Definition A.1 .

###### Definition A.2 .

###### Lemma A.3 .

###### Proof.

###### Definition A.4 .

###### Lemma A.5 .

###### Proof.

###### Definition A.6 .

###### Definition A.7 .

###### Lemma A.8 .

###### Proof.

###### Lemma A.9 .

###### Proof.

###### Proposition A.10 .

###### Proof.

###### Remark A.11 .

### A.3 Existence of a solution to the Poisson Equation

Now, we consider the Poisson equation (34) in a slightly more abstract manner:
Showing that a given Fisher-Rao flow $\rho_{t}$ is also a Wasserstein absolutely continuous curve requires solving the following (family of) PDEs

$$ $\displaystyle-\nabla\cdot\left(\rho_{t}\nabla s_{t}\right)=\left(\alpha_{t}- \overline{\alpha_{t}}\right)\rho_{t},\quad t\in[0,1],$ (52) $$

for the vector field $v_{t}=\nabla s_{t}$ satisfying a suitable integrability condition.
Here, $\alpha\in L_{2}(\boldsymbol{\rho})$, $\alpha_{t}\coloneqq\alpha(t,\cdot)$ and $\overline{\alpha_{t}}\coloneqq\mathbb{E}_{\rho_{t}}[\alpha_{t}]$ are given.

In order to do that, we look to find a weak solution to (52) in the space ${\mathcal{X}_{0}^{1}}$.

###### Definition A.12 .

An element $\phi\in{\mathcal{X}_{0}^{1}}$ is called a weak solution to (52) if it satisfies,

$$ $\displaystyle\int_{0}^{1}\int_{\mathbb{R}^{d}}\langle\nabla\phi_{t},\nabla\psi _{t}\rangle\,\mathrm{d}\rho_{t}\mathrm{d}t=\int_{0}^{1}\int_{\mathbb{R}^{d}} \psi_{t}\left(\alpha_{t}-\overline{\alpha_{t}}\right)\,\mathrm{d}\rho_{t} \mathrm{d}t,\quad\forall\psi\in{\mathcal{X}_{0}^{1}}.$ (53) $$

Note that for $\psi\in C_{c}^{\infty}((0,1)\times\mathbb{R}^{d})$ it holds for $g(t)\coloneqq\mathbb{E}_{\rho_{t}}[\psi(t,\cdot)]$ that $\psi-g\in{\mathcal{X}_{0}^{1}}$. Since $\nabla_{x}(g)=0=\int_{0}^{1}\int_{\mathbb{R}^{d}}g(t)(\alpha_{t}-\overline{
\alpha_{t}})\,\mathrm{d}\rho_{t}\mathrm{d}t$, we have that (53) is also fulfilled for all $\psi\in C_{c}^{\infty}((0,1)\times\mathbb{R}^{d})$.

We can now prove Theorem 1 from Section 2.2, which provides a solution to the Poisson problem (52) globally in time $t\in[0,1]$.

###### Theorem A.13 .

Assume that $\boldsymbol{\rho}$ from Definition A.1 satisfies (PPI) for some $K>0$, and that $\frac{1}{\boldsymbol{\rho}}\in L_{1,{\rm loc}}([0,1]\times\mathbb{R}^{d})$. Furthermore, assume that $\alpha_{t}-\overline{\alpha_{t}}\in L_{2}(\boldsymbol{\rho})$, i.e., $\|\alpha_{t}-\overline{\alpha_{t}}\|_{L_{2}(\boldsymbol{\rho})}^{2}\leq C<\infty$.
Then, there exists a unique weak solution $s\in{\mathcal{X}_{0}^{1}}$ to (52). Moreover, $s$ satisfies the following inequality

$$ $\displaystyle\int_{0}^{1}\int_{\mathbb{R}^{d}}\|\nabla s\|^{2}\,\mathrm{d}\rho _{t}\mathrm{d}t\leq K\int_{0}^{1}\int_{\mathbb{R}^{d}}\left(\alpha_{t}- \overline{\alpha_{t}}\right)^{2}\,\mathrm{d}\rho_{t}\mathrm{d}t.$ (54) $$

###### Proof.

Let $A:{\mathcal{X}_{0}^{1}}\to\mathbb{R}$ be defined as

$$ $\displaystyle A\psi:=\int_{0}^{1}\int_{\mathbb{R}^{d}}\psi_{t}\left(\alpha_{t} -\overline{\alpha_{t}}\right)\,\mathrm{d}\rho_{t}\mathrm{d}t.$ $$

Therefore, we can rewrite our system of PDEs (52) as

$$ $\displaystyle\langle s,\psi\rangle_{{\mathcal{X}_{0}^{1}}}=A\psi,\quad\forall \psi\in{\mathcal{X}_{0}^{1}}.$ (55) $$

Clearly, $A$ is a linear operator. Since $\boldsymbol{\rho}$ satisfies (PPI) the following holds:

$$ $\displaystyle\left|A\psi\right|^{2}$ $\displaystyle=\left|\int_{0}^{1}\int_{\mathbb{R}^{d}}\psi_{t}\left(\alpha_{t}- \overline{\alpha_{t}}\right)\,\mathrm{d}\rho_{t}\mathrm{d}t\right|^{2}=\left| \int_{[0,1]\times\mathbb{R}^{d}}\psi_{t}(x)\left(\alpha_{t}(x)-\overline{ \alpha_{t}}(x)\right)\,\mathrm{d}\boldsymbol{\rho}(t,x)\right|^{2}$ $\displaystyle\leq\Big{(}\int\left(\alpha_{t}-\overline{\alpha_{t}}\right)^{2} \,\mathrm{d}\boldsymbol{\rho}\Big{)}\cdot\Big{(}\int|\psi_{t}|^{2}\,\mathrm{d} \boldsymbol{\rho}\Big{)}\leq C\int_{0}^{1}\int_{\mathbb{R}^{d}}|\psi_{t}|^{2} \,\mathrm{d}\rho_{t}\mathrm{d}t$ $\displaystyle\leq C\,K\int_{0}^{1}\int_{\mathbb{R}^{d}}\|\nabla\psi_{t}\|^{2} \,\mathrm{d}\rho_{t}\mathrm{d}t=C\,K\|\psi\|^{2}_{{\mathcal{X}_{0}^{1}}}.$ $$

By Lemma A.8, the space ${\mathcal{X}_{0}^{1}}$ is a Hilbert space. Applying the Riesz Representation theorem, there exists a unique element $s\in{\mathcal{X}_{0}^{1}}$ satisfying (55).
∎

###### Remark A.14 .

Note that in the case of Fisher-Rao flows, $\alpha$ satisfies

$$ $\displaystyle\partial_{t}\rho_{t}=\left(\alpha_{t}-\overline{\alpha_{t}}\right )\rho_{t},\quad t\in[0,1].$ $$

Then, the Fisher-Rao action $\mathcal{A}_{\rm FR}(\rho_{t})$ of $\rho_{t}$ is given by $\int_{0}^{1}\|\alpha_{t}-\overline{\alpha_{t}}\|_{L_{2}(\rho_{t})}^{2}\,
\mathrm{d}t$.
Hence, the r.h.s. of (54) can be described by $K\,\mathcal{A}_{\rm FR}(\rho_{t})$.

###### Lemma A.15 .

Let $\boldsymbol{\rho}$ be such that $\rho_{t}$ is bounded (from above) for a.e. $t\in[0,1]$.
Then, for $s\in{\mathcal{X}_{0}^{1}}$ we have that $\nabla s_{t}\in T_{\rho_{t}}\mathcal{P}_{2}(\mathbb{R}^{d})$ for a.e. $t\in[0,1]$.

###### Proof.

For a.e. $t$, since $\rho_{t}$ is bounded from above by assumption, $C_{c}^{\infty}(\mathbb{R}^{d})$ is dense in $H^{1}(\rho_{t})$. Now, $s_{t}\in H^{1}(\rho_{t})$ immediately yields the claim
$\nabla s_{t}\in T_{\rho_{t}}\mathcal{P}_{2}(\mathbb{R}^{d})=\overline{\{\nabla
\phi:\phi\in C_{c}^{\infty}(\mathbb{R}^{d})\}}^{L_{2}(\rho_{t})}$.
∎

It follows the proof of Theorem 2 from Section 2.2, which says that certain Fisher-Rao curves are indeed also Wasserstein absolutely continuous curves.

###### Theorem A.16 .

Let $\rho_{0},\rho_{1}$ be probability densities on $\mathbb{R}^{d}$ and assume that $\rho_{t}$ is a Fisher-Rao curve defined by

$$ $\partial_{t}\rho_{t}=(\alpha_{t}-\overline{\alpha_{t}})\rho_{t},$ $$

such that $\boldsymbol{\rho}$ and $\alpha$ satisfy the assumptions of Theorem A.13.

Then, $\rho_{t}$ is a Wasserstein absolutely continuous curve with vector field $\nabla_{x}s$, where $s\in{\mathcal{X}_{0}^{1}}$ is the unique weak solution of (53). If $\rho_{t}$ is bounded for a.e. $t$, we have that $\nabla_{x}s_{t}\in T_{\rho_{t}}\mathcal{P}_{2}(\mathbb{R}^{d})$ for a.e. $t\in[0,1]$, i.e., for every other vector field $v_{t}$ satisfying the continuity equation for $\rho_{t}$, we have that $\|v_{t}\|_{L_{2}(\rho_{t})}\geq\|\nabla_{x}s_{t}\|_{L_{2}(\rho_{t})}$ for a.e. $t\in[0,1]$.

Furthermore, it holds that

$$ $\displaystyle\int_{0}^{1}\int_{\mathbb{R}^{d}}\|\nabla_{x}s_{t}\|^{2}\,\mathrm {d}\rho_{t}\mathrm{d}t\leq K\int_{0}^{1}\int_{\mathbb{R}^{d}}\left(\alpha_{t}- \overline{\alpha_{t}}\right)^{2}\,\mathrm{d}\rho_{t}\mathrm{d}t.$ (56) $$

Note that by the Benamou-Brenier formula, see (Ambrosio et al., , Equation 8.0.3), the left-hand side of (56) is an upper bound for the Wasserstein distance $W_{2}^{2}(\rho_{0},\rho_{1})$, while the double-integral on the right-hand side defines the Fisher-Rao action $\mathcal{A}_{\rm FR}$ of $\rho_{t}$.

###### Proof.

Using $(\alpha_{t}-\overline{\alpha_{t}})\rho_{t}=\partial_{t}\rho_{t}$ and Theorem A.13, we can conclude that $\rho_{t}$ satisfies the continuity equation

$$ $\partial_{t}\rho_{t}+\nabla\cdot(\rho_{t}v_{t})=0$ (57) $$

with a Borel vector field $v_{t}\coloneqq\nabla s_{t}$ in the sense of distributions.
Note that by construction, $s\in{\mathcal{X}_{0}^{1}}$, and hence, the gradient $\nabla_{x}s$ is already a Borel-measurable vector field on $[0,1]\times\mathbb{R}^{d}$ such that $\nabla_{x}s_{t}\in L_{2}(\rho_{t})$ for a.e. $t\in[0,1]$.
Further, (54) implies that $t\mapsto\|\nabla s_{t}\|_{L_{2}(\rho_{t})}\in L^{1}([0,T])$.

Together with the fact that $\rho_{t}$ is weakly continuous, we infer by (Ambrosio et al., , Theorem 8.3.1) that $\rho_{t}$ is Wasserstein absolutely continuous. By Lemma A.15, it holds that $\nabla s_{t}\in T_{\rho_{t}}\mathcal{P}_{2}(\mathbb{R}^{d})$, if $\rho_{t}$ is bounded for a.e. $t$. In this case, (Ambrosio et al., , Proposition 8.4.5) immediately implies that $\nabla s_{t}$ has minimal norm $\|\nabla s_{t}\|_{L_{2}(\rho_{t})}$ among all Borel vector fields $v_{t}$ fulfilling (57).
∎

Let us apply the above result to the curve given by linear interpolation of the energy functions $f_{0},f_{1}$.

###### Corollary A.17 .

Let $f_{0},f_{1}\in C^{2}(\mathbb{R}^{d})$ be lower bounded by some $\gamma\in\mathbb{R}$ such that the conditions (49) and (50) of Propostion A.10 are satisfied for some $\alpha,\eta>0$. Further, for simplicity, assume that $f_{0},f_{1}$ have at most polynomial growth outside a ball, i.e., there exists $R>0$ such that

$$ $|f_{i}(x)|\leq M\|x\|^{k}\quad\text{for all }\|x\|>R,~{}i=0,1.$ (58) $$

Then, Theorem A.16 can be applied to the Fisher-Rao curve $\rho_{t}$ given by $V_{t}=(1-t)f_{0}+tf_{1}$ and $\rho_{t}=Z_{t}^{-1}e^{-V_{t}}$ with $Z_{t}=\int_{\mathbb{R}^{d}}e^{-V_{t}}\,\mathrm{d}x$.

###### Proof.

First, by the drift condition (49) and linear growth (50), the measure $\boldsymbol{\rho}$ satisfies (PPI) by Proposition A.10. Furthermore, in Proposition A.10 we showed that $t\mapsto Z_{t}^{-1}$ is continuous on $[0,1]$, hence $(t,x)\mapsto\rho_{t}(x)=Z_{t}^{-1}e^{-V_{t}(x)}$ is strictly positive and continuous on $[0,1]\times\mathbb{R}^{d}$. Therefore, it immediately follows that $\frac{1}{\boldsymbol{\rho}}\in L_{1,{\rm loc}}([0,1]\times\mathbb{R}^{d})$.

We are left to check the $L_{2}$-condition on $\alpha_{t}=-\partial_{t}V_{t}=f_{0}-f_{1}$. First, since $t\mapsto Z_{t}^{-1}$ is continuous on $[0,1]$, hence bounded, it follows together with (50) and (58) that

$$ $\displaystyle|\overline{\alpha_{t}}|$ $\displaystyle\leq Z_{t}^{-1}\int_{\mathbb{R}^{d}}|f_{0}-f_{1}|e^{-V_{t}}\, \mathrm{d}x$ (59) $\displaystyle\leq Z_{t}^{-1}\int_{\|x\|\leq R}|f_{0}-f_{1}|e^{-V_{t}}\,\mathrm {d}x~{}+Z_{t}^{-1}\int_{\|x\|>R}|f_{0}-f_{1}|e^{-V_{t}}\,\mathrm{d}x$ (60) $\displaystyle\leq C_{1}\int_{\|x\|\leq R}|f_{0}-f_{1}|e^{-\gamma}\,\mathrm{d}x ~{}+C_{2}\,M\int_{\|x\|>R}\|x\|^{k}e^{-\alpha\|x\|}\,\mathrm{d}x$ (61) $\displaystyle\leq C<\infty,$ (62) $$

with a constant $C>0$ not depending on $t$. Hence, $\overline{\alpha_{t}}\in L_{2}(\boldsymbol{\rho})$, and we are left to prove ${\alpha_{t}}\in L_{2}(\boldsymbol{\rho})$. But this follows by the similar arguments

$$ $\displaystyle Z_{t}^{-1}\int_{\mathbb{R}^{d}}|f_{0}-f_{1}|^{2}e^{-V_{t}}\, \mathrm{d}x$ $\displaystyle\leq Z_{t}^{-1}\int_{\|x\|\leq R}|f_{0}-f_{1}|^{2}e^{-V_{t}}\, \mathrm{d}x~{}+Z_{t}^{-1}\int_{\|x\|>R}|f_{0}-f_{1}|^{2}e^{-V_{t}}\,\mathrm{d}x$ (63) $\displaystyle\leq C_{1}\int_{\|x\|\leq R}|f_{0}-f_{1}|^{2}e^{-\gamma}\,\mathrm {d}x~{}+C_{2}\,M^{2}\int_{\|x\|>R}\|x\|^{2k}e^{-\alpha\|x\|}\,\mathrm{d}x$ (64) $\displaystyle\leq C<\infty,$ (65) $$

again, with a constant $C>0$ not depending $t$. Integrating over $t\in[0,1]$ yields ${\alpha_{t}}\in L_{2}(\boldsymbol{\rho})$, and hence, all assumptions of Theorem A.16 are fulfilled.
∎

###### Remark A.18 .

Hence, the Fisher-Rao curves defined by linear interpolation $V_{t}=(1-t)f_{0}+tf_{1}$ are absolutely continuous in the Wasserstein sense, e.g., in the cases (see also Remark A.11)

- •
Gaussian distributions given by $f_{i}:=\frac{1}{2\sigma^{2}}\|x-m\|^{2}$,
- •
Laplacian distributions given by $f_{i}:=\frac{1}{\sigma}\|x-m\|$,
- •
Laplacian type distributions given by $f_{i}:=\min\{\|x\|,\|x-m\|\}$.

###### Definition A.12 .

###### Theorem A.13 .

###### Proof.

###### Remark A.14 .

###### Lemma A.15 .

###### Proof.

###### Theorem A.16 .

###### Proof.

###### Corollary A.17 .

###### Proof.

###### Remark A.18 .

### A.4 A teleportation issue with the linear interpolation

In this section we want to demonstrate that the norm $\|v_{t}\|_{L_{2}(\rho_{t})}$ of the (optimal) velocity field (5) belonging to the Fisher-Rao path given by linear interpolation $f_{t}:=(1-t)f_{0}+tf_{1}$, explodes for times $t$ close to $1$, when the measures $\rho_{0}$ and $\rho_{1}$ are asymmetrical to each other.

In the following, we restrict ourselves to the one-dimensional setting $\mathbb{R}^{1}$, in order to exploit the isometry from the Wasserstein space $\mathcal{P}_{2}(\mathbb{R}^{1})$ into the convex cone $\mathcal{C}\subset L_{2}(0,1)$ of *quantile functions*, also see (Duong et al., , Section 3). For the sake of self-containedness, we give a short introduction of the necessary tools:
For $\mu\in\mathcal{P}_{2}(\mathbb{R})$ we define its cumulative distribution function (CDF) by

$$ $F_{\mu}(x):=\mu(-\infty,x],\quad x\in\mathbb{R},$ (66) $$

and its left-continuous quantile function

$$ $Q_{\mu}(s):=\min\{x\in\mathbb{R}:F(x)\geq s\},\quad s\in(0,1).$ (67) $$

Then, it holds the isometric property

$$ $W_{2}^{2}(\mu,\nu)=\int_{0}^{1}|Q_{\mu}(s)-Q_{\nu}(s)|^{2}\,\mathrm{d}s\quad \text{for all }\mu,\nu\in\mathcal{P}_{2}(\mathbb{R}).$ (68) $$

Now, by (Ambrosio et al., , Theorem 8.3.1), the norm $\|v_{t}\|_{L_{2}(\rho_{t})}$ of the optimal field $v_{t}$ belonging to an absolutely continuous Wasserstein curve $\rho_{t}$ can be characterized by the *metric derivative* of the curve via

$$ $\|v_{t}\|_{L_{2}(\rho_{t})}=|\rho^{\prime}(t)|:=\lim_{s\to t}\frac{W_{2}(\rho_ {t},\rho_{s})}{|t-s|}\quad\text{for a.e. }t\in[0,1].$ (69) $$

Using the aforementioned isometry (68) and Komura’s theorem, it holds for a.e. $t\in[0,1]$

$$ $\lim_{s\to t}\frac{W_{2}(\rho_{t},\rho_{s})}{|t-s|}=\lim_{s\to t}\frac{\|Q_{ \rho_{t}}-Q_{\rho_{s}}\|_{L_{2}(0,1)}}{|t-s|}=\|\partial_{t}(Q_{\rho_{t}})\|_{ L_{2}(0,1)}.$ (70) $$

Hence, if we can explicitly compute the quantiles $Q_{\rho_{t}}$ (and its time derivative), we get an analytic representation of $\|v_{t}\|_{L_{2}(\rho_{t})}$ via

$$ $\|v_{t}\|_{L_{2}(\rho_{t})}^{2}=\int_{0}^{1}\Big{(}\partial_{t}(Q_{\rho_{t}})( s)\Big{)}^{2}\,\mathrm{d}s.$ (71) $$

In order to do so, we choose the prior and target measures the following way: Let $f_{0}(x):=|x|$ and $f_{1}(x):=2\min\{|x|,|x-m|\}$, where the second mode $m>0$ is fixed for the moment. By Corollary A.17 and Remark A.18, the Fisher-Rao curve $\rho_{t}$ given by linear interpolation of $f_{0},f_{1}$ is Wasserstein absolutely continuous. Now, it is straightforward (but tedious) to explicitly calculate the CDF and quantile function of $\rho_{t}$:

First, since $\rho_{t}\propto\rho_{0}^{1-t}\rho_{1}^{t}$ and $2\min\{|x|,|x-m|\}=|x|+|x-m|-\big{|}|x|-|x-m|\big{|}$, we obtain

$$ $\rho_{t}\propto e^{-|x|-t|x-m|+t||x|-|x-m||}\quad\text{for all }x\in\mathbb{R} ,~{}t\in[0,1],$ (72) $$

up to a normalization constant $Z_{t}$ which we calculate below.

#### Calculating the CDF of ρ t subscript 𝜌 𝑡 \rho_{t} italic_ρ start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT .

By dissolving the absolute values on suitable subintervals, we obtain for the CDF

$$ $Z_{t}\cdot F_{\rho_{t}}(x)~{}=~{}\frac{1}{1+t}e^{x(1+t)},\quad x\leq 0,$ (73) $$

$$ $Z_{t}\cdot F_{\rho_{t}}(x)~{}=~{}\frac{2}{1+t}-\frac{1}{1+t}e^{-x(1+t)},\quad 0 \leq x\leq\frac{m}{2},$ (74) $$

$$ $Z_{t}\cdot F_{\rho_{t}}(x)~{}=~{}\frac{2}{1+t}-\Big{(}\frac{1}{1+t}+\frac{1}{- 1+3t}\Big{)}e^{-\frac{m}{2}(1+t)}+\frac{1}{-1+3t}e^{x(-1+3t)-2tm},\quad\frac{m }{2}\leq x\leq m,$ (75) $$

$$ $\displaystyle Z_{t}\cdot F_{\rho_{t}}(x)~{}=~{}$ $\displaystyle\frac{2}{1+t}-\Big{(}\frac{1}{1+t}+\frac{1}{-1+3t}\Big{)}e^{- \frac{m}{2}(1+t)}$ (76) $\displaystyle+$ $\displaystyle\Big{(}\frac{1}{1+t}+\frac{1}{-1+3t}\Big{)}e^{-m(1-t)}-\frac{1}{1 +t}e^{-x(1+t)+2tm},\quad x\geq m,$ (77) $$

where the normalization constant is hence given by

$$ $\displaystyle Z_{t}=Z_{t}\cdot F_{\rho_{t}}(\infty)~{}=~{}$ $\displaystyle\frac{2}{1+t}-\Big{(}\frac{1}{1+t}+\frac{1}{-1+3t}\Big{)}e^{- \frac{m}{2}(1+t)}$ (78) $\displaystyle+$ $\displaystyle\Big{(}\frac{1}{1+t}+\frac{1}{-1+3t}\Big{)}e^{-m(1-t)}.$ (79) $$

#### Calculating the quantile function of ρ t subscript 𝜌 𝑡 \rho_{t} italic_ρ start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT .

Since $F_{\rho_{t}}(\cdot)$ is continuous and strictly increasing, its inverse is given by $F_{\rho_{t}}^{-1}=Q_{\rho_{t}}$. A simple calculation now yields

$$ $Q_{\rho_{t}}(s)~{}=~{}\frac{1}{1+t}\ln\big{(}Z_{t}(1+t)s\big{)},\quad 0<s\leq s _{1},$ (80) $$

$$ $Q_{\rho_{t}}(s)~{}=~{}-\frac{1}{1+t}\ln\big{(}-Z_{t}(1+t)s+2\big{)},\quad s_{1 }\leq s\leq s_{2},$ (81) $$

$$ $Q_{\rho_{t}}(s)~{}=~{}\frac{\ln\big{(}(-1+3t)(sZ_{t}-\frac{2}{1+t}+(\frac{1}{1 +t}+\frac{1}{-1+3t})e^{-\frac{m}{2}(1+t)})\big{)}+2tm}{-1+3t},\quad s_{2}\leq s \leq s_{3},$ (82) $$

$$ $Q_{\rho_{t}}(s)~{}=~{}\frac{-\ln\big{(}Z_{t}(1+t)(1-s)\big{)}+2tm}{1+t},\quad s _{3}\leq s<1.$ (83) $$

Here, the subintervals are given by

$$ $s_{1}=\frac{1}{Z_{t}(1+t)},$ (84) $$

$$ $s_{2}=\frac{1}{Z_{t}}\big{(}\frac{2}{1+t}-\frac{1}{1+t}e^{-\frac{m}{2}(1+t)} \big{)},$ (85) $$

$$ $s_{3}=\frac{1}{Z_{t}}\big{(}\frac{2}{1+t}-(\frac{1}{1+t}+\frac{1}{-1+3t})e^{- \frac{m}{2}(1+t)}+\frac{1}{-1+3t}e^{-m(1-t)}\big{)}.$ (86) $$

#### Calculating the time derivative ∂ t Q ρ t subscript 𝑡 subscript 𝑄 subscript 𝜌 𝑡 \partial_{t}Q_{\rho_{t}} ∂ start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT italic_Q start_POSTSUBSCRIPT italic_ρ start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT end_POSTSUBSCRIPT and the norm ‖ v t ‖ L 2 ⁢ ( ρ t ) 2 superscript subscript norm subscript 𝑣 𝑡 subscript 𝐿 2 subscript 𝜌 𝑡 2 \|v_{t}\|_{L_{2}(\rho_{t})}^{2} ∥ italic_v start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ∥ start_POSTSUBSCRIPT italic_L start_POSTSUBSCRIPT 2 end_POSTSUBSCRIPT ( italic_ρ start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) end_POSTSUBSCRIPT start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT .

Using standard differentiation rules, the time derivative of $Q_{\rho_{t}}$ can be computed separately on each subinterval $(s_{i},s_{i+1})$. Since the resulting derivatives become unbearably long, we only include them in the Supplementary Material. Nevertheless, they can be computed analytically as elementary functions, and be inserted into the desired integral (71). By numerical integration, we can finally approximate the norm $\|v_{t}\|_{L_{2}(\rho_{t})}^{2}$ arbitrarily close. The computed results are depicted in Figure 2 for different values of the mode $m>0$; roughly speaking, for late times $t\sim 1$, the norm of the velocity field grows *exponentially* with $m$.

On a side note, the integration of $\partial_{t}Q_{\rho_{t}}$ *only* on the subintervals $(s_{1},s_{2})$ and $(s_{2},s_{3})$ contributes to the explosion of $\|v_{t}\|_{L_{2}(\rho_{t})}^{2}$. Intuitively, this corresponds to the fact that mass gets shifted very abruptly *mainly* in the areas $x\in(0,\frac{m}{2})$ and $(\frac{m}{2},m)$ for late times $t$.

Figure: Figure 4: Evolution of probability densities $\rho_{t}$ for $m\in\{1,5,15,50\}$.
Refer to caption: extracted/5901341/imgs/teleportation/combined_density_row_m_1.png

### A.5 Additional Details on the Experiments

#### Implementation and evaluation details.

We implement the algorithms in Pytorch Paszke et al. (). For both experiments we use standard MLPs with ”swish” activation functions. For the linear and learned interpolations we parameterize the vector fields directly. The learned interpolation also uses the same network for $\psi$. The linear interpolation uses the same network for the time dependent function $C_{t}$. For the learned and gradient flow interpolations we implement $C_{t}$ using the same network with less parameters. The gradient flow interpolation also uses a smaller network for the time scheduling function $g$ as defined in Section 4. In case of normalized energy functions we can omit the parametrization of $C_{t}$ for the learned and gradient flow interpolations. We simulate the corresponding ODEs using the torchdiffeq Chen () package with the Runge-Kutta adaptive step size solver (”dopri5”). For the linear and learned interpolation we use $50$ time steps along which the loss is computed and the gradients are accumulated with a batch size of $256$. For the gradient flow interpolation we sample $4096$ particles at random uniform time points and therefore do not accumulate gradients. For the GMM example we train the linear and learned interpolations for $5*10^{4}$ iterations. For the many well example we train the linear interpolation for $5*10^{4}$ and the learned for $4*10^{4}$ iterations. The gradient flow interpolation is trained for $4*10^{5}$ iterations in both experiments.
We compute log weights of the generated samples by using the continuous change of variables formula and the corresponding implementation from Chen ().
We evaluate the methods by generating $5*10^{4}$ samples and log weight and computing the effective sample size, negative log likelihood and energy distance. We report mean and standard deviation over $10$ evaluation runs. The effective sample size is computed as in Midgley et al. (). We compute the energy distance using the GeomLoss library Feydy et al. (). For the target energy functions and plotting utilities we use the code provided in Midgley et al. ().

The gradient flow interpolation faced some difficulty in the many well experiment when, while sampling, particles entered regions of very low density. This could severely affect sampling times when using the adaptive step size solver. Therefore we used an Euler scheme and resampled a given batch if a NaN error is raised. This happens approximately for $0.01\%$ of the particles.

#### Choice of latent distribution.

The linear and learned interpolations depend significantly on the choice of initial distribution. In the following we visualize and compare numerically the results when starting in Gaussians with different standard deviations. Note that when using a small standard deviation not all modes are recovered. On the other hand choosing a very large standard deviation leads to the modes close to 0 0 not receiving enough mass.

Figure: Figure 5: Results for the linear interpolation for different values of $\sigma$.
Refer to caption: extracted/5901341/imgs/LINEAR_10.0.png

**Table 3: Comparison of effective sample size, negative log likelihood, and the energy distance for the linear interpolation with initial distribution being a Gaussian with standard deviation $\sigma$.**
|  | ESS $\uparrow$ | NLL $\downarrow$ | Energy Distance $\downarrow$ |
| --- | --- | --- | --- |
| $\sigma=10$ | 0.002 $\pm$ 0.003 | 7.071 $\pm$ 0.005 | 2.887 $\pm$ 0.021 |
| $\sigma=20$ | 0.063 $\pm$ 0.07 | 7.362 $\pm$ 0.005 | 0.592 $\pm$ 0.005 |
| $\sigma=30$ | 0.183$\pm$ 0.05 | 8.657 $\pm$ 0.012 | 0.247$\pm$ 0.016 |
| $\sigma=40$ | 0.092 $\pm$ 0.001 | 12.909 $\pm$ 0.023 | 0.554$\pm$ 0.008 |

Figure: Figure 6: Results for the learned interpolation for different values of $\sigma$.
Refer to caption: extracted/5901341/imgs/LEARNED_10.0.png

**Table 4: Comparison of effective sample size, negative log likelihood, and the energy distance for the learned interpolation with initial distribution being a Gaussian with standard deviation $\sigma$.**
|  | ESS $\uparrow$ | NLL $\downarrow$ | Energy Distance $\downarrow$ |
| --- | --- | --- | --- |
| $\sigma=10$ | 0.009 $\pm$ 0.011 | 7.275 $\pm$ 0.005 | 1.253 $\pm$ 0.017 |
| $\sigma=20$ | 0.605 $\pm$ 0.316 | 6.894$\pm$ 0.005 | 0.058$\pm$ 0.006 |
| $\sigma=30$ | 0.922$\pm$ 0.002 | 6.913 $\pm$ 0.003 | 0.121$\pm$ 0.009 |
| $\sigma=40$ | 0.139 $\pm$ 0.052 | 7.779 $\pm$ 0.012 | 0.531$\pm$ 0.014 |

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 4
questions: Am I correct in understanding that the linear and learned methods in your experiments are no different to the approach proposed in [5] ? thus from a methods standpoint the gradient flow interpolation is the novel algorithmic contribution ?
soundness: 3
strengths: The paper is rigorous and very well written and proposes a very sound approach for sampling from unnormalised densities based on the continuity equation, with strongly backed theoretical results and some modest numerical results; the paper also proposes a novel interpolation scheme (gradient flow interpolation) in the context of learning the vector field with promising numerical results.

In general I find the connections to Wasserstein gradient flows and the general formal machinery in the W_2 metric space to be a rather strong selling point of this work in contrast to prior works which explore learning the same vector field, during the review period I look forward to discussing/exploring this further with the authors, in order to improve the manuscript and highlight their contributions.
summary: This work contributions are twofold:

1. First, the work provides conditions under  which a probability flow admits a velocity field satisfying the continuity equation
2. Following from the theory this work then proposes learning the vector field via a PINN loss on the log of the continuity equation (PDE arising from taking the log of the probability flow) and explore several different interpolation schemes.

The work then presents some empirical results motivating the success of their proposed method.
weaknesses: Unfortunately, this work misses prior work, which already explores learning the exact same vector field / continuity equation in the exact same context/task.

1.  Prior work [1]  (ICLR 2024 , ArXived July  2023)  provides a family of objectives for learning the vector field of the continuity equation in the exact same continuity equation + vector field up to a divergence. Notice [1] also provides an existence result for the vector field in the context of sampling from boltzman distributions (See Proposition 3.2 Appendix D3), the authors should make it clear how their results are different (from what I can see this submission has stronger / more general results, however they should be more precise in comparing to prior work). Note that whilst [1] has formulated their work in the SDE setting it is easy to see that their SDE in Equation 21 satisfies the same continuity equation as your work (See Equation 50 in their appendix) when using the optimal drift / velocity field considered in both works. Finally note that concurrent work [3] (Appendix 5.6 ) provides an explicit connection between [1] and your PINN based objective 
2. Prior Work [2]  (workshop version [4] Published on March 2024 around 5 months before the ICLR deadline https://openreview.net/pdf?id=KwHPBIGkET) explores the exact same objective for learning (They call it ODE anneal, see [3] equation 26 for a more clear rewriting of the PINN objective in [2,4]). As you highlight [5] also explores this loss

So, unfortunately, I would say that works [1,2,4] cannot be deemed as concurrent and have already explored NN-based learning of the vector field in the continuity equation for sampling from Boltzmann densities.  Therefore, this work needs to both conceptually and empirically discuss/compare these prior / non-concurrent works as you have done with [5]. Finally, [2,4] have explored the **exact same PINN-based objective** for learning the vector field (and much more thoroughly so from an empirical standpoint), so I'm not sure from a methods standpoint what contribution on the methods side this work brings.

 It is a shame; the paper is very well written, and it is a good idea. However, it has indeed been already explored. I am happy to focus more on the theoretical contributions in the discussion period, and given a better understanding of their novelty I will be happy to increase my score; in particular, I suspect the assumptions in proposition 2 of [1] do seem (Assumption D1) like they might be more restrictive and their results seem to have a less broader scope than yours (i.e. you show the curves are a.c. in the Wasserstein 2 space, which feels like a stronger and more insightful result).

Another point that seems novel is the gradient flow interpolation from what I can see in works [1,2,4] they also explored learned interpolations too and in particular the $t(1-t) * \mathrm{NN}(x,t)$ parameterizations can be seen in [2,4,5], that said the gradient flow interpolation seems novel, If a more comprehensive comparison of this to prior works / motivation as to why this is an improvement over [1,2,4]  is provided I would also happily increase my score. 


[1] Nusken, N., Vargas, F., Padhy, S. and Blessing, D., 2024, May. Transport meets variational inference: Controlled monte carlo diffusions. In The Twelfth International Conference on Learning Representations: ICLR 2024.

[2] Sun, J., Berner, J., Richter, L., Zeinhofer, M., Müller, J., Azizzadenesheli, K. and Anandkumar, A., 2024. Dynamical measure transport and neural pde solvers for sampling. arXiv preprint arXiv:2407.07873.

[3] Albergo, M.S. and Vanden-Eijnden, E., 2024. NETS: A Non-Equilibrium Transport Sampler. arXiv preprint arXiv:2410.02711.

[4] Sun, J., Berner, J., Azizzadenesheli, K. and Anandkumar, A., Physics-informed neural networks for sampling. In ICLR 2024 Workshop on AI4DifferentialEquations In Science.

[5] Máté, B. and Fleuret, F., 2023. Learning interpolations between boltzmann densities. arXiv preprint arXiv:2301.07388.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: 1. General background section. What do you denote by $C_{c}^{\infty}$? Is it the set of (infinitely-many) differentiable functions?
2. Lines 261 - 267. I missed whether you introduce your method in these lines, or just explain the method by [Mate&Fleuret]. But in the latter case, this loss is different from [Mate&Flauret], because instead of sampling $x \in U[a, b]^d$ uniformly, as in eq. (17), they propagate $x$ along the learned vector field (treated as continuous normalizing flow).
3. Lines 359-360. I do not understand, what you mean by “[...] it is not clear if Boltzmann densities stay Boltzmann densities”

**Misprints**

1. Line 187 - no comma in $\nabla \cdot (\rho_t, v_t)$ in eq. (11).
2. Line 037-038. Velocity field $v: [0, 1]\times \mathbb{R}^d \rightarrow \mathbb{R}$ maps to scalar? Misprint?

**Minor comments**

A bit difficult to read the text (Introduction section, Related works subsection). Line 118, eq. (2). What is $s_t$ in eq. (2), what is $\alpha_t$, $\overline{\alpha}_t$ in eq. (2)? All of these quantities are introduced later, in the main text, but I think it is strange to introduce a formula in the introduction which will be understandable only after reading the main text. Some references/brief explanations should be given.

**Related works**

[Mate&Fleuret] Mate et. al., Learning Interpolations between Boltzmann Densities, TMLR’2023

[Xu] Xu et. al., Normalizing flow neural networks by JKO scheme, NeurIPS’2023

[Boffi] Boffi et. al., Probability flow solution of the fokker--planck equation
soundness: 2
strengths: The paper has some interesting theoretical results about absolute continuity and optimality of Fisher-Rao flows. Also, the paper demonstrates the connection between Fisher-Rao flows and Wasserstein Gradient flows (w.r.t KL functional).
summary: The paper deals with the problem of sampling from unnormalized density (with known energy function). This is done by means of learning a specific Fisher-Rao curve in the probability measures space, which transforms the data distribution to some simple distribution (Gaussian one is considered). Since the vector field of this curve is learned, it could be inverted and simulated with ODE for sampling.
weaknesses: I think that the paper has rather limited methodological and practical contributions.
1. In fact, the authors take the method from [Mate&Fleuret] (in particular, eq. 15 is exactly eqs. 11 and 12 from [Mate&Fleuret]), but with specific parameterization of the vector field $v_t$. In turn, the chosen vector field parameterization exactly corresponds to Wasserstein gradient flows (WGF) with KL divergence. Therefore, the proposed method - modeling WGF with KL from data to Gaussian (with subsequent inversion) using  techniques from [Mate&Fleuret]. At this point I want to note that there are other papers which also solve something similar to WGF inversion task with different techniques (e.g., [Boffi], JKO: [Xu]). So, what is the advantage of your proposed method? In particular, [Xu] manages to work with image data, while I am not sure if the presented method scales to high dims.
2. The scalability of the method is under question. The experimental illustrations do not stress test high-dimensional applications.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: Could you please explain in more detail the precise methodology used for solving the ODE (6)? Do you need a stiff ODE solver?

Both examples you've considered are 2 dimensional. Could you discuss the scalability of this methodology to higher dimensional problems?
soundness: 4
strengths: The theoretical results are quite fundamental, and show existence of unique solution to the Poisson equation that ensures Wasserstein absolute continuity. The idea of using a neural network to parameterize the interpolating distributions is also novel for Fisher-Rao flow, as far as I know. The numerical results sufficiently illustrate the methodology, and the choice of metrics makes sense.
summary: The paper considers the problem of sampling from probability densities. The key idea is to use a so-called Fisher-Rao flow, which is a general equation that transports a simple density to a more complex one, going through a sequence of interpolating densities that are only assumed to be known up to an unknown normalizing constant. The main theoretical results of this paper (Theorem 1 and 2) show under very general conditions the uniqueness of a weak solution of the Poisson equation of this flow for which the Fisher-Rao curve is Wasserstein absolute continuous. The solution of this Poisson equation is found by parametrizing it using neural networks, and minimizing a certain loss function (explained on pages 5 and 6). Given this solution, the ODE of equation (6) is solved backward to yield samples from the target distribution by transporting a simple initial distribution to the target. Multiple different paths (interpolating densities) are considered, namely linear interpolation, and learned interpolation, where the interpolation curve itself is also parameterized by a neural network, and certain loss function enforcing more smooth transitions along the Fisher-Rao curve is optimized.
weaknesses: There could be a more precise description of the ODE solver used to solve equation (6) backward, since this is the essence of the method for sampling from targets.
There is no theoretical guarantees for the neural network approximation that is used to solve (17) actually is able to solve it.
The scalability of this approach to high dimensional problems has not been demonstrated.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 2
questions: -
soundness: 3
strengths: - The paper introduces well the concepts and challenges it aims to address, for example using the analytical example of Figure 1 and 2.
- The paper highlights and fruitfully exploits connections between SDE/ODE sampling and gradient flows on probability spaces. 
- The paper draws its inspiration from prior literature, but draws new connections to propose amendments to the approach.
- The proposed sampling method appears to be superior to related ones as its initial distribution can be chosen as a standard Gaussian regardless of the variance/mode location in the target. 

I am not familiar enough with the formal literature on Wasserstein curves to judge the novelty of the theoretical results of the paper.
summary: The present paper considers the problem of sampling from a Boltzmann distribution known up to a normalization constant.

The author investigate theoretically the connection between the evolution of Boltzmann distributions with energy indexed by time, which give rise to Fisher Rao flow, and absolutely continuous curve in the Wasserstein space that obey a continuity equation. In particular, the author set the conditions under which a regular velocity field expressed as a gradient is solution to continuity equation associated with a the Boltzmann distribution curve.

Using this connection, the authors propose to learn a curve of Boltzmann densities bridging a base distribution to the target: assuming it follows a continuity equation for a velocity field expressed as a gradient a training objective is derived as in Mate & Fleuret (2023).  Going further, the authors take inspiration from the Fokker Plank equation of the Ornstein-Uhlenbeck process, recalling it’s interpretation as a Wassertein gradient flow on the reverse KL, to propose corresponding forms for the velocity field and learned energy. As such the learned curve is related to a variance preserving noising scheme of the target distribution.

Ultimately, the ODE associated with the velocity field can be integrated backward for sampling. The proposed method is compared against closely related methods that learn different interpolation schemes between a simple Gaussian base distribution and the target. The simple experiments in 2 and 8 dimensions suggest that the present method is less sensitive to the choice of the base distribution.
weaknesses: - The writing of the paper could be improved, in particular, some notations are regularly used before being introduced. 
- The discussion of related works should be extended to include:
	- Stochastic interpolants: https://arxiv.org/abs/2303.08797 
	- Non-parametric denoising based samplers: RDMC https://openreview.net/forum?id=kIPEyMSdFV & SLIPS https://proceedings.mlr.press/v235/grenioux24a.html
- The experimental validation is currently very limited, notably in terms of how the method scales with increasing dimension.

Minor:
- maybe misses the assumption that $\rho_1 = \rho_D$?
- Notations in Equation (2) are not defined, also lacking Equation (5)
- In (17), the space variable is sampled uniformly on an hypercube? How are the boundaries chosen?


- typos:
	- line 96: Kullback Leibner
	- extra comma between $\rho_t$ and $v_t$ in the divergence (11)

---

confidence: 2
contribution: 3
initial_rating: 6
presentation: 2
questions: - The experiments that are done in Máte-Fleuret, show that problems of mode-collapse occur more severely when the Gaussian distributions are unevenly distributed. While the dataset proposed by the authors is not homogeneous since the mean of the Gaussians are random, it might be better to include a test on a mixture with unequal weights and variances.
- Does there exist a case where the gradient flow method is clearly better than the learned one ?
- There very few experimental validations, I would like to see another dataset, maybe inspired from the one investigated in Máté and Fleuret.

It seems to me that the phenomena described by the use of the linear interpolation, as in fig 1, is very similar to what is called "first order transition" in physics. As a parameter is slightly change, the overlap between the distribution $\rho_t$ and $\rho_{t+\delta t}$ is very small. That being said, the condition for annealed importance sampling to work is that the overlap between two nearest neighbors of the interpolation remain high enough. Is there a way to interpret the diverging of the velocity as characterize by the authors with such a phenomenon ?

Minor remark:
- there is a tipo at "Leibner" line 96
soundness: 3
strengths: The paper make very interesting connections between Wasserstein Flows and the developed method. 
It also underlines its relation with Diffusion models.
summary: The paper deals with the problem of sampling configuration from an un-normalized Boltzmann density. The author consider an approach where the energy function is interpolated by a function $f_t$ that is learned, following an approach developed by Máté and Fleuret. The main scope of the paper is to prove that the curve of the Boltzmann densities are absolutely continuous Wasserstein curve. They also demonstrate on an analytical example that the approach of Máté and Fleuret lead to a discontinuity in the particle transport due to the explosion of the velocity field. They propose a new reparametrization for the interpolant where a single function has to be learned while the velocity field is fixed.
weaknesses: The article is particularly technical for those not-familiar with Wasserstein flows etc. While I find the results very interesting, I'm confused about the overuse of mathematical technicality, which tends to make the paper hard to grasp for a non-specialist. I'm wondering if the authors could ease part of it (eventually keeping a lot of details in the appendix) but putting an emphasis on the general direction and method. 

Among all these technicalities, I end up not understanding the details on how the function $\psi_t$ is learned, or the velocity in the "learned interpolation" case.

The results on the experiments are ok, but not super-convincing either.

Meta-review: The authors studied Fisher--Rao flows from both a theoretical and an algorithmic point of view. In particular, on a conceptual level, the paper has made clear contributions towards understanding when a flow satisfies the continuity equation, equivalently absolutely continuous in Wasserstein geometry. Furthermore, the authors provided a clear example towards understanding the problems with linear interpolation leading to an exploding norm, which has practical impacts in algorithm design. 

Despite the somewhat borderline scores, a very productive discussion period clarified the contributions of this work. More specifically, the goal of this paper was not to propose a new algorithm yielding directly performance improvements, but rather understanding a common issue for a large class of flow related algorithms. For these reasons, I would recommend accept for this paper.

Decision: Accept (Poster)