Year: 2019

Abstract:
We present a systematic framework for the Nesterov’s accelerated gradient flows in the spaces of probabilities embedded with information metrics. Here two metrics are considered, including both the Fisher-Rao metric and the Wasserstein-2 metric. For the Wasserstein-2 metric case, we prove the convergence properties of the accelerated gradient flows, and introduce their formulations in Gaussian families. Furthermore, we propose a practical discrete-time algorithm in particle implementations with an adaptive restart technique. We formulate a novel bandwidth selection method, which learns the Wasserstein-2 gradient direction from Brownian-motion samples. Experimental results including Bayesian inference show the strength of the current method compared with the state-of-the-art.

heading: 1 INTRODUCTION
text: Recently, optimization problems on the space of probability and probability models attract increasing attentions from machine learning communities. These problems include variational inference (Blei et al., 2017), Bayesian inference (Liu & Wang, 2016), Generative Adversary Networks (GAN, Goodfellow et al. (2014)), and policy optimizations (Zhang et al., 2018), etc. For instance, variational inference methods approximate a target density by minimizing the Kullback-Leibler (KL) divergence as the loss (objective) function.
Gradient descent methods with sampling efficient properties play essential roles to solve these optimization problems. Here the gradient descent direction often relies on the information metric over the probability space. This direction naturally reflects the change of the loss function with respect to the metric. In literature, two important information metrics, such as the Fisher-Rao metric and the Wasserstein-2 (in short, Wasserstein) metric, are of great interests (Amari, 1998; Otto, 2001; Lafferty, 1988). For the Fisher-Rao gradient, classical results including Adam (Kingma & Ba, 2014) and K-FAC (Martens & Grosse, 2015) demonstrate its effectiveness in probability models. For the Wasserstein gradient, many classical methods such as Markov chain Monte Carlo (MCMC) methods (Geman & Geman, 1987; Neal et al., 2011; Welling & Teh, 2011) and particle-based variational inference (ParVI) methods (Liu & Wang, 2016; Chen & Zhang, 2017; Chen et al., 2018) are based on this framework in the probability space. The strength of using the Wasserstein gradient is also shown in probability models such as GANs. (Arjovsky et al., 2017; Lin et al., 2018; Li et al., 2019).
The Nesterov’s accelerated method (Nesterov, 1983) is widely applied in accelerating the vanilla gradient descent under the Euclidean metric. It corresponds to a damped Hamiltonian flow, known as the accelerated gradient flow (Su et al., 2016). A natural question is whether there exists a counterpart of the accelerated gradient flow in the probability space under information metrics. For optimization problems on a Riemannian manifold, the accelerated gradient methods are studied by Liu et al. (2017); Zhang & Sra (2018). The probability space embedded with information metric can be viewed as a Riemannian manifold. Several previous works explore accelerated methods in this manifold under the Wasserstein metric. Liu et al. (2018; 2019) propose an acceleration framework of ParVI methods based on manifold optimization. Taghvaei & Mehta (2019) introduce the accelerated flow from an optimal control perspective. On the other hand, Cheng et al. (2017); Ma et al. (2019) explore and analyze the acceleration on MCMC, based on the underdamped Langevin dynamics.
In this paper, we present a unified framework of accelerated gradient flows in the probability space embedded with information metrics, named Accelerated Information Gradient (AIG) flows. From an information-differential-geometry perspective, we derive AIG flows by damping Hamiltonian flows, concerning both the Fisher-Rao metric and the Wasserstein metric. Then we focus on the Wasser-
stein metric with the KL divergence loss function. In Gaussian families, we verify the existence of the solution to AIG flows. Here we show that the AIG flow corresponds to a well-posed ODE system in the space of symmetric positive definite matrices. We rigorously prove the convergence rate of AIG flows based on the geodesic convexity of the loss function. Here we note that our proof removes the unnecessary technical assumption in (Taghvaei & Mehta, 2019, Theorem 1).
Besides, we handle two difficulties in numerical implementations of AIG flows. On the one hand, as pointed out by Taghvaei & Mehta (2019); Liu et al. (2019), the logarithm of density term (Wasserstein gradient of KL divergence) is hard to approximate in particle formulations. We propose a novel kernel selection method, whose bandwidth is learned by sampling from Brownian motions. We call it the BM method. On the other hand, we notice that the AIG flow can be a numerically stiff system, especially in high-dimensional sample spaces. This is because the solution of AIG flows can be close to the boundary of the probability space. To handle this issue, we propose an adaptive restart technique, which accelerates and stabilizes the AIG algorithm. Numerical results in toy examples, Gaussian measures and Bayesian Logistic regression indicate the validity of the BM method and the acceleration effects of the proposed AIG flow.
This paper is organized as follows. Section 2 briefly reviews the information metrics and their corresponding gradient flows and Hamiltonian flows in the probability space. In Section 3, we formulate various forms of AIG flows and analyze W-AIG flows in Gaussian measures. We theoretically prove the convergence rate of W-AIG flows in Section 4. Section 5 presents the discrete-time algorithm for W-AIG flows, including the BM method and the adaptive restart technique. Section 6 provides numerical experiments.

heading: 2 METRIC AND FLOWS IN THE PROBABILITY SPACE
text: Suppose that Ω is a region in Rn. Let F(Ω) denote the set of smooth functions on Ω. 〈·, ·〉 and ‖ · ‖ are the Euclidean inner product and norm in Rn. ∇, ∇· and ∆ represent the gradient, divergence and Laplacian operators in Rn. Denote the set of probability density
P(Ω) = { ρ ∈ F(Ω): ∫ Ω ρdx = 1, ρ ≥ 0 } .
The tangent space at ρ ∈ P(Ω) follows Tρ(Ω) = { σ ∈ F(Ω) : ∫ σdx = 0. } . The cotangent space at ρ, T ∗ρP(Ω), can be treated as the quotient space F(Ω)/R, which are functions in F(Ω) defined up to addition of constants. Definition 1 (Metric in the probability space) A metric tensor G(ρ) : TρP(Ω) → T ∗ρP(Ω) is an invertible mapping from the tangent space at ρ to the cotangent space at ρ. This metric tensor defines the metric (inner product) on the tangent space TρP(Ω). Namely, for σ1, σ2 ∈ TρP(Ω), we define
gρ(σ1, σ2) = ∫ σ1G(ρ)σ2dx = ∫ Φ1G(ρ) −1Φ2dx,
where Φi is the solution to σi = G(ρ)−1Φi, i = 1, 2.
We present two important examples of metrics on the probability space P(Ω): the Fisher-Rao metric from information geometry and the Wasserstein metric from optimal transport. Example 1 (Fisher-Rao metric) The inverse of the Fisher-Rao metric tensor is defined by
GF (ρ)−1Φ = ρ ( Φ− ∫ Φρdx ) , Φ ∈ T ∗ρP(Ω).
The Fisher-Rao metric on the tangent space is given by
gFρ (σ1, σ2) = ∫ Φ1Φ2ρdx− (∫ Φ1ρdx )(∫ Φ2ρdx ) , σ1, σ2 ∈ TρP(Ω),
where Φi is the solution to σi = ρ ( Φi − ∫ Φiρdx ) , i = 1, 2.
Example 2 (Wasserstein metric) The inverse of the Wasserstein metric tensor is defined by
GW (ρ)−1Φ = −∇ · (ρ∇Φ), Φ ∈ T ∗ρP(Ω).
The Wasserstein metric on the tangent space is given by
gWρ (σ1, σ2) = ∫ ρ 〈∇Φ1,∇Φ2〉 dx, σ1, σ2 ∈ TρP(Ω),
where Φi is the solution to σi = −∇ · (ρ∇Φi), i = 1, 2.

heading: 2.1 GRADIENT FLOWS
text: In learning, many problems can be formulated as the optimization problem in the probability space,
min ρ∈P(Ω) E(ρ).
Here E(ρ) is a divergence or metric loss functional between ρ and a target density ρ∗ ∈ P(Ω). One typical example of E(ρ) is the KL divergence from ρ to ρ∗,
E(ρ) = DKL(ρ‖ρ∗) = ∫ log ( ρ
ρ∗
) ρdx.
Another example is the maximum mean discrepancy (MMD, Gretton et al. (2012)),
E(ρ) = MMD(ρ, ρ∗) = ∫ ∫ (ρ(x)− ρ∗(x))K(x, y)(ρ(y)− ρ∗(y))dxdy,
where K(x, y) is a given kernel function. The gradient flow for E(ρ) in (P(Ω), gρ) takes the form
∂tρt = −G(ρt)−1 δE(ρt)
δρt .

heading: Here δE(ρt)δρt is the L
text: 2 first variation w.r.t. ρt. We formulate the gradient flow under either the Fisher-Rao metric or the Wasserstein metric. Example 3 (Fisher-Rao gradient flow) The Fisher-Rao gradient flow is given by
∂tρt = −ρt ( δE δρt − ∫ δE δρt ρtdx ) .
Example 4 (Wasserstein gradient flow) The Wasserstein gradient flow writes ∂tρt = ∇ · ( ρt∇ δE
δρt
) .

heading: 2.2 HAMILTONIAN FLOW
text: In this subsection, we briefly review the Hamiltonian flow in the probability space. By using the metric gρ in the probability space, we can define a Lagrangian by
L(ρt, ∂tρt) = 1
2 gρt(∂tρt, ∂tρt)− E(ρt).
The Euler-Lagrange equation for the Lagrangian follows
∂t
( δL
δ(∂tρt) ) = δL δρt + C(t), (1)
where C(t) is a spatially-constant function. Proposition 1 If we let Φt = δL/δ(∂tρt) = G(ρt)∂tρt, the equation (1) can be formulated as a system of (ρt,Φt), i.e.,  ∂tρt −G(ρt)−1Φt = 0,
∂tΦt + 1
2
δ
δρt
(∫ ΦtG(ρt) −1Φtdx ) + δE
δρt = 0,
(2)
where Φt is up to a spatially-constant function shrift. Here (2) is the Hamiltonian flow
∂t [ ρt Φt ] − [ 0 1 −1 0 ][ δ δρt H(ρt,Φt)
δ δΦt H(ρt,Φt)
] = 0, (3)
with respect to the HamiltonianH(ρt,Φt) = 12 ∫ ΦtG(ρt) −1Φtdx+ E(ρt).
Similarly, we can write the Hamiltonian flow under the Fisher-Rao metric or the Wasserstein metric. Example 5 (Fisher-Rao Hamiltonian flow) The Fisher-Rao Hamiltonian flow follows ∂tρt − Φtρt + ∫ Φtρtdx = 0, ∂tΦt + 1
2 Φ2t −
(∫ ρtΦtdx ) Φt + δE
δρt = 0. The corresponding Hamiltonian isHF (ρt,Φt) = 12 (∫ Φ2tρtdx− (∫ ρtΦtdx )2) + E(ρt). Example 6 (Wasserstein Hamiltonian flow) The Wasserstein Hamiltonian flow writes ∂tρt +∇ · (ρt∇Φt) = 0, ∂tΦt + 1
2 ‖∇Φt‖2 +
δE δρt = 0.
The corresponding Hamiltonian is HW (ρt,Φt) = 12 ∫ ‖∇Φt‖2ρtdx + E(ρt). This is identical to the Wasserstein Hamiltonian flow introduced by Chow et al. (2019).

heading: 3 ACCELERATED INFORMATION GRADIENT FLOW
text: Let αt ≥ 0 be a scalar function of t. We add a damping term αtΦt to the Hamiltonian flow (3).
∂t [ ρt Φt ] + [ 0 αtΦt ] − [ 0 1 −1 0 ][ δ δρt H(ρt,Φt)
δ δΦt H(ρt,Φt)
] = 0. (4)
This renders the Accelerated Information Gradient (AIG) flow ∂tρt −G(ρt)−1Φt = 0,
∂tΦt + αtΦt + 1
2
δ
δρt
(∫ ΦtG(ρt) −1Φtdx ) + δE
δρt = 0,
(AIG)
with initial values ρ0 = ρ0 and Φ0 = 0. The choice of αt depends on the geodesic convexity of E(ρ), which has an equivalent definition as follows. Definition 2 For a functional E(ρ) defined on the probability space, we say that E(ρ) has a βpositive Hessian (in short, Hess(β)) w.r.t. the metric gρ if there exists a constant β ≥ 0 such that for any ρ ∈ P(Ω) and any σ ∈ TρP(Ω), we have
gρ(HessE(ρ)σ, σ) ≥ βgρ(σ, σ). Here Hess is the Hessian operator w.r.t. gρ.
If E(ρ) is Hess(β) for β > 0, then αt = 2 √ β; if E(ρ) is Hess(0), then αt = 3/t.
Remark 1 The Nesterov’s accelerated method (Nesterov, 1983) is a first-order method to optimize f(x) in the Euclidean space. The corresponding accelerated gradient flow by Su et al. (2016) is equivalent to a damped Hamiltonian system[
ẋ ṗ
] + [ 0 αtp ] − [ 0 I −I 0 ] [ ∇xHE(x, p) ∇pHE(x, p) ] = 0, HE(x, p) = 1 2 ‖p‖2 + f(x).
with initial values x(0) = x0 and p(0) = 0. The choice of αt depends on the property of f(x). If f(x) is β-strongly convex, then αt = 2 √ β; if f(x) is convex, then αt = 3/t. We apply this Hamiltonian flow interpretation to construct (AIG) in the probability space with information metrics.
We give examples of AIG flows under either the Fisher-Rao metric or the Wasserstein metric. Example 7 (Fisher-Rao AIG flow) The Fisher-Rao AIG flow writes ∂tρt − Φtρt + (∫ Φtρtdx ) ρt = 0,
∂tΦt + αtΦt + 1
2 Φ2t −
(∫ ρtΦtdx ) Φt + δE
δρt = 0.
(F-AIG)
Example 8 (Wasserstein AIG flow) The Wasserstein AIG flow writes ∂tρt +∇ · (ρt∇Φt) = 0, ∂tΦt + αtΦt + 1
2 ‖∇Φt‖2 +
δE δρt = 0.
(W-AIG)
For the rest of this paper, we mainly focus on the Wasserstein metric. Here the AIG flow (Eulerian formulation in fluid dynamics) has a counterpart in the particle level (Lagrangian formulation). Proposition 2 Suppose that Xt ∼ ρt and Vt = ∇Φt(Xt) are the position and the velocity of a particle at time t. Then, the differential equation of the particle system corresponding to (W-AIG) writes
dXt = Vtdt, dVt = −αtVtdt−∇ ( δE
δρt
) (Xt)dt. (W-AIG-P)
If E(ρ) evaluates the KL divergence and ρ∗ ∝ exp(−f(x)), (W-AIG-P) is equivalent to dXt = Vtdt, dVt = −αtVtdt−∇f(Xt)dt−∇ log ρt(Xt)dt. (W-AIG-P-KL)
Remark 2 The∇ log ρt(Xt)dt term cannot be simply replaced by a Brownian motion dBt because ρt is the marginal distribution on Xt. Several previous works have studied the accelerated gradient flow of KL divergence in the probability space under the Wasserstein metric. Taghvaei & Mehta (2019) construct the accelerated gradient flow in the probability space based on Wibisono et al. (2016)’s variational formulation on the Nesterov’s accelerated method. Their flows coincide with (W-AIG-P-KL) with αt = 3/t after rescaling. The underdamped Langevin dynamics in (Cheng et al., 2017; Ma et al., 2019) damps the Hamiltonian flow of the particles, which is different from (W-AIG-P) as shown in (Taghvaei & Mehta, 2019). Liu et al. (2018; 2019) give the discrete-time accelerated algorithm from the perspective of manifold optimization.

heading: 3.1 WASSERSTEIN METRIC RESTRICTED TO GAUSSIAN
text: In this subsection, we demonstrate that (W-AIG) has an ODE formulation in Gaussian family. Denote N 0n to the multivariate Gaussian densities with zero means. Namely, if ρ0, ρ∗ ∈ N 0n , then we show that (W-AIG) has a solution (ρt,Φt) and ρt ∈ N 0n . Let Pn and Sn represent symmetric positive definite matrix and symmetric matrix with size n × n respectively. Each ρ ∈ N 0n can be uniquely expressed by its covariance matrix Σ ∈ Pn by ρ(x; Σ) = (2π)−n/2√
det(Σ) exp
( − 12x TΣ−1x ) . The Wasserstein metric on P(Rn) induces a metric onN 0n , which is a
totally-geodesic submanifold in P(Rn), see (Takatsu, 2008; Modin, 2016; Malagò et al., 2018). So there exists a Wasserstein metric on Pn, also known as the Bures metric. For Σ ∈ Pn, the tangent space and cotangent space follow TΣPn ' T ∗ΣPn ' Sn. Definition 3 (Wasserstein metric in Gaussian) For Σ ∈ Pn, the metric tensor G(Σ) : Sn → Sn is defined by G(Σ)−1S = 2(ΣS + SΣ). The Wasserstein metric on Sn is gΣ(A1, A2) = tr(A1G(Σ)A2) = tr(S1ΣS2), where Si ∈ Sn is the solution to Ai = ΣSi + SiΣ, i = 1, 2. Proposition 3 Given an energy function E(Σ), the Wasserstein gradient flow in Gaussian writes
Σ̇t = −2(Σt∇ΣtE(Σt) +∇ΣtE(Σt)Σt). Here∇Σt is the standard matrix derivative. The Hamiltonian flow is a system of (Σt, St), i.e.,{
Σ̇t − (StΣt + ΣtSt) = 0, Ṡt + S 2 t + 2∇ΣtE(Σt) = 0.
(5)
The corresponding Hamiltonian writes H(Σt, St) = tr(StΣtSt) + 2E(Σt).
Therefore, by adding the damping term αtSt, we obtain the Wasserstein AIG flow in Gaussian.{ Σ̇t − (StΣt + ΣtSt) = 0, Ṡt + αtSt + S 2 t + 2∇ΣtE(Σt) = 0, (W-AIG-G)
with initial values Σ0 = Σ0 and S0 = 0. For now, we consider E(Σ) to be the KL divergence.
E(Σ) , E(ρ( · ; Σ)) = 1 2
[ tr(Σ(Σ∗)−1)− log det(Σ(Σ∗)−1)− n ] , (6)
where Σ∗ is the covariance matrix of ρ∗. The following theorem proves the well-posedness of (W-AIG-G) and illustrates the connection between W-AIG flows in Pn and P(Rn). Theorem 1 Suppose that ρ0, ρ∗ ∈ Nn0 and their covariance matrices are Σ0 and Σ∗. E(Σ) defined in (6) evaluates the KL divergence from ρ to ρ∗. Let (Σt, St) be the solution to (W-AIG-G) with initial values Σ0 = Σ0 and S0 = 0. Then, for any t ≥ 0, Σt is well-defined and stays positive definite. Furthermore, we denote
ρt(x) = (2π)−n/2√
det(Σt) exp
( −1
2 xTΣ−1t x
) , Φt(x) = 1
2 xTStx+ C(t),
where C(t) = −t + 12 ∫ t 0 log det(Σs(Σ
∗)−1)ds. Then, (ρt,Φt) is the solution to (W-AIG) with initial values ρ0 = ρ0 and Φ0 = 0.

heading: 4 CONVERGENCE RATE ANALYSIS ON W-AIG FLOWS
text: In this section, we prove the convergence rate of (W-AIG).
Theorem 2 Suppose that E(ρ) satisfies Hess(β) for β > 0. The solution ρt to (W-AIG) with αt = 2 √ β satisfies
E(ρt) ≤ C0e− √ βt = O(e− √ βt).
If E(ρ) only satisfies Hess(0), then the solution ρt to (W-AIG) with αt = 3/t satisfies
E(ρt) ≤ C ′0t−2 = O(t−2).
Here the constants C0, C ′0 only depend on ρ0.
Remark 3 Here Hess(β) is equivalent to the β-geodesic convexity in the probability space w.r.t. gρ. For the Wasserstein metric, it is also known as β-displacement convexity; see (Villani, 2008, Chap 16). Consider the case where E(ρ) is the KL divergence and the target density takes the form ρ∗ ∝ exp(−f(x)). A sufficient condition for Hess(β) is that f(x) is β-strongly convex, see (Otto & Villani, 2000; Bakry & Émery, 1985). If E(ρ) satisfies Hess(β) for β > 0, then the classical analysis indicates that the solution to the Wasserstein gradient flow has an O(e−2βt) convergence rate. The W-AIG flow improves the convergence rate to O(e− √ βt), especially when β is close to 0.
Here we provide a sketch in the proof of Theorem 2. Given ρt, we can find the optimal transport plan Tt from ρt to ρ∗. Let T#ρ denote the push-forward density from ρ by the mapping T . The following proposition characterizes the inverse of the exponential map in probability space with the Wasserstein metric.
Proposition 4 Denote the geodesic curve γ(s) that connects ρt and ρ∗ by γ(s) = (sTt + (1 − s) Id)#ρt, s ∈ [0, 1]. Here Id is the identity mapping from Rn to itself. Then, γ̇(0) corresponds to a tangent vector −∇ · (ρt(x)(Tt(x)− x)) ∈ TρtP(Ω).
We first consider the case where E(ρ) satisfies Hess(β) for β > 0. Motivated by the Lyapunov function for Nesterov’s ODE in the Euclidean case, we construct the following Lyapunov function.
E(t) = e √ βt
( 1
2 ∫ ∥∥∥−√β(Tt(x)− x) +∇Φt(x)∥∥∥2 ρt(x)dx+ E(ρt)− E(ρ∗)) . (7) Proposition 5 Suppose that E(ρ) satisfies Hess(β) for β > 0. ρt is the solution to (W-AIG) with αt = 2 √ β. Then, E(t) defined in (7) satisfies Ė(t) ≤ 0. As a result,
E(ρt) ≤ e− √ βtE(t) ≤ e− √ βtE(0) = O(e− √ βt).
Note that E(0) only depends on ρ0. This proves the first part of Theorem 2. We now consider the case where E(ρ) satisfies Hess(0). We construct the following Lyapunov function.
E(t) = 1 2 ∫ ∥∥∥∥−(Tt(x)− x) + t2∇Φt(x) ∥∥∥∥2 ρt(x)dx+ t24 (E(ρt)− E(ρ∗)). (8)
Proposition 6 Suppose that E(ρ) satisfies Hess(0). ρt is the solution to (W-AIG) with αt = 3/t. Then, E(t) defined in (8) satisfies Ė(t) ≤ 0. As a result,
E(ρt) ≤ 4 t2 E(t) ≤ 4 t2 E(0) = O(t−2).
Because E(0) only depends on ρ0, we complete the proof. Remark 4 For the Hess(0) case, we obtain the same result in (Taghvaei & Mehta, 2019, Theorem 1). Their proof comes from the Lagrangian formulation (W-AIG-P) and our proof is based on the Eulerian formulation (W-AIG). However, their technical assumption E [( Xt + e −γtYt − T ρ∞ρt (Xt) ) · ddtT ρ∞ ρt (Xt) ] = 0 is only valid in 1-dimensional case. In Appendix C.4, we prove that this quantity is non-negative. This is due to the Hodge decomposition behind the optimal transport, see Lemma 1 in Appendix C.3.

heading: 5 DISCRETE-TIME ALGORITHM FOR W-AIG FLOWS
text: In this section, we present the discrete-time implementation of (W-AIG-P-KL). This implementation is simpler and more stable than the one in (Taghvaei & Mehta, 2019). Suppose that initial positions of a particle system {Xi0}Ni=1 are given and V i0 = 0. The time parameter t is related to the step size√ τ via t = √ τk. The update rule follows
V ik+1 = αkV i k − √ τ(∇f(Xik) + ξk(Xik)), Xik+1 = Xik + √ τV ik+1, (9)
for i = 1, 2 . . . N . If E(ρ) is Hess(β), then αk = 1− √ βτ
1+ √ βτ ; if E(ρ) is Hess(0) or β is unknown, then
αk = k−1 k+2 . Here ξk(x) is an approximation of∇ log ρt(x).
We review two common choices of ξk as follows. If Xik follows a Gaussian distribution, then
ξk(x) = −Σ−1k (x−mk), (10)
where mk and Σk are the mean and the covariance matrix of {Xik}Ni=1. For the non-Gaussian case, we use the kernel density estimation (KDE, Singh (1977)), ρ̃k(x) = 1N ∑N i=1K(x,X i k) to approximate ρt(x). Here K(x, y) is a kernel function. Then, ξk writes
ξk(x) = ∇ log ρ̃k(x) = ∑N i=1∇xK(x,Xik)∑N i=1K(x,X i k) . (11)
A common choice of K(x, y) is a Gaussian kernel with the bandwidth h, KG(x, y;h) = (2πh)−n/2 exp ( ‖x− y‖2/(2h) ) . There are two difficulties in the discretization. For one thing, the bandwidth h strongly affects the estimation of∇ log ρt, so we propose the BM method to adaptively learn the bandwidth from samples. For another, the second equation in (W-AIG) is the HamiltonJacobi equation, which usually has strong stiffness. We propose an adaptive restart technique to deal with this problem. Remark 5 Our numerical implementations of W-AIG flows can be viewed as a special case of ParVI methods. Compared to traditional MCMC methods, ParVI methods are more sample-efficient because make full use of a finite number of particles by taking particle interaction into account.

heading: 5.1 LEARN THE BANDWIDTH VIA BROWNIAN MOTION
text: SVGD uses a median (MED) method to choose the bandwidth, i.e.,
hk = 1
2 log(N + 1) median
( {‖Xik −X j k‖ 2}Ni,j=1 ) . (12)
Liu et al. (2018) propose a Heat Equation (HE) method to adaptively adjust bandwidth. Motivated by the HE method, we introduce the Brownian motion (BM) method to adaptively learn the bandwidth.
Given the bandwidth h, {Xi}Ni=1 and a step size s, we can compute two particle systems Y ik (h) = X i k − sξk(x;h), Zik = Xik + √ 2sBi, i = 1, . . . N, where Bi is the standard Brownian motion. We want to minimize MMD(ρ̂Y , ρ̂Z), the MMD between the empirical distribution ρ̂Y (x) = ∑N i=1 δ(x−Y ik (h)) and ρ̂Z(x) = ∑N i=1 δ(x−Zik). Here, the kernel K(x, y) for the MMD is the Gaussian kernel with a bandwidth of 1. So we optimize over h to minimize MMD(ρ̂Y , ρ̂Z), using the bandwidth hk−1 from the last iteration as the initialization. For simplicity we denote BM(hk−1, {Xik}Ni=1, s) as the output of the BM method. Remark 6 Besides KDE, there are other methods that approximate the ∇ log ρt(x) (compute ξk) via a kernel function, such as the blob method (Carrillo et al., 2019) and the diffusion map (Taghvaei & Mehta, 2019). The BM method can also select the kernel bandwidth for these methods.

heading: 5.2 ADAPTIVE RESTART
text: To enhance the practical performance, we introduce an adaptive restart technique, which shares the same idea of gradient restart in (Odonoghue & Candes, 2015; Wang et al., 2019) under the Euclidean case. Consider
ϕk = − N∑ i=1 〈 V ik+1,∇f(Xik) + ξk(Xik) 〉 , (13)
which can be viewed as discrete-time approximation of −gWρt (∂tρt, G W (ρt) −1 δE δρt ) = −∂tE(ρt). If ϕk < 0, then we restart the algorithm with initial values Xi0 = X i k and V i 0 = 0. This essentially keeps ∂tE(ρt) negative along the trajectory. The numerical results show that the adaptive restart accelerates and stabilizes the discrete-time algorithm. The overall algorithm is summarized below.
Algorithm 1 Particle implementation of Wasserstein AIG flow
Require: initial positions {Xi0}Ni=1, step size τ , number of iteration L. 1: Set k = 0, V i0 = 0, i = 1, . . . N . Set the bandwidth h0 by MED (12). 2: for l = 1, 2, . . . L do 3: Compute hl based on BM: hl = BM(hl−1, {Xik}Ni=1, √ τ).
4: Calculate ξk(Xik) by (10) or by (11) with the bandwidth hl. 5: Set αk based on whether E(ρ) is Hess(β) or Hess(0). For i = 1, 2, . . . N , update
V ik+1 = αkV i k − √ τ(∇f(Xik) + ξk(Xik)), Xik+1 = Xik + √ τV ik+1.
6: if RESTART then 7: Compute ϕk = − ∑N i=1 〈 V ik+1,∇f(Xik) + ξk(Xik) 〉 . 8: If ϕk < 0, set Xi0 = X i k and V i 0 = 0 and k = 0; otherwise set k = k + 1.
9: else 10: Set k = k + 1. 11: end if 12: end for

heading: 6 NUMERICAL EXPERIMENTS
text: In this section, we present several numerical experiments to demonstrate the validity of the BM method, the acceleration effect of the Wasserstein AIG flow, and the strength of the adaptive restart technique. Implementation details can be found in Appendix D.

heading: 6.1 TOY EXAMPLE
text: We first investigate the validity of the BM method in selecting the bandwidth. The target density ρ∗ is a toy bimodal distribution (Rezende & Mohamed, 2015). We compare two types of particle implementations of the Wasserstein gradient flow over the KL divergence:
Xik+1 = X i k − τ∇f(Xik) + √ 2τBik, X i k+1 = X i k − τ(∇f(Xik) + ξk(Xik)).
Here Bik ∼ N (0, 1) is the standard Brownian motion and ξk takes the form (11). The first method is known as the MCMC method and the second method is called the ParVI method. For the second method, the bandwidth h is selected by MED/HE/BM respectively. Figure 1 shows the distribution of 200 samples based on different methods. Samples from MCMC match the target distribution in a stochastic way; samples from MED collapse; samples from HE align tidily around the contour lines; samples from BM arrange neatly and are closer to samples from MCMC. This indicates that the BM method makes the particle system behave similar to MCMC, though in a deterministic way.

heading: 6.2 GAUSSIAN MEASURES
text: Next, we explore the effectiveness of (W-AIG) flow compared to the Wasserstein gradient flow and demonstrate the strength of the adaptive restart. The target density ρ∗ is a Gaussian distribution with zero mean on R100, the covariance matrix of ρ∗ is Σ∗ and W ∗ = (Σ∗)−1. Let L and β be the largest/smallest eigenvalue of W ∗. E(ρ) satisfies Hess(β) and the step size is τ = 1/(4L). The condition number of W ∗ is defined as κ = L/β. The large L indicates Σ∗ is close to be singular.
We first demonstrate the effectiveness of (W-AIG-G) in the ODE level. Detailed discretization is left in Appendix D.2. The initial value is set to be Σ0 = I . For now, WGF denotes the discretization of the Wasserstein gradient flow; AIG-(r)(s) denotes the discretization of the Wasserstein AIG flow. For letters in the parentheses, ‘r’ denotes using the adaptive restart and ‘s’ denotes utilizing β. Figure 2 presents the convergence of the KL divergence on two target distributions with small/large L. We observe that AIG converges faster compared to WGF, which verifies Theorem 2. The adaptive restart also accelerates the algorithm.
Then, we demonstrate the results in the particle level. The setting of ρ∗ is same as the previous experiment. The initial distribution of samples follows N (0, I) and the number of samples is N = 600. For a particle system {Xik}Ni=1, we record the KL divergence E(Σ̂k) (6) using the empirical covariance matrix Σ̂k. The left part of Figure 3 (small L) is almost identical to Figure 2, which verifies the acceleration effect of AIG flows. It also indicates that the adaptive restart helps to accelerate the convergence. From the right part of Figure 3 (large L), AIG and AIG-s diverge because of the ill target distribution, and the adaptive restart solves this problem.

heading: 6.3 BAYESIAN LOGISTIC REGRESSION
text: We perform the standard Bayesian logistic regression experiment on the Covertype dataset, following the same settings as Liu & Wang (2016). We compare our methods with MCMC, SVGD (Liu & Wang, 2016), WNAG (Liu et al., 2018) and WNes (Liu et al., 2019). We select the bandwidth using either the MED method or the proposed BM method. Figure 4 indicates that the BM method accelerates and stabilizes the performance of WGF and AIG. The performance of MCMC and WGF are similar and they achieve the best log-likelihood. In test accuracy, AIG-r converges faster than other methods and is more stable. The adaptive restart improves the overall performance of AIG.

heading: 7 CONCLUSION
text: In summary, we propose the framework of AIG flows by damping Hamiltonian flows with respect to certain information metrics in the probability space. AIG flows have been carefully studied in Gaussian families. Theoretically, we establish the convergence rate of W-AIG flows. Numerically, we propose the discrete-time algorithm and the adaptive restart technique to overcome the numerical stiffness of W-AIG flows. We introduce a novel kernel selection method by learning from Brownianmotion samples. The numerical experiments verify the acceleration effect of AIG flows and the strength of the adaptive restart. In future works, we intend to systematically explain the stiffness of AIG flows and the effect of the adaptive restart. We shall apply our results to general information metrics, especially for the Fisher-Rao metric. We expect to study the related sampling efficient optimization methods and discrete-time algorithms on general probability models.

heading: A THE PROOFS AND DERIVATIONS IN SECTION 2
text: In this section, we provide the derivation of the Euler-Lagrange equation, proofs in Section 2 and the Euler-Lagrange formulation of (AIG).

heading: A.1 DERIVATION OF THE EULER-LAGRANGE EQUATION
text: We derive the Euler-Lagrange equation (1) in this subsection. For a fixed T > 0 and two given densities ρ0, ρT , consider the variational problem
I(ρt) = inf ρt {∫ T 0 L(ρt, ∂tρt)dt ∣∣∣∣∣ ρ0 = ρ0, ρT = ρT } .
Let ht ∈ F(Ω) be the smooth perturbation function that satisfies ∫ htdx = 0, t ∈ [0, T ] and h0 = hT ≡ 0. Denote ρ t = ρt + ht. Note that
I(ρ t) = ∫ T 0 L(ρt, ∂tρt)dt+ ∫ T 0 ∫ ( δL δρt ht + δL δ(∂tρt) ∂tht ) dxdt+ o( ).
From dI(ρ t)
d ∣∣∣ =0
= 0, it follows that∫ T 0 ∫ ( δL δρt ht + δL δ(∂tρt) ∂tht ) dxdt = 0.
Note that h0 = hT ≡ 0. Perform integration by parts w.r.t. t yields∫ T 0 ∫ ( δL δρt − ∂t δL δ(∂tρt) ) htdxdt = 0.
Because ∫ htdx = 0, (1) holds with a spatially constant function C(t).

heading: A.2 THE PROOF OF PROPOSITION 1 IN SECTION 2
text: In this subsection, we derive the Hamiltonian flow in the probability space. First, we give a useful identity. Given a metric tensor G(ρ) : TρP(Ω)→ T ∗ρP(Ω), we have∫
σ1G(ρ)σ2dx = ∫ G(ρ)σ1σ2dx = ∫ Φ1G(ρ) −1Φ2dx = ∫ G(ρ)−1Φ1Φ2dx. (14)
Here Φ1 = G(ρ)−1σ1 and Φ2 = G(ρ)−1σ2. We then check that
δ
δρt
(∫ ∂tρtG(ρt)∂tρtdx ) = − δ
δρt
(∫ ΦtG(ρt) −1Φtdx ) . (15)
Let ρ̃t = ρt + h, where h ∈ TρtP(Ω). For all σ ∈ TρtP , it follows
G(ρt + h) −1G(ρt + h)σ = σ.
The first-order derivative w.r.t. of the left hand side shall be 0, i.e.,( ∂G(ρt) −1 ∂ρt · h ) G(ρt)σ +G(ρt) −1 ( ∂G(ρt) ∂ρt · h ) σ = 0.
Because ∂tρt = G(ρ)−1Φt, applying (14) yields∫ ∂tρt ( ∂G(ρt) ∂ρt · h ) ∂tρtdx = ∫ ΦtG(ρt) −1 ( ∂G(ρt) ∂ρt · h ) ∂tρtdx
=− ∫
Φt
( ∂G(ρt) −1 ∂ρt · h ) G(ρt)∂tρtdx = − ∫ Φt ( ∂G(ρt) −1 ∂ρt · h ) Φtdx.
(16)
Based on basic calculations, we can compute that∫ ∂tρtG(ρ̃t)∂tρtdx− ∫ ∂tρtG(ρt)∂tρtdx = ∫ ∂tρt ( ∂G(ρt) ∂ρt · h ) ∂tρtdx+ o( ), (17)
− ∫ ΦtG(ρ̃t) −1Φtdx+ ∫ ΦtG(ρt) −1Φtdx = − ∫ Φt ( ∂G(ρt) −1 ∂ρt · h ) Φtdx+ o( ). (18)
Combining (16), (17) and (18) yields (15). Hence, the Euler-Lagrange equation (1) is equivalent to
∂tΦt = 1
2
δ
δρt
(∫ ∂tρtG(ρt)∂tρtdx ) − δE δρt = −1 2 δ δρt (∫ ΦtG(ρt) −1Φtdx ) − δE δρt .
This equation combining with ∂tρt = G(ρ)−1Φt recovers the Hamiltonian flow (2). In short, the Euler-Lagrange equation (1) is from the primal coordinates (ρt, ∂tρt) and the Hamiltonian flow (2) is from the dual coordinates (ρt,Φt). Similar interpretations can be found in (Chow et al., 2019).

heading: A.3 THE EULER-LAGRANGIAN FORMULATION OF AIG FLOWS
text: We can formulate (AIG) as a second-order equation of ρt,
D2 Dt2 ρt + αt∂tρt +G(ρt) −1 δE δρt = 0.
Here D2/Dt2 is the covariant derivative in metric G. We can also explicitly write D 2
Dt2 ρt as follows.
D2 Dt2 ρt = ∂ttρt − (∂tG(ρt)−1)∂tρt + 1 2 G(ρt) −1 δ
δρt
(∫ ∂tρtG(ρt)∂tρtdx ) .

heading: B THE PROOFS IN SECTION 3
text: In this section, we present proofs of propositions and theorems in Section 3.

heading: B.1 THE PROOF OF PROPOSITION 2 IN SECTION 3
text: We start with an identity. For a twice differentiable Φ(x), we have
1 2 ∇‖∇Φ‖2 = ∇2Φ∇Φ = (∇Φ · ∇)∇Φ. (19)
From (W-AIG), it follows that ∂tρt +∇ · (ρt∇Φt) = 0. (20)
This is the continuity equation of ρt. Hence, on the particle level, Xt shall follows
dXt = ∇Φt(Xt)dt.
Let Vt = ∇Φt(Xt). Then, by the conservation of momentum and (W-AIG), we have
dVt = (∂t +∇Φt(Xt) · ∇)∇Φt(Xt)dt
= ( −αt∇Φt(Xt)− 1
2 ∇‖∇Φ‖2 −∇ δE
δρt
) dt+ (∇Φ · ∇)∇Φdt
=− αt∇Φt(Xt)dt−∇ δE
δρt (Xt)dt = −αtVtdt−∇
δE δρt (Xt)dt.

heading: B.2 THE PROOF OF PROPOSITION 3 IN SUBSECTION 3.1
text: In this subsection, we derive the Hamiltonian flow in Gaussian. For A ∈ Sn, we define the linear operator MA : Sn → Sn by
MAB = AB +BA, B ∈ Sn.
It is easy to verify that if A ∈ Pn, then M−1A is well-defined. For a flow Σt ∈ Pn, t ≥ 0, we define the Lagrangian L(Σt, Σ̇t) = 12gΣt(Σ̇t, Σ̇t) − E(Σt). The corresponding Euler-Lagrange equation writes
d
dt
dL dΣ̇t = dL dΣ . (21)
Let St = M−1Σt Σ̇t, i.e., Σ̇t = StΣt + ΣtSt. Then, it follows
gΣt(Σ̇t, Σ̇t) = tr(StΣtSt) = 1
2 tr((StΣt + ΣtSt)St) =
1 2 tr(Σ̇tSt) = 1 2 tr(Σ̇tM −1 Σt Σ̇t).
This leads to dL dΣ̇t = M−1Σt Σ̇t = St. For simplicity, we denote g = gΣt(Σ̇t, Σ̇t). First, we show that
dg
dΣt = −S2t .
Because St = M−1Σt Σ̇t. Given Σ̇t, St can be viewed as a continuous function of Σt. For anyA ∈ S n, define lA = tr((ΣtSt + StΣt)A).
0 = dlA dΣt = ∂St ∂Σt ∂LA ∂St + ∂lA ∂Σt = ∂St ∂Σt (AΣt + ΣtA) + (ASt + StA).
Here we view ∂ST /∂Σt as a linear operator on Sn. Let B = AΣt + ΣtA, then A = M−1Σt B. ∂St ∂Σt B +MStM −1 Σt B = 0 holds for all B ∈ Sn. Therefore, we have ∂St∂Σt = −MStM −1 Σt . Hence,
dg dΣt = ∂St ∂Σt ∂g ∂St + ∂g ∂Σt = −MStM−1Σt (StΣt + ΣtSt) + S 2 t = −MStSt + S2t = −S2t .
As a result, the Euler-Lagrange equation (21) is equivalent to
Ṡ 2 = dL dΣt = −S
2 t
2 −∇E(Σt). (22)
Combining (22) with Σ̇t = StΣt + ΣtSt renders the Hamiltonian flow (5).

heading: B.3 THE PROOF OF THEOREM 1 IN SUBSECTION 3.1
text: We first show that Σt stays in Pn. Suppose that Σt ∈ Pn for 0 ≤ t ≤ T . Define Ht = H(Σt, St) = tr(StΣtSt)/2 + E(Σt). We observe that (W-AIG-G) is equivalent to
Σ̇t = 2 ∂Ht ∂St , Ṡt = −αtSt − 2 ∂Ht ∂Σt . (23)
We show that Ht is decreasing with respect to t.
dHt dt = tr ( ∂Ht ∂St Ṡt + ∂Ht ∂Σt Σ̇t ) = tr ( ∂Ht ∂St ( −αtSt − 2 ∂Ht ∂Σt ) + 2 ∂Ht ∂Σt ∂Ht ∂St ) =− αt tr ( St ∂Ht ∂St ) = −αt 2 tr(St(ΣtSt + StΣt)) = −αt tr(StΣtSt) ≤ 0.
For simplicity, we denote W ∗ = (Σ∗)−1. Let λt be the smallest eigenvalue of Σt. Then, log det(ΣtW ∗) ≥ log det(Σt) = n log λt. Therefore,
−n 2 (log λt + 1) ≤ − 1 2 [log det(ΣtW ∗) + n] ≤ E(Σt) ≤ H(t) ≤ H(0),
which indicates that
λt ≥ exp ( − 2 n H(0)− 1 ) . (24)
This means that as long as Σt ∈ Pn, the smallest eigenvalue of Σt has a positive lower bounded. If there exists T > 0 such that ΣT /∈ Pn. Because Σt is continuous with respect to t, there exists T1 < T , such that Σt ∈ Pn, 0 ≤ t ≤ T1 and λT1 < exp (−2H(0)/n− 1), which violates (24).
We then reveal the relationship between (W-AIG) in P(Rn) and Pn. We can compute that
∂ ∂t det(Σt) = det(Σt) tr(Σ −1 t Σ̇t), ∂ ∂t Σ−1t = −Σ−1t Σ̇tΣ−1t .
Combining with Σ̇t = ΣtSt + StΣt, we obtain
tr(Σ−1t Σ̇t) = tr(St + Σ −1 t StΣt) = 2 tr(St),
tr(xΣ−1t Σ̇tΣ −1 t x) = tr(x TΣ−1t Stx+ x TStΣ −1 t x) = 2 tr(StΣ −1 t xx T ).
Therefore, it follows
∂tρt(x) = ∂
∂t
( 1√
det(Σt)
)√ det(Σt)ρt(x) + 1
2 tr(xTΣ−1t Σ̇tΣ −1 t x)ρt(x)
=− 1 2 tr(Σ−1t Σ̇t)ρt(x) + tr(StΣ −1 t xx T )ρt(x) = − tr(St(I − Σ−1t xxT ))ρt(x).
Note that∇Φt(x) = Stx. Hence,
−∇ · (ρt∇Φt) = − n∑ i=1 ∂i(ρt(x)Stx)i = − n∑ i=1 [ρt(x)∂i(Stx)i + (Stx)i∂iρt(x)]
=− ρt(x) [ tr(St) + (Stx) T (−Σ−1t x) ] = −ρt(x) tr(St(I − Σ−1t xxT )) = ∂tρt(x).
The first equation of (W-AIG) holds. Because ∂tΦt(x) = xT Ṡtx/2 + Ċ(t),
∂tΦt(x) + αtΦt(x) + 1
2 ‖∇Φt(x)‖2 =
1 2 xT Ṡtx+ αt 2 xTStx+ 1 2 xTS2t x+ Ċ(t)
=− xT∇ΣtE(Σt)x+ Ċ(t) = 1 2 xT (Σ−1t −W ∗)x+ Ċ(t).
Note that ρ∗ is the Gaussian density with the covariance matrix Σ∗. Because Ċ(t) = 1 2 log det(ΣtW ∗)− 1, we can compute
δE δρt = log ρt(x)− log ρ∗(x) + 1 = − 1 2 xT (Σ−1t −W ∗)x− 1 2 log det(ΣtW ∗) + 1
=− 1 2 xT (Σ−1t −W ∗)x− Ċ(t) = −(∂tΦt(x) + αtΦt(x) + 1 2 ‖∇Φt(x)‖2).
Therefore, the second equation of (W-AIG) holds. Because Σ0 = Σ0, S0 = 0 and C(0) = 0, we have ρ0 = ρ0 and Φ0 = 0. This completes the proof.

heading: C THE PROOFS IN SECTION 4
text: In this section, we briefly review some geometric properties of the probability space as a Riemannian manifold and present proofs of propositions in Section 4.

heading: C.1 A BRIEF REVIEW ON THE GEOMETRIC PROPERTIES OF THE PROBABILITY SPACE
text: Suppose that we have a metric gρ in the probability space P(Ω). Given two probability densities ρ0, ρ1 ∈ P(Ω), we define the distance as follows
d(ρ0, ρ1) = ( inf γ(s) {∫ 1 0 gγ(s)(γ̇(s), γ̇(s))ds : γ(0) = ρ 0, γ(1) = ρ1 })1/2 .
The minimizer γ(s) of the above problem is defined as the geodesic curve connecting ρ0 and ρ1. An exponential map at ρ0 ∈ P(Ω) is a mapping from the tangent space Tρ0P(Ω) to P(Ω). It requires that σ ∈ Tρ0P(Ω) is mapped to a point ρ1 ∈ P(Ω) such that there exists a geodesic curve γ(s) satisfying γ(0) = ρ0, γ̇(0) = σ and γ(1) = ρ1.

heading: C.2 THE PROOF OF PROPOSITION 4 IN SECTION 4
text: In this subsection, we characterize the inverse of the exponential map in the probability space with the Wasserstein metric. Let T st = (sTt + (1 − s) Id)−1, s ∈ [0, 1]. Then, based on the theory of optimal transport (Villani, 2003), we can write the explicit formula of the geodesic curve γ(s) by
γ(s) = T st #ρt = det(∇T st )ρt ◦ T st . Through basic calculations, we can compute that
d
ds T st ∣∣∣∣ s=0 = − d ds (sTt + (1− s) Id) ∣∣∣∣ s=0 = Id−Tt.
d
ds det(∇T st ) ∣∣∣∣ s=0 = d ds det(I + s(I −DTt) + o(s)) ∣∣∣∣ s=0 = tr(I −DTt).
Therefore, d
ds γ(s) ∣∣∣∣ s=0 (x) = tr(I −∇Tt)ρt(x) + 〈∇ρt(x), x− ϕt(x)〉
=∇ · (x− Tt(x))ρt(x) + 〈∇ρt(x), x− Tt(x)〉 = −∇ · (ρt(x)(Tt(x)− x)), which completes the proof.

heading: C.3 THE PROOF OF PROPOSITION 5 AND 6 IN SECTION 4
text: The main goal of this subsection is to prove the Lyapunov function E(t) is non-increasing. Preparations. We first give a better characterization of the optimal transport plan Tt. We can write Tt = ∇Ψt, where Ψt is a strictly convex function, see (Villani, 2003). This indicates that ∇Tt is symmetric. We then introduce the following proposition. Proposition 7 Suppose that E(ρ) satisfies Hess(β) for β ≥ 0. Let Tt(x) be the optimal transport plan from ρt to ρ∗, then
E(ρ∗) ≥ E(ρt) + ∫ 〈 Tt(x)− x,∇ δE
δρt
〉 ρdx+ β
2
∫ ‖Tt(x)− x‖2ρtdx.
This is a direct result of β-geodesic convexity of E(ρ) based on Proposition 4.
Next, let us denote ut = ∂t(Tt)−1 ◦ Tt. We show that ut satisfies ∇ · (ρt(ut −∇Φt)) = 0. (25)
Because (Tt)−1#ρ∗ = ρt, let ut = ∂t(Tt)−1 ◦Tt andXt = (Tt)−1X0, whereX0 ∼ ρ∗. This yields d dtXt = ut(Xt). The distribution of Xt follows ρt. By the Euler’s equation, ρt shall follows ∂tρt +∇ · (ρtut) = 0. Combining this with the continuity equation (20) yields (25).
Then, we formulate ∂tTt(x) with ut. By the Taylor expansion,
Tt+s(x) = Tt(x) + s∂tTt(x) + o(s).
Let y = (Tt)−1x. it follows
(Tt+s) −1(x) = (Tt) −1(x) + sut((Tt) −1(x)) + o(s) = y + sut(y) + o(s).
Therefore, 0 = Tt+s((Tt+s)
−1(x))− x = Tt+s(y + sut(y) + o(s))− x =Tt(y + sut(y)) + s∂tTt(y + sut(y))− x+ o(s) =Tt(y) + s∇Tt(y)ut(y) + s∂tTt(y)− x+ o(s) =s [∇Tt(y)ut(y) + ∂tTt(y)] + o(s).
We shall have ∇Tt(y)ut(y) + ∂tTt(y) = 0. Replacing y by x gives ∂tTt(x) = −∇Tt(x)ut(x). (26)
The following lemma illustrates two important properties of ut and ∂tTt.
Lemma 1 For ut satisfying (25), we have∫ 〈∇Φt − ut,∇Tt∇Φt〉 ρtdx ≥ 0, ∫ 〈∇Φt − ut,∇Tt(x)(Tt(x)− x)〉 ρt = 0.
PROOF We first notice that ut − ∇Φt is divergence-free in term of ρt. From −∇Ttut = ∂tTt = ∇∂tΨt, we observe that −∇Ttut is the gradient of ∂tΨt. Therefore,∫
〈∇Φt − ut,∇Ttut〉 ρt = − ∫ 〈∂tΨt,∇ · (ρt(∇Φt − ut))〉 = 0.
Based on our previous characterization on the optimal transport plan Tt,∇Tt = ∇2Ψt is symmetric positive definite. This yields that∫
〈∇Φt − ut,∇Tt∇Φt〉 ρtdx = ∫ 〈∇Φt − ut,∇Tt∇Φt〉 ρtdx− ∫ 〈∇Φt − ut,∇Ttut〉 ρt
= ∫ 〈∇Φt − ut,∇Tt(∇Φt − ut)〉 ρtdx ≥ 0.
The last inequality utilizes that ∇Tt is positie definite and ρt is non-negative. Then, we prove the equality in Lemma 1. Because ∇Tt(x)(Tt(x) − x) = 12∇(‖Tt(x) − x‖
2 + Tt(x) − ‖x‖2) is a gradient. Similarly, it follows∫
〈∇Φt − ut,∇Tt(x)(Tt(x)− x)〉 ρt = 0.
Lemma 1 and the relationship (26) gives − ∫ 〈∂tTt,∇Φt〉 ρtdx = ∫ 〈ut,∇Tt∇Φt〉 ρtdx ≤ ∫ 〈∇Φt,∇Tt∇Φt〉 ρtdx, (27)∫
〈∂tTt, Tt(x)− x〉 ρtdx = − ∫ 〈∇Φt,∇Tt(x)(Tt(x)− x)〉 ρtdx. (28)
Proof of Proposition 5. Based on the definition of the Wasserstein metric, we have ∂tE(ρt) = − ∫ δE
δρt ∇ · (ρt∇Φt)dx.
Differentiating E(t) (7) w.r.t. t renders
Ė(t)e− √ βt =β ∫ 〈∂tTt, Tt(x)− x〉 ρtdx− β
2
∫ ‖Tt(x)− x‖2∇ · (ρt∇Φt)dx
− √ β ∫ 〈∂tTt,∇Φt〉 ρtdx− √ β ∫ 〈Tt(x)− x, ∂t∇Φt〉 ρtdx + √ β ∫ 〈Tt(x)− x,∇Φt〉∇ · (ρt∇Φt)dx+ ∫ 〈∇Φt, ∂t∇Φt〉 ρtdx
− 1 2
∫ ‖∇Φt‖2∇ · (ρt∇Φt)− ∫ δE
δρt ∇ · (ρt∇Φt)dx
+
√ β
2
∫ ‖∇Φt‖2ρtdx− β ∫ 〈Tt(x)− x,∇Φt(x)〉 ρtdx
+
√ β3
2
∫ ‖Tt(x)− x‖2ρtdx+ √ β(E(ρt)− E(ρ∗)). (29)
For the part (29), Proposition 7 renders√ β3
2
∫ ‖Tt(x)− x‖2ρtdx+ √ βE(ρt) ≤ − √ β ∫ 〈 Tt(x)− x,∇ δE
δρt
〉 ρtdx. (30)
We first compute the terms with the coefficient β0 in Ė(t)e− √ βt. We observe that∫
〈∇Φt, ∂tΦt〉 ρtdx− 1
2
∫ ‖∇Φt‖2∇ · (ρt∇Φt)dx− ∫ δE
δρt ∇ · (ρt∇Φt)ρtdx
= ∫ 〈 ∂t∇Φt + 1
2 ∇‖∇Φt‖2 +∇
δE δρ ,∇Φt
〉 ρtdx = −2 √ β ∫ ‖∇Φt‖2ρtdx,
(31)
where the last equality uses (W-AIG) with αt = 2 √ β. Substituting (30) and (31) into the expression of Ė(t)e− √ βt yields
Ė(t)e− √ βt ≤β ∫ 〈∂tTt, Tt(x)− x〉 ρtdx− β
2
∫ ‖Tt(x)− x‖2∇ · (ρt∇Φt)dx
− β ∫ 〈Tt(x)− x,∇Φt〉 ρtdx− √ β ∫ 〈∂tTt,∇Φt〉 ρtdx − √ β ∫ 〈Tt(x)− x, ∂t∇Φt〉 ρtdx− √ β ∫ 〈 Tt(x)− x,∇ δE
δρt
〉 ρtdx
+ √ β ∫ 〈Tt(x)− x,∇Φt〉∇ · (ρt∇Φt)dx− 3 √ β
2
∫ ‖∇Φt‖2ρtdx.
(32)
Then, we deal with the terms with∇ · (ρt∇Φt). We have the following two identities∫ 〈Tt(x)− x,∇Φt〉∇ · (ρt∇Φt)dx = − ∫ 〈∇ 〈Tt(x)− x,∇Φt〉 ,∇Φt〉 ρtdx
=− ∫ 〈 ∇Φt,∇2Φt(x)(Tt(x)− x) + (∇Tt(x)− I)∇Φt 〉 ρtdx
=− 1 2
∫ 〈 Tt(x)− x,∇‖∇Φt‖2 〉 ρtdx− ∫ 〈∇Φt,∇Tt∇Φt〉 ρtdx+ ∫ ‖∇Φt‖2ρtdx. (33)
− 1 2
∫ ‖Tt(x)− x‖2∇ · (ρt∇Φt)dx = ∫ 〈(∇Tt(x)− I)(Tt(x)− x),∇Φt〉 ρtdx
= ∫ 〈Tt(x)− x,∇Tt∇Φt〉 ρtdx− ∫ 〈Tt(x)− x,∇Φt〉 ρtdx.
(34)
Hence, we can proceed to compute the terms with the coefficient √ β. (27) and (33) yields
− √ β ∫ 〈∂tTt,∇Φt〉 ρtdx− √ β ∫ 〈 Tt(x)− x, ∂t∇Φt +∇ δE
δρt
〉 ρtdx
− 3 √ β
2
∫ ‖∇Φt‖2ρtdx+ √ β ∫ 〈Tt(x)− x,∇Φt〉∇ · (ρt∇Φt)dx
=− √ β ∫ 〈∂tTt +∇Tt∇Φt,∇Φt〉 ρtdx− √ β
2
∫ ‖∇Φt‖2ρtdx
− √ β ∫ 〈 Tt(x)− x, ∂t∇Φt +∇ δE
δρ +
1 2 ∇‖∇Φt‖2
〉 ρtdx
≤− √ β
2
∫ ‖∇Φt‖2ρtdx+ 2β ∫ 〈Tt(x)− x,∇Φt〉 ρtdx.
(35)
Substituting (34) and (35) into (32) gives
Ė(t)e− √ βt +
√ β
2
∫ ‖∇Φt‖2ρtdx
≤β ∫ 〈∂tTt, Tt(x)− x〉 ρtdx− β
2
∫ ‖Tt(x)− x‖2∇ · (ρt∇Φt)dx
− β ∫ 〈Tt(x)− x,∇Φt〉 ρtdx+ 2β ∫ 〈Tt(x)− x,∇Φt〉 ρtdx
=β ∫ 〈∂tTt +∇Tt∇Φt, Tt(x)− x〉 ρtdx = 0,
where the last equality uses (28). In summary, we have
Ė(t)e− √ βt ≤ −
√ β
2
∫ ‖∇Φt‖2ρtdx ≤ 0.
Proof of Proposition 6. Differentiating E(t) (8) w.r.t. t, we compute that
Ė(t) = ∫ 〈∂tTt, Tt(x)− x〉 ρtdx− 1
2
∫ ‖Tt(x)− x‖2∇ · (ρt∇Φt)dx
− ∫ 〈 ∂tTt, t
2 ∇Φt
〉 ρtdx− ∫ 〈 Tt(x)− x, 1
2 ∇Φt +
t 2 ∂t∇Φt
〉 ρtdx
+ ∫ 〈 Tt(x)− x, t
2 ∇Φt
〉 ∇ · (ρt∇Φt)dx+ ∫ 〈 t
2 ∇Φt,
1 2 ∇Φt + t 2 ∂t∇Φt
〉 ρtdx
− 1 2 ∫ ∥∥∥∥ t2∇Φt ∥∥∥∥2∇ · (ρt∇Φt)dx− t24 ∫ δE δρt ∇ · (ρt∇Φt)dx+ t 2 (E(ρt)− E(ρ∗)).
(36)
Because E(ρ) is Hess(0), Proposition 7 yields E(ρt) = E(ρt)− E(ρ∗) ≤ − ∫ 〈 Tt(x)− x,∇ δE
δρt
〉 ρtdx. (37)
Utilizing the inequality (37) and substituting the expressions of terms involving ∂tTt and∇·(ρt∇Φt) in (36) with the expressions in (27) (28) and (33) (34), we obtain
Ė(t) ≤− ∫ 〈∇Φt,∇Tt(x)(Tt(x)− x)〉 ρtdx+ ∫ 〈Tt(x)− x,∇Tt∇Φt〉 ρtdx
− ∫ 〈Tt(x)− x,∇Φt〉 ρtdx+ t
2
∫ 〈∇Φt,∇Tt∇Φt〉 ρtdx
− 1 2
∫ 〈Tt(x)− x,∇Φt〉 ρtdx− t
2
∫ 〈∂t∇Φt, Tt(x)− x〉 ρtdx
− t 4
∫ 〈 Tt(x)− x,∇‖∇Φt‖2 〉 ρtdx− t
2
∫ 〈∇Φt,∇Tt∇Φt〉 ρtdx
+ t
2
∫ ‖∇Φt‖2ρtdx+ t
4
∫ ‖∇Φt‖2ρtdx+ t2
4
∫ 〈∇Φt, ∂t∇Φt〉 ρtdx
+ t2
8
∫ 〈 ∇Φt,∇‖∇Φt‖2 〉 ρtdx+ t2
4
∫ 〈 ∇Φt,∇ δE
δρt
〉 ρtdx
− t 2
∫ 〈 Tt(x)− x,∇ δE
δρt
〉 ρtdx.
(38)
The expression of (38) can be reformulated into
Ė(t) ≤− 3 2
∫ 〈Tt(x)− x,∇Φt〉 ρtdx+ 3t
4
∫ ‖∇Φt‖2ρtdx
− t 2
∫ 〈 Tt(x)− x, ∂t∇Φt + 1
2 ∇‖∇Φt‖2 +∇
δE
δρt
〉 ρtdx
+ t2
4
∫ 〈 ∇Φt, ∂t∇Φt + 1
2 ∇‖∇Φt‖2 +∇
δE
δρt
〉 ρtdx.
From (W-AIG) with αt = 3/t, we have the following equalities.
t2
4
∫ 〈 ∇Φt, ∂t∇Φt + 1
2 ∇‖∇Φt‖2 +∇
δE
δρt
〉 ρtdx = − 3t
4
∫ ‖∇Φt‖2ρtdx,
− t 2
∫ 〈 Tt(x)− x, ∂t∇Φt + 1
2 ∇‖∇Φt‖2 +∇
δE
δρt
〉 ρtdx = 3
2
∫ 〈Tt(x)− x,∇Φt〉 ρtdx.
As a result, Ė(t) ≤ 0. This completes the proof.

heading: C.4 COMPARISON WITH THE PROOF IN THE ACCELERATED FLOW
text: The accelerated flow in (Taghvaei & Mehta, 2019) is given by dXt dt = eαt−γtYt, dYt dt
= −eαt+βt+γt∇ ( δE
δρt
) (Xt). (39)
Here the target distribution satisies ρ∞(x) = ρ∗(x) ∝ exp(−f(x)). Suppose that we take αt = log p − log t, βt = p log t + logC and γt = p log t. Here we specify p = 2 and C = 1/4. Then the accelerated flow (39) is identical to (W-AIG-P-KL) if we replace Yt by 2t−3Vt. The Lyapunov function in (Taghvaei & Mehta, 2019) follows
V (t) = 1
2 E[‖Xt + e−γtYt − T ρ
∗
ρt (Xt)‖ 2] + eβt(E(ρ)− E(ρ∗))
= 1
2 E[‖Xt +
t 2 Vt − T ρ ∗ ρt (Xt)‖ 2] +
t2 4 (E(ρt)− E(ρ∗))
= 1
2 ∫ ∥∥∥∥−(Tt(x)− x) + t2∇Φt(x) ∥∥∥∥2 ρt(x)dx+ t24 (E(ρt)− E(ρ∗)).
The last equality is based on the fact that Vt = ∇Φt(Xt) and Tt = T ρ ∗
ρt is the optimal transport plan from ρt to ρ∗. This indicates that the Lyapunov function in (Taghvaei & Mehta, 2019) is identical to ours (8). The technical assumption in (Taghvaei & Mehta, 2019) follows
0 =E [( Xt + e −γtYt − T ρ ∗ ρt (Xt) ) · d dt T ρ ∗ ρt (Xt) ] =E [( Xt + t
2 Vt − Tt(Xt) ) · d dt Tt(Xt) ] =E [( Xt + t
2 Vt − Tt(Xt)
) · ((∂tTt)(Xt) +∇TtVt) ] = ∫ 〈 x− Tt(x) + t
2 ∇Φt(x), ∂tTt +∇Tt∇Φt
〉 ρtdx
Based on ∂tTt = −∇Ttut and Lemma 1, we have∫ 〈x− Tt(x), ∂tTt +∇Tt∇Φt〉 ρtdx = ∫ 〈x− Tt(x),∇Tt(∇Φt − ut)〉 ρtdx = 0.∫
〈∇Φt, ∂tTt +∇Tt∇Φt〉 ρtdx = ∫ 〈∇Φt,∇Tt(∇Φt − ut)〉 ρtdx
= ∫ 〈∇Φt − ut,∇Tt(∇Φt − ut)〉 ρtdx ≥ 0.
As a result, we have E [( Xt + e −γtYt − T ρ∞ρt (Xt) ) · d dt T ρ∞ρt (Xt) ] = t 2 ∫ 〈∇Φt − ut,∇Tt(∇Φt − ut)〉 ρtdx ≥ 0.
In 1-dimensional case, because ∇ · (ρt(ut −∇Φt)) = 0 indicates that ρt(ut − ∇Φt) = 0. For ρt(x) > 0, we have ut(x) − ∇Φt(x) = 0. So the technical assumption holds. In general cases, although ut = ∂t(Tt)−1 ◦ Tt satisfies ∇ · (ρt(ut −∇Φt)) = 0, but this does not necessary indicate that ut = ∇Φt. Hence, E [( Xt + e −γtYt − T ρ∞ρt (Xt) ) · ddtT ρ∞ ρt (Xt) ] = 0 does not necessary hold except for 1-dimensional case.
D IMPLEMENTATION DETAILS IN THE NUMERICAL EXPERIMENTS
In this section, we elaborate on the implementation details in the numerical experiments.

heading: D.1 DETAILS IN SUBSECTION 6.1
text: The initial distribution of the particle system follows the standard Gaussian N (0, I). The objective density function is ρ∗(x) ∝ exp(−2(‖x‖ − 3)2)(exp(−2(x1 − 3)2) + exp(−2(x1 + 3)2)), x ∈ R2. All methods run for 200 iterations using the same fixed step size τ = 0.1.

heading: D.2 DETAILS IN SUBSECTION 6.2
text: The discretization the Wasserstein gradient flow in Gaussian follows
Σk+1 = Σk − 2τ(Σk∇ΣkE(Σk) +∇ΣkE(Σk)Σk). (40)
The discretization of (W-AIG-G) is given by
Sk+1 = αkSk − √ τ(S2k + 2∇ΣkE(Σk)), Σk+1 = (I + √ τSk+1)Σk(I + √ τSk+1). (41)
The choice of αk is same as the one for (9), based on whether β is known. To ensure the update of Σk+1 won’t blow up, we check whether tr(τS2) > n. If this holds, we restart the algorithm. In fact, the update on Σk can be viewed as the exponential map from Σk with the direction ΣkSk+1 + Sk+1Σk. Nevertheless, this is an exponential map if and only if I + √ τSk+1 is positive definite.
The update rule of √ τSk+1 can rewrite into √ τSk+1 = αk √ τSk − ( √ τSk) 2 − 2τ∇ΣkE(Σk). (42) Because of the existence of−( √ τSk)
2, as long as the spectral radius (the eigenvalue with the largest absolute value) of √ τSk is greater than 1, then I + √ τSk+1 cannot maintain to be positive definite. So at this time we need to restart the algorithm. Nevertheless, to compute the spectral radius is computational costly. Instead, we use a weaker condition tr(τS2k) > n.
We also introduce the adaptive restart technique for (41). Consider
ϕk = − tr(Sk+1Σk∇ΣkE(Σk)).
This restarting criterion corresponds to the particle version of (13). Similarly, if ϕk < 0, we restart the algorithm using initial values Σ0 = Σk and S0 = 0.
For the particle implementation, the update rule of WGF writes
Xik = −τ(∇f(Xik)) + ξk(Xik)),
where ξk is computed based on KDE (11).

heading: D.3 DETAILS IN SUBSECTION 6.3
text: We follow the same setting as (Liu & Wang, 2016), which is also adopted by Liu et al. (2018; 2019). The data set is split into 80% for training and 20% for testing. The mini-batch size is taken as 50. For MCMC, the number of particles is N = 1000; for other methods, the number of particles is N = 100. The BM method is not applied to SVGD in selecting the bandwidth.
The initial step sizes for the compared methods are given in Table 1. The step size of SVGD is adjusted by Adagrad, which is same as (Liu & Wang, 2016). For WNAG, AIG and WRes, the step size is give by τl = τ0/l0.9 for l ≥ 1. The parameters for WNAG and Wnes are identical to (Liu et al., 2018) and (Liu et al., 2019). For MCMC, WGF and AIG-r, the step size is multiplied by 0.9 every 100 iterations. We record the cpu-time for each method in Table 2. The computational cost
of the BM method is much higher than the MED method because we need to evaluate the MMD of two particle systems several times in optimizing the subproblem.

Review summaries:
The paper attempts to develop a counterpart of the well-known Nesterov accelerated gradient method for gradient flows on the space of probability measures equipped with an information metric. This is an important problem which is useful for optimization on probability spaces. The accelerated gradient flow is developed by leveraging a damping Hamiltonian flow. The paper focuses mainly on the case with the Wasserstein metric and provides a convergence analysis. Practical considerations such as discretizing the accelerated flow and bandwidth selection are developed for the use of the method in practical problems.