Neural Sampling from Boltzmann Densities: Fisher-Rao Curves in the Wasserstein Geometry
Authors: Jannis Chemseddine
Christian Wald
Richard Duong
Gabriele Steidl
Venue: ICLR 2025 Poster (2025)
Keywords: Sampling
Boltzmann densities
Fisher-Rao Curves
Wasserstein Gradient Flows
Diffusion
Interpolations
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)

Abstract: We deal with the task of sampling from an unnormalized Boltzmann density $\rho_D$
by learning a Boltzmann curve given by energies $f_t$ starting in a simple density $\rho_Z$.
First, we examine conditions under which Fisher-Rao flows are absolutely continuous in the Wasserstein geometry.
Second, we address specific interpolations $f_t$ and  the learning of the related density/velocity pairs $(\rho_t,v_t)$.
It was numerically observed that the linear interpolation, 
which requires only a parametrization of the velocity field $v_t$,
suffers from  a "teleportation-of-mass" issue.
Using tools from the Wasserstein geometry,
we give an analytical example,
where we can precisely measure the explosion of the velocity field.
Inspired by Máté and Fleuret, who 
parametrize both $f_t$ and $v_t$, we propose an
interpolation which parametrizes only $f_t$ and fixes an appropriate $v_t$. 
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.

Published as a conference paper at ICLR 2025
NEURAL SAMPLING FROM BOLTZMANN DENSITIES: FISHER-RAO CURVES IN THE WASSERSTEIN GEOME- TRY
Jannis Chemseddine, Christian Wald, Richard Duong Gabriele Steidl
Institute of Mathematics
TU Berlin
Straße des 17. Juni 136 Berlin, Germany
[chemseddings ne, wald, duong, steidl] torpethhtu-bbrrinn
ABSTRACT
We deal with the task of sampling from an unnormalized Boltzmann density PD by learning Boltzmann curve given by energies fa starting in a simple density Pz- First, we examine conditions under which Fisher-Rat flows are absolutely contin- uous in the Wasserstein geometry. Second, we address specific interpolations and the learning of the related density/velocity pairs (pt. r2). It was numerically observed that the linear interpolation, which requires only a parametrization of the velocity field Dt, suffers from a "leleportation-of-maae 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 fi and Ut we propose an interpolation which parametrizes only Jt and fixes appropriate Ut. This corresponds to the Wasserstein gradient flow of the Kullback-Leibler divergence related to Langevin dynamics. We demon- strate by numerical examples that our model provides well-behaved flow field which successfully solves the above sampling task.
issue. Using tools from
INTRODUCTION
In this paper, we consider the problem of sampling from a Boltzmann density PD = f / Z with unknown normalizing constant ZD. Our approach is in the spirit of recent developments in generative modelling and aims to construct curve P = e /Z, interpolating between a simple density pz and the target PD. i.e. Po = pz and pr = PD. If such curve admits a velocity field it and there exists a solution y of
0,p:: (x(:), go Id where (p:))( p.dar,
then we can use 01 and pz to sample from PD. An important question is, whether for given family of functions ft, such velocity fields exists. While, for fixed time f, the existence of such velocity fields is well established, see e.g. Laggesn et (2014), global existence with integrability of [o, 1] x gd RRd in time and space addressed in this paper. This is directly related the question whether a large class of Fisher-Rao curves Pt is absolutely continuous in the Wasserstein geometry, and in particular, fulfills a continuity equation. This can be reduced 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, in time and space which admits moreover minimality property of (|v,)lL.(R<(,.))>
Determining 401 requires both finding a curve fe and the associated vector field us. One way is to choose a curve interpolating between fz and f' first and then to learn the velocity field. For the linear interpolation ft = (1 t)ʃz + ttfp 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 noted in (Mâté & Fleuret 2023 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 [|vellza px)- Another approach, recently proposed by Máté & Fleuret 2023). considers

---

Published as a conference paper at ICLR 2025
ft (1-t)) +tf++++++++ + -)) with unknown 1/2 and learns (v), vt) 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 = 7 fD + fx/2. Fixing the velocity field as U = Vft - fz). we only have to learn vt. The resulting PDE is the well-known Fokker- Planck equation related Langevin dynamics which has solution with Boltzmann densities pt. if, e.g., Pz is Gaussian. Then, the corresponding SDE is the Omstten- Ceeeenn-Uhlebeckk (OU) process. Here, the backward ODE of (T) 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 Boltz- mann densities is an absolutely continuous Wasserstein curve, see Theorems 1 and 2 and Appendix A In partieular, this includes the linear interpolation.
2. We verify by nontrivial analytic example that the velocity field corresponding to the linear interpolation can admit an exploding norm [[vellzz(R<...) leading to a "non-smooth" particle transport, see Figure I and Appendix B
3. We propose to learn curve using the parameterization ft f== TT-T- + tht which we call "gradient flow interpolation" by the following reason. With an appropriate velocity field which can be computed directly from ft, this resembles the Wasserstein gradient flow of the Kullback-Leible divergence with fixed second argument pz. This curve can be also described the Fokker-Planck equation of a Langevin SDE.
4. We learn the networks for the linear, learned and gradient flow interpolations and demon- strate the performance of these three methods on certain sampling problems.
Related Work. Classical approaches for sampling from unnormalized density are based on the Markov Chain Monte Carlo (MCMC) method, see Gilks et al. (1995), including variants as the Hamiltonian Monte Carlo (HMM) in Hoffman & Gelman (2014) or the Metropolis Adjusted Langevin Algorithm (MALA) in Girolami & Calderhead (2011) 1. Arising problems with the correct distribution of mass among different modes lead to the use of importance sampling Neal (2001) and
Neal
sequential Monte Carlo samplers Moral et al. (2006).
Newer approaches include the simulation of gradient flows of probability measures. In Chen let al.] (2024) Gaussian mixtures are used to approximate the Fisher-Rao gradient flow of the Kullback-Leibler (KL) divergence. This corresponds to a time rescaling of the linear interpolation. In Nüsken 2024| Maurais & Marzouk (2024) they approximate the latter curve by kernelizing the associated Poisson equation and simulating corresponding interacting particle system. Closely related Stein variational gradient descent Liu & Wangl (2016) kemeliies the Wasserstein gradient flow of the reverse KI. divergence. Recently in Chehab & Korba (2024) 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-Uhlenber process.
Recently, (stochastic) normalizing flows have been widely used to sample from unnormalized densities Noé et al. ) Wu al. 2 Hagemann et al. (2022). In Midgley et al. (2023) this problem is tackled without using a curve of probability measures by augmenting normalizing
flows, which are trained using an a-divergence loss, with annealed importance sampling. In Xu
et al. (022) the Wasserstein gradient flow of KL(,, PD) is approximated. They learn a network for approximating the corresponding JKO scheme. In similar approach in Hertrich & Gruhlke (2024) a functional F with minimum PD is considered. They then learn a vector field by approximating a JKO scheme for the Wasserstein gradient flow of F and combine with an importance based rejection method. They focus on F being the Kullback-Leible divergence with target distribution PD. Using KL for sampling is a natural choice and it is shown in Chen et al. (2023) that it is the only fddivergence such that the Wasserstein gradient flow does not depend on the normalization constant of PD-
Note that parametrizing the vector field using the score V Vlog P: is common practice for sampling via SDEs. A method that generates samples and scores from trajectories of the Fokker Planck equation of certain SDEs is presented in Boffi & Vanden-Fajjnden (2023). In Vargas et al. (2024)
2

---

Published as a conference paper at ICLR 2025
the vector field is learned by using the KL divergence on the path measures of forward and corresponding backward SDE assuming knowledge of V Vlggr, pt. In Zhang & Chen (202)) they formulate the sampling problem in terms of 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.. 2024, Section 3) where the KI. divergence between path measures associated to a controlled and an uncontrolled SDE is used to approximate the score of c.g. an OU process. A different approach to approximating the same process using the KL on path measures is formulated in Vargas et al. (2023a). Another way of approximating flow ending in PD is to consider the variance exploding diffusion SDE as done in Akhound -Sddegh et al. (2024), or to consider the time reversal of a pinned Brownian motion as in Vargas et al. (2023b) Reu et al. (0024)
The authors of Máté & Fleuret 2023) present a loss based on the continuity equation and the spe- ft cial form " of Boltzmann densities. This loss allows for using fixed interpolations fa and only learning Ut as well as learning ft and it simultaneously. Using a fixed interpolation and a loss based on the continuity equation also shows up in the physics informed neural networks (PINN) literature, (2024). When finishing this paper, we became aware of the recent preprint Sunn et al. (2024) and the workshop paper (Sun et al. comparing variety of PDE based sampling methods. They included a similar approach as the one proposed here, using the same parametrization to simulate
see e.g. AAlergoo & Sanden-Eijnder
Ornstein-Uhlenber SDE. However, coming from the Wasser- stein 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
V (peVs) = (aa-0a)) oz) Pt: I LE [0,
(2)
which relates Fisher- Rao curves and Wasserstein absolutely continuous curves. Here Vs, corre- sponds to the vector field of a Wasserstein absolutely continuous curve and o₁ - for or = Ep, [ax] corresponds to the gradient in the Fisher-Rao geometry. This equation was studied for fixed time t in various contexts. In nonlinear filtering it is used to study weak solutions of nonlinear filters, see e.g. (Laugesen al., 2014, Section 2..). Their theory for fixed time steps can be generalized the case of a time dependent equation as seen in Section A The close relationship between the Poincaré
inequality and fixed time solutions of [2) is discussed in Dhara & Kalamajska (0155).
2 WASSERSTEIN FLOWS OF BOLTZMANN DENSITIES
2.1 GENERAL BACKGROUND
Let P(RR)) denote the space of probability measures on Rdd For μ E P(Rd), let (P μ),, p e [L, oo), denotes the Banach space of (equivalence classes of) real-valued functions on Rd with finite
norm Ilfllz, .p) = (Jgg [fll du) 1/2 and (R Rdd 1n) be the Banach space of Borel vector fields Rd Rd with finite norm [|uIIr,, = (Sez |||i|| dy)) where || || is the Euclidean norm on Idd For the Lebesgue measure, we skip the μ in the notation. The space (2( of probability measures having finite second moments equipped with the Wasserstein(-2) metric
(Fe xE(()o.)))
W2 m1) = min
1/0 w1l² dh (1,
where T(pa, un) denotes the set of couplings plans क P(TRd IRd) having marginals μt P2(Rd4) = 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, weakly continuous curve μ [0, 1] + (2(")), t + Me is absolutely continuous, if there exists Borel vector field v: [0, 1] x d + IRd with fo || (B dl x such that the continuity equation
Oyp + V (mum)) 0
(3)
is fulfilled in the sense of distributions see, e.g. Ambrosio et al. 2005).. While for fixed μ there are many vector fields such that the continuity equation is fulfilled, there exists unique one such
'fo Sea d. + (Vzó, vx) dpedt 0 for all & E C ((0. 1) x Rdd)) where C ((0, 1) x Rd) denotes the space of smooth and compactly supported functions.
3

---

Published as a conference paper at ICLR 2025
that ( becomes minimal for almost every t E [o, 1]. In other words, if i fulfills the continuity equation for some vector field, then there exists a unique vector field that solves
arg min ty E ,2 (p) J
|| duu s.t. pep V (wmm) =0
(4)
Moreover, this optimal vector field is determined by the condition that us E Tp, P2(R4) for almost every t [0, 1]. with the regular tangent space
T.P2 (Vó: QE C Lx((
(5)
where the right-hand side is the closure in the Banach space L2 , 14), see (Ambrosio et al. 2005
$ 8). Finally, if je an absolutely continuous curve with Borel vector field it such that for every compact Borel set B c. R it holds So supæ ||eell + Lip(u,, B)
dt < oc, then there exists a solution
; [0. 1] x R + Rd of
Oep(t, x) ve(p(t, x)).
(00 x
(6)
which fulfills i4 Pt.p/l Mo Pr see e.g. (Ambrosio et al. 2005 Proposition 8.1.8).
2.2 WASSERSTEIN MEETS FISHER-RAO FLOWS
In the following, we are interested absolutely continuous probability measures with respect to the Lebesgue measure, i.e., the space (Pa(RR) {u pdr P L, p 0, fg= p will only consider
1} and will address such measures just by their densities. Moreover, Gallzmann/Gibbe densities depending on functions f [0, 1] x Rd + R with e-ht E L1 (I2) and x()xx)) E ((O, 1J) for all x E IRd and (t,, E (RRR,, pt) such that
e-
Ze = { c dr. Bd
Pe ==
Ze
In this case, straightforward computation gives
(8) where it maybe useful to note that E,, [Or = Ô log Ze. This differential equation is a Fisher-Ro Rao flow equation. More precisely, a weakly continuous curve p [0, 1] - pao pac(R. (R.4), 1. +> Pt is Fisher- Rao curve, if there exists measurable map [0, 1) IR + IR such that at,, -) E T [R. Pt) and we have a distibutional sense
0. = 0.p: = - (0.fe ; - EE, f+l Pt,
0.p: = (0;-aa))) ac)pts ā = E. (ar]-
Obviously, K delivers Fisher-Rao curve with 0₁ == -e For more information on Fisher- Rao flows see, c.g., Gallouët & Monsaingeon (2017); Wang & Li (2022). We are interested in
the question under which conditions a Fisher-Raa Rao curve determined by (8) is also an absolutely
continuous curve in the Wasserstein space. For later purposes, note that for Boltzmann densities the second summand of the continuity equation
(3) reads as
(14v)) = V (pert) = (-( V vz) Pt.
(10)
Nevertheless, using (8) for our Fisher-Rao curves the continuity equation becomes
(0. Ep, [a.xa)) P: (p.v.),
(11)
and in the ideal case that we can find a vector field 0t = Vs, compare 5). this can be rewritten as the family of Poisson equations
(0.S: Ep, pr V (p.Vs.),
(12)
which needs to be solved for 8t. For fixed time f, the solution of the Poisson equation has already been examined, e.g., in Laugesen et al. (2014) and for more recent work, see also Reich (20TT); Taghvaei & Mehtal (2023). However, the pointwise solution for each fixed does not ensure that (Pe. Vst) gives rise to an absolutely continuous Wasserstein curve, since by definition this requires additional properties of V,s : [0, 1] Rd Rd globally in t: i) V.s is Borel measurable on
[0, 1] x IR. and ii) + (t, ( E [(0, 1)). If also, iii) (as(t, E T, P2(24))) for a.e. t E [0, 1]. then it is known that the velocity field Vys is the optimal one in the sense of (4). To our best knowledge, the following two theorems, which are detailed in Appendix A tackle
rigorous solution of this problem for the first time.
4

---

Published as a conference paper at ICLR 2025
Theorem 1. Assune that p determined by [[o f(t, x) -o JR. f(t, 2)) dp,dt satisfies
so-called partial Poincaré inequality, see Definition A.7 for some K > O, and that 1 is locally integrable. Furthermore, suppose that 04 - aat [2( 1] x TR. p). Then, there exists a unique weak solution & in certain Hilbert space Xg the problem
p
J0 (V. Vete) dp = -fo f we (az - aa) dp for all db E Xg. SVL.
(13)
Theorem 2. Let Po,P1 11 E P(RR)) and assume that Pt is Fisher-Rao curve defined by D. Let p and a satisfy the assumptions of Theorem Then, Pt a Wasserstein absolutely continuous curve with vector field ss where X6 is the unique weak solution of (13). Pt is bounded from above, holds that V88t E TT P2(P2 for a.e i E (0, 1). Furthermore, it holds that
J gd
s dp,dt s K S(² [ (ar -a)) dp,dt.
(14)
Note that by the Benamou-Breniee formula, see (Ambrosio al. 2005 Equation 8.0.3), the left- hand side of (14) is an upper bound for the Wasserstein distance [ (1),, while the double- integral on the right-hand side defines the Fisher-Rao action of Pt.
3 NEURAL SAMPLING FROM BOLTZMANN DENSITIES
Our aim is to sample from a Boltzmann measure μán = Pn dæ, where we only have access to Sn : Rd R, but not the normalization constant Zp, by following measure flow starting in a simple to sample measure = pz dæ like the standard Gaussian one. We apply the continuity equation (TT) of Fisher-Rao curves, which can be rewritten by (10) and division by Pr. >> 0 as
a.f. Ep. (V vz) ve-
(15)
Indeed every spatially constant function C. which fulfills the above equation must be of the form C = Ep. [i. see Máté Fleuret (2023). This enables us to learn C,, so that (15) does no longer depend on normalization constants. We will deal with the following parameterizations of fr and it and how to learn them:
1. linear interpolation: ft = (1 t)fz tfD learn Ut:
2. learned interpolation by Máté and Fleuret: fe = (1- (1- + LlD + 6(1 f)z/ 4 learn (1h, Uz),
3. gradient flow interpolation: == T 14h vz V(Si -zz) - learn d/-
While by definition the first two settings interpolate between pz and pD and po = pz P1 PD. we will see that the newly proposed third variant interpolates between PD and pz by the choice of ut, i.e. Po PD and limy-120 Pr = pz- practice we reparametrize the latter curve into unit time,
see Remark 4
3.1 LINEAR & LEARNED INTERPOLATION
Linear interpolation.
The simplest interpolation is the linear one ft = (1 t)fz tfD. i.e.,
P: x P pf. Under suitable conditions on the end points this density fulfillss the assumptions of Theorem 2 see Corollary the Appendix, and therefore Wasserstein curve. By (8). also defines a Fisher-Rao curve, and the corresponding Fisher-Rao flow equation (9) is given by
a.p. (Jo Ep..fo f1)) P -
P00 Ep, Dog(
PD )
P1
P1
Note that, is even Fisher-Rao gradient flow of the negative log likelihood, see also Maurais & Marzouk 2024 given by ExTa((p) = -Ep log(

---

Published as a conference paper at ICLR 2025
There are variety of approaches for approximating the vector field corresponding the linear
interpolation. In Maurais Marzouk (2024); Nüsken (2024) a kernel-based method for solving
[2) proposed. Mäté & Fleuret 2023) they propose solve (15 which in case of the linear interpolation becomes
- C, (V V u 0. This equation has many solutions for the same Pt, and in order to enforce uniqueness we could demand that U Ss.
(16)
This approach leads to learning neural networks ( by minimizing the loss function
L(0)= =
[8(0, xz Q)]
(17)
E 2r, 1) = Si Jo C C (v, v. V v o= (2)
Remark 3. Since the approximation of high dimensional integrals is difficult Máté & Fleurer (2023)
propose to use an iterative approach by solving the ODE dest = of (4p)) ). sp = Ld for 8' from the previous minimization and solve
l(O) == BrEu(O.1))s-oo [E sp (z),
Heuristically this penalizes E00, xx,) stronger in regions where the intermediate solution has more mass. Finally, sampling can done assuming sufficient regularity of the learned uf by simu- lating the ODE 66
t-0.00
t-0.63
b-0.01
t-0.99
1-1.00
Figure 1: Evolution of the probability densities Ps x Po pi. where -2 2minf/=/./zmm// for = 50. For different values of m see
and P1 x
A teleportation issue with the linear interpolation. Máté & Fleuret (023) gave examples, where for an asymmetric constellation of target and prior measure, this path transports mass into distant modes only at very late times (sometimes referred mode switching) which may cause problems while learning the corresponding vector field. Mathematically, this corresponds to the fact that the vector field (xx) develops a singularity in x Rd for late times 6, (Máté & Fleuret 2023 Figure 6). In this work, we demonstrate that this phenomenon is more severe: even when averaging over the space z, the norm ||vellaz(e)) of the optimal field explodes for late times For an example in IR1 involving a Laplacian Po α e laxl centered at and a Laplacian-like distribution P1 & e 2in (in,([ar with a mode at and second mode at m > 0, see Figure [ we are able to give an explicit and analytic formula for the calculation of ||vy (2((:)) in the Appendix B Using this formula 37). demonstrate that ||v (p())) explodes when 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, see Figure 2
Learned Interpolation. cope with the above disadvantages of the linear interpolation, Máté & t)fo + tfi + t(1 +) with
Fleuret (2023) proposed parameterize also the flow by == (1 neural network viSi and to minimize according to (1) the above loss with
E00, = jārf (x) cfa + 6 0o
1.. == (0,, f2 0s).
(18)
and
L(O)= = Ereur [E op t)],
where o0 (2) is as in Remark trajectory of the intermediate solution. Note that in general there are infinitely many solutions of 16) and neither the parameterized curve nor the vector field v, is unique. Further of obtained this way may not be minimal in any sense. The proof of the validity of the interpolation with respect Theorem 2 is future work.

---

Published as a conference paper at ICLR 2025
Bellas
Ioa (olilii
m7
771 15
770 30
t
Figure 2: In the first three figures, the evolution of || La(p.) belonging to the path in Figure I de- picted for different values of the second mode m 0. For larger m, the norm ||0, (2)(())) approaches the limit ||0ilf2,((1)) at later times, hence with steeper slope. The last figure shows the log scale of I|v ((2)) depending on m. It demonstrates that || roughly grows exponentially with m.
3.2 GRADIENT FLOW INTERPOLATION
Recall that the linear interpolation of the energies Jo, leads multiplicative interpolation of densities (x( X pz(r))==pn(z) Note that this does not take account of any local behaviour of pz. PD neighbourhood of æ, resulting in irregular behavior of the vector field as shown in Figure and 2 In contrast, widely used paths in generative modelling [Song et al. (2021): Lipman
al. 2023 originate from an interpolation of densities via (spatial) convolutions Pt pz
PD b, for time schedules ur; be, corresponding to an interpolation X, = a + by Xo of associated independent random variables Z pz: Xo ~ PD on the spatial level. We aim to learn such a curve and associated vector field.
In this paper, we propose to use f; T-t fn + tiie and Vt. = V - fz) and to parameterize dh by neural network. Then, equation
15)
becomes
a.ft - Ep, [0. + (V,, V(fi fz)) Alf: - f2) =.
(19)
Indeed, will show in the paragraph below that this equation makes perfect sense, in particular, it has solution for pz NN 1). Then
compute (wit O.22) for f T n + tw by
minimizing the loss L in 177) with the function
E(0, 1) = 10. + ( V(S - fz)) - A( -z)) 0 = (0,, 02). (20)
Finally, sample from the target distribution, we simulate the ODE (6) backwards on [0, to obtain C [0, 7 x Rd + Rdd satisfying
0is(t, x)) =))(((,,)), ((T.2)
which fulfills pr- -1
Relation to Wasserstein gradient flows. stein gradient flows. A Wasserstein gradient flow is an absolutely continuous curve whose vector field i = grad, F T,, (P2(R))) is given by the negative Wasserstein gradient (or more gener- ally, subdifferential) of function F P2 (PRd) + R. Then, the continuity equation
Equation (19) is well-known in connection with Wasser-
O.p: V (pigrad,, F) =
(21)
produces a flow towards a minimizer F. A prominent example, see, e.g. (Ambrosio et al. 2005
Chapter 10.4), is F(p) = KL(plpz) arising from the Kullback-Leeibler divergence which has the unique minimizer pz- Since it is casy to show, see (Ambrosio et al. 2005 Lemma 10.4.1) I), that grad F = log lt,, pz
be rewrite (2
O.p. = (pr V log log
Po = PD,
(22)
which is also known as the Bokker-Plance' equation. Inserting Pr == = f. /Z, ZZ, results exactly in (9) The existence and uniqueness of solution of (22 is shown in (Ambrosio & Savaré 2007 Theorem
6.6). However, it is not clear if Boltzmann densities stay Boltzmann densities, Le. if the densities p strictly positive for all I. But in the special case, when pz is Gaussian, this is the case
7

---

Published as a conference paper at ICLR 2025
will see next.
There is close relation between 22 and SDEs, namely the Langevin SDE
dx, = Pz + vddB
(23)
corresponds to the Fokker-Planck equation (22 see e.g.. (Chewi et al. 2024 Chapter 6.2), i.e., X, ~ Pt- For pz x (0, Ta), (23) becomes dX, = - X,dt /2dB, which is Ornstein- Uhlenbeck (OU) process, and minimizer of (20) (up to a constant) log pr of the corresponding
Fokker-Planck equation. In (Sun et al.) 2024, Equation (18)) the loss (20) is derived from this
process. This Crnstein-Uhlenbed process has the closed form solution
X₁ = + e "Xo, Z~ N (0, (e - 1))a), Xo PD,
where Z is independent from Xy, see (Herrmann Massin 2019 Eq. (2.4)). From the closed form
solution it follows that ~ N (0, (1 - ==aaaa pD (e*.) > 0, and thus Pt. x e which implies
that 19h has a solution.
Remark 4. To obtain dynamics on the unit interval for the gradient flow interpolation apply an approach from the score based diffusion literature Song el al. (021). We use a linear time schedule B(1) Bmin 1 (Bmax (min) and the associated SDE
dx, 1 3(t)x,dt + 1 B(t)dBB Xo ~ PI-
(24)
The SDE (24) has a closed form solution given by
X = V g
Xo:
(25)
where g(t) = 1o (s)) ds and ~ (o), As done in Song et al. 2021 we choose Bmin = 0.1 and Bmax 20. These choices ensure that the distribution of X₁ is approximately N(0, Ia). Note that in contrast, finite-time interpolations as the linear one do not incur such a mixing error. We then aim to learn the solution of the associated Fokker- Planckk equation, this corresponds to setting 2t
3(0) (f. 11:11
The following example illustrates the behavior of the three interpolations, when learning the corre- sponding vector fields.
Example 5. Motivated by an example in Máté & Fleurei (023) and incorporating the idea of our teleportation example with asymmetrically diverging modes, we choose standard Gaussian source distribution and the target distribution
2 3 ((4.44)..) 3 (I-m,
1
For the linear and learned interpolations, we consider both sampling from uniform domains and sampling from the trajectories described in Remark 3 For the gradient flow interpolation, we sample points uniformly and interpolate them with samples from our latent distribution according to formula (25. The results are reported Figure Note that, in case of the learned interpolation, the mode collapse could be eventually alleviated by choosing different source distribution with larger support and/or heavier tails.
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 leam potential and obtain the vector field its gradient. In practice, choose to parameterize directly as in Máté & Fleuret (2023). Note that in practice for the gradient flow interpolation we apply Remark 4 Furthermore
improve performance by using learned schedul- ing function gt),, satisfying (()) = and (1)) = 0, and parametrizing the corresponding ft as
fo : fft to).
A major design choice at which points to evaluate the pointwise error E(0, xx t). For the gradient flow interpolation, chose the sampling strategy described in Example For the linear and learned
8

---

Published as a conference paper at ICLR 2025
linear (uniform)
linear (trj)
learned (uniform)
learned (trj)
gradient flow
linear (uniform)
learned (uniform)
gradient flow
Figure 3: Sampling results for m (top) and m = 15 (bottom) as in Example 5 source (black), target (red) and estimated (bluc). We cannot learn a meaningful vector field for the linear interpola- tion for both sampling strategies. Sampling along the trajectories ("trj") leads to mode collapse for the learned interpolation. While the learned interpolation with uniform sampling covers the modes for m 8, increasing m leads to mode collapse. In contrast, the gradient flow interpolation does mode collapse for both
interpolation we experimented both with sampling from uniform domains, or sampling along the trajectory as Máté Fleuret (2023). We found in our experiments that the method employed in Māté & Fleuret (2023) led better performance.
In the following, we evaluate the methods using the effective sample size, negative log likelihood and the energy distance Székely Rizzo (2013). See D for further details on the implementation and evaluation.
Gaussian Mixture Model. We use the experimental setup from Midgley et al.) (2023). The target distribution consists of mixture of 40 evenly weighted Gaussians in 2 dimensions. The means are distributed uniformly over [-40, 40)2. In Figure we plot the countour lines of the target distribu-
ground truth
linear
learned
gradient flow
Figure 4: Results for Gaussian Mixture Model with 40 modes. For the linear and learned interpo-
lation we showed the results for a = 30, for the gradient flow interpretation we used = 1.
tion as well as 1000 samples from the simulated flow and the ground truth, respectively. For this experiment an important factor the choice of latent distribution. We restricted ourselves to using Gaussian latent distribution. However as outlined in Section 3.1 the linear interpolation faces dif- ficulties when some modes are far away from the latent while others are close. order to mitigate this problem, experimented with using varying standard deviations σ; the impact of this choice visualized in D contrast, the gradient flow interpolation works well when starting in a standard Gaussian which what we do.
Many Well Distribution. As target distributions we consider a and 16-dimensional many well distribution similarly to the implementation of the target energy from Midgley al. (2023). More

---

Published as a conference paper at ICLR 2025
ESS
NLL
Energy
linear (a 30)
0.1831 0.05
8.657 0.122
0.247 0.016
learned (a = 20)
0.605 +: 0.316
6.8941 0.005
0.058 1 0.006
learned (a 30)
0.922 0.002
6.913 = 0.003
0.121 + 0.009
gradient flow (a = 1)
0.992 = 5e-5
6.864 + 0.005
0.0021 + 0.001
Table I: 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.
precisely, let
log (m)
Σ
-x
6z +
T;
Σ 2 constant.
2
1=m+1
The term -2 + 6z² + 24 creates two modes with unequal weights and the number of overall modes of Pm.d is Then our 8 dimensional many well distribution is p4.8 and the 16 dimensional target distribution P4,16- The results in Table 2 show that both the learned and gradient flow in- terpolations significantly outperform the linear interpolation for both dimensions. While the learned and gradient flow interpolations are competitive
= 8, the learned interpolation works better in d = 16. The latter behavior might be due to greater flexibility, since both the curve and the vector field are parametrized separately. Here the choice of latent distribution is not as relevant as starting in standard Gaussian is sufficient for all methods, since the modes of the many well distribution are concentrated close to zero.
8 8-dimensional P4.8
16-dimensional P4.16
Interpolation
ESS
NLL L
Energy
ESS
NLL
Energy |
Linear
0.985 0.016
9.689 i 0.01
0.278 a: 0.03
1395
21.85 =: 0.01
0,201 + 0.002
Learned
0.999 1e-5
6.876 0.01
9c-5 335
0.998#
3o-5
18.244
lc-4 + 3o-5
Gradient Flow
0.991 3e-4
6.989 0.01
8e-4 à: le-4
0.544 0.005
20.53 :: 0.17 0.009 a: 4e4
Table 2: Comparison of effective sample size (ESS), negative log likelihood (NLL), and energy dis- tance for different interpolations. evaluated for the 8- 8-dimensional and 16-dimensional experiments with m = 4. Note that the ESS is only meaningful in relation with the Energy.
5 CONCLUSIONS
Discussion. We studied different paths in the probability space in order to sample from unnormal- ized densities and compared their behaviour and computibility. The linear interpolation of energies suffers from 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 want to approximate and in turn conclusions about the regularity of the corresponding vector field can be drawn. In contrast, the proposed gradient flow approach allows for thorough mathematical description of the velocity field and the path itself, opening the doors for equations. We show numerical
rigorous analysis via Wasserstein gradient flows and Fokker- performance on sampling from Gaussian mixture models and 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 geome- try. 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 computation of divergences and Laplacians is expensive in high dimen- sions, requiring the use of e.g. Hutchinsons trace estimator. These issues can be addressed by developing more efficient strategies to sample points at which to evaluate the loss ε(0 x, t), such combining uniform sampling and trajectory based sampling. From theoretical point of view further investigation of the behaviour and regularity of the gradient flow path is needed, especially the behaviour of the norm (ell) would be of interest.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
J.C. acknowledges funding by the German Research Foundation (DFG) within the project SPP 2298 "Theoretical Foundations of Deep Learning". C. C.W. acknowledges funding by the DFG within the SFB "Tomography Across the Scales" (STE 571/19-1, project number: 495365311) R.D. acknowl- edges funding by the German Research Foundation (DFG) within the project STE 571/16-1.
REFERENCES
R. Adams and J. Fournier. Sobolev Spaces. Academic Press, 2003.
T. Akhound-Sadegh J. Rector-Brooks, A. J. Bose, S. Mittal, P. Lemos, C.-H. Liu, M. Sendera, S. Ravanbakhsh, G. Gidel, Bengio, al. Iterated denoising energy matching for sampling from Boltzmann densities. arXiv preprint arXiv:2402.061211 2024.
M. S. Albergo and F. Vanden-Fijnden. Nets: A non-equilibrium transport sampler. arXiv preprint
arXiv:2410.02711, 2024.
L. Ambrosio and G. Savaré. Gradient flows of probability measures. In Handbook of differential equations: evolutionary equations, volume 3, pp. 1-136. Elsevier, 2007.
L. Ambrosio, N. Gigli, and G. Savaré. Gradient Flows: in Metric Spaces and in the Space of Probability Measures. Springer Science & Business Media, 2005.
D. Bakry, F. Barthe, P. Cattiaux, and A. Guillin. A simple proof of the Poincaré inequality for a large class of probability measures. Electronic Communications in Probubility, 13:60-66, 2008
J. Berner, L. Richter, and K. Ullrich. An optimal control perspective on basfus generative modeling. Transactions on Machine Learning Research, 2024. ISSN 2835-8856. htts:
//openzeview. http://www....eerroooooaaaaaaaa1
N. M. Boffi and E. Vanden-Eijnden Probability flow solution of the Mlaker-Planck equation. Ma- chine Learning: Science and Technology, 4(3):035012, 2023.
O. Chehab and A. Korba. A practical diffusion path for sampling, 2024. URL httttt/ss///xx
org/ abs//b///// herg/abs/2466.0 2406. 14040
R. T. Q. Chen. torchdiffeq, 2018. URL
Y. Chen, D. Z. Huang, J. Huang, S. Reich, and A. M. Stuart. Gradient flows for sampling: Mcan- field models, Gaussian approximations and affine invariance. arXiv preprint arX1v:2302.11024, 2023.
Y. Chen, D. Z. Huang. J. Huang, S. Reich, and A. M. Stuart. Efficient, multimodal, and derivative- free Bayesian inference with Fisher-Rao gradient flows, 2024. URL htt.::://sxxxiv.ooo abs/2405 17263
S. Chewi, J. Niles-Weed, and P. Rigollet. vrXiv:2407.18166 2024.
Statistical optimal transport.
arXiv preprint
R. N. Dhara and A. Kalamajska. On equivalent conditions for the validity of Poincaré inequality on weighted Sobolev space with applications to the solvability of degenerated PDEs involving p-Lallacian. Journal of Mathematical Analysis and Applications, 432(():463 483, 2015.
R. Duong, V. Stein, R. Beinert, Hertrich, and G. Steidl. Wasserstein gradient flows of MMD functionals with distance kernel and Cauchy problems on quantile functions. arXiv preprint arXiv:2408.0744 07498, 2024.
J. Feydy, T. Séjouréé, F.-X. Vialard, S.-I. Amari, A. Trouve, and G. Peyré. Interpolating between optimal transport and MMD using Sinkhorn divergences. In The 22nd International Conference on Artificial Intelligence and Statistics, pp. 2681-2690, 2019.
T. O. Gallouët and L. Monsaingeon. A JKO splitting scheme for Kantorovich-Fist flows. SIAM Journal on Mathematical Analysis, 49(2):1
gradient
19(2):1100-1133) 100-1130. 2017.

---

Published as a conference paper at ICLR 2025
W. R. Gilks, S. Richardson, and D. Spiegelhalter. Markov Chain Monte Carlo in Practice. Springer, 1995.
M. Girolami and B. Calderhead. Riemann manifold Langevin and Hamiltonian Monte Carlo meth- ods. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 73:123 -214, 2011.
P. Hagemann, J. Hertrich, and G. Steidl. Stochastic normalizing flows for inverse problems: A Markov chain viewpoint. SIAM Journal on Uncertainty Quantification, 10:1162-1190, 2022.
S. Herrmann and N. Massin. Exit problem for OrstteinUullenbeck processes: random walk ap- proach; appear in Discrete Contin. Dyn. Syst. Ser. B, hal-02143409, 2019.
Hertrich and R. Gruhlke, arXiv:2407.20444, 2024.
Importance corrected neural JKO sampling. arXiv preprint
M. D. Hoffman and A. Gelman. The No-U-turn sampler: Adaptively setting path lengths in Hamil- tonian Monte Carlo. Journal of Machine Learning Research, 15:1593-1623, 2014.
A. Kufner and B. Opic. How to define reasonably weighted Sobolev spaces. Commentationes Mathematicue Universitatis Carolinae, 25(3):537-554, 1984.
R. S. Laugesen, P. G. Mehta, S. P. Meyn, and M. Raginsky. Poisson's 's equation in nonlinear filtering. arXiv preprini arXiv: 1412.5845, 2014.
Y. Lipman, R. Q. Chen, Ben- Ben-Hamu, M. Nickel, and M. Le. Flow matching for generative modeling. In The Eleventh International Conference Learning Representations, 2023, URL https:/ /ttp///ppeeeeeeeeee//e/e//. net
Q. Liu and D. Wang- Stein variational gradient descent: A general purpose Bayesian inference algorithm. Advances Neural Information Processing Systems, 29, 2016.
A. Maurais and Marzouk. Sampling in unit time with kernel Fisher-Rao flow. arXiv preprint arXiv:2401.03892, 2024.
I. I. Midgley, V. Stimper, G. N. C. Simm, B. Schölkopf, Jnd M. Hernández-Lobato. Flow annealed importance sampling bootstrap. arXiv preprint arXiv:2208.01883 2023.
P. Del Moral, A. Doucet, and A. Jasra. Sequential Monte Carlo samplers. Journal of the Royal Statistical Society Series Statistical Methodology, 2006.
B. Máté and F. Fleuret. Learning interpolations between Boltzmann densities. arXiv preprint 20X1v:2301.07388,, 2023.
R. M. Neal. Annealed importance sampling. Statistics and computing, 11:125-1399 2001.
M. Noble, L. Grenioux, M. Gabrié, and A. O. Durmus. Learned reference-based diffusion sampling for multi-modal distributions. arXiv preprint arXiv:2410.19449, 2024.
Noé, Olsson, J. Köhler, and H. Wu. Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning. Science, 2019.
N. Nüsken. Stein transport for Bayesian inference. arXiv preprint arXi::2409001464 2024. URL
htttps://arx: //arxiv/ wtg./bs/////.....eo..eo..ee.ee/e /btp///s..2.4.... 22tp:/2499...422
A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, A. Desmaison, A. Kopf, E. Yang, Z. DeVito, M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, L. Fang, J. Bai. and S. Chintala. PyTorch: an imperative style, high-performance deep learning library. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' Alché Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32, pp. 8024 8035, 2019.
A. Phillips, H. D. Dau, M. Hutchinson, De Bortoli, G. Deligiannidis, and Doucet. Particle denoising diffusion sampler. arXiv preprint arXiv:2402.06322 2024.
12

---

Published as a conference paper at ICLR 2025
S. Reich. A dynamical systems framework for intermittent data assimilation. BIT Numerical Math- ematics, 51:235-244 229,2449222 2011.
T. Reu, F. Vargas, A. Kerekes, and M. M. Bronstein. To smooth a cloud or to pin it down: Expres- siveness guarantees and insights on score matching in denoising diffusion models. The 40th Conference on Uncertainty in Artificial Intelligence, 2024.
Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and Poole. Score-based genera- tive modeling through stochastic differential equations. arXiv preprint arXiv:2011.13450 2021. 13456
URL htttttts///aaiiiii wwwpxxi..rggaa////0002
J. Sun, J. Berner, K. Azizzadenesheli, and A. Anandkumar. Physics-informed neural networks for sampling. In ICLR 2024 Workshop on A1DDiffereniaalEqquiinn In Science.
Sun, J. Berner, L. Richter, M. Zeinhofer, J. Müller, K. Azizzadenesheli, and A. Anandkumar. Dynamical measure transport and neural PDE solvers for sampling. arXi::2407.07873, 2024. URL httphttss///rxii.rrrrg.
27077 07873]
S. Syed, A. Bouchard-Côté, G. Deligiannidis, and A. Doucet. Non-reversible parallel tempering: a scalable highly parallel MCMC scheme. Journal of the Royal Statistical Society Series B: Statistical Methodology, 84:321-350 2022.
G. Székely and M. Rizzo. Energy statistics: a class of statistics based on distances. Journal of Statistical Planning and Inference, 1299-1222 2013.
A. Taghvaei and G. Mehta. A survey of feedback particle filter and related controlled interacting particle systems. Anmual Reviews in Control, 55:356-378, 2023.
F. Vargas, W. S. Grathwohl, and A. Doucet. Denoising diffusion samplers. In The Eleventh Interna tional Conference on Learning Representations, 2023a. URL Forum?id=8pvn f
net/
Vargas, A. Ovsianas, D. Fernandes, M. Girolami, N. D. Lawrence, and N. Nüsken. Bayesian learning via neural Schrödinger-Föllinne flows. Statisties and Computing, 33(1):3, 2023b.
Vargas, S. Padhy, D. Blessing, and Nüsken. Transport meets variational inference: Controlled Monte Carlo diffusions. In The Twelfth International Conference Learning Representations: ICLR 2024, 2024.
Y. Wang and W. Li. Accelerated information gradient flow. Journal of Scientific Computing, 90: 1-47, 2022.
D. Woodard, S. Schmidler, and M. Huber. Sufficient conditions for torpid mixing of parallel and simulated tempering. Electronic Journal of Probability, 14:780-804, 2009.
H. Wu, J. Köhler, and F. Noé. Stochastic normalizing flows. Advances in Neural Information Processing Systems, 33:5933-5944 2020.
C. Xu, X. Cheng, and Xie. Normalizing flow neural networks by JKO scheme. Advances in Neural Information Processing Systems, 36, 2024.
Q. Zhang and Y. Chen. Path integral sampler: A stochastic control approach for sampling. Proceed- ings of Machine Learning Research, 2022.
A WASSERSTEIN MEETS FISHER- RAO: SOLUTION OF A FAMILY OF POISSON EQUATIONS
We will show that under mild assumptions on f. the Fisher-Rao curve Pt from { is also Wasser- stein absolutely continuous curve with tangent vector field V In order to do that, we will
Overall, let V. denote differentiation in direction TR..
13

---

Published as a conference paper at ICLR 2025
describe as the unique solution of Poisson equation in suitable Hilbert space Xg More cisely, we show that there exists a unique & E [2( 1] x R2: p) (with given by Definition A.1 such that Vs, == Vs(t,,) fulfills the Poisson equation
(pr Vo) = a E, [a.s1)px
(26)
in distributional sense in X. For fixed time t, this has already been done, e.g., in Laugesen et al. (014) and Vargas al. (2024) But note that pointwise-in-time treatment does not allow to conclude Wasserstein absolute continuity of our curve Pt- Therefore, we generalize their methods to global-in- in-time ansatz; see also the discussion at the end of Section A In particular, under suitable assumptions on f, we find a unique weak solution of (26) such that
1. Es is Borel measurable on [0, 1] R.
2. t + st, (02)) E L1 (10, 11), i.e,, So / )ll dt < 00,
3. (t,(t, E Tp, P2((R) for a.. t E [0, 1].
These conditions ensure that the Fisher-Rao curve (8) is absolutely continuous in the Wasserstein sense, and that the velocity field vs is the optimal one fulfilling the continuity equation w.r.t. our curve Pt as described in (4).
As by-product, we obtain the estimate
dpttt S {K| f/² [ (a+ - az)) dp /// Ves(t, ) (vss(t
linking (an upper bound of) the Wasserstein-2 distance W2 P1) on the one side, and the Fisher- Rao action on the other.
Next will define the appropriate spaces for our analysis.
A.1 PREREQUISITES FOR SOLVING THE POISSON EQUATION
Definition A.1. Let L. + Pt be a weakly continuous curve [0. (Ro(). Then, for all Borel measurable B S. TRd the map ~> Pe (B) is measurable: Indeed, it is the pointwise limit of con- tinuous, hence measurable maps t + S. fv dpc with appropriate fr. Thus, defined by the disin-
tegration J1o.1 J XBS [(,, dp = S/' Jan f(,, x)) dp,dt, is a probability measure on [0, 1] x Rd, see e.g, (Ambrosio et al.) 2005 Chapter 5.3).
Definition A.2. Let p pdx E pac (TR.) with density p. We consider the Sobolev space
(T) = {1 L2 ,)) the weak derivative Jz. (W.Et L) exists and
0z, L2(p), Vi (l....,d}} d3}
with inner product (f. (f.9) H)) = (f.g) + (V Vo) Furthermore, define the space of Sobolev functions with zero mean by
(p) = (SEH / sdp= sdp=
Lemma A.3. Assume that pdx E (ac(RR)) and E (RRd). Then, (H), (: (())) is Hilbert space and Hg) is closed subspace.
Proof. The first claim is (Kufner & Opic 1984 Theorem 1.11). For the second claim note that for fn - f in H' (p) with fn E Hp))) we have that
f dp
= KV. 1)(0) S 115- - Snn 1)260) + 1[ L))))
Rd
< llf. fl +0 - 0.
14

---

Published as a conference paper at ICLR 2025
Definition A.4. Consider weakly continuous curve 0, 1] + pac (RR),, t Pt, and the associated measure p E (00 1] x Rd) given by Definition A.1 Define
H = { L2(10, RR p): the weak derivative f exists and
dz,f E L2(p), ViE d+1
with respect to the inner product (f (f.GHH1P)) == (f.9) + (V V. Veg) where f f(1,x). g g(1,x)). Now, consider the subspace
X = f
s(t,-) Hpt for almost every E [0.
For E Kg we will also write st. for denoting the function s(t," -)
Lemma A.5. Assume that 1 E (1,100 ((0, 1] x Rd). Then, the space (p) is a Hilbert space, and X is closed subspace of H (p).
Proof. The Hilbert space property of H. (p) again follows by Kufner & Opic 1984 sequence with Sp + s in H.
Theorem 1.11)
and we only prove the closedness X Now, let (s,) C Xg Then,
[[Jessstte e(t,-)) ((,, sts dppdt 0.
This implies pointwise convergence for a.c. t [0, 1] (along subsequence), i.c.
JeL (sn(tt.)) s(t, s(tt-)))s(tt))p-- VV
Vus(t, ) dp - +0 for a.c. t [0, 1].
Thus, have that sn (t, stt, converges in H' (pt) for almost all [0, 1]. Since H (pt) is closed, we have that the limit s(t, is in H (pt) for a.e. t.
Definition A.6. On My define bilinear form by setting
(U = (VaJ, v99) Lalp)-
Definition A.7. We say that probubility measure p fulfills the partial Poincaré inequality (PPI) if there exists K > O such that
l 11 I KllVilli
(PPI)
for all g E Kj. Note that usually formulated as Varplgl S K|[V 2 Aig but since we have Ep, [9t] = for Xy both notions coincide.
Note that the above definition is more general than the pointwise Poincaré inequality used in Vargas al. (2024), since we can allow single time points where the inequality may fail.
Lemma A.8. Assume that p fulfills the (PI) and that P L1,,oc 1] x RR). Then, (X) (, 1xg)
Hilbert space.
Proof. We have to show that (s))x; x3 is positive definite and X3 is complete. The first claim follows directly from the (PPI). The second follows from the fact that |lag(e) I VK I|| 11x S ([)))) and Xg is closed in (p).
VK 11|
This assumption PP[ is usually fulfilled our setting as we discuss in the following.
Lemma A.9. (Bakry et al. 2008 Corollary 1.6) Let p= e da be a probability measure on TRd and assume that V E C²(12") such that V is lower bounded. Assume furthermore that there exists n 0, R > 0 such that for all ||>||| >> R we have
(x. vv())) 2 n||=
15

---

Published as a conference paper at ICLR 2025
Then, there exists constant C C(d, n. (o 0. where R 2d R}, such that for all ó E Hj) it holds max ")
J ( SC |V
The constant C can be estimated as follows: Let == Bv.R == 2[( There exist C, C2.,, R. Ca, R > 0 (not depending on V) and a universal constant D >> 0 such that
4d (1 [1 (Ci, + Von. C T/2
Ca. D Ree))
Proof. We first show the claim for all smooth & E (RRR)) (J) with zero mean. Unfortunately in (Bakry et al. 2008 Corollary 1.6), the constants are not explicitly computed which is why we track them here. They consider for y n/2 the function given by W(r) W, (())( == ed|l=|l for
||=||| > R, and smooth continuation for ||0||| I R. They show that, if for L == A (VV, it holds that
LW S (x) + (XB(R):
then, the inequality [Ø] C SR V || du true with constant C (1 bsR). where the constant KR can be estimated by D R2e We are left to upper bound the constants and b. Let
||>||| < R. Then, we have that
ILW (x)) < (AW
VW(x)) < || 1o VW (B(R))
and for ||||| R. by choice of R,
LW (z) 11 d
1]
(xr, VV)
W(x) S 7) (d (d-1
+ W(x)
|2||
|
2
W
17
n(d
1)
nld
W(x) = - W(x),
n72
2d
2d
and thus the claim is true for o E C (R) (p) with zero mean, using the constants 0 = 4d
>
and C1 + + Cin,,,,Ct....llVVVVVVVVV VV
C3. where C1,y. = (IA,,,)ll=<<(BR)) C2, ft [[V R (l and C3,,,R F,,,,
In order to extend this result to Φ (1),,, note that (RR) is dense in H (TR))) see (Adams & Fournier 2003 Cor. 3.23). Since by assumption, V is lower bounded, i.c, V is upper bounded, C is also dense in (H) by Lebesgue's 's dominated convergence theorem. Now take a se- quence (ox) E C (IRd) converging to & in (n). Then, also on == Φn -fg on du converges to in H (n)), and on E (Rd) H H'()) (Jt) has mean. The first part now yields the claim. Li
For the specific path P: given by the linear potential interpolation V = (1 t)fo + tfi, the disin- tegration measure p fulfills the partial Poincaré inequality (PPT) given suitable assumptions on the end points fo, f1. as the next proposition shows.
Proposition A.10. Let fo, fi E (²(Rd) be lower bounded such that there exists n > 0 and R > O such that for all || 2 R it holds the drift condition
(x. ff(())) 2 n|[æ/|
i 0-0,1
(27)
and the linear growth condition
f((x) 2 a | i= 0,
(28)
Let V, (1 + and pt 7, -e- for 7e Sga e- dx, and p as Definition A.1 Then,
p satisfies the partial Poincaré inequality PPI). i.e. there exists > O such that for all Φ X3 holds
JONDe 8' d'dp 1..1xx
/V
16

---

p Published as a conference paper at ICLR 2025
Proof. First note that by the lower boundedness and linear growth assumption (28) normalization constant Ze is well-defined, and Pt is indeed probability measure for all t E [0, 1]. Further note that a exists if Pt is weakly continuous, see again Definition A.1 We now show the weak continuity of pt:
fu. fi, the
Let be a bounded and continuous function IR.. By the linear growth 288 of for and domi- nated convergence, the map & + Z₁ = .R e dæ is continuous on [0, 1]. Since it is also strictly positive on [0, 1], also its inverse + Z, is continuous
[0, 1]. By the same token, the map t t .g. ge VV dz is continuous, and altogether, t ++ Pt. is weakly continuous.
In order to prove the inequality, we conclude by (27) that (z, VV2(2)) 2 (1 - t)m||=|| + ty||=|| Z ||| for all ||ر||| 2 R. Let == max 2d R} and note that for (t, = V((x) + In())) is now and = [|
easy to check that -(t,-) (R()),, V VVvVll ((0.)
20. Then, by Proposition A.9 we know that
J o(t,)) dp sC sCLe /V (1
JRd
for all Φ Xg and all E [0, 1], with constant C which does not depend on t, hence the claim.
Remark A.11. (i) Gaussian distributions given by Si == 2 I| mll? satisfy the assumptions of Proposition A.10
(ii) In Proposition A.10 the regularity assumptions on the end points, i.e., for fi E (R can be
relaxed. enough assume that there exist sequences of functions (fo,n); (fi,) C² (RR))
which converge pointwise to fo and fi, respectively, such that the assumptions of Proposition A.10 are satisfied uniformly in n and such that [[V Va || I. are bounded in n. Hereby, also the Fisher-Rao curve given by linear interpolation of Jo == ||=[|| and fi
min[[[
[|2r]], [ll || m with m R2d, ie., the interpolation between two Laplacian (type) distribu-
tions, satisfies PPI..
A.2 EXISTENCE OF A SOLUTION TO THE POISSON EQUATION
Now, we consider the Poisson equation (26) in a slightly more abstract manner: Showing that a given Fisher-Rao flow pt is also Wasserstein absolutely continuous curve requires solving the following (family of) PDEs
-V (pt = (a: ar) Pt. E [0,
(29)
for the vector field 8/ Vs, satisfying a suitable integrability condition. Here, E L2(p), att,- -) and 0g == Ep, [ay] are given.
=
In order to do that, we look to find weak solution to (29 in the space Xg. Definition A.12. An element & E Xg is called a weak solution to (9 if satisfies,
[ the (as -or) dpidl, Vub E 23- Note that for vi E C (20 1) x R is holds for g(t) == :Bp;
[[ Vaha dpidi
(30)
E., (5,(t)))))) that ub - E Xo- Since Vx = 0 f Sae (())(a) an) dpıdl, have that (30) is also fulfilled for all 1/ C ((0, 1) x RR)..
We can now prove Theorem D from Section 2.2 which provides a solution to the Poisson problem
29h globally in time L [0, I].
Theorem A.13. Assume that p from Definition A.1 satisfies (PI) for some K > and that 1 Lo (10, 1) x Rd) Furthermore, assume that 0x
Le(p), Le,, lloe ore melli2a00) S C < co. Then, there exists a unique weak solution E X3 to 29 Moreover, satisfies the following inequality
L''
|[s] dp s K S/') / (a: (a)- dp,dt.
(31)
Jud
Proof. Let A: X3 + R be defined as
A := 1 / is (a, (a.- (a)) dp
17

---

Published as a conference paper at ICLR 2025
Therefore, rewrite system of PDEs (29
(s, x)x Ad,
Vb Ng.
(32)
Clearly, a linear operator. Since p satisfies
(PP]) the following holds:
|Anb I/' S 1/e (ar (an--aat an)) dpddt) 2
=
=I - (/x) (ar(x) - ax(x))) ap(t,,))
$ (ʃ (ac-ai)) dp) (ʃ l /eb)dd)) s sc Lo Jec dpydt dpydt
1/' fpc / dppdt = C
S SCK
By Lemma A.8 the space X is Hilbert space. Applying the Riesz Representation theorem, there exists a unique element & E Xg satisfying (32).
Remark A.14. Note that in the case of Fisher-Rac flows, a satisfies
dept
(at at) Pe: E [0.
Then, the Fisher-Rao action (FR(Pt) of pr is given by So | - or (ill2z(p)) dt. Hence, the rh.ss of 31 can be described by K AFR((1))
Lemma A.15. Let p be such that Pt is bounded (from above) for a.e. e [0, 1). Then, for s E X we have that ss E Tp. P2(R.) ) for a.e. t [0,
Proof. For a.e. I, since Pr is bounded from above by assumption, C (R4) is dense in H (pe). Now, E H (pt) immediately yields the claim Ss E Tp, P2( (Vo C (Rd)) L2(ps)
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 Po, P1 be probability densities on R2 and assume that is a Fisher-Rao curve defined by
O.pt = (a+ 04 Pt:
such that p and satisfy the assumptions of Theorem A.13
Then, Pt is a Wasserstein absolutely continuous curve with vector field V 8,, where 8 E X is the unique weak solution of (30). If Ps is bounded for a.e. t, we have that V T,, P2(R)) for a.e. t E [0,1], 1], i.e,, for every other vector field vr satisfying the continuity equation for Pt. we have that [[velLLo(P)) 2 Lolpe) for a.e. 0,
Furthermore, it holds that
S. [ 1 dp;dt K J. / (ar )t) dprdt.
(33)
Note that by the Benamou- Breenie--Beeereee formula, see (Ambrosio et al. 2005 Equation 8.0.3), the left-hand side of (33) is an upper bound for the Wasserstein distance W (Po. P1), while the double-integral on the right-hand side defines the Fisher-Rao action AFR of pt.
Proof. Using (a,-az)p: = 0.p: and Theorem A. equation
13 we can conclude that Pt satisfies the continuity
0ep: V (prv)) 0
(34)
with Borel vector field V Vi, in the sense of distributions. Note that by construction,
X)
and hence, the gradient v, is already a morel-measurable vector field on [o, 1] x Rdd such that
V.st L2(p)) for a.e. t E [0, 1). Further, (3T) implies that t + | L (10, T)).
Together with the fact that Pt is weakly continuous, we infer by (Ambrosio et al. 2005 Theorem 8.3.1) that P1 is Wasserstein absolutely continuous. By Lemmal LemmalA.15 it holds that V E T P2("), Pe bounded for a.c. t. In this case, (Ambrosio et al. 2005 Proposition 8.4.5) immediately implies that Vs, has minimal norm Vs; L2(e.) among all Borel vector fields v2 fulfilling (34). 11
18

---

Published as a conference paper at ICLR 2025
Let us apply the above result to the curve given by linear interpolation of the energy functions fo, fi-
Corollary A.17. Let fo. fi E (RR) be lower bounded by some E IR such that the conditions (27 and (8 of Propostion A.10 are satisfied for some a, 1) > Further, for simplicity, assume that Jo, J1 have at most polynomial growth outside ball, i.e., there exists R >> O such that
[f(z)| S M| for all ||]|| R 0,
(35)
Then, Theorem A.16 can be applied to the Fisher- Raa curve Pt given by V₁ (1 - t)fo tfi and P: Z, e Vs with Zt SR e dr.
Proof. First, by the drift condition 27 and linear growth 28 the measure P. satisfies (PP]) by Proposition A.10 Furthermore, in Proposition A.10 we showed that t + Z, is continuous on [0, 1]. hence x1,z) + (x) = Z - -Ve((t) strictly positive and continuous on [0, 1] Rd. There-
fore, it immediately follows that E L1,o0 [10 RR)..
p
We are left to check the Ly-condition on a = - = fo f1. First, since + Z continuous on [0, 1], hence bounded, follows together with (28 and (35) that
/ 15o -file dr
Rd
/
/ 1fo file -Vi dr
$ z
1fo -flle Vi dæ + Zz!
-ile-Ve
|00|/S
$ C1
1fo file dx + +C2 M I
|0||a -a |al dæ
|a|SR
||
< <
with constant C > 0 not depending on f. Hence, E T2(()), and we are left to prove an T2(p). But this follows by the similar arguments
Zi¹ | 1io filee dr s ZT¹
falee V. dr J ff file- Ve dx + Zz,' / |11122 Ifo
$ C1 /
Ifu fe dx C2
/
dr -l]=
|
||||>
< C < 00,
again, with constant C > 0 not depending Integrating over t E [0, yields as L2(p), and hence, all assumptions of Theorem A.16 are fulfilled.
Remark A.18. Hence, the Fisher-Rao curves defined by linear interpolation V, = (1 (1-t)fo + tfi are absolutely continuous the Wasserstein sense, e.g., the cases (see also Remark Ae
Gaussian distributions given by ft = 202 || mm
Laplacian distributions given by f;:= ||1 [ill-
Laplacian type distributions given by fi := min |z]|| || ml]).
Let us briefly discuss the differences between the theory developed here, and the prior results from Laugesen et al. 20144): Vargas al. 2024).
In Vargas al. 2024)) they assume that Pt E C and Poincaré inequality for each time t. We need much weaker conditions: we assume that + Pt weakly continuous, p = p,dt fulfills a (partial in æ) Poincaré inequality and Tp)( In partieular, allow single time points I, where the Poincaré inequality might fail for Pt-
((0, ( 1], R and E C ((0, 1] Rd)
Our solution Φ is an element of H' H' (p) C. = 1] x R4; p) and thus automatically mea- surable on the whole time domain & [0,1]. 1]. This makes it possible to show that Vol(
19

---

Published as a conference paper at ICLR 2025
is Borel vector field with regularity in time which implies absolute continuity of the flow Pr in the Wasserstein geometry. contrast, Vargas et al. (2024) only considers a solution for every single [0, 1] without any regularity in t. This Is result already shown in (Laugesen et al. 2014 Thm. 2.2) and motivated our investigation into a framework across the time t.
Furthermore, we give precise mild assumptions under which the partial Poincaré inequality is satisfied. In particular, cover the case of the linear interpolation. As result obtain that these curves are absolutely continuous in the Wasserstein geometry.
B TELEPORTATION ISSUE WITH THE LINEAR INTERPOLATION
In this section we want to demonstrate that the norm La(p)) the (optimal) velocity field (4) belonging to the Fisher-Rao path given by linear interpolation fr := (1 - fo + (fi, explodes for times / close to 1, when the measures Po and pi are asymmetrical to each other.
In the following, we restrict ourselves the one-dimensional setting TR: in order to exploit the isometry from the Wasserstein space P2 (R) into the convex cone C C (0, 1) of quantile func- tions, also see (Duong et al., 2024, Section 3). For the sake of self-containedness, we give a short introduction of the necessary tools: For μ (2R we define its cumulative distribution function (CDF) by
(a(2)) := ((())==(-o<)) m(000,0) DER,
and its left-continuous quantile function
(s(s) == min(x E R F(x) 2 s) SE (0,1).
Then, it holds the isometric property
(2, ( Qu())) for all μ, v E P2(R).
(36)
Now, by (Ambrosio et al. 2005 Theorem 8.3.1), the norm (a(e)) the optimal field e belong- ing to an absolutely continuous Wasserstein curve can be characterized by the metric derivative of the curve via
[[ellez(e))= = (p)( lim We(Pe, Po) for a.. t 1].
t
Using the aforementioned isometry 36) and Komura's 's theorem, it holds for a.e. t E [0, 1]
W2(Pr, Ps)
11p. Qp, (1)0
lim 8 8-
lim [-s)
t
Hence, if we can explicitly compute the quantiles Qp, (and its time derivative), we get an analytic representation of ||ve (ezee) via
| Ivellin(os) = S'' (o))) ds.
(37)
In order to do so, choose the prior and target measures the following way: Let fo(æ) |x|
and (x) : min æ/. Jr ml ml). where the second mode m > 0 is fixed for the moment. By
Corollary and7 and Remark A.18 the Fisher- curve P: given by linear interpolation of for fi is
Wasserstein absolutely continuous. Now, straightforward (but tedious) to explicitly calculate
the CDF and quantile function of P::
First, since Pt x pb p and 2 min { [z), Jx - m| = |a|| + | ml
||||| - |x ml/. we obtain
P: e x] -1--1--mm++(|aalll-----
for all x E R. t E [0, 1],
up normalization constant Z, which we calculate below.
20

---

Published as a conference paper at ICLR 2025
Calculating the CDF of pc- By dissolving the absolute values suitable subintervals, obtain for the CDF
1
Ze Z.FFo( (2)
(1++))
x ESO. O,
1
+
2
1
(x)
((++)
0 0 -
1+1
1
Ze Zp.Pp(()
2
11 +
1
(1+))
e (014333)) 2mm
m
S x S m,
1+t
-1+ 3t
1 + 3t
2
2
Ze 7. (x)
1 1 +
T (1+t)
1+t
-1 3t
e e-m(1 mm(111)
1
x(1+()+2tm
x 2 m,
1+t
-1+ + 3t
+
where the normalization constant is hence given by
Z: = Ze FF, = 1 I+1
E (7+()
1+1
-1 + 31
1
+ -1 3
mmm(1_t)
Calculating the quantile function of pt. Since Fx. is continuous and strictly increasing, its inverse given by F Pc =Qp.- simple calculation now yields
(s) I+1 In (2+(1+ 0<s581
1
(x)))) =
In (1 + ()8 2), 81 S 82,
1
(s) In (n(1++31)) (87,
2 + ( T- -143 ) x(+0)) + 21m
I+7
82 S 83
-1 31
Qp. (s) -In (Z((1 + t)(1 s)) + 2tm
83 3 s <
1tt
Here, the subintervals are given by
$1 =
Z₁(1 t)
1
-7 (1++)))
82
7,
1 +
1
1
1
83 = Z, 1+t
1+t + -1 3
)e -7 (1++)
e mm11 mm1--))
1 + 38
Calculating the time derivative a.Q,, and the norm I|0:1lE Using standard differentiation rules, the time derivative of Qp, can be computed separately on each subinterval (s;, 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 (37). By numerical integration, we can finally approximate the norm |/ellizaee)) arbitrarily close. The computed results are depicted Figure [ for different values of the mode m > 0; roughly speaking, for late times ~ 1, the norm of the velocity field grows exponentially with m.
On a side note, the integration of 0,Qp. only on the subintervals (s1, 82) and (s2, 8s) contributes to the explosion of [|ul222(e.) Intuitively, this corresponds to the fact that mass gets shifted very abruptly mainly in the areas E (0, 7 and (T) m) for late times t.
On related work. The teleportation issue described here has been discussed other works under the term mode switching label switching, see, e.g., Woodard al. (2009); Syed et al. (2022); Phillips et al. (202);; [hehab & Korbal (2024): Noblle et al. (2024). We complement these mostly empirical observations by theoretical results.
21

---

Published as a conference paper at ICLR 2025
t-0.0
t-0.26
1-0.52
t-0.7
1-1.00
5-0.00
t-0.30
b-0.53
t-0.82
1=1.00
b-0.00
b-0.10
t-0.91
1-1.00
b+0.00
0-0.3
b-0.99
b+1.00
Figure 5: Evolution of probability densities for {1,5, 15, 503.
22

---

Cl. Published as a conference paper at ICLR 2025
C ALGORITHMS
Algorithm Learning of, Cj in (17) for x² sampled from trajectories.
Parameters and Networks:
(0,, 02) and u Cf2
Functions: fr = fo tfi, uft Che
Loss: {00 x, 1) = [ fo - c+ + (Vir, -
Algorithm:
for 1 = 1,... do
grad = 0
for 1, ...,, 72 do
Draw x² pz, sjsm
S
ODEsolve o², solution of 0r. v (xt),, x² at time points & E [0,
Detach oy from computational graph
grad+= Vo E - 2(0,2%,11 t1)
end for Update v, of2 using grad end for Output: u/ C.2
Algorithm 2 Learning f( 02 C in (18) for x' sampled from trajectories.
Parameters and Networks: 0 (01, 02 03) and wit
02
Functions: f (1 fo Lfi (1 1) 01
Ut c.
02
Loss: (0,, .x,)) t) |0, f (x) C+ + (Vft 2t
Algorithm:
for i 1,...
grad = 0
1, l,.,, n do
Draw x pz, Sj
ODFsolve oy solution of au: = of (xt), xd at time points 4 E [0,
Detach oy from computational graph
grad+= Vo Em ET =(0, t1)
end for
22 ofa using grad Update ,0²² end for Output: 1ht 02 C...
Algorithm 3 Learning vi ch2 in 19 as done in Section 4
Networks: vb cC,²
Parameters and functions: 0 = (0,, 02) and = (1 (1- + tul CP2
Loss: E(0,
10: C²= ( v(f.h fz)) A =22)
Algorithm:
for i 1, I,.... N do
pz, 1 < $ and oif Ma,b, 1 و uniformly Draw L E [0, Sj S m uniformly and use 25 to set at,
Draw of
grad V, Σm +0,,
Update &et C.2 using grad
end for
Output: w/ C²²
23

---

Published as a conference paper at ICLR 2025
D ADDITIONAL DETAILS ON THE EXPERIMENTS
Implementation and evaluation details. We implement the algorithms in Pytorch Paszke et al. (201)). 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 16. The linear interpolation uses the same network for the time dependent function Ct. For the learned and gradient flow interpolations we implement C₁ using the same network with less parameters. The gradient flow interpolation also uses smaller network for the time scheduling function y as defined Section + In case of normalized energy functions we can omit the parametrization of C₁ for the learned and gradient flow interpolations. We simulate the corresponding ODEs using the torchdiffeq Chen (2018) package with the Runge-Kutte Kutta adaptive step size solver ("dopri5"). For the linear and learned interpolation the loss is computed and the gradients are accumulated with a batch size of 256. For the gradient sample 4096 particles at random uniform time points and therefore do not
use 50 time steps along which
flow interpolation accumulate gradients. We adjusted the number of iterations such that all methods ran approximately the same time the same hardware.
We compute log weights of the generated samples by using the continuous change of variables formula and the corresponding implementation from Chen (2018). We evaluate the methods by generating 5 10 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.] (2023). We compute the energy distance using the GeomLoss library Feydy et al. (2019). For the target energy functions and plotting utilities we use the code provided in Midgley al. (2023).
The gradient flow interpolation faced some difficulty 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 also used an Euler scheme and resampled a given batch if 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 Gaussians with different standard deviations. Note that when using small standard deviation not all modes are recovered. On the other hand choosing a very large standard deviation leads to the modes close to not receiving enough mass.
10
20
30
= 40
Figure 6: Results for the linear interpolation for different values of a.
ESS
NLL L
Energy |
10
0.002
7.071 1 0.005
2.887 0.021
σ 20
0.063 + 0.07
7.362 + 0.005
0.592 = 0.005
σ 30
0.183+ 0.05
8.657 + 0.012
0.247 0.016
σ 40
0.092 0.001
12.909 +: 0.023
0.554: 0.008
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 a.
24

---

Published as a conference paper at ICLR 2025
0 = 10
0= 20
30
0= 40
Figure 7:
Results for the learned interpolation for different values of a.
ESS 1
NLL 1
Energy +
T 10
0.009
0.011
7.275 + 0.005
1,253 0.017
a 20
0.605
0.316
6.894 0.005
0.058= 0.006
30
0.922 0.002
6.913 1 0.003
0.121 = 0.009
40
0.139 1 0.052
7.779 1: 0.012
0.531 = 0.014
Table 4: Comparison of effective sample size, negative log likelihood, and the energy distance for
the learned interpolation with initial distribution being Gaussian with standard deviation d.
25


Reviews:
[{"review_id": "uTdH0ibpEO", "reviewer": "Reviewer_m4s7", "summary": "This work contributions are twofold:\n\n1. First, the work provides conditions under  which a probability flow admits a velocity field satisfying the continuity equation\n2. 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.\n\nThe work then presents some empirical results motivating the success of their proposed method.", "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 ?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "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.\n\nIn 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.", "weaknesses": "Unfortunately, this work misses prior work, which already explores learning the exact same vector field / continuity equation in the exact same context/task.\n\n1.  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 \n2. 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\n\nSo, 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.\n\n 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).\n\nAnother 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. \n\n\n[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.\n\n[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.\n\n[3] Albergo, M.S. and Vanden-Eijnden, E., 2024. NETS: A Non-Equilibrium Transport Sampler. arXiv preprint arXiv:2410.02711.\n\n[4] Sun, J., Berner, J., Azizzadenesheli, K. and Anandkumar, A., Physics-informed neural networks for sampling. In ICLR 2024 Workshop on AI4DifferentialEquations In Science.\n\n[5] Máté, B. and Fleuret, F., 2023. Learning interpolations between boltzmann densities. arXiv preprint arXiv:2301.07388.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Nkb5yQJRVj", "reviewer": "Reviewer_rTNR", "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.", "questions": "1. General background section. What do you denote by $C_{c}^{\\infty}$? Is it the set of (infinitely-many) differentiable functions?\n2. 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).\n3. Lines 359-360. I do not understand, what you mean by “[...] it is not clear if Boltzmann densities stay Boltzmann densities”\n\n**Misprints**\n\n1. Line 187 - no comma in $\\nabla \\cdot (\\rho_t, v_t)$ in eq. (11).\n2. Line 037-038. Velocity field $v: [0, 1]\\times \\mathbb{R}^d \\rightarrow \\mathbb{R}$ maps to scalar? Misprint?\n\n**Minor comments**\n\nA 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.\n\n**Related works**\n\n[Mate&Fleuret] Mate et. al., Learning Interpolations between Boltzmann Densities, TMLR’2023\n\n[Xu] Xu et. al., Normalizing flow neural networks by JKO scheme, NeurIPS’2023\n\n[Boffi] Boffi et. al., Probability flow solution of the fokker--planck equation", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 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).", "weaknesses": "I think that the paper has rather limited methodological and practical contributions.\n1. 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.\n2. The scalability of the method is under question. The experimental illustrations do not stress test high-dimensional applications.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qPkMkFkrzq", "reviewer": "Reviewer_FHYM", "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.", "questions": "Could you please explain in more detail the precise methodology used for solving the ODE (6)? Do you need a stiff ODE solver?\n\nBoth examples you've considered are 2 dimensional. Could you discuss the scalability of this methodology to higher dimensional problems?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "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.", "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.\nThere is no theoretical guarantees for the neural network approximation that is used to solve (17) actually is able to solve it.\nThe scalability of this approach to high dimensional problems has not been demonstrated.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "SK2U1rAifp", "reviewer": "Reviewer_mKTY", "summary": "The present paper considers the problem of sampling from a Boltzmann distribution known up to a normalization constant.\n\nThe 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.\n\nUsing 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.\n\nUltimately, 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.", "questions": "-", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 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.\n- The paper highlights and fruitfully exploits connections between SDE/ODE sampling and gradient flows on probability spaces. \n- The paper draws its inspiration from prior literature, but draws new connections to propose amendments to the approach.\n- 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. \n\nI am not familiar enough with the formal literature on Wasserstein curves to judge the novelty of the theoretical results of the paper.", "weaknesses": "- The writing of the paper could be improved, in particular, some notations are regularly used before being introduced. \n- The discussion of related works should be extended to include:\n\t- Stochastic interpolants: https://arxiv.org/abs/2303.08797 \n\t- Non-parametric denoising based samplers: RDMC https://openreview.net/forum?id=kIPEyMSdFV & SLIPS https://proceedings.mlr.press/v235/grenioux24a.html\n- The experimental validation is currently very limited, notably in terms of how the method scales with increasing dimension.\n\nMinor:\n- maybe misses the assumption that $\\rho_1 = \\rho_D$?\n- Notations in Equation (2) are not defined, also lacking Equation (5)\n- In (17), the space variable is sampled uniformly on an hypercube? How are the boundaries chosen?\n\n\n- typos:\n\t- line 96: Kullback Leibner\n\t- extra comma between $\\rho_t$ and $v_t$ in the divergence (11)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FIVNncVlGr", "reviewer": "Reviewer_CXu4", "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.", "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.\n- Does there exist a case where the gradient flow method is clearly better than the learned one ?\n- There very few experimental validations, I would like to see another dataset, maybe inspired from the one investigated in Máté and Fleuret.\n\nIt 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 ?\n\nMinor remark:\n- there is a tipo at \"Leibner\" line 96", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper make very interesting connections between Wasserstein Flows and the developed method. \nIt also underlines its relation with Diffusion models.", "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. \n\nAmong 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.\n\nThe results on the experiments are ok, but not super-convincing either.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)