Title: ACCELERATED GRADIENT FLOW FOR PROBABILITY DISTRIBUTIONS

Abstract: This paper presents a methodology and numerical algorithms for constructing accelerated gradient flows on the space of probability distributions. In particular, we extend the recent variational formulation of accelerated gradient methods in Wibisono et al. (2016) from vector valued variables to probability distributions. The variational problem is modeled as a mean-field optimal control problem. The maximum principle of optimal control theory is used to derive Hamilton's equations for the optimal gradient flow. The Hamilton's equation are shown to achieve the accelerated form of density transport from any initial probability distribution to a target probability distribution. A quantitative estimate on the asymptotic convergence rate is provided based on a Lyapunov function construction, when the objective functional is displacement convex. Two numerical approximations are presented to implement the Hamilton's equations as a system of N interacting particles. The continuous limit of the Nesterov's algorithm is shown to be a special case with N = 1. The algorithm is illustrated with numerical examples.
Under review as a conference paper at ICLR 2019 The extension proposed in our work is based upon a generalization of the formula for the Lagrangian in Wibisono et al. ( 2016): (i) the kinetic energy term is replaced with the expected value of kinetic energy; and (ii) the potential energy term is replaced with a suitably defined functional on the space of probability distributions. The variational problem is to obtain a trajectory in the space of probability distributions that minimizes the action integral of the Lagrangian.
The variational problem is modeled as a mean-field optimal problem. The maximum principle of the optimal control theory is used to derive the Hamilton's equations which represent the first order optimality conditions. The Hamilton's equations provide a generalization of the Nesterov's ODE to the space of probability distributions. A candidate Lyapunov function is proposed for the convergence analysis of the solution of the Hamilton's equations. In this way, quantitative estimates on convergence rate are obtained for the case when the objective functional is displacement convex (McCann, 1997) . Table 1 provides a summary of the relationship between the original variational formulation in Wibisono et al. ( 2016) and the extension proposed in this paper.
We also consider the important special case when the objective functional is the relative entropy functional D(ρ|ρ ∞ ) defined with respect to a target probability distribution ρ ∞ . In this case, the accelerated gradient flow is shown to be related to the continuous limit of the Hamiltonian Monte-Carlo algorithm (Cheng et al., 2017) (Remark 1). The Hamilton's equations are finite-dimensional for the special case when the initial and the target probability distributions are both Gaussian. In this case, the mean evolves according to the Nesterov's ODE. For the general case, the Lyapunov function-based convergence analysis applies when the target distribution is log-concave.
As a final contribution, the proposed methodology is used to obtain a numerical algorithm. The algorithm is an interacting particle system that empirically approximates the distribution with a finite but large number of N particles. The difficult part of this construction is the approximation of the interaction term between particles. For this purpose, two types of approximations are described: (i) Gaussian approximation which is asymptotically (as N → ∞) exact in Gaussian settings; and (ii) Diffusion map approximation which is computationally more demanding but asymptotically exact for a general class of distributions.
The outline of the remainder of this paper is as follows: Sec. 2 provides a brief review of the variational formulation in Wibisono et al. (2016) . The proposed extension to the space of probability distribution appears in Sec. 3 where the main result is also described. The numerical algorithm along with the results of numerical experiments appear in Sec. 4. Comparisons with MCMC and Hamiltonian MCMC are also described. The conclusions appear in Sec. 5.
The gradient and divergence operators are denoted as ∇ and ∇• respectively. With multiple variables, ∇ z denotes the gradient with respect to the variable z. Therefore, the divergence of the vector field U is ∇ • U (x) = d n=1 ∇ xn U n (x). The space of absolutely continuous probability measures on R d with finite second moments is denoted by P ac,2 (R d ). The Wasserstein gradient and the Gâteaux derivative of a functional F is denoted as ∇ ρ F(ρ) and ∂F ∂ρ (ρ) respectively (see Appendix C for definition). The probability distribution of a random variable Z is denoted as Law(Z).

1. INTRODUCTION:

Optimization on the space of probability distributions is important to a number of machine learning models including variational inference (Blei et al., 2017) , generative models (Goodfellow et al., 2014; Arjovsky et al., 2017) , and policy optimization in reinforcement learning (Sutton et al., 2000) . A number of recent studies have considered solution approaches to these problems based upon a construction of gradient flow on the space of probability distributions (Zhang et al., 2018; Liu & Wang, 2016; Frogner & Poggio, 2018; Chizat & Bach, 2018; Richemond & Maginnis, 2017; Chen et al., 2018) . Such constructions are useful for convergence analysis as well as development of numerical algorithms.

In this paper, we propose a methodology and numerical algorithms that achieve accelerated gradient flows on the space of probability distributions. The proposed numerical algorithms are related to yet distinct from the accelerated stochastic gradient descent (Jain et al., 2017) and Hamiltonian Markov chain Monte-Carlo (MCMC) algorithms (Neal et al., 2011; Cheng et al., 2017) . The proposed methodology extends the variational formulation of (Wibisono et al., 2016) from vector valued variables to probability distributions. The original formulation of Wibisono et al. (2016) was used to derive and analyze the convergence properties of a large class of accelerated optimization algorithms, most significant of which is the continuous-time limit of the Nesterov's algorithm (Su et al., 2014) . In this paper, the limit is referred to as the Nesterov's ordinary differential equation (ODE).

Vector Probability distribution:

State-space

R d P 2 (R d ) Objective function f (x) F(ρ) := D(ρ|ρ ∞ )

Lagrangian e αt+γt 1 2 |e -αt u| 2 -e βt f (x) e αt+γt E 1 2 |e -αt U | 2 -e βt log( ρ(X) ρ∞(X) )

Lyapunov funct.

1 2 |x + e -γt y -x| 2 1 2 E[|X t + e -γt Y t -T ρ∞ ρt (X t )| 2 ] +e βt (f (x) -f (x)) +e βt (F(ρ t ) -F(ρ ∞ ))

Table 1 : Summary of the variational formulations for vectors and probability distributions.

2. REVIEW OF THE VARIATIONAL FORMULATION OF WIBISONO ET AL. (2016):

The basic problem is to minimize a C 1 smooth convex function f on R d . The standard form of the gradient descent algorithm for this problem is an ODE:

dX t dt = -∇f (X t ), t ≥ 0 (1)

Accelerated forms of this algorithm are obtained based on a variational formulation due to Wibisono et al. (2016) . The formulation is briefly reviewed here using an optimal control formalism. The Lagrangian L :

R + × R d × R d → R is defined as L(t, x, u) := e αt+γt 1 2 |e -αt u| 2 kinetic energy -e βt f (x) potential energy (2)

where t ≥ 0 is the time, x ∈ R d is the state, u ∈ R d is the velocity or control input, and the time-varying parameters α t , β t , γ t satisfy the following scaling conditions: α t = log p -log t, β t = p log t + log C, and γ t = p log t where p ≥ 2 and C > 0 are constants.

The variational problem is

EQUATION

The Hamiltonian function is

EQUATION

where y ∈ R d is dual variable and y • u denotes the dot product between vectors y and u.

According to the Pontryagin's Maximum Principle, the optimal control u *

t = arg max v H(t, X t , Y t , v) = e αt-γt Y t .

The resulting Hamilton's equations are

dX t dt = +∇ y H(t, X t , Y t , u t ) = e αt-γt Y t , X 0 = x 0 (5a) dY t dt = -∇ x H(t, X t , Y t , u t ) = -e αt+βt+γt ∇f (X t ), Y 0 = y 0 (5b)

The system (5) is an example of accelerated gradient descent algorithm. Specifically, if the parameters α t , β t , γ t are defined using p = 2, one obtains the continuous-time limit of the Nesterov's accelerated algorithm. It is referred to as the Nesterov's ODE in this paper.

For this system, a Lyapunov function is as follows:

EQUATION

where x ∈ arg min x f (x). It is shown in Wibisono et al. (2016) that upon differentiating along the solution trajectory, d dt V (t, X t , Y t ) ≤ 0. This yields the following convergence rate:

f (X t ) -f (x) ≤ O(e -βt ), ∀t ≥ 0 (7)

3 VARIATIONAL FORMULATION FOR PROBABILITY DISTRIBUTIONS

3.1. MOTIVATION AND BACKGROUND:

Let F : P ac,2 (R d ) → R be a functional on the space of probability distributions. Consider the problem of minimizing F(ρ). The (Wasserstein) gradient flow with respect to

EQUATION

where ∇ ρ F(ρ) is the Wasserstein gradient of F.

An important example is the relative entropy functional where

F(ρ) = D(ρ|ρ ∞ ) := R d log( ρ(x) ρ∞(x) )ρ(x) dx where ρ ∞ ∈ P ac,2 (R d

) is referred to as the target distribution. The gradient of relative entropy is given by

∇ ρ F(ρ) = ∇ log( ρ ρ∞ ). The gradient flow ∂ρ t ∂t = -∇ • (ρ t ∇ log(ρ ∞ )) + ∆ρ t (9)

is the Fokker-Planck equation (Jordan et al., 1998) . The gradient flow achieves the density transport from an initial probability distribution ρ 0 to the target (here, also equilibrium) probability distribution ρ ∞ ; and underlies the construction and the analysis of Markov chain Monte-Carlo (MCMC) algorithms. The simplest MCMC algorithm is the Langevin stochastic differential equation (SDE):

dX t = -∇f (X t ) dt + √ 2 dB t , X 0 ∼ ρ 0

where B t is the standard Brownian motion in R d .

The main problem of this paper is to construct an accelerated form of the gradient flow (8). The proposed solution is based upon a variational formulation. As tabulated in Table 1 , the solution represents a generalization of Wibisono et al. (2016) from its original deterministic finite-dimensional to now probabilistic infinite-dimensional settings.

The variational problem can be expressed in two equivalent forms: (i) The probabilistic form is described next in the main body of the paper; and (ii) The partial differential equation (PDE) form appears in the Appendix. The probabilistic form is stressed here because it represents a direct generalization of the Nesterov's ODE and because it is closer to the numerical algorithm.

3.2. PROBABILISTIC FORM OF THE VARIATIONAL PROBLEM:

Consider the stochastic process {X t } t≥0 that takes values in R d and evolves according to:

dX t dt = U t , X 0 ∼ ρ 0

where the control input {U t } t≥0 also takes values in R d , and ρ 0 ∈ P ac,2 (R d ) is the probability distribution of the initial condition X 0 . It is noted that the randomness here comes only from the random initial condition.

Suppose the objective functional is of the form F(ρ) = F (ρ, x)ρ(x) dx. The Lagrangian L :

R + × R d × P ac,2 (R d ) × R d → R is defined as L(t, x, ρ, u) := e αt+γt 1 2 |e -αt u| 2 kinetic energy -e βt F (ρ, x)

potential energy (10) This formula is a natural generalization of the Lagrangian (2) and the parameters α t , β t , γ t are defined exactly the same as in the finite-dimensional case. The stochastic optimal control problem is:

EQUATION

where ρ t = Law(X t ) ∈ P ac,2 (R d ) is the probability density function of the random variable X t .

The Hamiltonian function

H : R + × R d × P ac,2 (R d ) × R d × R d → R

for this problem is given by (Carmona & Delarue, 2017, Sec. 6.2.3 ):

EQUATION

where y ∈ R d is the dual variable.

3.3. MAIN RESULT:

Theorem 1. Consider the variational problem (11).

(i) The optimal control U * t = e αt-γt Y t where the optimal trajectory {(X t , Y t )} t≥0 evolves according to the Hamilton's equations:

dX t dt = U * t = e αt-γt Y t , X 0 ∼ ρ 0 (13a) dY t dt = -e αt+βt+γt ∇ ρ F(ρ t )(X t ), Y 0 = ∇φ 0 (X 0 ) (13b)

where φ 0 is any convex function and ρ t := Law(X t ). (ii) Suppose also that the functional F is displacement convex and ρ ∞ is its minimizer. Define the energy along the optimal trajectory

EQUATION

where the map T ρ∞ ρt : R d → R d is the optimal transport map from ρ t to ρ ∞ . Suppose also that the following technical assumption holds:

E[(X t + e -γt Y t -T ρ∞ ρt (X t )) • d dt T ρ∞ ρt (X t )] = 0. Then dV dt (t) ≤ 0.

Consequently, the following rate of convergence is obtained along the optimal trajectory:

F(ρ t ) -F(ρ ∞ ) ≤ O(e -βt ), ∀t ≥ 0 (15)

Proof sketch. The Hamilton's equations are derived using the standard mean-field optimal control theory Carmona & Delarue (2017) . The Lyapunov function argument is based upon the variational inequality characterization of a displacement convex function (Ambrosio et al., 2008, Eq. 10.1.7) .

The detailed proof appears in the Appendix. We expect that the technical assumption is not necessary. This is the subject of the continuing work.

3.4. RELATIVE ENTROPY AS THE FUNCTIONAL:

In the remainder of this paper, we assume that the functional F(ρ) = D(ρ|ρ ∞ ) is the relative entropy where ρ ∞ ∈ P ac,2 (R d ) is a given target probability distribution. In this case the Hamilton's equations are given by

EQUATION

where ρ t = Law(X t ) and f = -log(ρ ∞ ). Moreover, if f is convex (or equivalently ρ ∞ is logconcave), then F is displacement convex with the unique minimizer at ρ ∞ and the convergence estimate is given by D(ρ t |ρ ∞ ) ≤ O(e -βt ).

Remark 1. The Hamilton's equations ( 16) with the relative entropy functional is related to the under-damped Langevin equation (Cheng et al., 2017) . The difference is that the deterministic term ∇ log(ρ t ) in ( 16) is replaced with a random Brownian motion term in the under-damped Langevin equation. More detailed comparison appears in the Appendix D.

3.5. QUADRATIC GAUSSIAN CASE:

Suppose the initial distribution ρ 0 and the target distribution ρ ∞ are both Gaussian, denoted as N (m 0 , Σ 0 ) and N (x, Q), respectively. This is equivalent to the objective function f (x) being quadratic of the form f

(x) = 1 2 (x -x) Q -1 (x -x).

Therefore, this problem is referred to as the quadratic Gaussian case. The following Proposition shows that the mean of the stochastic process (X t , Y t ) evolves according to the Nesterov ODE (5):

Proposition 1. (Quadratic Gaussian case) Consider the variational problem (11) for the quadratic Gaussian case. Then (i) The stochastic process (X t , Y t ) is a Gaussian process. The Hamilton's equations are given by:

dX t dt = e αt-γt Y t , dY t dt = -e αt+βt+γt (Q -1 (X t -x) -Σ -1 t (X t -m t ))

where m t and Σ t are the mean and the covariance of X t .

(ii) Upon taking the expectation of both sides, and denoting

n t := E[Y t ] dm t dt = e αt-γt n t , dn t dt = -e αt+βt+γt Q -1 (m t -x) ∇f (mt)

which is identical to Nesterov ODE (5).

4. NUMERICAL ALGORITHM:

The proposed numerical algorithm is based upon an interacting particle implementation of the Hamilton's equation ( 16). Consider a system of N particles {(X i t , Y i t )} N i=1 that evolve according to:

dX i t dt = e αt-γt Y i t , X i 0 ∼ ρ 0 dY i t dt = -e αt+βt+γt (∇f (X i t ) + I (N ) t (X i t ) interaction term ), Y i 0 = ∇φ 0 (X i 0 )

The interaction term I (N ) t is an empirical approximation of the ∇ log(ρ t ) term in ( 16). We propose two types of empirical approximations as follows:

1. Gaussian approximation: Suppose the density is approximated as a Gaussian N (m t , Σ t ). In this case, ∇ log(ρ t (x)) = -Σ t -1 (x -m t ). This motivates the following empirical approximation of the interaction term:

EQUATION

Algorithm 1 Interacting particle implementation of the accelerated gradient flow 18) or ( 19)

Input: ρ 0 , φ 0 , N , t 0 , ∆t, p, C, K Output: {X i k } N,K i=1,k=0 Initialize {X i 0 } N i=1 i.i.d ∼ ρ 0 , Y i 0 = ∇φ 0 (X i 0 ) Compute I (N ) 0 (X i 0 ) with (

for k = 0 to K -1 do t k+ 1 2 = t k + 1 2 ∆t Y i k+ 1 2 = Y i k -1 2 Cpt 2p-1 k+ 1 2 (∇f (X i k ) + I (N ) k (X i k ))∆t X i k+1 = X i k + p t p+1 k+ 1 2 Y i k ∆t Compute I (N )

k+1 (X i k+1 ) with ( 18)or ( 19)

Y i k+1 = Y i k+ 1 2 -1 2 Cpt 2p-1 k+ 1 2 (∇f (X i k+1 ) + I (N ) k+1 (X i k+1 ))∆t t k+1 = t k+ 1 2 + 1 2 ∆t end for where m (N ) t := N -1 N i=1 X i t is the empirical mean and Σ (N ) t := 1 N -1 N i=1 (X i t -m (N ) t )(X i t - m (N ) t

) is the empirical covariance.

Even though the approximation is asymptotically (as N → ∞) exact only under the Gaussian assumption, it may be used in a more general settings, particularly when the density ρ t is unimodal. The situation is analogous to the (Bayesian) filtering problem, where an ensemble Kalman filter is used as an approximate solution for non-Gaussian distributions (Evensen, 2003) .

2.. Diffusion map approximation::

This is based upon the diffusion map approximation of the weighted Laplacian operator (Coifman & Lafon, 2006; Hein et al., 2007) . For a C 2 function f , the weighted Laplacian is defined as ∆ ρ f := 1 ρ ∇ • (ρ∇f ). Denote e(x) = x as the coordinate function on R d . It is a straightforward calculation to show that ∇ log(ρ) = ∆ ρ e. This allows one to use the diffusion map approximation of the weighted Laplacian to approximate the interaction term as follows:

(DM)

EQUATION

where the kernel k (x, y) = g (x,y)

√ N i=1 g (y,X i )

is constructed empirically in terms of the Gaussian kernel g (x, y) = exp(-|x -y| 2 /(4 )). The parameter is referred to as the kernel bandwidth. The approximation is asymptotically exact as ↓ 0 and N ↑ ∞. The approximation error is of order

O( ) + O( 1 √ N d/4

) where the first term is referred to as the bias error and the second term is referred to as the variance error (Hein et al., 2007) . The variance error is the dominant term in the error for small values of , whereas the bias error is the dominant term for large values of (see Figure 3(d) ).

The resulting interacting particle algorithm is tabulated in Table 1 . The symplectic method proposed in (Betancourt et al., 2018) is used to carry out the numerical integration. The algorithm is applied to two examples as described in the following sections.

Remark 2. For the case where there is only one particle ( N = 1), the interaction term is zero and the system (17) reduces to the Nesterov ODE (5). approximation, the interaction term is approximated as

EQUATION

Despite the apparent similarity of the two formulae, (19) for diffusion map approximation and (20) for density estimation, the nature of the two approximations is different. The difference arises because the kernel k (x, y) in ( 19) is data-dependent whereas the kernel in (20) is not. While both approximations are exact in the asymptotic limit as N ↑ ∞ and ↓ 0, they exhibit different convergence rates. Numerical experiments presented in Figure 3 

4.3. COMPARISON WITH MCMC AND HMCMC:

This section contains numerical experiment comparing the performance of the accelerated algorithm 1 using the diffusion map (DM) approximation ( 19) and the density estimation (DE)-based approximation ( 20) with the Markov chain Monte-Carlo (MCMC) algorithm studied in Durmus & Moulines (2016) and the Hamiltonian MCMC algorithm studied in Cheng et al. (2017) .

We consider the problem setting of the mixture of Gaussians as in example 4.2. All algorithms are simulated with a fixed step-size of ∆t = 0.1 for K = 1000 iterations. The performance is measured by computing the mean-squared error in estimating the expectation of the function ψ(x) = x1 x≥0 . The mean-square error at the k-th iteration is computed by averaging the error over M = 100 runs:

m.s.e k = 1 M M m=1 1 N N i=1 ψ(X i,m t k ) -ψ(x)ρ ∞ (x) dx 2 (21)

The numerical results are depicted in Figure 3 . Figure 3 (a) depicts the m.s.e as a function of N . It is observed that the accelerated algorithm 1 with the diffusion map approximation admits an order of magnitude better m.s.e for the same number of particles. It is also observed that the m.s.e decreases rapidly for intermediate values of N before saturating for large values of N , where the bias term dominates (see discussion following Eq. 19). 

× N matrix [k (X i , X j )] N i,j=1

, while the computational cost of the MCMC and HMCMC algorithms scale as O(N ). The computational complexity may be improved by (i) exploiting the sparsity structure of the N × N matrix ; (ii) sub-sampling the particles in computing the empirical averages; (iii) adaptively updating the N × N matrix according to a certain error criteria. of , where bias and variance dominates the error, respectively, the two algorithms have similar m.s.e. However, for intermediate values of , the diffusion map approximation has smaller variance, and thus lower m.s.e.

5. CONCLUSION AND DIRECTIONS FOR FUTURE WORK:

The main contribution of this paper is to extend the variational formulation of Wibisono et al. (2016) to obtain theoretical results and numerical algorithms for accelerated gradient flow in the space of probability distributions. In continuous-time settings, bounds on convergence rate are derived based on a Lyapunov function argument. Two numerical algorithms based upon an interacting particle representation are presented and illustrated with examples. As has been the case in finite-dimensional settings, the theoretical framework is expected to be useful in this regard. Some direction for future include: (i) removing the technical assumption in the proof of the Theorem 1; (ii) analysis of the convergence under the weaker assumption that the target distribution satisfies only a spectral gap condition; and (iii) analysis of the numerical algorithms in the finite-N and in the finite ∆t cases.

A PDE FORMULATION OF THE VARIATIONAL PROBLEM:

An equivalent pde formulation is obtained by considering the stochastic optimal control problem (11) as a deterministic optimal control problem on the space of the probability distributions. Specifically, the process {ρ t } t≥0 is a deterministic process that takes values in P ac,2 (R d ) and evolves according to the continuity equation

∂ρ t ∂t = -∇ • (ρ t u t )

where u t : R d → R d is now a time-varying vector field. The Lagrangian L : R

EQUATION

The optimal control problem is:

EQUATION

The Hamiltonian function H : R

EQUATION

where φ ∈ C(R d ; R) is the dual variable and the inner-product ∇φ,

u L 2 (ρ) := R d ∇φ(x) • u(x)ρ(x) dx

B RESTATEMENT OF THE MAIN RESULT AND ITS PROOF:

We restate Theorem 1 below which now includes the pde formulation as well.

Theorem 2. Consider the variational problem (11)-( 23).

(i) For the probabilistic form (11) of the variational problem, the optimal control U * t = e αt-γt Y t , where the optimal trajectory {(X t , Y t )} t≥0 evolves according to the Hamilton's odes:

dX t dt = U * t = e αt-γt Y t , X 0 ∼ ρ 0 (25a) dY t dt = -e αt+βt+γt ∇ ρ F (ρ t )(X t ), Y 0 = ∇φ 0 (X 0 ) (25b)

where φ 0 is a convex function, and ρ t = Law(X t ). (ii) For the pde form (23) of the variational problem, the optimal control is u * t = e αt-γt ∇φ t (x), where the optimal trajectory {(ρ t , φ t )} t≥0 evolves according to the Hamilton's pdes:

EQUATION

(iii) The solutions of the two forms are equivalent in the following sense:

Law(X t ) = ρ t , U t = u t (X t ), Y t = ∇φ t (X t )

(iv) Suppose additionally that the functional F is displacement convex and ρ ∞ is its minimizer. Define

EQUATION

where the map T ρ∞ ρt : R d → R d is the optimal transport map from ρ t to ρ ∞ . Suppose also that the following technical assumption holds: E[(X t + e -γt Y t -T ρ∞ ρt (X t )) • d dt T ρ∞ ρt (X t )] = 0. Then dV dt (t) ≤ 0. Consequently, the following rate of convergence is obtained along the optimal trajectory F (ρ t ) -F (ρ ∞ ) ≤ O(e -βt ), ∀t ≥ 0

Proof.

(i) The Hamiltonian function defined in ( 12) is equal to

H(t, x, ρ, y, u) = y • u -e γt-αt 1 2 |u| 2 + e αt+γtβt F (ρ, x)

after inserting the formula for the Lagrangian. According to the maximum principle in probabilistic form for (mean-field) optimal control problems (see (Carmona & Delarue, 2017, Sec. 6.2.3 )), the optimal control law U * t = arg min v H(t, X t , ρ t , Y t , v) = e αt-γt Y t and the Hamilton's equations are

dX t dt = +∇ y H(t, X t , ρ t , Y t , U * t ) = U * t = e αt-γt Y t dY t dt = -∇ x H(t, X t , ρ t , Y t , U * t ) -Ẽ[∇ ρ H(t, Xt , ρ t , Ỹt , Ũ * t )(X t )]

where Xt , Ỹt , Ũ * t are independent copies of X t , Y t , U * t . The derivatives

∇ x H(t, x, ρ, y, u) = e αt+βt+γt ∇ x F (ρ, x) ∇ ρ H(t, x, ρ, y, u) = e αt+βt+γt ∇ ρ F (ρ, x) It follows that dY t dt = -e αt+βt+γt ∇ x F (ρ t , X t ) + Ẽ[∇ ρ F (ρ t , Xt )(X t )] = -e αt+βt+γt ∇ ρ F(ρ)(X t )

where we used the definition F(ρ) = F (x, ρ)ρ(x) dx and the identity (Carmona & Delarue, 2017, Sec. 5.2.2 Example 3)

∇ ρ F(ρ)(x) = ∇ x F (ρ, x) + ∇ ρ F (ρ, x)(x)ρ(x) dx

(ii) The Hamiltonian function defined in ( 24) is equal to

H(t, ρ, φ, u) = ∇φ(x) • u(x) - 1 2 e γt-αt |u(x)| 2 ρ(x) dx + e αt+γt+βt F(ρ)

after inserting the formula for the Lagrangian. According to the maximum principle for pde formulation of mean-field optimal control problems (see (Carmona & Delarue, 2017, Sec. 6.2.4 )) the optimal control vector field is u * t = arg min v H(t, ρ t , φ t , v) = e αt-γt ∇φ t and the Hamilton's equations are:

∂ρ t ∂t = + ∂H ∂φ (t, ρ t , φ t , u t ) = -∇ • (ρ t ∇u * t ) ∂φ t ∂t = - ∂H ∂ρ (t, ρ t , φ t , u t ) = -(∇φ • u * -e γt-αt 1 2 |u * t | 2 + e αt+γt+βt ∂F ∂ρ (ρ t ))

inserting the formula u * t = e αt-γt ∇φ t concludes the result.

(iii) Consider the (ρ t , φ t ) defined from ( 26). The distribution ρ t is identified with a stochastic process Xt such that d Xt dt = e αt-γt ∇φ t ( Xt ) and Law( Xt ) = ρ t . Then define Ỹt = ∇φ t ( Xt ). Taking the time derivative shows that

d Ỹt dt = d dt ∇φ t ( Xt ) = ∇ 2 φ t ( Xt ) d Xt dt + ∇ ∂φ t ∂t (X t ) = e αt-γt ∇ 2 φ t ( Xt )∇φ t ( Xt ) -e αt-γt ∇ 2 φ t ( Xt )∇φ t (X t ) -e αt+βt+γt ∇ ∂F ∂ρ (ρ t )( Xt ) = -e αt+βt+γt ∇ ∂F ∂ρ (ρ t )( Xt ) = -e αt+βt+γt ∇ ρ F(ρ t )( Xt )

with the initial condition Ỹ0 = ∇φ 0 ( X0 ), where we used the identity ∇ , 2017, Prop. 5.48) . Therefore the equations for Xt and Ỹt are identical. Hence one can identify (X t , Y t ) with ( Xt , Ỹt ).

x ∂F ∂ρ (ρ) = ∇ ρ F(ρ) (Car- mona & Delarue

(iv) The energy functional

V (t) = 1 2 E |X t + e -γt Y t -T ρ∞ ρt (X t )| 2 first term + e βt (F(ρ) -F(ρ ∞ )) second term

Then the derivative of the first term is

E (X t + e -γt Y t -T ρ∞ ρt (X t )) • (e αt-γt Y t -γt e -γt Y t -e αt+βt ∇ ρ F(ρ t )(X t ) + ξ(T ρ∞ ρt (X t )))

where ξ(T ρ∞ ρt (X t )) := d dt T ρ∞ ρt (X t ). Using the scaling condition γt = e αt the derivative of the first term simplifies to E (X t + e -γt Y t -T ρ∞ ρt (X t )) • (-e αt+βt ∇ ρ F(ρ t )(X t ) + ξ(T ρ∞ ρt (X t )))

Example: Let F(ρ) = log( ρ(x) ρ∞(x) )ρ(x) dx be the relative entropy functional. Consider a path ρ t in P ac,2 (R d ) such that ∂ρt ∂t = -∇ • (ρ t u) with ρ 0 = ρ ∈ P ac,2 (R d ). where {B t } t≥0 is the standard Brownian motion.

Consider next, the the accelerated flow (16). Denote v t := e αt-γt Y t . Then, with an appropriate choice of scaling parameters (e.g. α t = 0, β t = 0 and γ t = -γt ):

EQUATION

The scaling parameters are chosen here for the sake of comparison and do not satisfy the ideal scaling conditions of Wibisono et al. (2016) .

The sdes (28) and ( 29) are similar except that the stochastic term √ 2 dB t in ( 28) is replaced with a deterministic term -∇ x log(ρ t (X t )) in (29). Because of this difference, the resulting distributions are different. Let p t (x, v) denote the joint distribution on (X t , v t ) of ( 28) and let q t (x, v) denote the joint distribution on (X t , v t ) of ( 29). Then the corresponding Fokker-Planck equations are:

∂p ∂t (x, v) = -∇ x • (p t (x, v)v) + ∇ v • (p t (x, v)(γv + ∇f (x))) + ∆ v p t (x, v) ∂q ∂t (x, v) = -∇ x • (q t (x, v)v) + ∇ v • (q t (x, v)(γv + ∇f (x))) + ∇ v • (q t (x, y)∇ x log(ρ t (x)))

where ρ t (x) = q t (x, v) dv is the marginal of q t (x, y) on x. The final term in the Fokker-Planck equations are clearly different. The joint distributions are different as well.

The situation is in contrast to the first order Langevin equation, where the stochastic term √ 2 dB t and -∇ log(ρ t (X t )) are equivalent, in the sense that the resulting distributions have the same marginal distribution as a function of time. To illustrate this point, consider the following two forms of the Langevin equation:

EQUATION

)

EQUATION

:



Figures and Tables:

Figure 1: Simulation result for the Gaussian case (Example 4.1): (a) The time traces of the particles; (b) The KL-divergence as a function of time.

(a)-(d)  show that the diffusion map approximation has a much smaller variance for intermediate values of N . Theoretical understanding of the difference is the subject of continuing work.4.1 GAUSSIAN EXAMPLEConsider the Gaussian example as described in Sec. 3.5. The simulation results for the scalar (d = 1) case with initial distribution ρ 0 = N (2, 4) and target distribution N (x, Q) where x = -5.0 and Q = 0.25 is depicted in Figure1-(a)-(b). For this simulation, the numerical parameters are as follows: N = 100, φ 0 (x) = 0.5(x -2), t 0 = 1, ∆t = 0.1, p = 2,C = 0.625, and K = 400. The result numerically verifies the O(e -βt ) = O( 1 t 2 ) convergence rate derived in Theorem 1 for the case where the target distribution is Gaussian.4.2 NON-GAUSSIAN EXAMPLEThis example involves a non-Gaussian target distribution ρ ∞ = 1 2 N (-m, σ 2 ) + 1 2 N (m, σ 2 ) which is a mixture of two one-dimensional Gaussians with m = 2.0 and σ 2 = 0.8. The simulation results are depicted in Figure2-(a)-(b). The numerical parameters are same as in the Example 4.1. The interaction term is approximated using the diffusion map approximation with = 0.01. The numerical result depicted in Figure2-(a)show that the diffusion map algorithm converges to the mixture of Gaussian target distribution. The result depicted in Figure2-(b) suggests that the convergence rate O(e -βt ) also appears to hold for this non-log-concave target distribution. Theoretical justification of this is subject of continuing work.

Figure 2: Simulation result for the non-Gaussian case (Example 4.2): (a) The time traces of the particles; (b) The KL-divergence as a function of time.

Figure 3(b) depicts the m.s.e as a function of the number of iterations for a fixed number of particles N = 100. It is observed that the accelerated algorithm 1 displays the quickest convergence amongst the algorithms tested.

Figure 3(c) depicts the average computational time per iteration as a function of the number of samples N . The computational time of the diffusion map approximation scales as O(N 2 ) because it involves computing a N × N matrix [k (X i , X j )] N i,j=1, while the computational cost of the MCMC and HMCMC algorithms scale as O(N ). The computational complexity may be improved by (i) exploiting the sparsity structure of the N × N matrix ; (ii) sub-sampling the particles in computing the empirical averages; (iii) adaptively updating the N × N matrix according to a certain error criteria.

Figure 3: Simulation-based comparison of the performance of accelerated algorithm 1 using the diffusion map (DM) approximation (19), the density estimation (DE)-based approximation (20) with the MCMC and HMCMC algorithms: (a) the mean-squared error (m.s.e) (21) as a function of the number of samples N ; (b) the m.s.e as a function of the number of iterations; (c) the average computational time per iteration as a function of the number of samples; (d) m.s.e comparison between the diffusion map and the density estimation-based approaches as a function of the kernel bandwidth .

D RELATIONSHIP WITH THE UNDER-DAMPED LANGEVIN EQUATIONA basic form of the under-damped (or second order) Langevin equation is given inCheng et al. (2017) dX t = v t dt dv t = -γv t dt -∇f (X t ) dt +

Table content: <table><tr><td/><td/><td/><td/><td/><td/><td/><td/><td>Then</td></tr><tr><td>d dt</td><td colspan="2">F(ρ t ) = log(</td><td colspan="3">ρ t (x) ρ ∞ (x)</td><td>)</td><td colspan="2">∂ρ t ∂t</td><td>(x) dx +</td><td>∂ρ t ∂t</td><td>(x) dx</td></tr><tr><td/><td colspan="3">= -log(</td><td colspan="4">ρ t (x) ρ ∞ (x)</td><td>)∇ • (ρ t (x)u(x)) dx</td></tr><tr><td/><td colspan="4">= ∇ x log(</td><td colspan="4">ρ t (x) ρ ∞ (x)</td><td>) • u(x) ρ t (x) dx</td></tr><tr><td colspan="9">where the divergence theorem is used in the last step. The definitions of the Gâteaux derivative and</td></tr><tr><td colspan="2">Wasserstein gradient imply</td><td/><td/><td/><td/><td/><td/></tr><tr><td/><td>∂F ∂ρ</td><td colspan="7">(ρ)(x) = log(</td><td>ρ(x) ρ ∞ (x)</td><td>)</td></tr><tr><td/><td colspan="8">∇ ρ F(ρ)(x) = ∇ x log(</td><td>ρ(x) ρ ∞ (x)</td><td>)</td></tr><tr><td/><td/><td/><td/><td/><td/><td/><td/><td>√</td><td>2 dB t</td><td>(28)</td></tr></table>

Peer review (iclr2019):
 This paper derives accelerated gradient flow formula in the space of probability measures from the view of optimal control formalism. The generalization of variational formulation from finite space to the space of probability measures seems new, but the resulting PDE seems to be a known result, which is the Fokker-Planck equation (with some minor modifications) for the 2nd order Langevin dynamic. From this point of view, the resulting algorithm from the derived PDE seems not having much practical advantage over SGHMC (a stochastic version of 2nd order Langevin dynamics). Actually, I think the derivation of accelerated gradient flow formula from the view of optimal control formalism does not seem necessary. One can get the same formula by deriving it from Wasserstein gradient flows. When considering the functional as relative entropy, one can derive the formula simply from the Fokker-Planck equation of 2nd order Langevin dynamics. As a result, the proposed methods seems to be a new way to derive the Wasserstein gradient flow (or Fokker-Planck equation), which does not make impact the algorithm, e.g., both ways result in the same algorithm. Besides, I found the writing needs to be improved. There are a lot of background missing, or the descriptions are not clear enough. For example: 1. Page 2: the divergence operator is not defined, though I think it is a standard concept, but would be better to define it. 2. Page 2: the Wasserstein gradient and Gateaux derivative are not defined, what are the specific meanings of .nabla_.rho F(.rho) and .partial F / .partial .rho? 3. 1st line in Section 2: convex function f of d real variables seems odd, I guess the author means argument of f is d-dimensional variable. 4. Section 2, the authors directly start with the variational problem (3) without introducing the problem. Why do we need to variational problem? It would be hard to follow for some one who does not have such background. 5. Similarly, what is the role of Lyapunov function here in (6)? Why do we need it? 6. Why do you define the Lagrangian L in the form of (10)? What is the relation between (10) and (2)? 7. It is not clear what "The stochastic process (X_t, Y_t) is Gaussian" means in Proposition 1? It might need to be rephrased. 8. Second last line in page 5: I guess .nabla .log(.rho) should be .nabla.log(.rho_t). For the theory, I think eq.15 only applies when the PDE, e.g. (13), is solved exactly, thus there is not too much practical impact, as it is well known from the Wasserstein gradient theory that the PDE decays exponentially, as stated in the theorem. When considering numerical solutions, I think this results is useless. For the relation with SGHMC, let*s look at eq.16. Actually, the derivative of the log term .nabla .log .rho_t(X_t)) is equivalent to a brownian motion term. This can be seen by considering the Fokker-Planck equation for Brownian motion, which is exactly d .rho_t = .Delta .rho_t. Consequently, instead of using the numerical approximations proposed later, one cane simply replacing this term with a Brownian motion term, which reduces to SGHMC (with some constant multipliers in front). The authors then shows empirically that the proposed method is better than SGHMC, which I think only comes from the numerical methods. For the kernel approximation, it makes the particles in the algorithm interactive. This resembles other particle optimization based algorithms such as SVGD, or the latest particle interactive SGLD proposed in [1] or [2[. I think these methods need to be compared. [1] Chen et al (2018), A Unified Particle-Optimization Framework for Scalable Bayesian Sampling. [2] Liu et al (2018), https://arxiv.org/pdf/1807.01750.pdf To sum up, though the derivation of accelerated gradient flow formula seems interesting, the resulting algorithm does not seem benefit from this derivation. The algorithm seems to be able to derived from a more direct way of using Wasserstein gradient flows, which results in a Wasserstein gradient flow for 2nd order Langevin dynamics, and is thus well known. The experiments are not convincing, and fail to show the advantage of the proposed method. The proposed method needs to be compared with other related methods.