Year: 2018

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.

text: 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.

heading: 1 INTRODUCTION
text: 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).
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 MonteCarlo 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.
Notation: 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∇xnUn(x). The space of absolutely continuous probability measures on Rd with finite second moments is denoted by Pac,2(Rd). 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).

heading: 2 REVIEW OF THE VARIATIONAL FORMULATION OF WIBISONO ET AL. (2016)
text: The basic problem is to minimize a C1 smooth convex function f on Rd. The standard form of the gradient descent algorithm for this problem is an ODE:
dXt dt = −∇f(Xt), 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+ × Rd × Rd → R is defined as
L(t, x, u) := eαt+γt ( 1
2 |e−αtu|2︸ ︷︷ ︸ kinetic energy − eβtf(x)︸ ︷︷ ︸ potential energy
) (2)
where t ≥ 0 is the time, x ∈ Rd is the state, u ∈ Rd 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+ logC, and γt = p log t where p ≥ 2 and C > 0 are constants.
The variational problem is
Minimize u J(u) = ∫ ∞ 0 L(t,Xt, ut) dt Subject to dXt dt = ut, X0 = x0
(3)
The Hamiltonian function is
H(t, x, y, u) = y · u− L(t, x, u) (4)
where y ∈ Rd 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,Xt, Yt, v) = e
αt−γtYt. The resulting Hamilton’s equations are
dXt dt = +∇yH(t,Xt, Yt, ut) = eαt−γtYt, X0 = x0 (5a)
dYt dt = −∇xH(t,Xt, Yt, ut) = −eαt+βt+γt∇f(Xt), Y0 = y0 (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:
V (t, x, y) = 1
2
∣∣x+ e−γty − x̄∣∣2 + eβt(f(x)− f(x̄)) (6)
where x̄ ∈ arg minx f(x). It is shown in Wibisono et al. (2016) that upon differentiating along the solution trajectory, ddtV (t,Xt, Yt) ≤ 0. This yields the following convergence rate:
f(Xt)− f(x̄) ≤ O(e−βt), ∀t ≥ 0 (7)

heading: 3 VARIATIONAL FORMULATION FOR PROBABILITY DISTRIBUTIONS

heading: 3.1 MOTIVATION AND BACKGROUND
text: Let F : Pac,2(Rd)→ R be a functional on the space of probability distributions. Consider the problem of minimizing F(ρ). The (Wasserstein) gradient flow with respect to F(ρ) is
∂ρt ∂t = ∇ · (ρt∇ρF(ρt)) (8)
where ∇ρF(ρ) is the Wasserstein gradient of F. An important example is the relative entropy functional where F(ρ) = D(ρ|ρ∞) :=∫ Rd log( ρ(x) ρ∞(x)
)ρ(x) dx where ρ∞ ∈ Pac,2(Rd) 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):
dXt = −∇f(Xt) dt+ √
2 dBt, X0 ∼ ρ0 where Bt is the standard Brownian motion in Rd.
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.

heading: 3.2 PROBABILISTIC FORM OF THE VARIATIONAL PROBLEM
text: Consider the stochastic process {Xt}t≥0 that takes values in Rd and evolves according to: dXt dt = Ut, X0 ∼ ρ0 where the control input {Ut}t≥0 also takes values in Rd, and ρ0 ∈ Pac,2(Rd) is the probability distribution of the initial condition X0. 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 :

heading: R+ × Rd × Pac,2(Rd)× Rd → R is defined as
text: L(t, x, ρ, u) := eαt+γt ( 1
2 |e−αtu|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:
Minimize J(u) = E [∫ ∞
0
L(t,Xt, ρt, Ut) dt ] Subject to
dXt dt
= Ut, X0 ∼ ρ0 (11)
where ρt = Law(Xt) ∈ Pac,2(Rd) is the probability density function of the random variable Xt.
The Hamiltonian function H : R+ × Rd × Pac,2(Rd) × Rd × Rd → R for this problem is given by (Carmona & Delarue, 2017, Sec. 6.2.3):
H(t, x, ρ, y, u) := u · y − L(t, x, ρ, u) (12)
where y ∈ Rd is the dual variable.

heading: 3.3 MAIN RESULT
text: Theorem 1. Consider the variational problem (11).
(i) The optimal control U∗t = e αt−γtYt where the optimal trajectory {(Xt, Yt)}t≥0 evolves
according to the Hamilton’s equations: dXt dt = U∗t = e αt−γtYt, X0 ∼ ρ0 (13a)
dYt dt = −eαt+βt+γt∇ρF(ρt)(Xt), Y0 = ∇φ0(X0) (13b)
where φ0 is any convex function and ρt := Law(Xt). (ii) Suppose also that the functional F is displacement convex and ρ∞ is its minimizer. Define
the energy along the optimal trajectory
V (t) = 1
2 E[|Xt + e−γtYt − T ρ∞ρt (Xt)| 2] + eβt(F(ρ)− F(ρ∞)) (14)
where the map T ρ∞ρt : R d → Rd is the optimal transport map from ρt to ρ∞. Suppose also that the following technical assumption holds: E[(Xt + e−γtYt − T ρ∞ρt (Xt)) · d dtT ρ∞ ρt (Xt)] = 0. Then dVdt (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.

heading: 3.4 RELATIVE ENTROPY AS THE FUNCTIONAL
text: In the remainder of this paper, we assume that the functional F(ρ) = D(ρ|ρ∞) is the relative entropy where ρ∞ ∈ Pac,2(Rd) is a given target probability distribution. In this case the Hamilton’s equations are given by
dXt dt = eαt−γtYt, X0 ∼ ρ0 (16a)
dYt dt = −eαt+βt+γt(∇f(Xt) +∇ log(ρt(Xt)), Y0 = ∇φ0(X0) (16b)
where ρt = Law(Xt) 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.

heading: 3.5 QUADRATIC GAUSSIAN CASE
text: Suppose the initial distribution ρ0 and the target distribution ρ∞ are both Gaussian, denoted as N (m0,Σ0) and N (x̄, Q), respectively. This is equivalent to the objective function f(x) being quadratic of the form f(x) = 12 (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 (Xt, Yt) 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 (Xt, Yt) is a Gaussian process. The Hamilton’s equations are given by:
dXt dt = eαt−γtYt, dYt dt = −eαt+βt+γt(Q−1(Xt − x̄)− Σ−1t (Xt −mt))
where mt and Σt are the mean and the covariance of Xt.
(ii) Upon taking the expectation of both sides, and denoting nt := E[Yt]
dmt dt = eαt−γtnt, dnt dt = −eαt+βt+γt Q−1(mt − x̄)︸ ︷︷ ︸ ∇f(mt)
which is identical to Nesterov ODE (5).

heading: 4 NUMERICAL ALGORITHM
text: The proposed numerical algorithm is based upon an interacting particle implementation of the Hamilton’s equation (16). Consider a system of N particles {(Xit , Y it )}Ni=1 that evolve according to:
dXit dt = eαt−γtY it , X i 0 ∼ ρ0 dY it dt = −eαt+βt+γt(∇f(Xit) + I (N) t (X i t)︸ ︷︷ ︸
interaction term
), Y i0 = ∇φ0(Xi0)
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 GaussianN (mt,Σt). In this case,∇ log(ρt(x)) = −Σt−1(x−mt). This motivates the following empirical approximation of the interaction term:
I (N) t (x) = −Σ (N) t −1 (x−m(N)t ) (18)
Algorithm 1 Interacting particle implementation of the accelerated gradient flow Input: ρ0, φ0, N , t0, ∆t, p, C, K Output: {Xik} N,K i=1,k=0
Initialize {Xi0}Ni=1 i.i.d∼ ρ0, Y i0 = ∇φ0(Xi0) Compute I(N)0 (X i 0) with (18) or (19) for k = 0 to K − 1 do tk+ 12 = tk + 1 2∆t
Y i k+ 12 = Y ik − 12Cpt 2p−1 k+ 12 (∇f(Xik) + I (N) k (X i k))∆t Xik+1 = X i k + p
tp+1 k+1
2
Y ik∆t
Compute I(N)k+1(X i k+1) with (18)or (19) Y ik+1 = Y i k+ 12 − 12Cpt 2p−1 k+ 12 (∇f(Xik+1) + I (N) k+1(X i k+1))∆t tk+1 = tk+ 12 + 1 2∆t
end for
where m(N)t := N −1∑N i=1X 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 C2 function f , the weighted Laplacian is defined as ∆ρf := 1ρ∇ · (ρ∇f). Denote e(x) = x as the coordinate function on Rd. 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) I(N)t (X i t) =
1 ∑N j=1 k (X i t , X j t )(X
j t −Xit)∑N
j=1 k (X i t , X j t )
(19)
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).
Remark 3. (Comparison with density estimation) The diffusion map approximation algorithm is conceptually different from an explicit density estimation-based approach. A basic density estimation is to approximate ρ(x) ≈ 1N ∑N i=1 g (x,X i t) where g (x, y) is the Gaussian kernel. Using such an
approximation, the interaction term is approximated as
(DE) I(N)t (X i t) =
1
2
∑N j=1 g (X i t , X j t )(X
j t −Xit)∑N
j=1 g (X i t , X j t )
(20)
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(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.

heading: 4.1 GAUSSIAN EXAMPLE
text: Consider 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 Figure 1-(a)-(b). For this simulation, the numerical parameters are as follows: N = 100, φ0(x) = 0.5(x − 2), t0 = 1, ∆t = 0.1, p = 2,C = 0.625, and K = 400. The result numerically verifies the O(e−βt) = O( 1t2 ) convergence rate derived in Theorem 1 for the case where the target distribution is Gaussian.

heading: 4.2 NON-GAUSSIAN EXAMPLE
text: This example involves a non-Gaussian target distribution ρ∞ = 12N (−m,σ 2) + 12N (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 Figure 2-(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 Figure 2-(a) show that the diffusion map algorithm converges to the mixture of Gaussian target distribution. The result depicted in Figure 2-(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.

heading: 4.3 COMPARISON WITH MCMC AND HMCMC
text: 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) = x1x≥0. The mean-square error at the k-th iteration is computed by averaging the error over M = 100 runs:
m.s.ek = 1
M M∑ m=1
( 1
N N∑ i=1 ψ(Xi,mtk )− ∫ ψ(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).
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(N2) because it involves computing a N ×N matrix [k (Xi, Xj)]Ni,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 theN×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.
Finally, we provide comparison between diffusion map approximation (20) and the density-based approximation (20): Figure 3(d) depicts the m.s.e for these two approximations as a function of the kernel-bandwidth for a fixed number of particles N = 100. For very large and for very small values
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.

heading: 5 CONCLUSION AND DIRECTIONS FOR FUTURE WORK
text: 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.

heading: A PDE FORMULATION OF THE VARIATIONAL PROBLEM
text: 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 Pac,2(Rd) and evolves according to the continuity equation
∂ρt ∂t = −∇ · (ρtut)
where ut : Rd → Rd is now a time-varying vector field. The Lagrangian L : R+ × Pac,2(Rd) × L2(Rd;Rd)→ R is defined as:
L(t, ρ, u) := eαt+γt [∫
Rd
1 2 |e−αtu(x)|2ρ(x) dx− eβtF(ρ)
] (22)
The optimal control problem is: Minimize ∫ ∞ 0 L(t, ρt, ut) dt
Subject to ∂ρt ∂t
+∇ · (ρtut) = 0 (23)
The Hamiltonian functionH : R+ × Pac,2(Rd)× C(Rd;R)× L2(Rd;Rd)→ R is
H(t, ρ, φ, u) := 〈∇φ, u〉L2(ρ) − L(t, ρ, u) (24)
where φ ∈ C(Rd;R) is the dual variable and the inner-product 〈∇φ, u〉L2(ρ) := ∫ Rd ∇φ(x) · u(x)ρ(x) dx

heading: B RESTATEMENT OF THE MAIN RESULT AND ITS PROOF
text: 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−γtYt,
where the optimal trajectory {(Xt, Yt)}t≥0 evolves according to the Hamilton’s odes: dXt dt = U∗t = e αt−γtYt, X0 ∼ ρ0 (25a)
dYt dt = −eαt+βt+γt∇ρF (ρt)(Xt), Y0 = ∇φ0(X0) (25b)
where φ0 is a convex function, and ρt = Law(Xt). (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:
∂ρt ∂t = −∇ · (ρt eαt−γt∇φt︸ ︷︷ ︸ u∗t ), initial condn. ρ0 (26a) ∂φt ∂t = −eαt−γt |∇φt| 2 2 − eαt+γt+βt∇ρF (ρ) (26b)
(iii) The solutions of the two forms are equivalent in the following sense:
Law(Xt) = ρt, Ut = ut(Xt), Yt = ∇φt(Xt)
(iv) Suppose additionally that the functional F is displacement convex and ρ∞ is its minimizer. Define
V (t) = 1
2 E(|Xt + e−γtYt − T ρ∞ρt (Xt)| 2) + eβt(F (ρ)− F (ρ∞)) (27)
where the map T ρ∞ρt : R d → Rd is the optimal transport map from ρt to ρ∞. Suppose also that the following technical assumption holds: E[(Xt + e−γtYt − T ρ∞ρt (Xt)) · d dtT ρ∞ ρt (Xt)] = 0. Then dVdt (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 minv H(t,Xt, ρt, Yt, v) = e
αt−γtYt and the Hamilton’s equations are
dXt dt = +∇yH(t,Xt, ρt, Yt, U∗t ) = U∗t = eαt−γtYt dYt dt = −∇xH(t,Xt, ρt, Yt, U∗t )− Ẽ[∇ρH(t, X̃t, ρt, Ỹt, Ũ∗t )(Xt)]
where X̃t, Ỹt, Ũ∗t are independent copies of Xt, Yt, U ∗ t . The derivatives
∇xH(t, x, ρ, y, u) = eαt+βt+γt∇xF̃ (ρ, x) ∇ρH(t, x, ρ, y, u) = eαt+βt+γt∇ρF̃ (ρ, x)
It follows that dYt dt = −eαt+βt+γt ( ∇xF̃ (ρt, Xt) + Ẽ[∇ρF̃ (ρt, X̃t)(Xt)] ) = −eαt+βt+γt∇ρF(ρ)(Xt)
where we used the definition F(ρ) = ∫ F̃ (x, ρ)ρ(x) dx and the identity (Carmona & Delarue, 2017, Sec. 5.2.2 Example 3)
∇ρF(ρ)(x) = ∇xF̃ (ρ, 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+βtF(ρ)
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 minvH(t, ρt, φt, v) = eαt−γt∇φt and the Hamilton’s equations are:
∂ρt ∂t = + ∂H ∂φ (t, ρt, φt, ut) = −∇ · (ρt∇u∗t ) ∂φt ∂t = −∂H ∂ρ (t, ρt, φt, ut) = −(∇φ · 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 X̃t such that dX̃tdt = e
αt−γt∇φt(X̃t) and Law(X̃t) = ρt. Then define Ỹt = ∇φt(X̃t). Taking the time derivative shows that
dỸt dt = d dt ∇φt(X̃t) = ∇2φt(X̃t) dX̃t dt +∇∂φt ∂t (Xt)
= eαt−γt∇2φt(X̃t)∇φt(X̃t)− eαt−γt∇2φt(X̃t)∇φt(Xt)− eαt+βt+γt∇ ∂F
∂ρ (ρt)(X̃t)
= −eαt+βt+γt∇∂F ∂ρ (ρt)(X̃t) = −eαt+βt+γt∇ρF(ρt)(X̃t)
with the initial condition Ỹ0 = ∇φ0(X̃0), where we used the identity∇x ∂F∂ρ (ρ) = ∇ρF(ρ) (Carmona & Delarue, 2017, Prop. 5.48). Therefore the equations for X̃t and Ỹt are identical. Hence one can identify (Xt, Yt) with (X̃t, Ỹt).
(iv) The energy functional
V (t) = 1 2 E [ |Xt + e−γtYt − T ρ∞ρt (Xt)| 2 ]
︸ ︷︷ ︸ first term + eβt(F(ρ)− F(ρ∞))︸ ︷︷ ︸ second term
Then the derivative of the first term is E [ (Xt + e −γtYt − T ρ∞ρt (Xt)) · (e αt−γtYt − γ̇te−γtYt − eαt+βt∇ρF(ρt)(Xt) + ξ(T ρ∞ρt (Xt))) ] where ξ(T ρ∞ρt (Xt)) := d dtT ρ∞ ρt (Xt). Using the scaling condition γ̇t = e
αt the derivative of the first term simplifies to
E [ (Xt + e −γtYt − T ρ∞ρt (Xt)) · (−e αt+βt∇ρF(ρt)(Xt) + ξ(T ρ∞ρt (Xt))) ]
Upon using the technical assumption, E[(Xt + e−γtYt − T ρ∞ρt (Xt)) · ξ(T ρ∞ ρt (Xt))] = 0 the derivative of the first term simplifies to
E [ (Xt + e −γtYt − T ρ∞ρt (Xt)) · (−e αt+βt∇ρF (ρt)(Xt)) ] The derivative of the second term is
d dt (second term) = β̇teβt(F(ρt)− F(ρ∞)) + eβt d dt F(ρt)
= eαt+βt(F(ρt)− F(ρ∞)) + eβtE[∇ρF(ρt)(Xt)eαt−γtYt]
where we used the scaling condition β̇t = eαt and the chain-rule for the Wasserstein gradient (Ambrosio et al., 2008, Ch. 10, E. Chain rule). Adding the derivative of the first and second term yields:
dV
dt (t) = eαt+βt
( F(ρt)− F(ρ∞)− E [ (Xt − T ρ∞ρt (Xt)) · ∇ρF(ρt)(Xt) ]) which is negative by variational inequality characterization of the displacement convex function F(ρ) (Ambrosio et al., 2008, Eq. 10.1.7).
We expect that the technical assumption can be removed. This is the subject of the continuing work.

heading: C WASSERSTEIN GRADIENT AND GÂTEAUX DERIVATIVE
text: This section contains definitions of the Wasserstein gradient and Gâteaux derivative (Ambrosio et al., 2008; Carmona & Delarue, 2017).
Let F : Pac,2(Rd)→ R be a (smooth) functional on the space of probability distributions.
Gâteaux derivative: The Gâteaux derivative of F at ρ ∈ Pac,2(Rd) is a real-valued function on Rd denoted as ∂F∂ρ (ρ) : R d → R. It is defined as a function that satisfies the identity
d dt F(ρt) ∣∣∣∣ t=0 = ∫ Rd ∂F ∂ρ (ρ)(x)(−∇ · (ρ(x)u(x))) dx
for all path ρt in Pac,2(Rd) such that ∂ρt∂t = −∇ · (ρtu) with ρ0 = ρ ∈ Pac,2(R d).
Wasserstein gradient: The Wasserstein gradient of F at ρ is a vector-field on Rd denoted as∇ρF(ρ) : Rd → Rd. It is defined as a vector-field that satisfies the identity
d dt F(ρt) ∣∣∣∣ t=0 = ∫ Rd ∇ρF(ρ)(x) · u(x) ρ(x) dx
for all path ρt in Pac,2(Rd) such that ∂ρt∂t = −∇ · (ρtu) with ρ0 = ρ ∈ Pac,2(R d).
The two definitions imply the following relationship (Carmona & Delarue, 2017, Prop. 5.48):
∇ρF(ρ)(·) = ∇x ∂F
∂ρ (ρ)(·)
Example: Let F(ρ) = ∫
log( ρ(x)ρ∞(x) )ρ(x) dx be the relative entropy functional. Consider a path ρt in
Pac,2(Rd) such that ∂ρt∂t = −∇ · (ρtu) with ρ0 = ρ ∈ Pac,2(R d). Then
d dt F(ρt) =
∫ log( ρt(x)
ρ∞(x) ) ∂ρt ∂t (x) dx+ ∫ ∂ρt ∂t (x) dx
= − ∫ log( ρt(x)
ρ∞(x) )∇ · (ρt(x)u(x)) dx
= ∫ ∇x log( ρt(x)
ρ∞(x) ) · u(x) ρt(x) dx
where the divergence theorem is used in the last step. The definitions of the Gâteaux derivative and Wasserstein gradient imply
∂F ∂ρ (ρ)(x) = log( ρ(x) ρ∞(x) )
∇ρF(ρ)(x) = ∇x log( ρ(x)
ρ∞(x) )

heading: D RELATIONSHIP WITH THE UNDER-DAMPED LANGEVIN EQUATION
text: A basic form of the under-damped (or second order) Langevin equation is given in Cheng et al. (2017)
dXt = vt dt
dvt = −γvt dt−∇f(Xt) dt+ √ 2 dBt (28)
where {Bt}t≥0 is the standard Brownian motion.
Consider next, the the accelerated flow (16). Denote vt := eαt−γtYt. Then, with an appropriate choice of scaling parameters (e.g. αt = 0, βt = 0 and γt = −γt ):
dXt = vt dt
dvt = −γvt dt−∇f(Xt) dt−∇x log(ρt(Xt)) (29)
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 dBt in (28) is replaced with a deterministic term −∇x log(ρt(Xt)) in (29). Because of this difference, the resulting distributions are different. Let pt(x, v) denote the joint distribution on (Xt, vt) of (28) and let qt(x, v) denote the joint distribution on (Xt, vt) of (29). Then the corresponding Fokker-Planck equations are: ∂p
∂t (x, v) = −∇x · (pt(x, v)v) +∇v · (pt(x, v)(γv +∇f(x))) + ∆vpt(x, v) ∂q ∂t (x, v) = −∇x · (qt(x, v)v) +∇v · (qt(x, v)(γv +∇f(x))) +∇v · (qt(x, y)∇x log(ρt(x)))
where ρt(x) = ∫ qt(x, v) dv is the marginal of qt(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 dBt and −∇ log(ρt(Xt)) 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:
dXt = −∇f(Xt) dt+ √ 2 dBt (30)
dXt = −∇f(Xt) dt−∇ log(ρt(Xt)) (31)
Let pt(x) denote the distribution of Xt of (30) and let qt(x) denote the distribution of Xt of (31). The corresponding Fokker-Planck equations are as follows
∂p ∂t (x) = −∇ · (pt(x)∇f(x)) + ∆pt(x) ∂q ∂t (x) = −∇ · (qt(x)∇f(x)) +∇ · (qt(x)∇ log(ρt(x)))
= −∇ · (qt(x)∇f(x)) +∇ · (qt(x)∇ log(qt(x)))
= −∇ · (qt(x)∇f(x)) + ∆qt(x)
where we used ρt(x) = qt(x). In particular, this implies that the marginal probability distribution of the stochastic process Xt are the same for first order Langevin sde (30) and (31) .

Review summaries:
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).

This paper introduces a functional extension of the Bregman Lagrangian framework of Wibisono et al. 2016. The basic idea is to define accelerated gradient flows on the space of probability distribution. Because the defined flows include a term depending on the current distribution of the system, which is difficult to compute in general, the authors introduce an interacting particle approximation as a practical numerical approximation. The experiments are a proof-of-concept on simple illustrative toy examples.