Title: ACCELERATED INFORMATION GRADIENT FLOW

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.

1. INTRODUCTION:

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.

2. METRIC AND FLOWS IN THE PROBABILITY SPACE:

Suppose that Ω is a region in R n . Let F(Ω) denote the set of smooth functions on Ω. •, • and • are the Euclidean inner product and norm in R n . ∇, ∇• and ∆ represent the gradient, divergence and Laplacian operators in R n . 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 ) = σ 1 G(ρ)σ 2 dx = Φ 1 G(ρ) -1 Φ 2 dx, 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

G F (ρ) -1 Φ = ρ Φ -Φρdx , Φ ∈ T * ρ P(Ω).

The Fisher-Rao metric on the tangent space is given by

g F ρ (σ 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

G W (ρ) -1 Φ = -∇ • (ρ∇Φ), Φ ∈ T * ρ P(Ω).

The Wasserstein metric on the tangent space is given by

g W ρ (σ 1 , σ 2 ) = ρ ∇Φ 1 , ∇Φ 2 dx, σ 1 , σ 2 ∈ T ρ P(Ω),

where Φ i is the solution to σ i = -∇ • (ρ∇Φ i ), i = 1, 2.

2.1. GRADIENT FLOWS:

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(ρ) = D KL (ρ ρ * ) = 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 .

Here δE(ρt) δρt is the L 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 ρ t dx .

Example 4 (Wasserstein gradient flow) The Wasserstein gradient flow writes

∂ t ρ t = ∇ • ρ t ∇ δE δρ t .

2.2. HAMILTONIAN FLOW:

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

EQUATION

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.,

EQUATION

where Φ t is up to a spatially-constant function shrift. Here (2) is the Hamiltonian flow

EQUATION

with respect to the Hamiltonian

H(ρ t , Φ t ) = 1 2 Φ t G(ρ t ) -1 Φ t dx + 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 ρ t dx = 0, ∂ t Φ t + 1 2 Φ 2 t - ρ t Φ t dx Φ t + δE δρ t = 0.

The corresponding Hamiltonian is

H F (ρ t , Φ t ) = 1 2 Φ 2 t ρ t dx - ρ t Φ t dx 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

H W (ρ t , Φ t ) = 1 2 ∇Φ t 2 ρ t dx + E(ρ t )

. This is identical to the Wasserstein Hamiltonian flow introduced by Chow et al. (2019) .

3. ACCELERATED INFORMATION GRADIENT FLOW:

Let α t ≥ 0 be a scalar function of t. We add a damping term α t Φ t to the Hamiltonian flow (3).

EQUATION

This renders the Accelerated Information Gradient (AIG) flow

   ∂ t ρ t -G(ρ t ) -1 Φ t = 0, ∂ t Φ t + α t Φ t + 1 2 δ δρ t Φ t G(ρ t ) -1 Φ t dx + δ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 ρ (Hess E(ρ)σ, σ) ≥ β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

ẋ ṗ + 0 α t p - 0 I -I 0 ∇ x H E (x, p) ∇ p H E (x, p) = 0, H E (x, p) = 1 2 p 2 + f (x).

with initial values x(0) = x 0 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 ρ t dx ρ t = 0, ∂ t Φ t + α t Φ t + 1 2 Φ 2 t - ρ t Φ t dx Φ 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 X t ∼ ρ t and V t = ∇Φ t (X t ) 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

dX t = V t dt, dV t = -α t V t dt -∇ δE δρ t (X t )dt. (W-AIG-P)

If E(ρ) evaluates the KL divergence and ρ * ∝ exp(-f (x)), (W-AIG-P) is equivalent to

dX t = V t dt, dV t = -α t V t dt -∇f (X t )dt -∇ log ρ t (X t )dt. (W-AIG-P-KL)

Remark 2 The ∇ log ρ t (X t )dt term cannot be simply replaced by a Brownian motion dB t because ρ t is the marginal distribution on X t . 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.

3.1. WASSERSTEIN METRIC RESTRICTED TO GAUSSIAN:

In this subsection, we demonstrate that (W-AIG) has an ODE formulation in Gaussian family. Denote N 0 n to the multivariate Gaussian densities with zero means. Namely, if ρ 0 , ρ * ∈ N 0 n , then we show that (W-AIG) has a solution (ρ t , Φ t ) and ρ t ∈ N 0 n . Let P n and S n represent symmetric positive definite matrix and symmetric matrix with size n × n respectively. Each ρ ∈ N 0 n can be uniquely expressed by its covariance matrix Σ ∈ P n by ρ(x; Σ) =

(2π) -n/2 √ det(Σ) exp -1 2 x T Σ -1

x . The Wasserstein metric on P(R n ) induces a metric on N 0 n , which is a totally-geodesic submanifold in P(R n ), see (Takatsu, 2008; Modin, 2016; Malagò et al., 2018) . So there exists a Wasserstein metric on P n , also known as the Bures metric. For Σ ∈ P n , the tangent space and cotangent space follow

T Σ P n T * Σ P n S n . Definition 3 (Wasserstein metric in Gaussian) For Σ ∈ P n , the metric tensor G(Σ) : S n → S n is defined by G(Σ) -1 S = 2(ΣS + SΣ). The Wasserstein metric on S n is g Σ (A 1 , A 2 ) = tr(A 1 G(Σ)A 2 ) = tr(S 1 ΣS 2 ), where S i ∈ S n is the solution to A i = ΣS i + S i Σ, i = 1, 2.

Proposition 3 Given an energy function E(Σ), the Wasserstein gradient flow in Gaussian writes

Σt = -2(Σ t ∇ Σt E(Σ t ) + ∇ Σt E(Σ t )Σ t ).

Here ∇ Σt is the standard matrix derivative. The Hamiltonian flow is a system of (Σ t , S t ), i.e.,

Σt -(S t Σ t + Σ t S t ) = 0, Ṡt + S 2 t + 2∇ Σt E(Σ t ) = 0.

(5)

The corresponding Hamiltonian writes

H(Σ t , S t ) = tr(S t Σ t S t ) + 2E(Σ t ).

Therefore, by adding the damping term α t S t , we obtain the Wasserstein AIG flow in Gaussian.

Σt -

(S t Σ t + Σ t S t ) = 0, Ṡt + α t S t + S 2 t + 2∇ Σt E(Σ t ) = 0, (W-AIG-G)

with initial values Σ 0 = Σ 0 and S 0 = 0. For now, we consider E(Σ) to be the KL divergence.

EQUATION

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 P n and P(R n ).

Theorem 1 Suppose that ρ 0 , ρ * ∈ N n 0 and their covariance matrices are Σ 0 and Σ * . E(Σ) defined in (6) evaluates the KL divergence from ρ to ρ * . Let (Σ t , S t ) be the solution to (W-AIG-G) with initial values Σ 0 = Σ 0 and S 0 = 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 x T Σ -1 t x , Φ t (x) = 1 2 x T S t x + C(t),

where

C(t) = -t + 1 2 t 0 log det(Σ s (Σ * ) -1 )ds. Then, (ρ t , Φ t )

is the solution to (W-AIG) with initial values ρ 0 = ρ 0 and Φ 0 = 0.

4. CONVERGENCE RATE ANALYSIS ON W-AIG FLOWS:

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 ) ≤ C 0 e - √ β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 0 t -2 = O(t -2 ).

Here the constants C 0 , 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 & É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 T t 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)

= (sT t + (1 - s) Id)#ρ t , s ∈ [0, 1].

Here Id is the identity mapping from R n to itself. Then, γ(0) corresponds to a tangent vector

-∇ • (ρ t (x)(T t (x) -x)) ∈ T ρt P(Ω).

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.

EQUATION

)

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 Ė(t) ≤ 0. As a result,

E(ρ t ) ≤ e - √ βt E(t) ≤ e - √ βt E(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.

EQUATION

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 Ė(t) ≤ 0. As a result,

E(ρ t ) ≤ 4 t 2 E(t) ≤ 4 t 2 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 X t + e -γt Y t -T ρ∞ ρt (X t ) • d dt T ρ∞ ρt (X t ) = 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.

5. DISCRETE-TIME ALGORITHM FOR W-AIG FLOWS:

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 {X i 0 } N i=1 are given and V i 0 = 0. The time parameter t is related to the step size √ τ via t = √ τ k. The update rule follows

EQUATION

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 X i k follows a Gaussian distribution, then

EQUATION

where m k and Σ k are the mean and the covariance matrix of {X i k } N i=1 . For the non-Gaussian case, we use the kernel density estimation (KDE, Singh (1977) 

EQUATION

A common choice of K(x, y) is a Gaussian kernel with the bandwidth h, K G (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 Hamilton-Jacobi 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.

5.1. LEARN THE BANDWIDTH VIA BROWNIAN MOTION:

SVGD uses a median (MED) method to choose the bandwidth, i.e.,

EQUATION

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, {X i } N i=1 and a step size s, we can compute two particle systems

Y i k (h) = X i k -sξ k (x; h), Z i k = X i k + √ 2sB i , i = 1, . . . N

, where B i 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 i k (h)) and ρZ (x) = N i=1 δ(x -Z i k ).

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 h k-1 from the last iteration as the initialization. For simplicity we denote BM(h k-1 , {X i k } N i=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.

5.2. ADAPTIVE RESTART:

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

EQUATION

which can be viewed as discrete-time approximation of

-g W ρt (∂ t ρ t , G W (ρ t ) -1 δE δρt ) = -∂ t E(ρ t ). If ϕ k < 0, then we restart the algorithm with initial values X i 0 = X i k and V i 0 = 0.

This essentially keeps ∂ t E(ρ 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. Compute h l based on BM:

h l = BM(h l-1 , {X i k } N i=1 , √ τ ).

4:

Calculate ξ k (X i k ) by (10) or by ( 11) with the bandwidth h l .

5::

Set α k based on whether E(ρ) is Hess(β) or Hess(0). For i = 1, 2, . . . N , update

V i k+1 = α k V i k - √ τ (∇f (X i k ) + ξ k (X i k )), X i k+1 = X i k + √ τ V i k+1 . 6: if RESTART then 7: Compute ϕ k = - N i=1 V i k+1 , ∇f (X i k ) + ξ k (X i k ) . 8: If ϕ k < 0, set X i 0 = 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

6. NUMERICAL EXPERIMENTS:

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.

6.1. TOY EXAMPLE:

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:

X i k+1 = X i k -τ ∇f (X i k ) + √ 2τ B i k , X i k+1 = X i k -τ (∇f (X i k ) + ξ k (X i k )).

Here B i k ∼ 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. 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. The setting of target densities is identical to the ones in Figure 2 .

6.3. BAYESIAN LOGISTIC REGRESSION:

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. 

7. CONCLUSION:

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.

A THE PROOFS AND DERIVATIONS IN SECTION 2:

In this section, we provide the derivation of the Euler-Lagrange equation, proofs in Section 2 and the Euler-Lagrange formulation of (AIG).

A.1 DERIVATION OF THE EULER-LAGRANGE EQUATION:

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 h t ∈ F(Ω) be the smooth perturbation function that satisfies h t dx = 0, t ∈ [0, T ] and

h 0 = h T ≡ 0. Denote ρ t = ρ t + h t . Note that I(ρ t ) = T 0 L(ρ t , ∂ t ρ t )dt + T 0 δL δρ t h t + δL δ(∂ t ρ t ) ∂ t h t dxdt + o( ).

From

dI(ρ t ) d =0 = 0, it follows that T 0 δL δρ t h t + δL δ(∂ t ρ t ) ∂ t h t dxdt = 0.

Note that h 0 = h T ≡ 0. Perform integration by parts w.r.t. t yields

T 0 δL δρ t -∂ t δL δ(∂ t ρ t ) h t dxdt = 0.

Because h t dx = 0, (1) holds with a spatially constant function C(t).

A.2 THE PROOF OF PROPOSITION 1 IN SECTION 2:

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

EQUATION

)

Here Φ 1 = G(ρ) -1 σ 1 and Φ 2 = G(ρ) -1 σ 2 . We then check that

EQUATION

)

Let ρt = ρ t + h, where h ∈ T ρt P(Ω). For all σ ∈ T ρt P, it follows

G(ρ t + h) -1 G(ρ t + h)σ = σ.

The first-order derivative w.r.t. of the left hand side shall be 0, i.e.,

EQUATION

Under review as a conference paper at ICLR 2020

Based on basic calculations, we can compute that

EQUATION

EQUATION

Combining ( 16), ( 17) and ( 18) yields (15). Hence, the Euler-Lagrange equation ( 1) is equivalent to

∂ t Φ t = 1 2 δ δρ t ∂ t ρ t G(ρ t )∂ t ρ t dx - δE δρ t = - 1 2 δ δρ t Φ t G(ρ t ) -1 Φ t dx - δ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) .

A.3 THE EULER-LAGRANGIAN FORMULATION OF AIG FLOWS:

We can formulate (AIG) as a second-order equation of ρ t ,

D 2 Dt 2 ρ t + α t ∂ t ρ t + G(ρ t ) -1 δE δρ t = 0.

Here D 2 /Dt 2 is the covariant derivative in metric G. We can also explicitly write D 2 Dt 2 ρ t as follows.

D 2 Dt 2 ρ t = ∂ tt ρ t -(∂ t G(ρ t ) -1 )∂ t ρ t + 1 2 G(ρ t ) -1 δ δρ t ∂ t ρ t G(ρ t )∂ t ρ t dx .

B THE PROOFS IN SECTION 3:

In this section, we present proofs of propositions and theorems in Section 3.

B.1 THE PROOF OF PROPOSITION 2 IN SECTION 3:

We start with an identity. For a twice differentiable Φ(x), we have

EQUATION

From (W-AIG), it follows that

EQUATION

This is the continuity equation of ρ t . Hence, on the particle level, X t shall follows

dX t = ∇Φ t (X t )dt. Let V t = ∇Φ t (X t ).

Then, by the conservation of momentum and (W-AIG), we have

dV t = (∂ t + ∇Φ t (X t ) • ∇)∇Φ t (X t )dt = -α t ∇Φ t (X t ) - 1 2 ∇ ∇Φ 2 -∇ δE δρ t dt + (∇Φ • ∇)∇Φdt = -α t ∇Φ t (X t )dt -∇ δE δρ t (X t )dt = -α t V t dt -∇ δE δρ t (X t )dt.

B.2 THE PROOF OF PROPOSITION 3 IN SUBSECTION 3.1:

In this subsection, we derive the Hamiltonian flow in Gaussian. For A ∈ S n , we define the linear operator

M A : S n → S n by M A B = AB + BA, B ∈ S n .

It is easy to verify that if A ∈ P n , then M -1 A is well-defined. For a flow Σ t ∈ P n , t ≥ 0, we define the Lagrangian L(Σ t , Σt ) = 1 2 g Σt ( Σt , Σt ) -E(Σ t ). The corresponding Euler-Lagrange equation writes d dt

EQUATION

Let S t = M -1 Σt Σt , i.e., Σt = S t Σ t + Σ t S t . Then, it follows

g Σt ( Σt , Σt ) = tr(S t Σ t S t ) = 1 2 tr((S t Σ t + Σ t S t )S t ) = 1 2 tr( Σt S t ) = 1 2 tr( Σt M -1 Σt Σt ).

This leads to dL

d Σt = M -1 Σt Σt = S t .

For simplicity, we denote g = g Σt ( Σt , Σt ). First, we show that

dg dΣ t = -S 2 t .

Because S t = M -1 Σt Σt . Given Σt , S t can be viewed as a continuous function of Σ t . For any

A ∈ S n , define l A = tr((Σ t S t + S t Σ t )A). 0 = dl A dΣ t = ∂S t ∂Σ t ∂L A ∂S t + ∂l A ∂Σ t = ∂S t ∂Σ t (AΣ t + Σ t A) + (AS t + S t A).

Here we view ∂S T /∂Σ t as a linear operator on

S n . Let B = AΣ t + Σ t A, then A = M -1 Σt B. ∂St ∂Σt B + M St M -1 Σt B = 0 holds for all B ∈ S n . Therefore, we have ∂St ∂Σt = -M St M -1 Σt . Hence, dg dΣ t = ∂S t ∂Σ t ∂g ∂S t + ∂g ∂Σ t = -M St M -1 Σt (S t Σ t + Σ t S t ) + S 2 t = -M St S t + S 2 t = -S 2 t .

As a result, the Euler-Lagrange equation ( 21) is equivalent to

EQUATION

Combining ( 22) with Σt = S t Σ t + Σ t S t renders the Hamiltonian flow (5).

B.3 THE PROOF OF THEOREM 1 IN SUBSECTION 3.1:

We first show that Σ t stays in P n . Suppose that Σ t ∈ P n for 0 ≤ t ≤ T . Define H t = H(Σ t , S t ) = tr(S t Σ t S t )/2 + E(Σ t ). We observe that (W-AIG-G) is equivalent to

EQUATION

We show that H t is decreasing with respect to t.

dH t dt = tr ∂H t ∂S t Ṡt + ∂H t ∂Σ t Σt = tr ∂H t ∂S t -α t S t -2 ∂H t ∂Σ t + 2 ∂H t ∂Σ t ∂H t ∂S t = -α t tr S t ∂H t ∂S t = - α t 2 tr(S t (Σ t S t + S t Σ t )) = -α t tr(S t Σ t S t ) ≤ 0.

For simplicity, we denote W * = (Σ * ) -1 . Let λ t be the smallest eigenvalue of Σ t . Then, log det(Σ t W * ) ≥ log det(Σ t ) = n log λ t . Therefore,

EQUATION

)

This means that as long as Σ t ∈ P n , the smallest eigenvalue of Σ t has a positive lower bounded.

If there exists T > 0 such that Σ T / ∈ P n . Because Σ t is continuous with respect to t, there exists T 1 < T , such that Σ t ∈ P n , 0 ≤ t ≤ T 1 and λ T1 < exp (-2H(0)/n -1), which violates (24).

We then reveal the relationship between (W-AIG) in P(R n ) and P n . We can compute that

∂ ∂t det(Σ t ) = det(Σ t ) tr(Σ -1 t Σt ), ∂ ∂t Σ -1 t = -Σ -1 t Σt Σ -1 t .

Combining with Σt = Σ t S t + S t Σ t , we obtain

tr(Σ -1 t Σt ) = tr(S t + Σ -1 t S t Σ t ) = 2 tr(S t ), tr(xΣ -1 t Σt Σ -1 t x) = tr(x T Σ -1 t S t x + x T S t Σ -1 t x) = 2 tr(S t Σ -1 t xx T ). Therefore, it follows ∂ t ρ t (x) = ∂ ∂t 1 det(Σ t ) det(Σ t )ρ t (x) + 1 2 tr(x T Σ -1 t Σt Σ -1 t x)ρ t (x) = - 1 2 tr(Σ -1 t Σt )ρ t (x) + tr(S t Σ -1 t xx T )ρ t (x) = -tr(S t (I -Σ -1 t xx T ))ρ t (x).

Note that ∇Φ t (x) = S t x. Hence,

-∇ • (ρ t ∇Φ t ) = - n i=1 ∂ i (ρ t (x)S t x) i = - n i=1 [ρ t (x)∂ i (S t x) i + (S t x) i ∂ i ρ t (x)] = -ρ t (x) tr(S t ) + (S t x) T (-Σ -1 t x) = -ρ t (x) tr(S t (I -Σ -1 t xx T )) = ∂ t ρ t (x).

The first equation of (W-AIG) holds. Because

∂ t Φ t (x) = x T Ṡt x/2 + Ċ(t), ∂ t Φ t (x) + α t Φ t (x) + 1 2 ∇Φ t (x) 2 = 1 2 x T Ṡt x + α t 2 x T S t x + 1 2 x T S 2 t x + Ċ(t) = -x T ∇ Σt E(Σ t )x + Ċ(t) = 1 2 x T (Σ -1 t -W * )x + Ċ(t).

Note that ρ * is the Gaussian density with the covariance matrix Σ * . Because Ċ(t) = 1 2 log det(Σ t W * ) -1, we can compute

δE δρ t = log ρ t (x) -log ρ * (x) + 1 = - 1 2 x T (Σ -1 t -W * )x - 1 2 log det(Σ t W * ) + 1 = - 1 2 x T (Σ -1 t -W * )x -Ċ(t) = -(∂ t Φ t (x) + α t Φ t (x) + 1 2 ∇Φ t (x) 2 ).

Therefore, the second equation of (W-AIG) holds. Because Σ 0 = Σ 0 , S 0 = 0 and C(0) = 0, we have ρ 0 = ρ 0 and Φ 0 = 0. This completes the proof.

C THE PROOFS IN SECTION 4:

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.

C.1 A BRIEF REVIEW ON THE GEOMETRIC PROPERTIES OF THE PROBABILITY SPACE:

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 ρ 0 P(Ω) to P(Ω). It requires that σ ∈ T ρ 0 P(Ω) is mapped to a point ρ 1 ∈ P(Ω) such that there exists a geodesic curve γ(s) satisfying γ(0) = ρ 0 , γ(0) = σ and γ(1) = ρ 1 .

Lemma 1 For u t satisfying (25), we have

∇Φ t -u t , ∇T t ∇Φ t ρ t dx ≥ 0, ∇Φ t -u t , ∇T t (x)(T t (x) -x) ρ t = 0.

PROOF We first notice that u t -∇Φ t is divergence-free in term of ρ t . From -∇T t u t = ∂ t T t = ∇∂ t Ψ t , we observe that -∇T t u t is the gradient of ∂ t Ψ t . Therefore,

∇Φ t -u t , ∇T t u t ρ t = - ∂ t Ψ t , ∇ • (ρ t (∇Φ t -u t )) = 0.

Based on our previous characterization on the optimal transport plan T t , ∇T t = ∇ 2 Ψ t is symmetric positive definite. This yields that

∇Φ t -u t , ∇T t ∇Φ t ρ t dx = ∇Φ t -u t , ∇T t ∇Φ t ρ t dx - ∇Φ t -u t , ∇T t u t ρ t = ∇Φ t -u t , ∇T t (∇Φ t -u t ) ρ t dx ≥ 0.

The last inequality utilizes that ∇T t is positie definite and ρ t is non-negative. Then, we prove the equality in Lemma 1. Because ∇T

t (x)(T t (x) -x) = 1 2 ∇( T t (x) -x 2 + T t (x) -x 2 ) is a gradient. Similarly, it follows ∇Φ t -u t , ∇T t (x)(T t (x) -x) ρ t = 0.

Lemma 1 and the relationship (26) gives

EQUATION

EQUATION

Proof of Proposition 5. Based on the definition of the Wasserstein metric, we have

∂ t E(ρ t ) = - δE δρ t ∇ • (ρ t ∇Φ t )dx.

Differentiating E(t) (7) w.r.t. t renders

EQUATION

For the part (29), Proposition 7 renders

EQUATION

We first compute the terms with the coefficient β 0 in Ė(t)e - √ βt . We observe that

EQUATION

where the last equality uses (W-AIG) with α t = 2 √ β. Substituting ( 30) and ( 31) into the expression of Ė(t)e - √ βt yields

EQUATION

Then, we deal with the terms with ∇ • (ρ t ∇Φ t ). We have the following two identities

EQUATION

EQUATION

Hence, we can proceed to compute the terms with the coefficient √ β. ( 27) and (33) yields

EQUATION

Substituting ( 34) and ( 35) into (32) gives

Ė(t)e - √ βt + √ β 2 ∇Φ t 2 ρ t dx ≤β ∂ t T t , T t (x) -x ρ t dx - β 2 T t (x) -x 2 ∇ • (ρ t ∇Φ t )dx -β T t (x) -x, ∇Φ t ρ t dx + 2β T t (x) -x, ∇Φ t ρ t dx =β ∂ t T t + ∇T t ∇Φ t , T t (x) -x ρ t dx = 0,

where the last equality uses (28). In summary, we have

Ė(t)e - √ βt ≤ - √ β 2 ∇Φ t 2 ρ t dx ≤ 0.

Proof of Proposition 6. Differentiating E(t) (8) w.r.t. t, we compute that

EQUATION

Because E(ρ) is Hess(0), Proposition 7 yields As a result, Ė(t) ≤ 0. This completes the proof.

E(ρ t ) = E(ρ t ) -E(ρ * ) ≤ - T t (

C.4 COMPARISON WITH THE PROOF IN THE ACCELERATED FLOW:

The accelerated flow in (Taghvaei & Mehta, 2019) Here the target distribution satisies ρ ∞ (x) = ρ * (x) ∝ exp(-f (x)). Suppose that we take α t = log p -log t, β t = p log t + log C 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 Y t by 2t -3 V t . The Lyapunov function in (Taghvaei & Mehta, 2019) follows

V (t) = 1 2 E[ X t + e -γt Y t -T ρ * ρt (X t ) 2 ] + e βt (E(ρ) -E(ρ * )) = 1 2 E[ X t + t 2 V t -T ρ * ρt (X t ) 2 ] + t 2 4 (E(ρ t ) -E(ρ * )) = 1 2 -(T t (x) -x) + t 2 ∇Φ t (x) 2 ρ t (x)dx + t 2 4 (E(ρ t ) -E(ρ * )).

The last equality is based on the fact that V t = ∇Φ t (X t ) and T t = 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 In 1-dimensional case, because ∇ • (ρ t (u t -∇Φ t )) = 0 indicates that ρ t (u t -∇Φ t ) = 0. For ρ t (x) > 0, we have u t (x) -∇Φ t (x) = 0. So the technical assumption holds. In general cases, although u t = ∂ t (T t ) -1 • T t satisfies ∇ • (ρ t (u t -∇Φ t )) = 0, but this does not necessary indicate that u t = ∇Φ t . Hence, E X t + e -γt Y t -T ρ∞ ρt (X t ) • d dt T ρ∞ ρt (X t ) = 0 does not necessary hold except for 1-dimensional case.

0 =E X t + e -γt Y t -T ρ * ρt (X t ) • d dt T ρ * ρt (X t ) =E X t + t 2 V t -T t (X t ) • d dt T t (X t ) =E X t + t 2 V t -T t (X t ) • ((∂ t T t )(X t ) + ∇T t V t ) = x -T t (

D IMPLEMENTATION DETAILS IN THE NUMERICAL EXPERIMENTS:

In this section, we elaborate on the implementation details in the numerical experiments.

D.1 DETAILS IN SUBSECTION 6.1:

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(x 1 -3) 2 ) + exp(-2(x 1 + 3) 2 )), x ∈ R 2 . All methods run for 200 iterations using the same fixed step size τ = 0.1.

:



Figures and Tables:

Particle implementation of Wasserstein AIG flow Require: initial positions {X i 0 } N i=1 , step size τ , number of iteration L. 1: Set k = 0, V i 0 = 0, i = 1, . . . N . Set the bandwidth h 0 by MED (12). 2: for l = 1, 2, . . . L do 3:

Figure 1: The effect of the BM method. Samples are plotted as blue dots. Left to right: MCMC, MED, HE and BM. All methods are run for 200 iterations with the same initialization.

Figure 2: The acceleration effect of W-AIG flow and the strength of adaptive restart (ODE level). The target density is a Gaussian distribution with zero mean on R 100 . Left: L = 1, κ ≈ 3.8 × 10 3 . Right: β = 1, κ ≈ 4.0 × 10 3 . 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 {X i k } N i=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.

Figure 3: The acceleration effect of W-AIG flow and the strength of adaptive restart (particle level). The setting of target densities is identical to the ones in Figure 2.

Figure 4: Comparison of different methods on Bayesian logistic regression, averaged over 10 independent trials. The shaded areas show the variance over 10 trials. Top: BM; Bottom: MED. Left: Test accuracy; Right: Test log-likelihood.

x) -x, ∇ δE δρ t ρ t dx. (37)Utilizing the inequality (37) and substituting the expressions of terms involving ∂ t T t and ∇•(ρ t ∇Φ t ) in (36) with the expressions in (27) (28) and (33) (34), we obtainĖ(t) ≤ -∇Φ t , ∇T t (x)(T t (x) -x) ρ t dx + T t (x) -x, ∇T t ∇Φ t ρ t dx -T t (x) -x, ∇Φ t ρ t dx + t 2 ∇Φ t , ∇T t ∇Φ t ρ t dx ∇Φ t , ∇T t ∇Φ t ρ t dxFrom (W-AIG) with α t = 3/t, we have the following equalities.

Table content: <table><tr><td/><td/><td/><td colspan="3">-T + 1 2 t 2 ∇Φ t</td><td colspan="5">2 ρ t dx +</td><td>t 4</td><td>∇Φ t</td><td>2 ρ t dx +</td><td>t 2 4</td><td>∇Φ t , ∂ t ∇Φ t ρ t dx</td><td>(38)</td></tr><tr><td/><td/><td/><td>+</td><td>t 2 8</td><td colspan="6">∇Φ t , ∇ ∇Φ t</td><td>2 ρ t dx +</td><td>t 2 4</td><td>∇Φ t , ∇</td><td>δE δρ t</td><td>ρ t dx</td></tr><tr><td/><td/><td/><td>-</td><td>t 2</td><td colspan="5">T t (x) -x, ∇</td><td>δE δρ t</td><td>ρ t dx.</td></tr><tr><td colspan="11">The expression of (38) can be reformulated into</td></tr><tr><td/><td/><td/><td colspan="2">Ė(t) ≤ -</td><td>3 2</td><td colspan="5">T t (x) -x, ∇Φ t ρ t dx +</td><td>3t 4</td><td>∇Φ t</td><td>2 ρ t dx</td></tr><tr><td/><td/><td/><td/><td>-</td><td>t 2</td><td colspan="5">T t (x) -x, ∂ t ∇Φ t +</td><td>1 2</td><td>∇ ∇Φ t</td><td>2 + ∇</td><td>δE δρ t</td><td>ρ t dx</td></tr><tr><td/><td/><td/><td/><td>+</td><td>t 2 4</td><td/><td colspan="4">∇Φ t , ∂ t ∇Φ t +</td><td>1 2</td><td>∇ ∇Φ t</td><td>2 + ∇</td><td>δE δρ t</td><td>ρ t dx.</td></tr><tr><td/><td/><td>t 2 4</td><td colspan="4">∇Φ t , ∂ t ∇Φ t +</td><td colspan="2">1 2</td><td colspan="2">∇ ∇Φ t</td><td>2 + ∇</td><td>δE δρ t</td><td>ρ t dx = -</td><td>3t 4</td><td>∇Φ t</td><td>2 ρ t dx,</td></tr><tr><td>-</td><td>t 2</td><td colspan="5">T t (x) -x, ∂ t ∇Φ t +</td><td>1 2</td><td colspan="3">∇ ∇Φ t</td><td>2 + ∇</td><td>δE δρ t</td><td>ρ t dx =</td><td>3 2</td><td>T</td></tr></table>

is given by dX t dt = e αt-γt Y t , dY t dt = -e αt+βt+γt ∇ δE δ ρt (X t ).

Table content: <table><tr><td>(39)</td></tr></table>

x) + t 2 ∇Φ t (x), ∂ t T t + ∇T t ∇Φ t ρ t dxBased on ∂ t T t = -∇T t u t and Lemma 1, we havex -T t (x), ∂ t T t + ∇T t ∇Φ t ρ t dx = x -T t (x), ∇T t (∇Φ t -u t ) ρ t dx = 0. ∇Φ t , ∂ t T t + ∇T t ∇Φ t ρ t dx = ∇Φ t , ∇T t (∇Φ t -u t ) ρ t dx = ∇Φ t -u t , ∇T t (∇Φ t -u t ) ρ t dx ≥ 0.As a result, we haveE X t + e -γt Y t -T ρ∞ ρt (X t ) •

Table content: <table><tr><td>d dt</td><td>T ρ∞ ρt (X t ) =</td><td>t 2</td><td>∇Φ</td></tr></table>

Peer review (iclr2020):
I acknowledge reading the rebuttal of the authors. Thank you for your clarifications and explanation. My point was this paper would make a good submission to ICLR if it was better motivated presented and explained to a wider audience. Unfortunately in its current form it can only reach a limited audience. #### Summary of the paper: The paper proposes accelerated information flows under Wasserstein or Fisher rao metric . i.e a method for solving minimization of probability functional , where the probability space is either endowed with the Wasserstein or the Fisher distance. Gradient descent methods in euclidian spaces can be accelerated using a type of momentum , this paper extends this to gradient flows, using similar formalism of Hamiltonian that appeals to the dynamic of a the particles and the velocity (or momentum, ). Writing down the lagrangian one obtains two co-evolving PDE one of the dynamic of the density and one for the potential . The PDEs are specified for both Fisher Rao, and Wasserstein distance. The hamiltonian for example for the Wasserstein distance . and PDEs amounts a continuity equation for the density evolving with drift , the evolution of the momentum is also given by a PDE. Proposition 2 of the paper gives the particles differential equation corresponding to the system of PDEs. For the energy being the KL divergence an explicit expression is given , this expression remains difficult in practice since it needs the knowledge of the density . Authors propose in the application section to use gaussian approximation , or using a kernel density estimators. The Bandwidth of the kernel is then choosen using a heuristic proposed in the paper. The paper then focuses on deriving expression for flows when the densities are centered gaussians, and this amounts to an ODE on the covariance , the ODE is discretized in Appendix D.2 to lead to computational method. Then convergence of the flow is analyzed for the wasserstein accelerated flows, under *_x0008_eta- convextiy* in the wasserstein sense of the functional. Some experiments of the particle based method are shown on synthetic experiments and in bayesian logistic regression. Review: Contribution/ Clarity: The main contribution of the paper is in deriving the accelerated gradient flow for the wasserstein distance this was also addressed in a recent paper [Taghevia and Mettha 2019]. The technical contribution is interesting but given that this field of flows in probability space is still not very well spread in the ML community, I wonder if ICLR is the best fit for this type of work. I support good theoretical work, but I think the authors could have done a better job in exposing the ideas how they extend form euclidean space, to manifolds, to probability spaces gradually. Simple derivations of euclidean space Hamiltonian will help the reader that is not exposed to such literature. I think the paper will benefit from a less technical writing in introducing the ideas coming from euclidean space and in conveying the intuitions. Comments: - In the proof of Proposition 2 you give the expression of evolution of by conservation of the momentum. Could you please elaborate more how you obtain this expression, and where you proved the conservation of momentum? - In term of damping if ones uses the Wasserstein Fisher Rao flows , one obtain also accelerartion , maybe you can comment on that ? since you analyze both flows , would be interesting to discuss the relation to Global convergence of neuron birth-death dynamics, that shows that an acceleration is obtained via WFR flows, since it will introduce a damping as well. - since MCMC and BM method lead to similar result what is the advantage of the wasserstein accelerated flow? one could also implement also an accelerated langevin dynamic