An Explicit Expansion of the Kullback-Leibler Divergence along its Fisher-Rao Gradient Flow
Authors: Carles Domingo-Enrich
Aram-Alexandre Pooladian
Venue: Accepted by TMLR ()

Abstract: Let $V_* : \mathbb{R}^d \to \mathbb{R}$ be some (possibly non-convex) potential function, and consider the probability measure $\pi \propto e^{-V_*}$. When $\pi$ exhibits multiple modes, it is known that sampling techniques based on Wasserstein gradient flows of the Kullback-Leibler (KL) divergence (e.g. Langevin Monte Carlo) suffer poorly in the rate of convergence, where the dynamics are unable to easily traverse between modes. In stark contrast, the work of Lu et al. (2019; 2022) has shown that the gradient flow of the KL with respect to the Fisher-Rao (FR) geometry exhibits a convergence rate to $\pi$ is that \textit{independent} of the potential function. In this short note, we complement these existing results in the literature by providing an explicit expansion of $\text{KL}(\rho_t^{\text{FR}}\|\pi)$ in terms of $e^{-t}$, where $(\rho_t^{\text{FR}})_{t\geq 0}$ is the FR gradient flow of the KL divergence. In turn, we are able to provide a clean asymptotic convergence rate, where the burn-in time is guaranteed to be finite. Our proof is based on observing a similarity between FR gradient flows and simulated annealing with linear scaling, and facts about cumulant generating functions. We conclude with simple synthetic experiments that demonstrate our theoretical findings are indeed tight. Based on our numerical findings, we conjecture that the asymptotic rates of convergence for Wasserstein-Fisher-Rao gradient flows are possibly related to this expansion in some cases.


Published in Transactions on Machine Learning Research (06/2023)
An Explicit Expansion of the Kullback-Leibler Divergence along its Fisher-Rao Gradient Flow
Carles Domingo-Enrich
cd2754 @myu. edu
Courant Institute of Mathematical Sciences New York University
-ham-Alexandro Pooladian Center for Data Science New York University
owan-aleeandre.oooaaaieee Any.. edu
Reviewed on OpenReview: https: wwtp///vew...eee////ee formm? id= 9p
Abstract
Let V. : Rd + R be some (possibly non-convex) potential function, and consider the prob- ability measure x e-v. When exhibits multiple modes, it is known that sampling techniques based on Wasserstein gradient flows of the Kullback-Leibler (KL) divergence (e.g. Langevin Monte Carlo) suffer poorly in the rate of convergence, where the dynamics are unable casily traverse between modes. In stark contrast, the work of [Lu et al. (2019) 2022) has shown that the gradient flow of the KL with respect to the Fisher- (FR) ge- ometry exhibits a convergence rate to that is independent of the potential function. In this short note, complement these existing results in the literature by providing explicit expansion of KL(p, FR [[T) in terms of e where (oPR)1z) is the FR gradient flow of the KL divergence. In turn, we are able to provide a clean asymptotic convergence rate, where the burn-in time is guaranteed to be finite. Our proof is based on observing a similarity between FR gradient flows and simulated annealing with linear scaling, and facts about cumulant generating functions. We conclude with simple synthetic experiments that demonstrate our theoretical findings are indeed tight. Based the asymptotic rates of convergence for SonseerstiinFissere-ee gradient flows are possibly related to this expansion in some cases.
numerical findings, we conjecture that
1 Introduction
Sampling from distribution with an unknown normalization constant is a widespread task in several sci- entific domains. Namely, the goal is to generate samples from probability measure
TΦt) x e V.(())
where Va : Rdd R is some (possibly non-convex) potential function that is available for queries. most cases, the target measure æ is only known up to the normalization constant. Applications of sampling from π include Bayesian statistics, high-dimensionall integration, differential privacy, statistical physics and uncertainty quantification; see Gelman al.
(995)] Robert al. (1999 MacKay (2003); Johannes & Polson (2010g Von Toussaint 2011 Kobyzev et al. (2020) Chewi
2022 for thorough treatments.
Recent interest in the task of sampling stems from the following paradigm: sampling is nothing but optimiza-
tion over the space of probability measures (Wibisono) 2018). This interpretation is due to the connection
between the celebrated work of Jordan, Kinderleher, and Otto (Jordan et al. 1998) and the Langevin diffu-
sion dynamics given by
dX, vV, dt v2dB..
(1)

---

Published in Transactions on Machine Learning Research (06/2023)
where dB, is Brownian motion Indeed, the work of Jordan et al. (1998) demonstrates that the path in the space of proabability measures given by the law of Eq. (1 is the same as the Wasserstein gradient flow (i.e. steepest descent curve in the Wasserstein metric) of the Kullback-Leible (KL) divergence
(L(p||T) (L(oll))=
log L dp.
We write (ot )c20 E P(RR) for the law of the path given by Eq.
1 (see Section 2.2.1 for precise definition).
A central problem in this area has been to bound the convergence rate of W to т in certain similarity metries (e.g. the KT. divergence itself, or the Wasserstein distance) under different conditions on л. These
bounds translate to convergence rates for the Langevin Monte Carlo (LMC) sampling algorithm (Dalalyan
& Tsybakov 2012; Vempala & Wibissnoo 2019: Durmus et al., 2021 Chewi et al. 2022 upon accounting
for discretization errors.
The classical result is follows: assuming that satisfies Log-Sobolev inequality (LSI) with constant
CLSI > O, we obtain the following convergence rate (Stam 1959
Gross 1975
Markowich & Villani 1999)
KL W (|x) s KL W )))
21.
(2)
which holds for all / Recall that ı satisfies LST if for all smooth test functions g,
ent= s 2C
where (g) == (g Eng log Eng. For example, when V, is a-strongly convex, an LSI with CLsT 1/a holds. LSI hold more generally, but sometimes with very large constants CLSI- Indeed, for multimodal distributions such as mixtures of Gaussians, CIsI scales exponentially the height of the potential barrier between modes (Holley Stroock 1987 Arnold et al.] 2000). This impacts convergence at the discrete-time level, and thus hinders our ability lo generate samples using LMC.
Another geometry that gives rise to gradient flows over probability measures is the Fisher-Rao (FR) geometry; see Section 2.2.2 for definitions. Similar to the case of Wasserstein gradient flows, we let (p. be the FR gradient flow of the KL divergence. Recent work by Lu and collaborators has shown that the convergence Pi FR + π occurs at a rate that is independent of the potential function V.. This is in stark contrast to the case of Wasserstein gradient flows, where the rate of convergence is intimately related to the structure V. through the LSI constant. In their first work, Lu et al. 2019) show that for any 8 E (0, }l there exists a t, 2 log(5*) such that for all t 2 t,,
KL (T)) S KL PR ([)) (.)
where they require warm-start condition (L [[m] S 1, and assumption (B) (see Section 3 In Lu et al. (2022), the authors show that the KL divergence is always contracting under (e(")) even in the absence of warm-start, though with worse rate. Combined, these two results provide the first contimous-time convergence rates of the gradient flow of the KL divergence under the FR geometry to π.
(WFR) geometry. The WFR geometry has recently been used analyse the convergence dynamics of parameters of neural networks (Ciiatt 2022). mean-field games (Rotskoff et al. 2019), and has shown to be useful in statistical tasks such Gaussian variational inference [Lambert et al. 2022), and identifying parameters of a Gaussian mixture model Yan et al. 2023). In the context of sampling, particle-based methods that follow dynamicc governed by WFR gradient flow of the KL, written (pWFR) )220 are known to escape the clutches of slow-convergence that plague the Wasserstein geometry. simple observation [Lu et al. 2022 Remark 2.4) gives the following continuous-time convergence rate for & 2 t;:
Merging both these geometries gives rise to the well-defined CnnserrstiinFissher
KL WPR () S KL
WFR
[[)) min {e Cuss (0-00s,,-11--3000--
-(2-351(e-2a)
(5)
This equation is be understood from the perspective of Itô caleulus.
2

---

Published in Transactions on Machine Learning Research (06/2023)
where 6 and t, are as in the FR convergence rate (4) Loosely speaking, this "decoupled rate" is consequence of the Wasserstein and FR geometries being orthogonal to one another; this is made precise in Gallouët & Monsaingeon (2017).
As elegant as this last connection may seem, the convergence rate Eq. 4 and consequently Eq. 15 should appear somewhat unsatisfactory to the reader. Tt raises the natural question of whether or not the factor of appearing in the rate is avoidable, and whether the upper bound in Eq. is tight.
1.1 Main contributions
We close this gap for the KL divergence and any q-Rényi divergence. Using different proof technique than existing work, prove the following asymptotic rate of convergence for the flow (p, 122, namely,
(p.R)
P π
KL [|
[T)) Varx log
+(E -
(6)
and a similar result. holds for all q-Rényi divergences. Our assumptions are weaker to that of prior work, and given that this is a tight asymptotic convergence rate, we conjecture that the assumptions are likely unavoidable in the large t regime. Our proof technique provides an explicit expansion of (L(p)"1[])) (and q- (-Rényi) in terms of et We supplement our finding with simulations for all three geometries, indicating that our convergence rate in fact tight for Fisher-Rao gradient flows, and sheds light on possible conjectures for the convergence rate of WFR gradient flows.
Notation
For probability measure p (((R))) and a function f: TRdd IR, we sometimes use the shorthand (f)p f We let log(.) denote the natural logarithm, and we use the standard shorthand notation f Og)), meaning there exists constant C> 0 such that f S Cg.
2
Background
2.1
Definitions
The study of gradient flows has a rich history in both pure and applied mathernatics. The development of the relevant calculus to understand gradient flows is not the purpose of this note, and we instead provide a barcbones introduction. However, we strongly recommend the interested reader consult standard textbooks
the topic, namely Ambrosio et. al. (2005). and the first chapter of Chewi (022)
Let ((R)) be the space of probability measures over gd. A functional F P(RR) R is defined the space of probability mcasures, with p F(p) R. We call 5F((p) the first variation of F at if for signed
measure 17 such that dr = 0, it holds that
F(p En) - F(p) = 1 -6FFl))
lim 6-+
(7)
The Kullback-Leibler (KL) divergence of measure P with respect to some fixed target mcasure is defined as KL(pl(T) log & dp for absolutely continuous with respect to For x eV, the first variation of the KL, divergence is given by
+ {KL(-I+)())()) = log p(z) log + V-(x) log Zi
(8)
T(x)
where Z₁ is the normalizing constant for
A more general notion of dissimilarity between probability measures is the q-Rényi divergence: for E [1, oo). we define (pl(T)))) to be the q-Rényi divergence with respect to ग, given by
1 (pl(pl --1 log (e)) dm
(9)
3

---

Published in Transactions on Machine Learning Research (06/2023)
for measures p that are absolutely continuous with respect to Rg recovers the KL divergence in the limit + 1, and when 2, R2((((z) log(x2 + 1), where xe is the chi-squared divergence, written explicitly as
(el((T) = Var L - (G)
dm
2.2
Gradient flows of the Hullback-Leibles divergence
2.2.1
Wasserstein gradient flow
In its dynamic formulation, the 2- Wasserstein distance between two probability measures Po: P1 with bounded second moments can be written as (Villani 2008 Benamou & Brenier 2000
(.' [[x(()))) dx) de s.t. depr + V (pees) = 0,
Wa p1) = inf
(10)
(Per))
where (er))) is curve of probability densities over Rd, and (v)) a curve of (²(Rd) vector fields. The constraint is known the contimuity equation, with endpoints Po and P1. For functional F: P(RT) the Wasserstein gradient flow is the curve of measures (p)) )c20 that satisfies the continuity equation with the vector field replaced by the steepest descent under the Wasserstein geometry,
Ft v(, VoF(p!)
where the last equation is simply the (standard) spatial gradient of the first variation of F. Plugging in the expression for the first variation of the KL divergence (8, we see that the law of the Langevin diffusion is given by @ which satisfies
0ep. W V (p. (V log P + VV.))
(11)
This equation may be rewritten as 0,p, = (VV.p,") + Ap!" which readily identifies the Fokker- Planck equation for the potential V. The equation describes the evolution of the distribution of particle that moves according to the stochastic differential equation [ At the particle level, the key aspect of Wasserstein gradient flows that they model partiele transport, and that makes them useful for high- dimensional applications such as LMC. what follows, we will sometimes abbreviate Wasserstein gradient flow to W-GF.
2.2.2 Fisher-Rao gradient flow
The Fisher-Rao distance, or Hellinger-Kakutnn distance, between probability measures has a long history in statistics and information theory (Hellinger 1909 Kakutani 1948). It can be defined as (Bogachey 2007 Gallouët & Monsaingeon 2017)
(pere) ( (1) s.t. Oep: rePt.
(Po Pi) inf
where (Pe)cefo.i) is again curve of probability measures, and (rz)eejo.1 is a curve of (T2() functions. Together, they satisfy the prescribed equation, with endpoints equal to Po and Pr. The Fisher-Rao gradient flow of the KL divergence, also known as Birth-Dealh dynamics, is the curve of measures (0f )(20 that satisfies (Gallouët & Monsaingeon
2017 Lu al. 2019
0xp.
PR.
FR az,
Or == log
Pi
KL(P) PR (T)..
The first term adjusts mass (i.e. gives birth to or kills mass) according the log-ratio of p FR and the target measure The last term preserves the total mass, that /f. P(Rd) for all time.

---

Published in Transactions on Machine Learning Research (06/2023)
Expanding this equation, we have
aep (x) (log(pp (x)) + V.(2) (log(p!") + V.) p. )p (x))
(12)
We henceforth omit the superscript FR for the Fisher-Rao gradient flow of the KL divergence unless the notation becomes ambiguous. For short-hand, we make use of the abbreviation FR-GF for Fisher-Rao gradient flows.
The FR-GF may be simulated using a system of weighted particles (see Appendix B Unlike for the W-GF, in this case the positions of the particles are fixed; only the weights change over time. Hence, simulate the FR-GF one is forced to grid the underlying space R., This is feasible only for small dimensions d. Consequently, FR-GFs cannot be simulated in high dimensions, which makes them impractical for sampling applications.
2.2.3 Wasserstein-Fishees Rao geometry gradient flow
The Wasserstein-Fisher Raoo distance between probability measures arises combination of the Wasserstein
and the Fisher-Rao distances (Chizal et al. 2018 2015 Kondralyev el al. 2016 Liero el. al. 2016 2018
Tt is defined as
(peried LC ((lee())))))))))))))))) dr s.t. Ospe + V (pere) = rtPt.
(PR
inf
rex))
where, for each / E [0, 1], the triple (Pe:Ve, U/s T2) lives in P(R4) x (2( x (²(R),, and they simultaneously satisfy the constraint equation, which has endpoints Po and (l, as well. Similarly, the Wasserstein-Fishee Fiiisteiii--sss--Raaa gradient flow of the KL divergence is the solution of PDE that incorporates the terms in the Wasserstein and Fisher-Rar gradient flows (Eq. [11 and Eq. (12)
WFR O.p. O.P.ER
(p. WER (V log WFR VV.)) (log(Pt WFR + V. -
VV. (log(p) (FR)) WFR) + V.) won) WFR
(13)
Similar the other geomctrics, we write WFR-GF shorthand for Chasserstein-Fishezz
gradient flow At the particle level, WFR-GFs are able to captur ure both Iransport and weight updales, which is why they enjoy convergence rate that at least matches the better rate between W- and FR-GFs (recall Eq. (
and is clearly superior in practice in some instances. Hence, any improvement in the convergence analysis of either W- or FR-GEs translates to improving our understanding of WFR-GFs.
2.3 Simulated annealing dynamics
Simulated annealing is a technique seen in several works when allempting to either optimize a function or sample from a multimodal probability distribution, and has a long history (Pincus) 1970 Kirkpatrick et al. 1983), and plays a crucial role our analysis. what follows, we introduce the annealing path with linear scaling, and conclude with a proposition.
Consider the time-dependent measure (μT)re(o.1) corresponding to the annealing path, with linear scaling, initialized at the measure ni Po e eVo, By definition, Per admits the density
p+(x) (o)
7 7,
VoV dx,
(14)
Z+
for E 10. Note that indeed, µ = To this end, it will be convenient to rewrite Eq. the log-density of pr. Remark that
14 in terms of
log(u+(4)) = (V.((x) Vo(x)) - Vo(x:) - log Zr
(15)
One can check that the pointwise derivative of the density ne (with respect to r) is
0z) = -(V.(2)) Vo(x) (V. Vo)n.)
(16)
5

---

Published in Transactions on Machine Learning Research (06/2023)
From this, obtain that
log(ur())) + (x( - (lg((u)) + Vo)
-V.((x) - Vo(=)) Vo(x) (-T(V. VV)) - 1o+ V.), + +V((2) -(--)
--](V(<x) - Vo(x)) Vo(x) +V(()) (-T(V, -V) V/ +V=)
(17)
(1 (V)()) -
- (1 --, - Volm
= -1 (v - (V. - VT)
Note that in the first equality, we used that the log-partition is constant and gets cancelled out by the difference of the two terms. Consequently, Eq. (16) can be rewritlen, for (0, 1), as
0-p+(x) = - (log(u+())) + V.(x)) (log(u>) + +v.), (v,) 1
(18)
A first observation is that that the linear schedule т in the exponent of Eq. 14 results in dynamics that resemble the Fisher-Rao gradient flow of the KL divergence, up to a reparamctorization that can be made explicit. Indeed, if one compares Eq. (18) with Eq. (12), the only difference is the factor r in the right- hand side of Eq. 18). Since the solution of the Fisher-Rao gradient flow of the KL divergence is unique (see
Proposition 4 in Appendix A appropriate time reparameterization of the annealed dynamics (14) will 12 We summarize this observation in the following proposition, which were unable to find a citation for in the literature.
yield the solution
Proposition 1. Let (μ+)re(0.1] be as defined Eq. 14h. The Fisher-Run gradient flow (pt)ezo of KT(pl(z) (i.e. solving Eq. 12 is given by P = Pi-
Proof. If we write l as a function of we have that
dt (T) - dt O.P((T) = depe(t) dr
(())))+(()()()()(00)00) (T1g(()))))))) E(x) (log(Pe(r)) + E) ((rs))))
(19)
Identifying Pi(r) with Pr. and establishing direct comparison with Eq. 18 we obtain that for Eq.
L9
hold, t(r) must fulfill dt (r) = =__-- With the initial condition that =(0) 0, this differential equation has dr
--r- the following unique solution:
((T) ['
1 1-8 ds logg1 7).
(20)
That is, we have that t(7) = log(11 - -), or equivalently, T(t) = - -t
2.4 Cumulants and their power series
Our argument hinges on observing a relation between the above gradient flows and their connection to cunulants of random variable. Recall that for random variable Y, its commulant-generatiin function to be Ky(=)= = log V*.. The th cumulant 6n. of the random variable Y is defined as the th derivative of Ky- ovaluated 0, that is, K, = x<f")()) (0). Similar to moment-generating functions, (y finite in some neighborhood of (-co,en) en), then holds that Ky smooth (in fact, holomorphie) (see e.g. (Shiryaev 1984 Section 11.12.8). Moreover, (x)z) admits the following infinite series expansion
Ky(z)
Σ Kn
n!
In particular, one easily check that xn E and Aa Var(Y).

---

Published in Transactions on Machine Learning Research (06/2023)
3
Main result
The goal of this section is to prove our main result, which is an explicit expansion of the KL divergence in terms of log-cumulants of the random variable log (o(X) throughout, and we will make their uses explicit when necessary.
=(X) where X ~ We make the following assumptions
(A1) V4 LT(),,
po(w) (A2) There exists E R2, such that inf, m(a)tta 0.
Assumption (A1) ensures that æ has finite differential entropy, and is a relatively weak condition. (A2) asks that at least some mass initially placed along the support of w. (A2) is, however, a much weaker
assumption that what is currently used in the literature. To be precise, Lu al. (2019 2022) assume a particular case of (A2); namely
(B) There exists M > 0 such that inf,
palw)
2 M m(a))
This is the same as (A2) when a is constrained to be O, and they make explicit use of the constant M > in their rate. Note that. (A2) is weaker the larger is, (x>)+aa decreases faster. As a comparison, if Po and are Gaussians, (A2) covers the setting where both have arbitrary means and covariances, while constraining = 0 only covers the cases in which the covariance matrix of po is strictly larger than the one of π in the positive definite order.
The following theorem is our main contribution. While here we have stated an asymptotic expression, in lact a more general expression available as an infinite power series for times large enough, and appears explicitly in the proof (Eq. (99 and Eq. (30)).
Theorem 1. Suppose (A1) and (A2) hold. Then for any q E (1,
KL(Pe|[T) 22 e- + O(e and (p(p(( = q82 e-2 +0y(E-31)
2022--2
(21)
where Var, (log ) The remainder terms (ce-3t) and Og(e-31) depend on the initialization Po and on the target π.
Note that the result (4] by Lu ct al. 2019) implics an asymptotic rate very close to e e-2t, but there are significant. differences bet ween both: beyond the fact that our result holds under much weaker assumptions, we characterize exactly the asymptotic decay of KT(p(]))), while they only provide an upper-bound that becomes less tight as 6 goes zero (because the constant increases).
Remark The coefficient Na nothing more than the variance under a of the first-voriation of the KL divergence at pu (recall Eq. 8
3.1 Proof
Given potentials V, and Vo such that -V and pn eVV,, define the random variable
Y= - Vo(X) where X N~T.
(22)
Note that we can set V.(x) - log (αr and Vo(x) log Po(x), which means that Y log Po(X). adding any constant term to V, and Vo (or solely Vo) also yields a valid construction of Y. Proposition 2. Let Y be as in Eg. (22) Let (pr)rejo.1) be follow the simulated annealing dynamics from Eq. 14). It holds that
but
=(X)
KL(U, (T] (1 1)-r)Ký(( (1 - Kr T),
(23)
1
R,((111T)= 9-1
Kr(g(1 -)) -
Ky -17)
(24)
q 9-1

---

Published in Transactions on Machine Learning Research (06/2023)
Proof. We first identify the following relationship, which arises from simple manipulation of Eq. 14
Ky(1 -T) =log log (fe" (f()-)(V(())v((e)))) vo(z))
e V-(x)
Z1 dx log //e- dr
log 7- log Z1.
Z1
(25)
Using this expression, we can expand the KL divergence between 14, and as follows:
KL(u+1 [|T] = log |r M4 - |l08 (e (e- (--tV----))-------- Z,
(1005 e-V.Zi" '/'
djl
e
log Z1 llog 2, + (1- (1- Voolur =1 +V. Vo)n (T))
Another fact about cumulant generaling functions that we can exploit is the following differential relationship
-V. - Volu, = J7 -
(26)
Altogether, this gives
(T((++1=)== = (1-))) (1- Ky(1 -)
(27)
The general q-Rényi case is deferred to the appendix, where the computation is similar.
The following lemma uses both (A1) and (A2) to establish that (r(=) is finite in some neighborhood of Beg (0), which implies that Ky admits the series expansion will require in the sequel. The proof is deferred to the appendix.
Proposition 3. Suppose (A1) and (A2) are satisfied. Then there crists some constant 6p > 0 such that the cumulant generuting function of Y, Ky(z)= = log E]ex=] is finite on some neighborhood of E B (0). Moreover, inside this neighborhood, Ky(2. (2) is holomorphic and we have the series expansion
Ky(z)= =
Kn
(28)
n!
m21
We conclude with the proof of our main result.
Proof of Theorem 7 We begin with the expression of the KT, divergence. Note that since Ky(2) is smooth for sufficiently close to the origin, it holds that
Ky(=) = Σ
Kn
n-1
n21
Using the parameterization of Eq. 27 and the series expansion for K2 T), our expression for KL(u>][=) reads
(L(u>1|T) (1 T)
Kin
(1 - -)-
Kn
(1 7)" n!!
n
1)!
n
n
Kn
11
1
(1 - 7)
n!
n.
m>
Rin
n(n 2)} (1 7)
n22
Expanding the relation and replacing 7(1) = 1 11 gives
KL(p)][T)
= 122 e
-2t 2t +
Σ
Kn
nt
(29)
2
nln - 2)
n23
8

---

Published in Transactions on Machine Learning Research (06/2023)
We now do the same manipulations for Rg (1=11
(1T(L+) =
1
Ken
(q(1 - 7))
(1 - 7)
9
n!!
9
n!
n²1
n
1
Kl T) + Σ
q
Ki
- 7)
q
ki -7)
Σ
Kp
(1 7)
q -
n!
-
n22 n!!
n22
9" Kn
(1 T)
n22 1 n!
Substituting r(1) = - e-
and expanding out the first term yields
12 2x Rg(Pel(= = 1
9" - g.n
e
(30)
n!
n23
Lemma 2
in the appendix justifies that the higher-order terms of Eq. concludes the proof.
29
and Eq.
are O(e-3), which 30
4 Numerical simulations
We present simple numerical simulations that demonstrates our asymptotic convergence rate of the KL divergence the FR gradient flows, well comparison with the WFR- and W-GFs, We consider two target distributions over the set [- T), each with two initializations:
1. Target distribution mi: We set ग. & e with (z 2.5 cos (2x) + +inns)n()). This distribution has two modes with different weights and has been studied previously by Lu et al. (2019) We consider two initial distributions:
(a) V with V. = =,,,, which has two modes in locations where has little mass.
(b) my x e with Vi = 2.5 cos(2x), which has two modes in almost the same positions as m but with equal weight.
Target distribution æ2: We set 22 x - with V2(x)) = -6 cos(2). This distribution has one mode. We consider two initial distributions:
(c) Te xe-V. with Ve V22 which has one mode in location where æ has little mass.
(d) Va with Va 0, which is the uniform distribution.
Fig. shows the target energies Vi, 122 and the initial energies Va, Vn, Va, Va2 introduced above. Fig. 2 shows the evolution of the KL divergence along the FR, WFR and W gradient flows. It also contains plots of the dominant term X2 -2t of the approximation of the KL divergence docay for FR flows (sco Thcorem displayed as dotted lines. Table shows the slopes of each curve from Fig. 2 at large times (see Appendix for details the computation of slopes).
Some observations are in order:
As predicted by Theorem the curves (L [|()) approach the curves - as t grows.
For æı, the curves (L (()) and KL(p, (|)) initialized at wa are very close for small times. The reason that VV1 and VV, are very close the regions where Tı and π. have most of the mass. Consequently, the term V (p) WFR log WPR + VV))), which is the difference between the FR and the WFR PDEs, is small at initialization.
The curves (I()WW)x) behave very differently for कं and n2 (sce Table Indeod, since æj is bimodal Cis((TT) is quite large (thus convergence is slow), whereas Ta is unimodal, with much smaller log-Sobolev constant.
9

---

Published in Transactions on Machine Learning Research (06/2023)
Energies V1, and Vo and
Energies V2. and Ve and
(targel)
(targel)
-2
Figure 1: Energios of the target and initial distributions.
Flows of (a(m(m
Flows De
10
10
10"
10
Du(((n))
10
10
10"
FR
WFR
Itc
RL
10
Th WFR Rb
WFR IT
FR Rd
W fl
W
0.0 0.5
1.5
2.0
25
3.0
Evolution of the KL divergence with respect to (left) and (right) along their respective FR (solid lines), WFR (dash-dotted lines) and W (dashed lines) gradient flows. Each plot contains flows initialized at two probability measures: in the left plot these are Ta (blue, top curves at t = 0) and m (orange); in the right plot, *e (blue, top curves at 0) and æ. (orange). The dolled lines show the curves ge (for the appropriate values K2), introduced in Theorem
Figure 2:
The curves ((T also behave differently for both target distributions. For Tı, it decays only slightly faster than KL(pFR|mmg while for T2 it goes down much faster than both KL(p and KL(pWPR|| Interestingly, looking at Table 1 we observe that the asymplotic slopes of the WFR are very close to the sum of the slopes for FR and W. This seems to indicate that at large times, the KL divergence decays like 2 oun i.e. that the W and FR terms act more or less independently.
Target T1
Target T22
Init. va
Tnit. TB
Tnit. Re
Tnit. Wd
FR
-2.0016
-2.0002
-2.0028
-2.0014
WFR
-2.0771
-2.0759
-12.8190
-12.8632
W
-0.0811
-0.0811
-10.7784
-10.8538
Table 1: Large-time slopes of the KL divergence vs. time curves in a semi-logarithmic plot (Fig. for the three flows. Soe Appendix B for details on the computation of the slopes.
10

---

Published in Transactions on Machine Learning Research (06/2023)
5 Conclusion
In this work, using relatively simple proof technique, showed that the Kullback-Leibler divergence along its Fisher-Rao gradient flow )(20 can be written as a power-series expansion, resulting in tight asymptotic convergence rate for large times. A similar expansion holds for R, where Rg is any q-Rényi divergence. Our findings were verified with simple numerical experiments, where we also simulated Wasserstein and Wassersleir--Fishe--Re gradient flows. Our simulations indicated that, in some cases, the convergence rate of the WFR gradient flow scales like e (2+(2/Cta))) an observation that we hope can be made precise in future work. A second direction is to extend our proof technique from the KL divergence to general Bregman divergences.
Acknowledgments
AAP thanks Anna Korba for introducing the main WFR references to him. The authors collectively thank Joan Bruna, Jonathan Niles-Weed, Sinho Chewi, and Andre Wibisono for helpful discussions, and the anony- mous reviewers who helped improve the quality of the presentation. CD acknowledges Meta Al Research funding source. AAP acknowledges NSF Award 1922658 and Meta AT Research.
References
Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient flows: in metric spaces and in the space of probabilily measures. Springer Science & Business Media, 2005.
Anton Arnold, Peter Markowich, and Andreas Unterreiter. On convex Sobolev incqualities and the rate of convergence to equilibrium for Pokker-Planck type equations. Communications in Partial Differential Equations, 26, 05 2000.
Jean-David Benamou and Yann Brenier. A computational fluid mechanics solution to the Kantorvhatoo mass transfer problem. Numerische Mathematik, 84(3):375-30) 2000.
V.I. Bogachev. Measure Theory. Number 1 in Measure Theory. Springer Berlin Heidelberg, 2007.
Sinho Chewi. Log-concave sampling. 2022.
Sinho Chewi, Murat A Erdogdu, Mufan Li, Ruoqi Shen, and Shunshi Zhang, Analysis of Langevin Monte Carlo from Poincare to Log-Soboley, In Proceedings of Thirty Fifth Conference on Learning Theory, volume 178 of Proceedings of Machine Learning Research. PMLR, 02 05 Jul 2022.
Lenaic Chizat. Sparse optimization on measures with pvrrpaaraeeterized gradient descent. Mathematical Programmmngg, 194(1-2):487 532, 2022.
Lenaic Chizat, Bernhard Schmitzer, Gabriel Peyré, and François-Xaval Xavier Vialard. An interpolating distance between optimal transport and Fisher-Rao. Foundations of Compulational Mathematics, 18, 06 2015.
Lénaïc Chizat, Gabriel Peyré, Bernhard Schmitzer, and François- Xavier Vialard. Unbalanced optimal trans- port: dynamic and Kantorovich formulations. Journal of Funclional Analysis, (7)301) 012000313,,000 2018.
A.S. Dalalyan and A.B. Tsybakov. Sparse regression learning by aggregation and Langevin Monte-Carlo Journal of Computer and System Sciences, 1853) 2012.
Alain Durmus, Szymon Majewski, and Blažej Miasojedow. Analysis of Langevin Monte Carlo via convex optimization. J. Mach. Learn. Res., 20(1):2666-277 2711, 2021.
Thomas O Gallouët and Leonard Monsaingeon. A JKO splitting scheme for Kantorovich Fisher Rao gra- dient flows. SIAM Journal on Mathematical Analysis, 49(2):1100 1130, 2017.
Andrew Gelman, John B Carlin, Hal S Stern, and Donald B Rubin. Boyesion data analysis. Chapman and HaR//CRC, 1995.
11

---

Published in Transactions on Machine Learning Research (06/2023)
Leonard Gross. Logarithmic Sobolev inequalities. American Journal of Mathematies, 9744):1011-1083, 1975.
E. Hellinger. Neue Begrindung der Theorie quadratischer Formen von unendlichviclen Veränderlichen. Journal für die reine und angewandle Malhemalik, (136):210-2)) 271, 1909.
Richard Holley and Daniel Stroock. Logarithmic Sobolev inequalities and stochastic Ising models. Journal of Statistical Physics, 46(5):1159-11999) Mar 1987. ISSN 1572-9613.
Michael Johannes and Nicholas Polson. MCMC methods for contimuous- time financial econometrics. Handbook of Financial Econometries: Applications, pp. 1-72. Elsevier, 2010.
Richard Jordan, David Kinderlehrer, and Felix Otto. The variational formulation of the Fokkor-Planck equation. SIAM journal on mathematical analysis, 29(1):1-17 1998.
Shizuo Kakutani. On equivalence of infinite product measures. Annals of Mathematics, 49(1):2144 -224, 1948.
S. Kirkpatrick, C. D. Gelatt, and M. P. Vecchi. Optimization by simulated annealing. Science, 220(4598); 671-680, 1983.
Ivan Kobyzev, Simon JD Prince, and Marcus A Brubaker. Normalizing flows: An introduction and review of current methods. IEEE transactions on pattern analysis and machine intelligence, (3111::36443977) 2020.
Stanislay Kondratyev, Léonard Monsaingeon, and Dmitry Vorotnikov. new optimal transport distance on the space of finite Radon measures. Advances in Differential Equations, (1 1164, 2016.
Mare Lambert, Sinho Chewi, Francis Bach, Silvère Bonnabel, and Philippe Rigollet. Variational inference via Wasserstein gradient flows. arXiv preprint arXiv:2205. 15902, 2022.
Matthias Liero, Alexander Mielke, and Giuseppe Savaré. Optimal transport in competition with reaction: The Hellinger-Kantorovt distance and geodesic curves. SIAM Journal on Mothematical Analysis, 48(4): 2869 2911, 2016.
Matthias Liero, Alexander Mielke, and Giuseppe Savaré. Optimal entropy-transport problems and a new Hellingen-Kantrrovich distance between positive measures. Inventiones mathemalicae, 211, 03 2018.
Yulong Lu, Jianfeng Lu, and James Nolen. Accelerating Langevin sampling with birth-death. arXiv preprinl arXiv:1905.098 2019.
Yulong Lu, Dejan Slepčev, and Lihan Wang. Birth-death dynamics for sampling: Global convergence, approximations and their asymptotics. arXiv preprint arXiv:2211.00450, 2022.
David JC MacKay. Information theory, inference and learning algorithms. Cambridge university press, 2003.
P. A. Markowich and C. Villani. On the trend to equilibrium for the Fokker-Planck equation: An interplay between physics and functional analysis. In Physics and Functional Analysis, Matematica Contemporanea (SBM) 19, pp. 1-29, 1999.
Martin Pincus. A Monte Carlo method for the approximate solution of certain types of constrained opti- mization problems. Operations Research, 12(6):1225-1228, 1970.
Christian P Robert, George Casella, and George Casella. Monte Carlo statistical methods, volume 2. Springer, 1999.
Grant Rotskoff, Samy Jelassi, Joan Bruna, and Eric Fajnen-Eijnden. Global convergence of neuron birth- death dynamics. arXiv preprinl 1902:1190,,01843, 2019.
Al'bert Nikolaevich Shiryaev. Probability, Graduate texts in mathematics 95. Springer- Verlag, New York, 1984. ISBN 0781489900188
12

---

Published in Transactions on Machine Learning Research (06/2023)
A.J. Stam. Some inequalities satisfied by the quantities of information of Fisher and Shannon. Information and Control, 1(2):101-112, 1959.
Santosh Vempala and Andre Wibisono. Rapid convergence of the unadjusted Langevin algorithm: Isoperime- try suffices. In Advonces in Newral Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
C. Villani, Optimal Transport: Old and New. Grundlehren der mathematischen Wissenschaften, Springer Berlin Heidelberg, 2008.
Udo Von Toussaint. Bayesian inference in physics. Reviews of Modern Physics, 83(3):943, 2011.
Andre Wibisono. Sampling as optimization in the space of measures: The Langevin dynamics as composite optimization problem. In Conference on Learning Theory, pp. 2093 3027. PMLR, 2018.
Yuling Yan, Kaizheng Wang, and Philippe Rigollet. Learning Gaussian mixtures using the Wasserstein- Fisher-Rao gradient flow. arXiv preprint arXiv:2301.00 01662,,,..1166, 2023.
A Remaining proofs
Proposition (Uniqueness of the Fisher-Rao gradient flow of the KL divergence). Given a target potential V. and initial measure po, the solution of Eq. unique.
Proof. Consider the PDE
O(pe(x) -1)g(((())) (x)))
Ho = Po
(31)
Note that this is fact an ODE for each point x, that we can rewrite as o log ue(x) = - = (log((((())++V+ The unique solution of this ODE with initial condition (og:)(a)) log (xx) = (oog))((z)) V. ( Thus, we conclude that Eq. (31) has a unique solution.
Now, given solution Pt of Eq. 12 with initial condition Po, define 0t as
log (x() log (o()) ''' c'=<(lo((p)) + Vo) ds
(32)
Note that R₂ fo c' (og((()))) ds satisfics the ODE - (log(())) +V)),, with the initial
condition Rg = O. Using this, we observe that òr is solution of Eq. 31
o, log (x)) = 0, lt + O.R. -(log(((())) + V.(x) - (log(pt) + v.)) + R.
==(log(p(x)) + R, V.(x)) = - ((og((((())) V..)i)
Also, note that the map (pt):20 + (pt) defined by Eq. 332 is invertible, as Pc(x)) Pt(x)/ pv(y) dy. This follows from the fact that and गt are proportional to each other, and that Pt integrates to 1.
Finally, suppose that of and p are two solutions of Eq. (12) with initial condition Pp- Via the construction Eq. (32). they yield solutions of and p of Eq. (31) with initial condition Po- The uniqueness of the solution of Eq. (31h implies that oe = of. Since the map (pe)ezo + (6:)20 is invertible, we obtain that pf p which concludes the proof.
13

---

Published in Transactions on Machine Learning Research (06/2023)
Proof of Proposition (Continued). We perform similar manipulations in the case with the KL diver-
gence:
-VV(V--Vo--pp) (2x)-
Rg R1((4=11T)=
q - log
dm
(-.V -9
Zr
log
(Q(1-T)(V6-Vo)
dı
21
Ky (1 -T)
q (log Z, log 21)
q
q Ky 7) - Ky 7)
where the last line we again used Eq.
25 This completes the proof.
Proof of Proposition 3 By (A1), the partition function F(I) Sga (VV, de differentiable at This is because F* V dr(x). Hence, F(t) is finite on an interval (1 2e, 1] for some e1.
Note that the assumption (A2) be written equivalently as & == inf, aV, Vo(x) > -oc. We obtain that for all E [0, ())))
-vV V.(x)) -())))=) (ea (ea-))VV
(33)
- 1)V.(s) + 1)V)(x)
Equivalently,
exp(Kr(-c))- [ --VVVV dx e-5 e-(1-ei)V.(0) de e -1) +oc.
{
(34)
Also, for all c C 1), using the convexity of the exponential function we have that
exp(Ky(e)) f cc(Va())-Vo(e))-vkk dr / ga
((L-)))cc())-cooo0 dr
(35)
(1 c)- V-(x) (Vv()))) + cc Volat dr = (1 6) + 67o < +00. Rd
(36)
Hence, the cumulant-generatinn function (y(t)= log Ee is finite on neighborhood (-co, co) with co min L, e1/a]. Applying Lemma we conclude that there exists > 0 such that for (.), we have that (2 = E 2"
The following lemma, which make explicit, is well-known fact in probability theory. In short, since the moment-generatio function is analytic in some neighborhood, and is non-negative, taking the logarithm is safe as everything is analytic. The interested reader can consult e.g. (Shiryaev 1984 Section 11.12.8) which dissects this in detail.
Lemma 1. Assume that the cunnulant-generutinn function Kr = log EY is fovile on neighborhood (-Eo, so) of zero. Then, Ky(2)= = log Ee=Y as function on the complex plane is holomorphic on the open ball (0) of radius e centered at zero, for some e > O. Moreover, for 2 E (0), we have that
Ky (2)
Σ
(37)
nl
Lemma 2 (End of the proof of Theorem
We have that
(L(p)
K2 c = O -3),
9K2 22 RR((oll
c
0(e
(38)
2
2
14

---

Published in Transactions on Machine Learning Research (06/2023)
Proof. Lemma implies that the series for Ky centered zero has convergence radius c, for some >> 0. Since the derivative of a series has the same radius of convergence, we obtain that
H(z) = =x)) (x) = Σ
Kn
2)n
n22
has convergence radius e as well. Hence, by the Cauchy-Hadamarr theorem, ! 2 lim sup co where Kn Cn n(n 2]
This implics that for all < d < c, there exists constant Ce 0 such that for all 0, lcal Ccl(d)". Consequently, for all C with |z| < 1/e',
(2)
K2 31
+00
Kn
21 S Ce'
+000
[2]
Ce (12 3
(39)
n(n 2)
=
Using Eq. 23 we get that for any constant 7 > 0, ift - loo + y (or equivalently, e de
(KL(P) K2 - S Ce 2
e-31 Ce (e)((1e-)) = =(e))))
(40)
which concludes the proof for the KL divergence. For the Rényi divergence, the proof is analogous (note that in that case the series q-1 Ky (gz) - Ky (2) has convergence radius e/g).
Details on the numerical simulations
To run the simulations in Section 4 we discretized the interval [-x, TT) in 2000 equispaced points. Let à 2Tr /n. For each algorithm and initialization, we construct sequences (rx)) where 25 R represents the normalized log-density at each point. We let v, E IR" be the (non-normalized) energy of the larget distribution ion, obtained by evaluating V. at the discrelization points. Similarly, Vv., Ar, E IR" are the evaluations of VV, and AV, at the points (note that VV, is a scalar because the distributions are one-dimensional).
We used the following discretizations for the Fisher-Rao, Wasserstein and Wasserstein-Fisher gradient flows:
ii Fisher-Rao GF: We use mirror descent in log-space. The update reads:
TR+1 3h. e(-vv- - -Ik);
Th+1 f- Tst1 - log Ee
(ii) Wasserstein GF: We approximate numerically the gradient and the laplacian of the log-density:
Vi E [m], Vi E [n], (Arx)* 3 + (o + + 2x/)//² Jk+1 + Ib + e(Au + Ara + (Vv, + VZL)
(Vxx)² + (xtt1 - /
(41)
We use periodic boundary conditions, so that the first discretization point is adjacent to the last one for the purposes computing derivatives.
iii) Wasserstein-Fiste Fisher- Rao GF: We combine the two previous updates. Letting Vry and Ark be as in Eq. (41) we have
To+1 +- Tk + e(-Ve-
+ Av. + Ar + (Vv, + VIL)
Eb+1 +- 8x+1 - log Σ
E t
i=1
15

---

Published in Transactions on Machine Learning Research (06/2023)
We used stepsizes 2.5 10-6 and e 1 x 10 for the experiments target distributions (1) and (2).
respectively. The slopes in Table 1 are obtain by taking 0 < l1 < l2 and computing
log(KL(Prslla) ll))) -
l|x))
-
We different values for l and 22 for cach target distribution; and 22 must be large enough to capture the asymptotic slope of the curve, but not too large to avoid numerical errors. For all the curves corresponding to target we take t = 7.0 and to = 7.5. For target we take: for FR, t = 6.875 and t2 7.0; for WFR, 4 = 1.875 and b = 2.0; for W, t = 2.75 and t2 = 2.755.
16


Reviews:
[{"review_id": "Ni5WpxYXmp", "reviewer": "Reviewer_sQ4p", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper is written in a concise way and the proof seems straightforward. The problem studied here is also of great importance. However, clarity over previous work is lacking. In particular, the comparison of the results obtained in this paper with the immediately relevant work on this topic seems not explicit, which makes it hard to judge the contribution of the current paper.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "In this paper, the authors improve the results of the convergence of the Wasserstein Fisher-Rao gradient flow that was recently established by Lu et al. (2019, 2022). The key technique in this paper is the use of an infinite power series expansion related to the cumulant-generating function of random variables. \n", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "In the introduction of the convergence rates for Langevin Monte Carlo sampling algorithms, you might need to elaborate a little bit more on the differences among these works. In particular, it is worth discussing other representing papers such as [1] that directly derive the convergence of discrete-time LMC and [2] that shows the convergence of the stochastic version of LMC with an arbitrary choice of mini-batch size.\n[1] Global convergence of Langevin dynamics-based algorithms for nonconvex optimization. Advances in Neural Information Processing Systems, 2018.\n[2] Faster convergence of stochastic gradient Langevin dynamics for non-log-concave sampling. In Uncertainty in Artificial Intelligence, 2021.\n\nThe authors claim that the results presented in this paper are tight and require fewer assumptions than previous work. It would be helpful if the authors could provide more detail on this point, including the lower bound of the problem and the relaxation of assumptions from previous work. In addition, it would be valuable to discuss how the assumptions are essential to their paper.\n\nIn Section 2.4, it is unclear whether $\\kappa_n$ represents the n-th moment of Y.\n\nThere is a lack of comprehensive comparison between the results presented in this paper (Theorem 1) and those of existing results, particularly Lu et al. (2022). It is noted that the results in Lu et al. (2022) hold for any $\\delta\\in(0,½]$. Furthermore, if a value is arbitrarily chosen for $\\delta$, the result in their paper is also in the same order as that of Theorem 1. Moreover, the term $t^*$ would also be absorbed by the constant $C_1$. Hence, it is suggested that the authors provide a more detailed comparison between their results and those of existing work.\n\nRegarding the proof, it is mentioned that even though the decomposition of the KL divergence only involves terms like $e^{-2t}$ and $e^{-3t}$, the $O(e^{-3t})$ term actually hides problem-dependent constants related to the infinite series expansion. Therefore, it is not directly comparable with Lu et al. (2022).\n\nMinor questions:\n\n* In the abstract: **$\\pi$ is that independent of** -> **$\\pi$ that is independent of**\n* Around Eq (3): **satisfies a LSI** -> **satisfies an LSI**\n* Below Eq (3): **when $V^∗$ $\\alpha$-strongly convex** -> **when $V^∗$ ``is`` $\\alpha$-strongly convex**\n", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "ZSlmeZBYCO", "reviewer": "Reviewer_ekbs", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths \n\n1. As far as I know, the connection to simulated annealing is novel, and honestly somewhat surprising. I believe this observation alone is significant. \n2. Using the cumulant-generating function to compute a power series expansion is also a new technique based on my knowledge, at least it is not well known and therefore welcomed. The authors were able to use this expansion to show the first order term does indeed decay at the desired rate $e^{-2t}$, which is much cleaner than choosing $\\delta,t_*$ in previous works. \n\nWeaknesses \n\nOnly several minor clarifications and questions in the next section. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors drew a very interesting connection between simulated annealing and Wasserstein--Fisher--Rao gradient flow. Using this connect, the authors were able to compute an explicit expansion of the Fisher--Rao gradient flow decay of KL and Renyi divergences. The authors also provided convincing numerical evidence supporting their theoretical results. ", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I have several minor questions for clarity. \n1. The authors briefly commented on Wasserstein and Fisher--Rao geometry being essentially orthogonal. Can I interpret the WFR Riemannian metric as a sum of Wasserstein and FR separately? \n2. If we can take the above interpretation, does the connection to simulated annealing still hold if we vary the coefficients in the linear combination of the two Riemannian metrics? \n3. Can we compute a higher order term in the $O(e^{-3t})$ expansion, so that we get a type of Taylor remainder term? \n4. In your numerical simulations, can you describe how the gradient flow evolution and KL were calculated? In particular, you chose the support $[-\\pi,pi)$, but is the boundary condition period, i.e. making it a torus? ", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "a6sZNRIKD3", "reviewer": "Reviewer_5Xmo", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The proof is short, succinct, and well-written.\n- Assumption A2 allows the consideration of many new cases, which were unreachable with previous results.\n\nWeaknesses:\n- First and foremost, the result is not algorithmic (and probably cannot be made, in full generality). Thus, while this is a fine paper in Wasserstein metric geometry, it is tangentially related to the main topics of TMLR. Still, there is a lot of interest in the Wasserstein gradient flows in the ML community, so I would not hold this point against the paper.\n- The improvement in the convergence rate is somewhat incremental compared to previous results. More on this in the next section.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper studies the convergence rate of gradient flow for the KL divergence along the Fisher-Rao dynamics towards its equilibrium at a predetermined target measure. \nThe main result is that, under very mild assumptions on the target measure and the initial point of the dynamics, the KL divergence decays exponentially fast. Remarkably, the exponent does not depend on either measure, initial, or target. This result improves upon previous works by relaxing the necessary assumptions and by a slight improvement to the convergence rate.\nAt the technical level, the authors establish an asymptotic expansion of the KL divergence along the dynamics, in terms of the cumulants of the score function.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Before outlining my requested changes, let me note that I enjoyed reading the paper and found the result interesting. Since I found the claims made in the submission to be supported by accurate by convincing proofs and since, as indicated above, I believe some individuals in TMLR's audience will be interested in knowing the findings of this paper, I am inclined to recommend acceptance, subject to addressing the issues I raise below.\n\n- Currently, Theorem 1 gives the impression that the error term in the expansion of the KL divergence is bounded by Ce^{-3t}, where C is a *universal* constant. This is also the way the big O notation is presented in the notations paragraph.\nAs far as I could see from the proof of Theorem 1, this is not the case. C is a constant that depends on the cumulants and on the parameter \\alpha, introduced in Assumption A2. In particular, C depends on both the initial and terminal points of the dynamics. I think that this point should be emphasized and discussed.\n\n- In line with my above comment, the discussion about Assumption B is a bit awkward. The existence of the constant M is equivalent to Assumption A2 (when \\alpha = 1). The point is that the convergence rate in the cited paper depends explicitly on M. In the present paper, the convergence rate will also depend on M (and \\alpha), but only implicitly through the cumulants.\n\n- I found the introduction to be a bit confusing. A lot of emphasis is given to the Wasserstein/Langevin dynamics flow and the Fisher-Rao-Wasserstein flow when the results only deal with the Fisher-Rao flow. While I think that the addition of background and relevant results is great and well-written, it may perhaps be better to separate it from the main focus of the paper. I propose to add a new section about the 'bigger picture' in which the relevant material and corollaries can be presented. (This comment is a matter of personal taste and I leave it to the authors' discretion) \n\n- Line below Equation 11. The superscript W is missing from some of the rhos.\n\n- Statement of theorem 1. I'm not sure what the sentence for \"...for t large enough...\" is trying to convey. Isn't the expansion valid for any t? The big O notation is supposed to take care if smaller t's. Also I'm not sure why at the end of the proof of Theorem 1, t needs to be taken to infinity to complete the proof. What is needed to show, and what the authors do show, is that the coefficients in the expansion do not grow too fast.\n\n- Equation 32. Is the equation missing the normalizing constants?", "broader_impact_concerns": "No concerns", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: Three expert reviewers reviewed the paper and they all had positive evaluation, which AC agrees with. This is a nice theoretical paper with great presentation style which makes it accessible to a wider ML audience. AC recommends accepting this paper for publication.
Decision: Accept (poster)