NATURAL GRADIENT VIA OPTIMAL TRANSPORT I

arXiv:1803.07033v1 [math.OC] 16 Mar 2018

WUCHEN LI AND GUIDO MONTÚFAR

Abstract. We study a natural Wasserstein gradient flow on manifolds of probability
distributions with discrete sample spaces. We derive the Riemannian structure for the
probability simplex from the dynamical formulation of the Wasserstein distance on a
weighted graph. We pull back the geometric structure to the parameter space of any given
probability model, which allows us to define a natural gradient flow there. In contrast to
the natural Fisher-Rao gradient, the natural Wasserstein gradient incorporates a ground
metric on sample space. We discuss implementations following the forward and backward
Euler methods. We illustrate the analysis on elementary exponential family examples.

1. Introduction
The statistical distance between histograms plays a fundamental role in statistics and
machine learning. It provides the geometric structure on statistical manifolds [3]. Learning problems usually correspond to minimizing a loss function over these manifolds. An
important example is the Fisher-Rao metric on the probability simplex, which has been
studied especially within the field of information geometry [3, 6]. A classic result due to
Chentsov [10] characterizes this Riemannian metric as the only one, up to scaling, that
is invariant with respect to natural statistical embeddings by Markov morphisms (see
also [8, 18, 28]). Using the Fisher-Rao metric, a natural Riemannian gradient descent
method is introduced [2]. This natural gradient has found numerous successful applications in machine learning (see, e.g., [1, 32, 37, 23, 31]).
Optimal transport provides another statistical distance, named Wasserstein or Earth
Mover’s distance. In recent years, this metric has attracted increasing attention within
the machine learning community [5, 15, 27, 33]. One distinct feature of optimal transport
is that it provides a distance among histograms that incorporates a ground metric on
sample space. In particular, the L2 -Wasserstein distance has a dynamical formulation,
which exhibits a metric tensor structure. The set of probability densities with this metric
forms an infinite-dimensional Riemannian manifold, named density manifold [17]. The
gradient descent method in the density manifold, called Wasserstein gradient flow, has
been widely studied in the literature; see [30, 35] and references.
A question intersecting optimal transport and information geometry arises: What is
the natural Wasserstein gradient descent method on the parameter space of a statistical
model? In optimal transport, the Wasserstein gradient flow is studied on the full space
of probability densities, and shown to have deep connections with the ground metrics on
Key words and phrases. Optimal transport; Information geometry; Wasserstein statistical manifold;
Machine learning.
1

2

LI, MONTÚFAR

sample space deriving from physics [29], fluid mechanics [9] and differential geometry [21].
We expect that these relations also exist on parametrized probability models, and that
the Wasserstein gradient flow can be useful in the optimization of objective functions that
arise in machine learning problems. By incorporating a ground metric on sample space,
this method can serve to implement useful priors in the learning algorithms.
We are interested in developing synergies between the information geometry and optimal
transport communities. In this paper we take a natural first step in this direction. We
introduce the Wasserstein natural gradient flow on the parameter space of probability
models with discrete sample spaces. The L2 -Wasserstein metric on discrete states was
introduced in [11, 22, 25]. Following the settings from [12, 13, 16, 19], the probability
simplex forms the Riemannian manifold called Wasserstein probability manifold. The
Wasserstein metric on the probability simplex can be pulled back to the parameter space
of a probability model. This metric allows us to define a natural Wasserstein gradient
method on parameter space.
In the literature one finds several formulations of optimal transport for continuous
sample spaces. On the one hand, there is the static formulation, known as Kantorovich’s
linear programming [35]. Here, the linear program is to find the minimal value of a
functional over the set of joint measures with given marginal histograms. The objective
functional is given as the expectation value of the ground metric with respect to a joint
probability density measure. On the other hand, there is the dynamical formulation,
known as the Benamou-Brenier formula [7]. This dynamic formulation gives the metric
tensor for measures by lifting the ground metric tensor of sample spaces.
Both static and dynamic formulations are equivalent in the case of continuous state
spaces. However, the two formulations lead to different metrics in the simplex of discrete
probability distributions. The major reason for this difference is that the discrete sample
space is not a length space.1 Thus the equivalence result in classical optimal transport is no
longer true in the setting of discrete sample spaces. We note that for the static formulation,
it defines the distance for the discrete probability simplex, in which the gradient operator is
a metric derivative depending on the dual variable of linear programming. In this setting,
there is no Riemannian metric tensor for the discrete probability simplex. See [13, 22] and
the Appendix for a more detailed discussion.
The connection between optimal transport and information geometry has been studied in [4, 36]. These works mainly use the static formulation of optimal transport on
discrete sample spaces. Many connections between static optimal transport and information geometry have been found. In contrast to these works, we focus on the dynamical
formulation of optimal transport in order to define a Riemannian metric tensor on the discrete probability simplex. Other works have studied the Gaussian family of distributions
with L2 -Wasserstein metric in continuous sample space [26, 34]. Compared to them, we
consider discrete state spaces and arbitrary parametric models.
This paper is organized as follows. In Section 2 we briefly review the Riemannian
manifold structure in probability space introduced by optimal transport in the cases of
1A length space is one in which the distance between points can be measured as the infimum length of

continuous curves between them.

NATURAL GRADIENT VIA OPTIMAL TRANSPORT

3

continuous and discrete sample spaces. In Section 3 we introduce Wasserstein statistical
manifolds by isometric embedding into the probability manifold, and derive the corresponding gradient flows. In Section 4 we discuss a few examples.
2. Optimal transport on continuous and discrete sample spaces
In this section, we briefly review the results of optimal transport. In particular, we
derive the corresponding Riemannian manifold structure for probability distributions with
a discrete support set.
2.1. Optimal transport on continuous sample space. We start with a review of
the optimal transport problem on continuous spaces. This will guide our discussion of
the finite state case. For related studies, we refer the reader to [17, 35] and the many
references therein.
Denote the sample space by (Ω, g Ω ). Here Ω is a finite dimensional smooth Riemannian
manifold, for example Rd or the open unit ball therein. Its inner product is denoted by
g Ω and its volume form by dx. Denote the geodesic distance of Ω by dΩ : Ω × Ω → R+ .
Consider the set P2 (Ω) of Borel measurable probability density functions on Ω with
finite second moment. Given ρ0 , ρ1 ∈ P2 (Ω), the L2 -Wasserstein distance (also called
Wasserstein metric in the optimal transport literature) between ρ0 and ρ1 is denoted by
W : P(Ω) × P(Ω) → R+ . There are two equivalent ways of defining this distance. On
one hand, there is the static formulation. This refers to the following linear programming
problem:
Z
W (ρ0 , ρ1 )2 =

inf

π∈Π(ρ0 ,ρ1 ) Ω×Ω

dΩ (x, y)2 π(dx, dy) ,

(1)

where the infimum is taken over the set Π of joint probability measures on Ω × Ω that
have marginals ρ0 , ρ1 .
On the other hand, the Wasserstein distance W can be written in a dynamical formulation, where a probability path ρ : [0, 1] → P2 (Ω) connecting ρ0 , ρ1 is considered. This
refers to a variational problem known as the Benamou-Brenier formula:
Z 1Z
0 1 2
W (ρ , ρ ) = inf
gxΩ (∇Φ(t, x), ∇Φ(t, x))ρ(t, x)dxdt ,
(2a)
Φ

0

Ω

where the infimum is taken over the set of Borel potential functions [0, 1] × Ω → R. Each
potential function Φ determines a corresponding density path ρ as the solution of the
continuity equation
∂ρ(t, x)
+ div(ρ(t, x)∇Φ(t, x)) = 0 , ρ(0, x) = ρ0 (x) , ρ(1, x) = ρ1 (x) .
(2b)
∂t
Here div and ∇ are the divergence and gradient operators in Ω. The continuity equation
is well known in physics.
The equivalence of the static (1) and dynamical (2) formulations is well known. For
the reader’s convenience we give a sketch of proof in the appendix. In this paper we
focus on the variational formulation (2). In fact, this formulation entails the definition

4

LI, MONTÚFAR

of a Riemannian structure as we now discuss. For simplicity, we only consider the set of
smooth and strictly positive probability densities
Z
n
o
∞
P+ (Ω) = ρ ∈ C (Ω) : ρ(x) > 0 ,
ρ(x)dx = 1 ⊂ P2 (Ω) .
Ω

C ∞ (Ω)

Denote F(Ω) :=
of P+ (Ω) is given by

the set of smooth real valued functions on Ω. The tangent space

Z
n
o
Tρ P+ (Ω) = σ ∈ F(Ω) :
σ(x)dx = 0 .
Ω

Given Φ ∈ F(Ω) and ρ ∈ P+ (Ω), define
VΦ (x) := −div(ρ(x)∇Φ(x)) .
We assume the zero flux condition
Z
VΦ (x)dx = 0 .
Ω

In
of the continuity equation, the zero flux condition is equivalent to requiring that
R view
∂ρ
dx
= 0, which means that the space integral of ρ is always 1. When Ω is compact
Ω ∂t
without boundary, this is automatically satisfied. This is also true when Ω = Rd and ρ
has finite second moment. Thus VΦ ∈ Tρ P+ (Ω). The elliptic operator ∇ · (ρ∇) identifies
the function Φ on Ω modulo additive constants with the tangent vector VΦ of the space of
densities (for more details see [17, 21]). This gives an isomorphism
F(Ω)/R → Tρ P+ (Ω);

Φ 7→ VΦ .

Define the Riemannian metric (inner product) on the tangent space of positive densities
g W : Tρ P+ (Ω) × Tρ P+ (Ω) → R by
Z
W
gρ (VΦ , VΦ̃ ) =
gxΩ (∇Φ(x), ∇Φ̃(x))ρ(x)dx ,
Ω

where Φ(x), Φ̃(x) ∈ F(Ω)/R. The inner product endows P+ (Ω) with an infinite dimensional Riemannian metric tensor. In other words, the variational problem (2) is a
geometric action energy [7, 21] in (P+ (Ω), g W ). In literature [17], (P+ (Ω), g W ) is called
density manifold.
2.2. Dynamical optimal transport on discrete sample spaces. We translate the
dynamical perspective from the previous section to discrete state spaces, i.e., we replace
the continuous space Ω by the discrete space I = {1, · · · , n}.
To encode the metric tensor of discrete states, we first need to introduce a ground
metric notion on sample space. We do this in terms of a graph with weighted edges,
G = (I, E, ω), where I is the vertex set, E is the edge set, and ω = (ωij )i,j∈I ∈ Rn×n are
the edge weights. These weights satisfy
(
ωji > 0, if (i, j) ∈ E
ωij =
.
0,
otherwise

NATURAL GRADIENT VIA OPTIMAL TRANSPORT

5

As mentioned above, the weights encode the ground metric on the discrete state space.
More precisely, we write
1
ωij = G 2 , if (i, j) ∈ E ,
(3)
(dij )
where dG
ij represents the distance between states i and j. The set of neighbors or adjacent
vertices of i is denoted by N (i) = {j ∈ V : (i, j) ∈ E}.
The probability simplex supported on the vertices of G is defined by
n
n
o
X
n
P(I) = (p1 , · · · , pn ) ∈ R :
pi = 1 , pi ≥ 0 .
i=i

Here p = (p1 , . . . , pn ) is a probability vector with coordinates pi corresponding to the
probabilities assigned to each node i. We denote the interior of the probability simplex by
P+ (I). This consists of the strictly positive probability distributions, with pi > 0, i ∈ I.
Next we introduce the variational problem (2) on discrete states. First we need to
define the “metric tensor” on graphs. A vector field v = (vij )i,j∈V ∈ Rn×n on G is a
skew-symmetric matrix:
(
−vji , if (i, j) ∈ E
.
vij =
0,
otherwise
A potential function Φ = (Φi )ni=1 ∈ Rn defines gradient vector field ∇G Φ = (∇G Φij )i,j∈V ∈
Rn×n on the graph G by the finite differences
(√
ωij (Φi − Φj ) if (i, j) ∈ E
.
∇G Φij =
0
otherwise
√
Here we use ω rather than the more natural 1/d. This is for the simplicity of notations.
So that we can represent the gradient, divergence and Laplacian matrix in a multiplicity
of weight, instead of dividing the ground metric.
We define an inner product of tangent vectors at each state i ∈ I by
1 X 2
giI (v, v) :=
vij .
2
j∈N (i)

In particular, the gradient vector field ∇G Φ defines a kinetic energy at each state i ∈ I by
1 X
giI (∇G Φ, ∇G Φ) :=
(Φi − Φj )2 ωij
2
j∈N (i)

We next define the expectation value of kinetic energy with respect to the probability
distribution p:
X
pi + pj
1 X
(∇G Φ, ∇G Φ)p :=
pi giI (∇G Φ, ∇G Φ) =
ωij (Φi − Φj )2
.
2
2
i∈I

(i,j)∈E

The latter can also be written as
n
X
X

pi + pj
(∇G Φ, ∇G Φ)p =
Φi
ωij (Φi − Φj )
= ΦT − divG (p∇G Φ) ,
2
i=1

j∈N (i)

6

LI, MONTÚFAR

where
− divG (p∇G Φ) :=

 X

ωij (Φi − Φj )

j∈N (i)

pi + pj 
.
2
i∈I

(4)

There are two definitions hidden in (4). First, divG : Rn×n → Rn maps any given vector
field m on the graph to
X √
divG (m) = (
ωij mji )ni=1 .
j∈N (i)

Second, the probability weight m = p∇G Φ of the gradient vector field ∇G Φ is defined by
(
√
pi +pj
2 (Φi − Φj ) ωij , if (i, j) ∈ E ,
mij =
0,
otherwise
where

pi +pj
2

represents the probability weight on the edge (i, j) ∈ E.

Now we are ready to introduce the L2 -Wasserstein metric on P+ (I).
Definition 1. For any p0 , p1 ∈ P+ (I), define the Wasserstein distance W : P+ (I) ×
P+ (I) → R by
nZ 1
o
0 1 2
W (p , p ) := inf
(∇G Φ(t), ∇G Φ(t))p(t) dt .
p(t),Φ(t)

0

Here the infimum is taken over pairs (p(t), Φ(t)) with p ∈ H 1 ((0, 1), Rn ) and Φ : [0, 1] → Rn
measurable, satisfying
d
p(t) + divG (p(t)∇G Φ(t)) = 0, p(0) = p0 , p(1) = p1 .
dt
Remark 1. It is worth mentioning that the metric given in Definition 1 is different from
the metric defined by linear programming. In other words, denote the distance dG (i, j)
between two vertices i and j is the length of a shortest (i, j)-path. If (i, j) ∈ E, then
dG (i, j) is same as the ground metric defined in (3). Then
n
n
n X
o
X
X

2
0
0 1 2
W (p , p ) 6≡ min
dG (i, j) πij :
πij = pj ,
πij = p1i , πij ≥ 0 . (5)
π

1≤i,j≤n

i=1

j=1

The reason for this in-equivalence is that discrete sample space I is not a length space. In
other words, there is no continuous path in I connecting two nodes in I. For more details
see discussions in the appendix.
2.3. Wasserstein manifold of discrete probability distributions. In this section we
introduce the primal coordinates of the discrete probability simplex with L2 Wasserstein
Riemannian metric. Our discussion follows the recent work [19]. The probability simplex
P(I) is a manifold with boundary. To simplify the discussion, we focus on the interior
P+ (I). The geodesic properties on the boundary ∂P(I) have been studied in [16].
Let us focus on the Riemannian structure. In the following we introduce an inner
product on the tangent space
n
n
o
X
Tp P+ (I) = (σi )ni=1 ∈ Rn :
σi = 0 .
i=1

NATURAL GRADIENT VIA OPTIMAL TRANSPORT

7

Denote the space of potential function on I by F(I) = Rn . Consider the quotient space
F(I)/R = {[Φ] | (Φi )ni=1 ∈ Rn } ,
where [Φ] = {(Φ1 +c, · · · , Φn +c) | c ∈ R} are functions defined up to addition of constants.
We introduce an identification map via (4)
V : F(I)/R → Tp P+ (I) ,

VΦ = −divG (p∇G Φ) .

In [11] it is shown that VΦ : F(I)/R → Tp P+ (I) is a well defined map, linear, and one
to one. I.e., F(I)/R ∼
= Tp∗ P+ (I), where Tp∗ P+ (I) is the cotangent space of P+ (I). This
identification induces the following inner product on Tp P+ (I).
We first present this in a dual formulation, which is known in the literature [21].
Definition 2 (Inner product in dual coordinates). The inner product gpW : Tp P+ (I) ×
Tp P+ (I) → R takes any two tangent vectors σ1 = VΦ1 and σ2 = VΦ2 ∈ Tp P+ (I) to
gpW (σ1 , σ2 ) = σ1T Φ2 = σ2T Φ1 = (∇G Φ, ∇G Φ)p .

(6)

The above is written in the dual coordinates. We shall now give the inner product
in primal coordinates. The following matrix operator will be the key to the Riemannian
metric tensor of (P+ (I), g W ).
Definition 3 (Linear weighted Laplacian matrix). Given a weighted graph G = (I, E, ω),
I = {1, . . . , n}, the matrix function L(·) : Rn → Rn×n is defined by
L(a) = DT Λ(a)D ,

a = (ai )ni=1 ∈ Rn ,

where
• D ∈ R|E|×n is the discrete gradient operator
√

if i = k, i > j
 ωij ,
√
D(i,j)∈E,k∈V = − ωij , if j = k, i > j ,


0,
otherwise
• −DT ∈ Rn×|E| is the discrete divergence operator, also named oriented incidence
matrix [14], and
• Λ(a) ∈ R|E|×|E| is a weight matrix depending on a,
(
ai +aj
if (i, j) = (k, l) ∈ E
2
Λ(a)(i,j)∈E,(k,l)∈E =
.
0
otherwise
Consider some p ∈ P+ (I). From spectral graph theory [14], we know that L(p) can be
decomposed as


0
 λ1 (p)



L(p) = U (p) 
 U (p)T .
..


.
λn−1 (p)

8

LI, MONTÚFAR

Here 0 < λ1 (p) ≤ · · · ≤ λn−1 (p) are the eigenvalues of L(p) in ascending order, and
U (p) = (u0 (p), u1 (p), · · · , un−1 (p)) is the corresponding orthogonal matrix of eigenvectors
with
1
u0 = √ (1, · · · , 1)T .
n
We write L(p)† for the pseudo-inverse of L(p), i.e.,


0
1


λ1 (p)


†
L(p) = U (p) 
 U (p)T .
..


.
1
λn−1 (p)

With σi = L(p)Φi , i = 1, 2, we see that
†
T
σ1 T L(p)† σ2 = ΦT
1 L(p)L(p) L(p)Φ2 = Φ1 L(p)Φ2 = (∇G Φ1 , ∇G Φ2 )p .

Now we are ready to give the inner product in primal coordinates.
Definition 4 (Inner product in primal coordinates). The inner product gpW : Tp P+ (I) ×
Tp P+ (I) → R is defined by
gpW (σ1 , σ2 ) := σ1 T L(p)† σ2 ,

for any σ1 , σ2 ∈ Tp P+ (I) .

In other words, the variational problem from Definition 1 is a minimization of geometry
energy functional in P+ (I), i.e.,
nZ 1
o
W (p0 , p1 )2 =
inf
ṗ(t)T L(p(t))† ṗ(t)dt : p(0) = p0 , p(1) = p1 .
p(t)∈P+ (I)

0

This defines a Wasserstein Riemannian strucure on the probability simplex. For more
details of Riemannian formulas see [19]. Following [17] we could call (P+ (I), g W ) discrete
density manifold. However, this could be easily confused with other notions from information geometry, and hence we will use the more explicit terminology Wasserstein discrete
probability manifold, or Wasserstein manifold for short.
3. Natural gradient for the Wasserstein manifold
In this section we study parametric probability models endowed with the L2 -Wasserstein
Riemannian metric. We define this in the natural way, by pulling back the Riemannian
structure from the Wasserstein manifold that we discussed in the previous section. This
allows us to introduce a natural gradient flow on the parameter space of statistical models.
3.1. Wasserstein statistical manifold. Consider a statistical model defined by a triplet
(Θ, I, p). Here, I = {1, · · · , n} is the sample space, Θ is the parameter space, which is an
open subset of Rd , d ≤ n − 1, and p : Θ → P+ (I) is the parametrization function,
p(ξ) = (pi (ξ))ni=1 ,

ξ∈Θ.

In the sequel we will assume that rank(Jξ p(ξ)) = d, so that the parametrization is locally
injective.

NATURAL GRADIENT VIA OPTIMAL TRANSPORT

9

We define a Riemannian metric g on Θ as the pull-back of metric g W on P+ (I). In
other words, we require that p : (Θ, g) → (P+ (I), g W ) is an isometric embedding:
W
(dp(ξ)(a), dp(ξ)(b))
gξ (a, b) :=gp(ξ)
T

= dp(ξ)(a) L(p(ξ))† dp(ξ)(b) .
Pn ∂pi (ξ) n
Since dp(ξ)(a) =
j=1 ∂ξj aj i=1 = Jξ p(ξ)a, we arrive at the following definition.

Definition 5. For any pair of tangent vectors a, b ∈ Tξ Θ = Rd , define
T
gξ (a, b) := aT Jξ p(ξ) L(p(ξ))† (Jξ p(ξ))b ,
where Jξ (p(ξ)) = ( ∂p∂ξi (ξ)
)1≤i≤n,1≤j≤d ∈ Rn×d is the Jacobi matrix of the parametrization
j
p, and L(p(ξ)) ∈ Rn×n is the pseudo-inverse of the linear weighted Laplacian matrix.
This inner product is consistent with the restriction of the Wasserstein metric g W to
p(Θ). For this reason, we call p(Θ), or (Θ, I, p), together with the induced Riemannian
metric g, Wasserstein statistical manifold.
We need to make sure that the embedding procedure is valid, because the metric tensor
L(p)† is only of rank n − 1. The next lemma shows that (Θ, g) is a well defined ddimensional Riemannian manifold.
Lemma 6. For any ξ ∈ Θ, we have
λmin (ξ) =

inf

a∈Rd ,kak2 =1

gξ (a, a) > 0 .

In addition, gξ is smooth as a function of ξ, so that (Θ, g) is a smooth Riemannian
manifold.
T

Proof. We only need to show that Jξ p(ξ) L(p(ξ))† (Jξ p(ξ) ∈ Rd×d is a positive definite
matrix. Consider
T

aT Jξ p(ξ) L(p(ξ))† (Jξ p(ξ) a = 0 ,
where 0 ∈ Rn−1 . Since L(p) only has one simple eigenvalue 0 with eigenvector u0 , then
for some constant c ∈ R1 .
Pn
Pn
T ∂p(ξ)
√1
√1
Since uT
0 p(ξ) = n
i=1 pi (ξ) = 0, we have that u0 ∂ξj = n
i=1
Jξ p(ξ)a = cu0 ,

(7)
∂pi (ξ)
∂ξj

= 0, i.e.,

uT
0 Jξ p(ξ) = 0 .
Left multiply u0 into (11), we obtain
T
0 = uT
0 Jξ p(ξ)a = cu0 u0 = c .

Thus c = 0, and (11) forms
Jξ p(ξ)a = 0 .
Since rank(Jξ p(ξ)) = d < n, we have a = 0, which finishes the proof.



10

LI, MONTÚFAR

We briefly illustrate some geometric calculations on (Θ, g). Given ξ0 , ξ1 ∈ Θ, the
Riemannian distance on (Θ, g) is given by
nZ 1
o
2
T
˙
˙
Dist(ξ0 , ξ1 ) =
inf
ξ G(ξ)ξ : ξ(0) = ξ0 , ξ(1) = ξ1 .
ξ(t)∈C 1 ((0,1);Θ)

0

where the metric tensor is given by
G(ξ) = (Jξ p(ξ))T L(p(ξ))† (Jξ p(ξ)) .
Following the standard Riemannian geometry, the cotangent geodesic flow forms
(
ξ˙ − G(ξ)−1 S = 0
.
∂ T
S G(ξ)−1 S = 0
Ṡ + 12 ∂ξ

(8)

(9)

˙ It is worth recalling the following facts. If p
Here S is the Legendre transformation of ξ.
is an identity map, then (9) translates to
(
ṗ + divG (p∇G S) = 0
P
.
Ṡ + 41 j∈N (i) (∇G S)2 = 0
In addition, if I = Ω and we replace i by x and pi (t) by ρ(t, x), the above becomes
(
∂ρ(t,x)
+ div(ρ(t, x)∇S(t, x)) = 0
∂t
,
∂S(t,x)
+ 12 (∇S(t, x))2 = 0
∂t
which are the standard continuity and Hamilton-Jacobi equations on Ω. For these reasons,
we call the two equations in (9) the continuity equation and the Hamilton-Jacobi equation
in parameter space.
3.2. Gradient flow on Wasserstein statistical manifold. In this section we introduce
the natural Riemannian gradient flow on (Θ, g).
Consider a loss function F : P+ (I) → R. Thus we focus on the composition F ◦ p : Θ →
R. The Riemannian gradient of F (p(ξ)) is defined as follows. Given some ∇g F (p(ξ)) ∈
Tξ Θ, we have
gξ (∇g F (p(ξ)), a) = dξ F (p(ξ)) · a , for any a ∈ Tξ Θ ,
P
where dξ F (p(ξ)) · a = di=1 ∂ξ∂ i F (p(ξ))ai . The gradient flow satisfies

(10)

dξ
= −∇g F (p(ξ)) .
dt
The next theorem establishes an explicit formulation of the gradient flow.
Theorem 7 (Wasserstein gradient flow). The gradient flow of a functional F : P+ (I) → R
is given by
dξ
= −G(ξ)−1 ∇ξ F (p(ξ)) ,
dt
where ∇ξ is the Euclidean gradient of F (p(ξ)) w.r.t. ξ. More explicitly,

†
dξ
= − (Jξ p(ξ))T L(p(ξ))† Jξ p(ξ) (Jξ p(ξ))T ∇p F (p(ξ) ,
(11)
dt
where ∇p is the Euclidean gradient of F (p) w.r.t. p.

NATURAL GRADIENT VIA OPTIMAL TRANSPORT

11

Proof. The proof follows directly from (10). Notice that
T

gξ (∇g F (p(ξ)), a) = ∇g F (p(ξ))T Jξ p(ξ) L(p(ξ))† (Jξ p(ξ) a = ∇ξ F (p(ξ))T a ,
T

and Jξ p(ξ) L(p(ξ))† (Jξ p(ξ) is an invertible matrix. Hence
†
∇g F (p(ξ)) = (Jξ p(ξ))T L(p(ξ))† Jξ p(ξ) dξ F (p(ξ)) .
We compute dξ F (p(ξ)) as
dξ F (p(ξ)) =

n
X
n
∂pj (ξ) n
∂
∂
F (p(ξ)) i=1 =
F (p(ξ)) ·
= (Jξ p(ξ))T ∇p F (p(ξ)) .
∂ξi
∂pj
∂ξi i=1
j=1

This concludes the proof of (11).



Equation (11) is the generalization of the Wasserstein gradient flow the probability
simplex to the Wasserstein gradient flow on parameter space. If p is an identity map
p(ξ) = ξ, ξ = (pi )ni=1 , with the parameter space Θ equal to the entire probability simplex,
then (11) is
dp
= −L(p)∇p F (p) = divG (p∇G dp F (p)) ,
dt
which is the Wasserstein gradient flow on the discrete probability simplex. In particular,
if I = Ω, then it represents
∂ρ(t, x)
= divx (ρ(t, x)∇x dρ F(ρ)) ,
∂t
which is the Wasserstein gradient flow on Ω. From now on, we call (11) the Wasserstein
gradient flow in parameter space.
In fact, the definition of Wasserstein gradient flow shares many similarities with the
steepest gradient descent defined as follows. Consider
arg min F (p(ξ + ∆ξ))
∆ξ

s.t.

1
W (p(ξ), p(ξ + ∆ξ))2 =  .
2

(12)

By taking the second-order Taylor approximation of the Wasserstein distance at ξ, we get
W (p(ξ), p(ξ + ∆ξ))2 = ∆ξ T G(ξ)∆ξ + o(∆ξ 2 ) ,
where G(ξ) is the metric tensor of (Θ, g) defined in (8), inherited from Wasserstein manifold. We take the first-order approximation of F (p(ξ + ∆ξ)) in (12) by
arg min F (p(ξ)) + ∆ξ T ∇ξ F (p(ξ))
∆ξ

s.t.

1 T
∆ξ G(ξ)∆ξ =  .
2

By the Lagrangian method with Lagrange multiplier λ, we have
∆ξ = λG(ξ)−1 ∇ξ F (p(ξ)) .
The above derivations lead to the Wasserstein natural gradient
∇g F (p(ξ)) = G(ξ)−1 ∇ξ F (p(ξ)) .

12

LI, MONTÚFAR

Remark 2. In the standard natural gradient [2], we replace (12) by
arg min F (p(ξ + ∆ξ))
∆ξ

s.t.

KL(p(ξ)kp(ξ + ∆ξ)) =  ,

where KL represents the Kullback-Leibler divergence (relative entropy) from p(ξ) to
p(ξ + ∆ξ). Our definition just changes the KL-divergence by the Wasserstein metric
on a specified graph.
3.3. Numerical methods. Given the gradient flow (11), there are two standard choices
of time discretization, namely the forward Euler scheme and the backward Euler scheme.
Denote the step size by h > 0. The forward Euler method of (11) computes a discretized
trajectory by
ξ k+1 = ξ k − h∇g F (p(ξ k )) ,
while the backward Euler method computes
ξ k+1 = arg min F (p(ξ)) +
ξ∈Θ

Dist(ξ, ξ k )2
,
2h

where Dist is the geodesic distance in parameter space (Θ, g).
In the information geometry literature, the forward Euler method is often referred to as
natural gradient method. In Wasserstein geometry, the backward Euler method in often
called the Jordan-Kinderlehrer-Otto (JKO) scheme. In the following we give pseudo code
for both numerical methods.

Natural Wasserstein gradient method
1.
2.

for k = 1, 2, . . . while not converged
Choose a suitable step size hk > 0 ;
†
ξ k+1 = ξ k − hk (Jξ p(ξ k ))T L(p(ξ k ))† Jξ p(ξ k ) (Jξ p(ξ k ))T ∇p F (p(ξ k )) ;
end

Natural Jordan-Kinderlehrer-Otto scheme
1.
2.

for k = 1, 2, . . . while not converged
Choose a suitable adaptive step size hk > 0 ;
k )2
ξ k+1 = arg minξ∈Θ F (p(ξ)) + Dist(ξ,ξ
;
2hk
end

In practice, the forward Euler method is usually easier to implement than the backward
Euler method. We would also suggest to implement the natural Wasserstein gradient using
this method for minimization problems in machine learning.

NATURAL GRADIENT VIA OPTIMAL TRANSPORT

13

As known in optimization, the JKO scheme can also be useful for non-smooth objective
functions. Moreover, the backward Euler method is usually unconditionally stable, which
means that one can choose a large step size h for computations.
4. Examples
In this section, we discuss the Wasserstein metric and gradient flow for two exponential
families of probabilities.
Example 1 (Wasserstein geodesics). Consider the sample space I = {1, 2, 3} with an
unweighted graph 1 − 2 − 3. The probability simplex for this sample space is a triangle
in R3 :
3
n
o
X
3
3
P(I) = (pi )i=1 ∈ R :
pi = 1, pi ≥ 0 .
i=1

Following Definition 1, the L2 -Wasserstein distance is given by
Z
p1 (t) + p2 (t)
p2 (t) + p3 (t)
1 1
0 1 2
(Φ1 (t) − Φ2 (t))2
+ (Φ2 (t) − Φ3 (t))2
dt,
W (p , p ) := inf
2
2
Φ(t) 2 0
(13)
where the infimum is taken over paths Φ : [0, 1] → R3 . Each Φ defines p : [0, 1] → R3 as
the solution of the differential equation

p1 +p2

ṗ1 = (Φ1 − Φ2 ) 2
2
3
ṗ2 = (Φ2 − Φ1 ) p1 +p
+ (Φ2 − Φ3 ) p2 +p
2
2


p2 +p3
ṗ3 = (Φ3 − Φ2 ) 2
with boundary condition p(0) = p0 , p(1) = p1 .
Consider local coordinates in (13). We parametrize a probability vector as p = (p1 , 1 −
p1 − p3 , p3 ), with parameters (p1 , p3 ). Then (13) can be written as
n Z 1 ṗ (t)2
o
ṗ3 (t)2
1
0 1 2
W (p , p ) :=
inf
+
dt .
(14)
1 − p1 (t)
p(t) : p(0)=p0 , p(1)=p1
0 1 − p3 (t)
This can be solved numerically2for different choices of the boundary conditions.
We fix three distributions
1
1
1
q 1 = (6, 1, 1), q 2 = (1, 6, 1), q 3 = (1, 1, 6)
8
8
8
and solve (14) for three choices of the boundary conditions:
p0 = q 1 , p1 = q 2 ;

p0 = q 1 , p1 = q 3 ;

p0 = q 2 , p1 = q 3 .

(15)

(16)

This gives us a geodesic triangle between q 1 , q 2 , q 3 , which is illustrated in Figure 1. It
can be seen that (P+ (G), W ) has a non Euclidean geometry. Moreover, we see that the
geodesics depend on the graph structure on sample space, where state 2 is qualitatively
differently from states 1 and 3.
2We use the direct method, which is a standard technique in optimal control. here the time is discretized, and the sum replacing the integral is minimized by means of gradient descent with respect to
(p(t)i )i=1,3,t∈{t1 ,...,tN } ∈ R2×N . A standard reference for these techniques is [20].

14

LI, MONTÚFAR

1

1

1

q

q

p

θ

q3
3

q2

q3

q2
2

3

2

1

2

3
2

2

1.5

1.5

1

1

0.5

0.5

0

0

2

1

-0.5

-0.5

-1

-1

-1.5
-2
2

3

3

1

1

-1.5
2

-2

2

-2 -1.5 -2-1 -1.5
-0.5 -10 -0.5
0.5 01 0.5
1.5 12 1.5
1

2

1

Figure 1. The Wasserstein geodesic triangle from Example 1 plotted in
the probability simplex (left) and in exponential parameter space (right).
The path connecting q 1 and q 3 bends towards q 2 ; something that does not
happen for the other two paths. This illustrates how, as a result of the
ground metric on sample space, state 2 is treated differently from 1 and 3.
We can make the same derivations in terms of an exponential parametrization. Consider
the parameter space Θ = {θ = (θ1 , θ2 ) ∈ R2 } and the parametrization p : Θ → P+ (I) with
p1 (θ) =

eθ1

e θ1
,
+ eθ2 + 1

p3 (θ) =

eθ1

eθ2
,
+ eθ2 + 1

p2 (θ) = 1−p1 (θ)−p3 (θ) =

eθ1

1
.
+ eθ2 + 1

We rewrite the Wasserstein metric (14) in terms of θ. Denote p(θk ) = pk , k = 0, 1. Then
the Wasserstein metric in the coordinate system θ is
!
1
o
nZ 1
0
Dist(θ0 , θ1 )2 =
inf
θ̇T Jθ (p1 , p3 )T 1−p3 (θ)
J
(p
,
p
)
θ̇dt
.
1
3
θ
1
0
θ(t) : θ(0)=θ0 , θ(1)=θ1
0
1−p1 (θ)
The resulting geodesic triangle in Θ is plotted in the right panel of Figure 1.
For comparison, we compute the exponential geodesic triangle between the same distributions q 1 , q 2 , q 3 . This is shown in Figure 2. In this case, there is no distinction between
the states 1, 2, 3 and the three paths are symmetric.
P The exponential geodesic between two
distributions p0 and p1 is given by (p0 )1−t (p1 )t / x (p0 )1−t (p1 )t , t ∈ [0, 1].
Example 2 (Wasserstein gradient flow on an independence model). We next illustrate
the Wasserstein gradient flow over the independence model of two binary variables. The
sample space is I = {−1, +1}2 . For simplicity, we denote the states by a = (−1, −1),
b = (−1, +1), c = (+1, −1), d = (+1, +1). We consider the square graph
b−d
|
|
a−c

NATURAL GRADIENT VIA OPTIMAL TRANSPORT

1

1

q

q

p

θ

q3
3

q2

q3

q2
2

3

2

2

1.5

1.5

1

1

0.5

0.5

0

0

2

1

2

1

-0.5

-0.5

-1

-1

-1.5

15

3

3

1

1

-1.5
2

-2
2

-2

2

-2 -1.5 -2-1 -1.5
-0.5 -10 -0.5
0.5 01 0.5
1.5 12 1.5
1

2

1

Figure 2. Exponential geodesic triangle plotted in the probability simplex
(left) and in exponential parameter space (right). Exponential geodesics
correspond to straight lines in exponential parameter space.
with vertices I, edges E = {{a, b}, {b, d}, {a, c}, {c, d}}, and weights ω = (ωab , ωbd , ωac , ωcd ) ∈
RE attached to the edges. The edge weights correspond to the inverse squared ground metric that we assign to the sample space I. The probability simplex for this sample space is
the tetrahedron
n
o
X
P(I) = (p(x))x∈I ∈ R4 :
p(x) = 1, p(x) ≥ 0 .
x∈I

Following Definition 4, the Wasserstein metric tensor is given by gpW = L(p)† , which is
the inverse of the linear weighted Laplacian metric L from Definition 3. In this example
the latter is
 pa +pb
pa +pc
pa +pb
pa +pc
ωab


L(p) = 

2

+ ωac

b
−ωab pa +p
2
c
−ωac pa +p
2

0

2

−ωab 2
b
d
ωab pa +p
+ ωbd pb +p
2
2
0
d
ωbd pb +p
2

−ωac 2
0
c
d
ωac pa +p
+
ωcd pc +p
2
2
pc +pd
−ωcd 2

0



d
−ωbd pb +p

2

d
−ωcd pc +p
2
pb +pd
pc +pd
ωbd 2 + ωcd 2

The independence model consist of the joint distributions that satisfy p(x1 , x2 ) = p(x1 )p(x2 ).
This can be parametrized in terms of Θ = {ξ = (ξ1 , ξ2 ) ∈ [0, 1]2 }, where ξ1 = p1 (x1 = +1),
ξ2 = p2 (x2 = +1) describe the marginal probability distributions. The parametrization
p : Θ → P(I) is then

(1 − ξ1 )(1 − ξ2 ) if (x1 , x2 ) = (−1, −1)



(1 − ξ )ξ
if (x1 , x2 ) = (−1, +1)
1 2
p(ξ)(x1 , x2 ) =
.

ξ1 (1 − ξ2 )
if (x1 , x2 ) = (+1, −1)



ξ1 ξ2
if (x1 , x2 ) = (+1, +1)
The model p(Θ) ⊂ P(I) is a two dimensional manifold. The parameter space Θ inherits
the Riemannian structure g W from P(I), which is computed as follows. Denote the Jacobi

.

16

LI, MONTÚFAR

matrix of the parametrization by



−(1 − ξ2 ) −(1 − ξ1 )
 −ξ2
1 − ξ1 
 ∈ R4×2 .
Jξ p(ξ) = 
 1 − ξ2
−ξ1 
ξ2
ξ1

Then g W induces a metric tensor on Θ given by

G(ξ) = Jξ p(ξ)T L(p(ξ))† Jξ (p(ξ)) ∈ R2×2 .

We now consider a discrete optimization problem via stochastic relaxation and illustrate
the gradient flow. Consider following potential function on I, taken from [24]:


0



−2
f (x1 , x2 ) = x1 + 2x2 + 3x1 x2 =

−4



6

if
if
if
if

(x1 , x2 ) = (−1, −1)
(x1 , x2 ) = (−1, +1)
.
(x1 , x2 ) = (+1, −1)
(x1 , x2 ) = (+1, +1)

We are to minimize F (p) = Ep [f ], i.e.,

F (p(ξ)) =

X

f (x1 , x2 )p1 (x1 )p2 (x2 ) = −4ξ1 − 2ξ2 + 12ξ1 ξ2 .

(x1 ,x2 )∈I

By Theorem 7, the Wasserstein gradient flow is

ξ˙ = −G(ξ)−1 ∇ξ F (p(ξ)) .

For our function, the standard Euclidean gradient is ∇ξ F (p(ξ)) = (−4+12ξ2 , −2+12ξ1 )T .
The inverse matrix G is computed numerically from J and L.

NATURAL GRADIENT VIA OPTIMAL TRANSPORT

bd

1

−2
b

17

= 0.1

6
ωbd

0.8

d
2

0.6

ωab

ωcd

0.4

0.2

a

c

ωac

0
0

−4

0
bd

1

0.2

0.4

0.6

0.8

1

0.6

0.8

1

1

=1

bd

1

0.6

0.6
2

0.8

2

0.8

= 10

0.4

0.4

0.2

0.2

0

0
0

0.2

0.4

0.6
1

0.8

1

0

0.2

0.4
1

Figure 3. Negative Wasserstein gradient on the parameter space [0, 1]2 of
the two-bit independence model from Example 2. We fix the state graph
shown on the top left, and a function f with values shown in gray next to
the state nodes. We evaluate the gradient flow for three different choices
of the graph weight ωbd . When the weight ωbd is small, the flow from d
towards b (a local minimum) is suppressed. A large weight has the opposite
effect. The contours are for the objective function F (p(ξ)) = Ep(ξ) [f ].

In Figure 3 we plot the negative Wasserstein gradient vector field in the parameter
space Θ = [0, 1]2 . As can be seen, the Wasserstein gradient direction depends on the
ground metric on sample space (encoded in the edge weights). If b and d are far away,
there is higher tendency to to go c, rather than b. This reflects the intuition that, the more
ground distance between b and d, the harder for the probability distribution to move from
its concentration place b to d. We observe that the the attraction region of the two local
minimizers changes dramatically as the ground metric between b and d changes, i.e., as
ωbd varies from 0.1, 1, 10.
This is different in the Fisher-Rao gradient flow, plotted in Figure 4, which is independent of the ground metric on sample space.

18

LI, MONTÚFAR

Fisher-Rao

1

0.8

2

0.6

0.4

0.2

0
0

0.2

0.4

0.6

0.8

1

1

Figure 4. Fisher-Rao gradient vector field for the same objective function
of Figure 3.

5. Discussion
We introduced the Wasserstein statistical manifolds, which are submanifolds of the
probability simplex with the L2 -Wasserstein Riemannian metric tensor. With this, we
defined an optimal transport natural gradient flow on parameter space.
We hope that this paper contributes to strengthening the emerging interactions between
information geometry and optimal transport, in particular in relation to machine learning
problems. We suggest that many studies from information geometry will have a natural
analog or extension in the Wasserstein statistical manifold. Some questions to consider
include the following.
Is it possible to characterize the Wasserstein metric on probability manifolds through an
invariance requirement of Chentsov type? For instance, the work [28] formulates extensions
of Markov embeddings for polytopes and weighted point configurations. Is there a weighted
graph structure for which the corresponding Wasserstein metric recovers the Fisher metric?
One of the key limitations of natural gradient methods in machine learning applications is the computational cost. This difficulty will likely be inherited in the Wasserstein
gradient flow. However, the key innovation coming from the Wasserstein gradient is that
it incorporates a ground metric in sample space. This could be fixed or trained. This
could have a positive effect in terms of generalization performance, as it provides means
to introduce preferences in the hypothesis space. The specific form of such a regularization
still needs to be developed and investigated. In this regard, a natural question is how to
define natural ground metric notions.
The Wasserstein distance has already been shown to be useful in machine learning, for
instance in training restricted Boltzmann machines and generative adversarial networks.
We think that the Wasserstein natural gradient can be a useful method for a variety of
minimization problems. We will continue to work on the above problems.

NATURAL GRADIENT VIA OPTIMAL TRANSPORT

19

Appendix
In this appendix we review the equivalence of static and dynamical formulations of the
metric formally. For more details see [35].

L2 -Wasserstein

Consider the duality of linear programming.
Z
Z
o
nZ Z 1
1
0 1 2
2
0
W (ρ , ρ ) = inf
dΩ (x, y) π(x, y)dxdy :
πdy = ρ (x),
πdx = ρ1 (y)
π≥0
2
Ω Ω 2
Ω
Ω
Z
o
nZ
1
= sup
Φ1 (y)ρ1 (y)dy −
Φ0 (x)ρ0 (x)dx : Φ1 (y) − Φ1 (x) ≤ dΩ (x, y)2 .
2
Ω
Ω
Φ1 ,Φ0
(17)
By standard considerations, the supremum in the last formula is attained when
1
Φ1 (y) = sup Φ0 (x) + dΩ (x, y)2 .
2
x∈Ω

(18)

This means that Φ1 , Φ0 are related to the viscosity solution of the Hamilton-Jacobi equation on Ω:
∂Φ(t, x) 1 Ω
(19)
+ gx (∇Φ(t, x), ∇Φ(t, x)) = 0 ,
∂t
2
with Φ0 (x) = Φ(0, x), Φ1 (x) = Φ(1, x). Hence (17) becomes
nZ
o
1
∂Φ(t, x) 1 Ω
0 1 2
W (ρ , ρ ) = sup
Φ1 (x)ρ1 (x)−Φ0 (x)ρ0 (x)dx :
+ gx (∇Φ(t, x), ∇Φ(t, x)) = 0 .
2
∂t
2
Φ
Ω
By the duality of above formulas, we can obtain variational problem (1). In other words,
consider the dual variable of Φt = Φ(t, x) by the density path ρt = ρ(t, x), then
1
W (ρ0 , ρ1 )2
2
Z
Z
1 1
0 0
= sup inf
Φ ρ − Φ ρ dx −

1Z



1
ρt ∂t Φt + gxΩ (∇Φt , ∇Φt )dx dt
2
Φt ρt
Ω
0
Ω
Z
Z 1Z
Z 1Z
1 Ω
1 1
0 0
= sup inf
Φ ρ − Φ ρ dx −
ρt ∂t Φt dxdt −
gx (∇Φt , ∇Φt )ρt dxdt
ρ
t
Φt
Ω
0
Ω
0
Ω 2
Z 1Z
Z 1Z
1 Ω
Ω
= sup inf
∂t ρt Φt − gx (∇Φt , ∇Φt )ρt dxdt +
gx (∇Φt , ∇Φt )ρt dxdt
Φt ρt
0
Ω
0
Ω 2
Z 1Z
Z 1Z
1 Ω
gx (∇Φt , ∇Φt )ρt dxdt
= inf sup
Φt (∂t ρt + div(ρ∇Φt ))dt +
ρt Φt
0
Ω
0
Ω 2
nZ 1Z 1
o
= inf
gxΩ (∇Φt , ∇Φt )ρt dxdt : ∂t ρt + div(ρ∇Φt ) = 0, ρ0 = ρ0 , ρ1 = ρ1 .
ρt
0
Ω 2
The third equality is derived by integration by parts w.r.t. t and the fourth equality is by
switching infimum and supremum relations and integration by parts w.r.t. x.
In the above derivations, the relation of Hopf-Lax formula (18) and Hamilton-Jacobi
equation (19) plays a key role for the equivalence of static and dynamic formulations of

20

LI, MONTÚFAR

the Wasserstein metric. This is also a consequence of the fact that the sample space Ω is
a length space, i.e.,
o
nZ 1
2
Ω
dΩ (x, y) = inf
gγ(t)
(γ̇, γ̇)dt : γ(0) = x, γ(1) = y .
γ(t)

0

However, in a discrete sample space I, there is no path γ(t) ∈ I connecting two discrete
points. Thus the relation between (18) and (19) does not hold on I. This indicates that
in discrete sample spaces, the Wasserstein metric in Definition (1) can be different from
the one defined by linear programming (5). See many related discussions in [11, 22].
Acknowledgement The authors would like to thank Prof. Luigi Malagò for his inspiring
talk at UCLA in December 2017.

References
[1] S. Amari. Neural learning in structured parameter spaces - natural Riemannian gradient. In M. C.
Mozer, M. I. Jordan, and T. Petsche, editors, Advances in Neural Information Processing Systems 9,
pages 127–133. MIT Press, 1997.
[2] S. Amari. Natural gradient works efficiently in learning. Neural Computation, 10(2):251–276, 1998.
[3] S. Amari. Information Geometry and Its Applications. Number volume 194 in Applied mathematical
sciences. Springer, Japan, 2016.
[4] S. Amari, R. Karakida, and M. Oizumi. Information geometry connecting Wasserstein distance and
Kullback-Leibler divergence via the Entropy-Relaxed Transportation Problem. arXiv:1709.10219 [cs,
math], 2017.
[5] M. Arjovsky, S. Chintala, and L. Bottou. Wasserstein GAN. arXiv:1701.07875 [cs, stat], 2017.
[6] N. Ay, J. Jost, H. Lê, and L. Schwachhöfer. Information Geometry. Ergebnisse der Mathematik und
ihrer Grenzgebiete. 3. Folge / A Series of Modern Surveys in Mathematics. Springer International
Publishing, 2017.
[7] J.-D. Benamou and Y. Brenier. A computational fluid mechanics solution to the Monge-Kantorovich
mass transfer problem. Numerische Mathematik, 84(3):375–393, 2000.
[8] L. Campbell. An extended Čencov characterization of the information metric. Proceedings of the
American Mathematical Society, 98:135–141, 1986.
[9] E. A. Carlen and W. Gangbo. Constrained Steepest Descent in the 2-Wasserstein Metric. Annals of
Mathematics, 157(3):807–846, 2003.
[10] N. N. Čencov. Statistical decision rules and optimal inference, volume 53 of Translations of Mathematical Monographs. American Mathematical Society, Providence, R.I., 1982. Translation from the
Russian edited by Lev J. Leifman.
[11] S.-N. Chow, W. Huang, Y. Li, and H. Zhou. Fokker–Planck Equations for a Free Energy Functional
or Markov Process on a Graph. Archive for Rational Mechanics and Analysis, 203(3):969–1008, 2012.
[12] S.-N. Chow, W. Li, and H. Zhou. A discrete Schrodinger equation via optimal transport on graphs.
arXiv:1705.07583 [math], 2017.
[13] S.-N. Chow, W. Li, and H. Zhou. Entropy dissipation of Fokker-Planck equations on graphs.
arXiv:1701.04841 [math], 2017.
[14] F. R. K. Chung. Spectral Graph Theory. Number no. 92 in Regional conference series in mathematics.
Published for the Conference Board of the mathematical sciences by the American Mathematical
Society, Providence, R.I, 1997.
[15] C. Frogner, C. Zhang, H. Mobahi, M. Araya-Polo, and T. Poggio. Learning with a Wasserstein Loss.
arXiv:1506.05439 [cs, stat], 2015.
[16] W. Gangbo, W. Li, and C. Mou. Geodesic of minimal length in the set of probability measures on
graphs. arXiv:1712.09266 [math], 2017.
[17] J. D. Lafferty. The density manifold and configuration space quantization. Transactions of the American Mathematical Society, 305(2):699–741, 1988.

NATURAL GRADIENT VIA OPTIMAL TRANSPORT

21

[18] G. Lebanon. Axiomatic geometry of conditional models. IEEE Transactions on Information Theory,
51(4):1283–1294, 2005.
[19] W. Li. Riemannian structures of Probability simplex set via optimal transport on graphs. Preprint.
[20] W. Li, P. Yin, and S. Osher. Computations of Optimal Transport Distance with Fisher Information
Regularization. Journal of Scientific Computing, 2017.
[21] J. Lott. Some Geometric Calculations on Wasserstein Space. Communications in Mathematical
Physics, 277(2):423–437, 2007.
[22] J. Maas. Gradient flows of the entropy for finite Markov chains. Journal of Functional Analysis,
261(8):2250–2292, 2011.
[23] L. Malagò, M. Matteucci, and G. Pistone. Towards the geometry of estimation of distribution algorithms based on the exponential family. In Proceedings of the 11th Workshop Proceedings on Foundations of Genetic Algorithms, FOGA ’11, pages 230–242, New York, NY, USA, 2011. ACM.
[24] L. Malagò and G. Pistone. Natural Gradient Flow in the Mixture Geometry of a Discrete Exponential
Family. Entropy, 17(12):4215–4254, 2015.
[25] A. Mielke. A gradient structure for reaction–diffusion systems and for energy-drift-diffusion systems.
Nonlinearity, 24(4):1329–1346, 2011.
[26] K. Modin. Geometry of Matrix Decompositions Seen Through Optimal Transport and Information
Geometry. Journal of Geometric Mechanics, 9(3):335–390, 2017.
[27] G. Montavon, K.-R. Müller, and M. Cuturi. Wasserstein Training of Restricted Boltzmann Machines.
In D. D. Lee, M. Sugiyama, U. V. Luxburg, I. Guyon, and R. Garnett, editors, Advances in Neural
Information Processing Systems 29, pages 3718–3726. Curran Associates, Inc., 2016.
[28] G. Montúfar, J. Rauh, and N. Ay. On the Fisher metric of conditional probability polytopes. Entropy,
16(6):3207–3233, 2014.
[29] E. Nelson. Quantum Fluctuations. Princeton series in physics. Princeton University Press, Princeton,
N.J, 1985.
[30] F. Otto. The geometry of dissipative evolution equations: The porous medium equation. Communications in Partial Differential Equations, 26(1-2):101–174, 2001.
[31] R. Pascanu and Y. Bengio. Revisiting natural gradient for deep networks. In International Conference
on Learning Representations 2014 (Conference Track), Apr. 2014.
[32] J. Peters, S. Vijayakumar, and S. Schaal. Natural actor-critic. In J. Gama, R. Camacho, P. B. Brazdil,
A. M. Jorge, and L. Torgo, editors, Machine Learning: ECML 2005, pages 280–291, Berlin, Heidelberg,
2005. Springer Berlin Heidelberg.
[33] Gabriel Peyre, Marco Cuturi Computational Optimal Transport. arXiv:1803.00567, 2018.
[34] A. Takatsu. Wasserstein geometry of Gaussian measures. Osaka Journal of Mathematics, 48(4):1005–
1026, 2011.
[35] C. Villani. Optimal Transport: Old and New. Number 338 in Grundlehren der mathematischen Wissenschaften. Springer, Berlin, 2009.
[36] T.-K. L. Wong. Logarithmic divergences from optimal transport and R\’enyi geometry.
arXiv:1712.03610 [cs, math, stat], 2017.
[37] S. Yi, D. Wierstra, T. Schaul, and J. Schmidhuber. Stochastic search using the natural gradient.
In Proceedings of the 26th Annual International Conference on Machine Learning, ICML ’09, pages
1161–1168, New York, NY, USA, 2009. ACM.

