A CCELERATED G RADIENT F LOW FOR P ROBABILITY DISTRIBUTIONS
Anonymous authors
Paper under double-blind review
A BSTRACT
This paper presents a methodology and numerical algorithms for constructing accelerated gradient ﬂows 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-ﬁeld optimal control problem.
The maximum principle of optimal control theory is used to derive Hamilton’s equations for the optimal gradient ﬂow.
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.
I NTRODUCTION
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 ﬂow 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 ﬂows 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 signiﬁcant 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 deﬁned 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-ﬁeld optimal problem.
The maximum principle of the optimal control theory is used to derive the Hamilton’s equations which represent the ﬁrst 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.
and the extension proposed in this paper.
We also consider the important special case when the objective functional is the relative entropy functional D ( ρ | ρ ∞ ) deﬁned with respect to a target probability distribution ρ ∞ .
In this case, the accelerated gradient ﬂow 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 ﬁnite-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 ﬁnal 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 ﬁnite but large number of N particles.
The difﬁcult 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.
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 ﬁeld U is ∇ · U ( x ) =
(cid:80) dn =1 ∇ x n U n ( x ) . The space of absolutely continuous probability measures on R d with ﬁnite 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 deﬁnition).
The probability distribution of a random variable Z is denoted as Law ( Z ) .
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
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:
[Equation]
Accelerated forms of this algorithm are obtained based on a variational formulation due to Wibisono et al.
The formulation is brieﬂy reviewed here using an optimal control formalism.
The
Lagrangian L : R + × R d × R d → R is deﬁned as
[Equation]
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
Minimize u
[Equation]
Subject to
[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
[Equation]
t
[Equation]
t [Equation] t Y t .
The resulting Hamilton’s equations are
[Equation]
The system (5) is an example of accelerated gradient descent algorithm.
Speciﬁcally, if the parameters α t , β t , γ t are deﬁned 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, dd t V ( t, X t , Y t ) ≤ 0 .
This yields the following convergence rate:
[Equation]
V ARIATIONAL FORMULATION FOR PROBABILITY DISTRIBUTIONS 3. M OTIVATION 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 ﬂow with respect to F ( ρ ) is [Equation]
[Equation]
where ∇ ρ F ( ρ ) is the Wasserstein gradient of F .
An important example is the relative entropy functional where F ( ρ ) =
[Equation]
:= (cid:82) R d log( ρ ( x ) ρ ∞ ( x ) )
ρ ( x ) d x where ρ ∞ ∈ P ac , 2 ( R d ) is referred to as the target distribution.
The gradient of relative entropy is given by ∇ ρ F ( [Equation] = ∇ log( ρρ
[Equation]
The gradient ﬂow
[Equation]
(9) is the Fokker-Planck equation (Jordan et al., 1998).
The gradient ﬂow 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):
[Equation]
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 ﬂow (8).
The proposed solution is based upon a variational formulation.
As tabulated in Table 1, the solution represents a generalization of Wibisono et al.
from its original deterministic ﬁnite-dimensional to now probabilistic inﬁnite-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. P ROBABILISTIC FORM OF THE VARIATIONAL PROBLEM
Consider the stochastic process { X t } t ≥ 0 that takes values in R d and evolves according to:
[Equation]
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 ( ρ )
= (cid:82) ˜ F ( ρ, x ) ρ ( x ) d x .
The Lagrangian L :
R + × R d × P ac , 2 ( R d ) × R d → R is deﬁned as
[Equation]
This formula is a natural generalization of the Lagrangian (2) and the parameters α t , β t , γ t are deﬁned exactly the same as in the ﬁnite-dimensional case.
The stochastic optimal control problem is:
[Equation]
to
[Equation]
11) 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. M AIN 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 ≥ evolves according to the Hamilton’s equations:
[Equation]
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.
Deﬁne 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 )) · dd t T ρ ∞ ρ t ( X t )] = 0 .
Then d V d t ( t ) ≤ 0 .
Consequently, the following rate of convergence is obtained along the optimal trajectory:
[Equation]
Proof sketch.
The Hamilton’s equations are derived using the standard mean-ﬁeld 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. R ELATIVE 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. Q UADRATIC G AUSSIAN 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 ) =
( x − ¯ x ) (cid:62) 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:
[Equation]
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 ]
[Equation]
which is identical to Nesterov ODE (5) [Equation]
N UMERICAL 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 it , Y it ) } Ni =1 that evolve according to:
[Equation]
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:
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]
Input: ρ 0 , φ 0 , N , t
[Equation]
∆
[Equation]
, p ,
[Equation]
K
Output: { X ik } N,Ki =1 ,k =
Initialize { X i 0 } Ni =1 i.i.d ∼ ρ 0 , Y i
= ∇ φ 0 ( X i 0 )
Compute I ( N )( X i 0 ) with (18) or (19) for k = 0 to K − 1 do
[Equation]
Compute I ( N ) k +
( X ik +1 ) with (18)or (19)
[Equation]
X
[Equation]
+
[Equation]
X
[Equation]
∆
[Equation]
t
[Equation]
=
[Equation]
t end for where m ( N ) t :=
N − 1 (cid:80) Ni =
X it is the empirical mean and Σ ( N ) t :=
[Equation]
X
[Equation]
m
[Equation] [Equation]
X
[Equation]
− m ( N ) t ) (cid:62) 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) ﬁltering problem, where an ensemble Kalman ﬁlter is used as an approximate solution for non-Gaussian distributions (Evensen, 2003).
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 deﬁned 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:
[Equation]
) where the kernel k (cid:15) ( x, y ) =
[Equation]
( y,X i ) is constructed empirically in terms of the Gaussian kernel g (cid:15) ( x, y ) =
exp ( −| x − y | 2 / (4 (cid:15) )) .
The parameter (cid:15) is referred to as the kernel bandwidth.
The approximation is asymptotically exact as (cid:15) ↓ 0 and N ↑ ∞ .
The approximation error is of order
O ( (cid:15) )
+ O ( 1 √ N(cid:15) d/ 4 ) where the ﬁrst 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 (cid:15) , whereas the bias error is the dominant term for large values of (cid:15) (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 ) ≈ 1 N (cid:80) Ni =1 g (cid:15) ( x, X it ) where g (cid:15) ( x, y ) is the Gaussian kernel.
Using such an
[Figure]
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 (cid:15) ( 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 (cid:15) ↓ 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.
4. G AUSSIAN E XAMPLE
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 .
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) , t 0 = 1 , ∆ t = 0 .
1 , p = 2 , C = 0 .
625 , and K = 400 .
The result numerically veriﬁes the O ( e − βt ) convergence rate derived in Theorem 1 for the case where the target distribution is Gaussian.
4. N ON -G AUSSIAN EXAMPLE
This example involves a non-Gaussian target distribution ρ ∞ = 12 N ( − m, σ 2 ) + 12 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 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 (cid:15) = 0 .
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 justiﬁcation of this is subject of continuing work.
t=
tt=tt=tttt(a) 1010 1 t
KL ( t | )
O ( e t ) (b)
4. C OMPARISON WITH MCMC AND HMCMC
This section contains numerical experiment comparing the performance of the accelerated algorithmusing 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.
We consider the problem setting of the mixture of Gaussians as in example 4.2.
All algorithms are simulated with a ﬁxed 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 ) = x 1 x ≥ 0 .
The mean-square error at the k -th iteration is computed by averaging the error over M = 100 runs:
[Equation]
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 ﬁxed 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 (cid:15) ( X i , X j )]
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 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.
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 (cid:15) for a ﬁxed number of particles N = 100 .
For very large and for very small values
[Figure]
of (cid:15) , where bias and variance dominates the error, respectively, the two algorithms have similar m.s.e.
However, for intermediate values of (cid:15) , the diffusion map approximation has smaller variance, and thus lower m.s.e.
C ONCLUSION 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 ﬂow 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 ﬁnite-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 satisﬁes only a spectral gap condition; and (iii) analysis of the numerical algorithms in the ﬁnite- N and in the ﬁnite ∆ 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.
Speciﬁcally, 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
[Equation]
where u t : R d → R d is now a time-varying vector ﬁeld.
The Lagrangian L : R + × P ac , 2 ( R d ) ×
L 2 ( R d ; R d ) → R is deﬁned as:
[Equation]
The optimal control problem is:
Minimize
[Equation]
Subject to
[Equation]
The Hamiltonian function H : R + × P ac , 2 ( R d ) × C ( R d ; R ) ×
L 2 ( R d ; R d ) → R is H
[Equation]
(24) where φ ∈ C ( R d ; R ) is the dual variable and the inner-product (cid:104)∇ φ, u (cid:105)
L 2 ( ρ ) :
= (cid:82) R d ∇ φ ( x ) · u ( x ) ρ ( x ) d x
B R ESTATEMENT OF THE MAIN RESULT AND ITS PROOF
We restate Theorem 1 below which now includes the pde formulation as well.
Theorem
[Equation]
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:
[Equation]
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:
[Equation]
(iv) Suppose additionally that the functional F is displacement convex and ρ ∞ is its minimizer.
Deﬁne
[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 )) · dd t T ρ ∞ ρ t ( X t )] = 0 .
Then d V d t ( t ) ≤ 0 .
Consequently, the following rate of convergence is obtained along the optimal trajectory
[Equation]
Proof.
(i)
The Hamiltonian function deﬁned in (12) is equal to
[Equation]
after inserting the formula for the Lagrangian.
According to the maximum principle in probabilistic form for (mean-ﬁeld) 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
[Equation]
where ˜ X t , ˜ Y t , ˜
U ∗ t are independent copies of X t , Y t , U ∗ t .
The derivatives
[Equation]
It follows that
[Equation]
where we used the deﬁnition F ( ρ ) = (cid:82) ˜ F ( x, ρ )
ρ ( x ) d x and the identity (Carmona & Delarue, 2017, Sec.
5.2.2 Example 3)
[Equation]
(ii)
The Hamiltonian function deﬁned in (24) is equal to
[Equation]
after inserting the formula for the Lagrangian.
According to the maximum principle for pde formulation of mean-ﬁeld optimal control problems (see (Carmona & Delarue, 2017, Sec.
the optimal control vector ﬁeld is u ∗ t = arg min v H ( t, ρ t , φ t , v ) =
e α t − γ t ∇ φ t and the
Hamilton’s equations are:
[Equation]
inserting the formula u ∗ t =
e α t − γ t ∇ φ t concludes the result.
(iii) Consider the ( ρ t , φ t ) deﬁned from (26).
The distribution ρ t is identiﬁed with a stochastic process ˜ X t such that d ˜ X t d t =
e α t − γ t ∇ φ t ( ˜ X t ) and Law ( ˜ X t ) =
ρ t .
Then deﬁne ˜
Y t = ∇ φ t ( ˜ X t ) .
Taking the time derivative shows that
[Equation]
with the initial condition
˜ Y
= ∇ φ 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 ( X t , Y t ) with ( ˜ X t , ˜ Y t ) .
(iv) The energy functional
[Equation]
Then the derivative of the ﬁrst term is
[Equation] [Equation]
where ξ ( T ρ ∞ ρ t ( X t )) :
= dd t T ρ ∞ ρ t ( X t ) .
Using the scaling condition ˙ γ t
=
e α t the derivative of the ﬁrst term simpliﬁes to
[Equation]
Upon using the technical assumption, E [( X t + e − γ t Y t − T ρ ∞
[Equation]
t ( X t )) · ξ ( T ρ ∞
[Equation]
t ( X t ))]
=
the derivative of the ﬁrst term simpliﬁes to
[Equation]
The derivative of the second term is
[Equation]
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 ﬁrst and second term yields:
[Equation]
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.
C W ASSERSTEIN GRADIENT AND G ÂTEAUX DERIVATIVE
This section contains deﬁnitions of the Wasserstein gradient and Gâteaux derivative (Ambrosio et al., 2008; Carmona & Delarue, 2017).
Let F : P ac , 2 ( R d ) → R be a (smooth) functional on the space of probability distributions.
Gâteaux derivative:
The Gâteaux derivative of F at ρ ∈ P ac , 2 ( R d ) is a real-valued function on R d denoted as ∂ F ∂ρ ( ρ ) :
R d → R .
It is deﬁned as a function that satisﬁes the identity
[Equation]
x )
[Equation]
x for all path ρ t in P ac , 2 ( R d ) such that ∂ρ t ∂t = −∇ · ( ρ t u ) with ρ 0 = ρ ∈ P ac , 2 ( R d ) .
Wasserstein gradient: The Wasserstein gradient of F at ρ is a vector-ﬁeld on R d denoted as ∇ ρ F ( ρ ) :
R d → R d .
It is deﬁned as a vector-ﬁeld that satisﬁes the identity
[Equation]
for all path ρ t in P ac , 2 ( R d ) such that ∂ρ t ∂t = −∇ · ( ρ t u ) with ρ 0 = ρ ∈ P ac , 2 ( R d ) .
The two deﬁnitions imply the following relationship (Carmona & Delarue, 2017, Prop. 5.48):
[Equation]
Example: Let F ( ρ ) =
(cid:82) log( ρ ( x ) ρ ∞ ( x ) )
ρ ( x ) d x 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 ) .
Then
[Equation]
where the divergence theorem is used in the last step.
The deﬁnitions of the Gâteaux derivative and
Wasserstein gradient imply
[Equation]
D R ELATIONSHIP WITH THE UNDER - DAMPED L ANGEVIN EQUATION
A basic form of the under-damped (or second order)
Langevin equation is given in Cheng et al.
[Equation]
(28) where { B t } t ≥ 0 is the standard Brownian motion.
Consider next, the the accelerated ﬂow (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.
The sdes (28) and (29) are similar except that the stochastic term √ 2 d B 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:
[Equation]
where ρ t ( x ) = (cid:82) q t ( x, v ) d v is the marginal of q t ( x, y ) on x .
The ﬁnal term in the Fokker-Planck equations are clearly different.
The joint distributions are different as well.
The situation is in contrast to the ﬁrst order Langevin equation, where the stochastic term √ 2 d B 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]
Let p t ( x ) denote the distribution of X t of (30) and let q t ( x ) denote the distribution of X t of (31).
The corresponding Fokker-Planck equations are as follows
[Equation]
where we used ρ t ( x )
=
q t ( x ) .
In particular, this implies that the marginal probability distribution of the stochastic process X t are the same for ﬁrst order Langevin sde (30) and (31) .