A Mini-Course on Optimization and Dynamics
From Euclidean Gradient Descent to Wasserstein Gradient Flow
Jia-Jie Zhu
Weierstrass Institute for Applied Analysis and Stochastics, Berlin

August 15, 2023

Euclidean gradient descent

Optimization in Rd
min f (x ).

x ∈Rd

We optimize using the gradient descent algorithm
xk+1 = xk − τk · ∇f (xk )using the “variational principle”
xk+1 ∈ argminx f (xk ) + h∇f (xk ), x − xk i +1kx − xk k22
2τkFrom GD to Mirror descent

We define the Bregman divergence associated with φ asDφ (x , y ) = φ(x ) − φ(y ) − ∇φ(y )> (x − y ).
Mirror descent update (with quadratic term replaced by Bregman)xk+1 = argminx f (xk ) + h∇f (xk ), x − xk i +1
Dφ (x , xk )τkExample of MD: Euclidean norm

Mirror map φ(x ) = 12 kx k22 .
The resulting mirror descent algorithm
xk+1 = argminx f (xk ) + h∇f (xk ), x − xk i +This is equivalent to gradient descent with stepsize τk .

1kx − xk k22 .
2τkExample of MD: negative entropy

Mirror map φ(x ) = di=1 x (i) log x (i). The resulting Bregman divergence is theP
KL-divergence Dφ (x , y ) = di=1 x (i) log yx (i)(i) .
P

If we restrict x to a (discrete probability) simplex 1> x = 1, then the MD updatexk+1 = argmin1> x =1 f (xk ) + h∇f (xk ), x − xk i +1
Dφ (x , xk )τkhas the closed-form as exponentiated gradientx (i)k e −τk gkx (i)k+1 = Pn,−τk gk
j=1 x (j)k e

i = 1, 2, . . . , nEuclidean gradient descent as discretization of ODE gradient flow

1kx − x k k2
2τis the explicit Euler scheme for the ODE (for simplicity, we take constant time step τ )x k+1 ∈ argminx h∇f (x k ), x iEd +ẋ (t) = −∇f (x (t)).
The solution x (t) is an ODE gradient flow and the ODE is the gradient flow equation(GFE). In GF terms, the solution x (t) is also called a curve of maximal slope (steepestdescent).

Gradient flow dynamics: (nonlinear) ODEẋ (t) = −∇f (x (t))
ẋ (t) ∈ X provides the rate (or velocity) (we can see)−∇f (x (t)) ∈ X ∗ provides the (thermodynamic) force (can’t see; shadow price)
The equation should be written in the force-balance form
IR ẋ (t) = −∇f (x (t)) ∈ X ∗ ,IR : X → X ∗ is the Riesz isomorphism.

If, in the non-Euclidean setting, X  X ∗ , then we have both force space and rate spaceGFE.

Energy dissipation balance (equality)
22≥ ab)Fenchel-Young For convex ψ, (proof is trivial; a +b2
ψ(x ) + ψ ∗ (ξ) ≥ hx , ξi, ∀(x , ξ) ∈ X × X ∗ .
Furthermore, if ψ is proper, lsc, and convex, (x ∗ , ξ ∗ ) is optimal.
By Fenchel(-Young) duality and optimalityd1
1f (x (t)) =X ∗ h∇f (x (t)), ẋ iX = −k∇f (x (t))k2 = − kẋ k2 + k∇f (x )k2
dt
22

Energy does not necessarily decrease along non-solutions, i.e., only inequality
d1
1f (z(t)) ≥ − kżk2 + k∇f (z(t))k2 .
dt22

Evolutionary variational inequality (EVI)λ : ODE
Suppose the energy functional f is proper, upper semicontinuous, λ-convex for someλ ∈ R , i.e., f can be non-convex, ∀s ∈ [0, 1], ∀u0 , u1 ∈ Rdλf ((1 − s)u0 + su1 ) ≤ (1 − s)f (u0 ) + sf (u1 ) − s(1 − s)ku0 − u1 k2 .
2and has compact sublevel sets. Then for any initial condition in the x (0) ∈ Rd , there
exists a unique solution at time t, x (t) ∈ Rd .
Furthermore, the ODE solution x (t) satisfies (EVI)λ , for t, s ∈ [0, T ].
11kx (t) − νk2 ≤ e−λ(t−s) kx (s) − νk2 + Mλ (t − s)(f (ν) − f (x (t))),22ZτMλ (τ ) =e −λ(τ −s) ds,0∀ν ∈ dom(F ) ⊂ Rd .

Using (EVI)λ , we can effortlessly extract convergence results. Suppose a minimizer ofthe energy exists x ∗ ∈ arginf x ∈Rd f (x ), we set ν = x ∗ , s = 0 in (EVI)λkx (t) − x ∗ k2 ≤ e−λt kx (0) − x ∗ k2 + 2Mλ (t − s)

inf f (x ) − f (x (t))

x ∈Rd≤ e−λt kx (0) − x ∗ k2

Gradient flow convergence without (strong) convexity: ODE
Impose the Polyak-Łojasiewicz inequality, suppose an optimizing x ∗ existsk∇f (x (t))k2 ≥ c · (f (x ) − f (x ∗ )).
Starting from EDB
df (x (t)) = −k∇f (x (t))k2 ≤ −c · (f (x ) − f (x ∗ )) ≤ 0,dt
implies exponential convergence of the gradient flowf (x (t)) − f (x ∗ ) ≤ e −c·t (f (x (0)) − f (x ∗ )) .
We can also recover the log-Sobolev inequality by setting the energy f as theKL-divergence.

Optimization over probability measures

inf F (µ)µ∈PI P: set of probability measures; the probability “simplex”
I We will work with two types of (probability) measuresdµ(x ) = ρ(x ) dx ,µ=Xi∈II M + ⊇ P: non-negative measures; “cone”
I F : objective function; “energy”

αi δxi α ∈ ∆Optimization over probability measures

What can’t we just do gradient descent?
µk+1 = µk − τk · ∇F (µk )I ∇F (µk ) is undefinedI µk+1 must be a probability measure; care needs to be taken
I What can we do instead?

A variational approach

Recall the “variational” formulation of gradient descentx k+1 ∈ argminx h∇f (x k ), x iRd +1kx − x k k2 ⇐⇒ xk+1 = xk − τ · ∇f (xk )2τ

for a suitable τ . This is the variational principle.
Can we do the same for probability measures?
1µk+1 ∈ arginf µ∈P F (µ) + D2 (µ, µk )τfor some “distance” measure D. This is sometimes called the Minimizing Movement
Scheme (MMS).

Variational approach and MMS

1µk+1 ∈ arginf µ∈P F (µ) + D(µ, µk )
τWe must specify the important ingredients
Energy :FGeometry :DThe merit of the right gradient flow formulation of a dissipative evolution equation isthat it separates energetics and kinetics: The energetics endow the state space with afunctional, the kinetics endow the state space with a (Riemannian) geometry via themetric tensor. [Otto 2001]

Geometry: Wasserstein distanceDefinition. The p-Wasserstein distance** between probability measures P, Q on Rd(with p-th finite moments, p ≥ 1) is defined through the following Kantorovich problem
Wpp (P, Q) := infnZ(1)(2)|x − y |p dΠ(x , y ) π# Π = P, π# Π = QoDual Kantorovich problem
Wpp (P, Q) = supnZZψ1 (x ) dP(x ) +ψ2 (y ) dQ(y ) ψ1 (x ) + ψ2 (y ) ≤ |x − y |pDynamic formulation: Benamou–Brenier
W22 (P, Q) = infnZ 1Z|vt |2 dµt dt µ0 = P, µ1 = Q,0odµt + div(vt µt ) = 0dt

Entropy regularization (Sinkhorn divergence)ZinfΠc(x , y )dΠ(x , y ) + λDφ (ΠkP ⊗ Q)oGeometry: (Csizsar) φ-divergence
Relative entropy is defined as
Dφ (µ|ν) =R  φ dµ dν

if µ  ν+∞otherwisedνWe can choose the φ functions from the following table to obtain: identity (trivial),Kullback, Hellinger, χ2(table: J. Zhu)

Preliminary: first variation over measures and subdifferentials
The first variation of a functional F at µ ∈ P is defined as a function δFδµ [µ]dF (µ +  · v )|=0 =dZδF[µ](x ) dv (x )
δµfor any perturbation in measure v such that µ +  · v ∈ P.
The variational principledF (µ +  · v )|=0 = 0d
for all variation v , states the “optimality condition”.
We also summon the Fréchet differential on a Banach space X as a set in the dual space
DF := {ξ ∈ X ∗ | F (µ) ≥ F (ν) + hξ, µ − νiX + o (kµ − νkX ) for µ → ν}Three types of energy functionals

Suppose µ(x ) = ρ(x ) dx WLOG,F(%) =Zf (%(x ))dx ,V(%) =ZV (x )d%,1W(%) =2ZZ

W (x − y )d%(x )d%(y )We calculate the first variations (by following the definition)
δF(%) = f 0 (%),
δ%δV(%) = V ,δ%δW(%) = W ∗ %δ%Back to (discrete-time) gradient flow

Optimization
inf F (µ)µ∈P

Recall MMS

1µk+1 ∈ arginf µ∈P F (µ) + D(µ, µk )
τWe have specified the important ingredients
Energy :F ,Geometry :DWe can construct a concrete instance of MMS for gradient flow by “mix-and-match”.Wasserstein-MMS: Jordan-Kinderlehrer-Otto (JKO) scheme

a.k.a. Minimizing Movement Scheme (MMS):µk+1 ∈ inf F (µ) +µ∈P

1 2W (µ, µk )2τ 2This formulation is very general in the sense that it includes nonlinear-in-measure F .
We should think of this as the gradient descent algorithm for prob. measures.

Otto’s Gradient flow equation in the Wasserstein spaceµk+1 ∈ inf F (µ) +µ∈P

1 2W (µ, µk )2τ 2

Continuous-time limit τ → 0, we have (non-trivially) the gradient flow equation (GFE)∂t µ − ∇ · (µ∇δF[µ]) = 0
δµwhich describes the dissipation of energy F in (Prob(X̄ ), W2 ). [Otto et al 90s-2000s,Ambrosio 2005]
In a different flavor, we can write it just like ODE ẋ = −∇f (x ) (in the rate form; primalvs. dual force-balance)∂t µ = −KOtto (µ) DF = ∇ · (µ∇DF ).

Example: WGF of (Boltzmann/KL/relative) Entropy
nonlinear (in measure) energy (e.g., in variational inference)
ZF (µ) = DKL (µkπ) =log(δµ(x ))ρ(x ) dxδπδF[µ] = log ρ − log π,
δµdµdensity ρ := dLThe Fokker-Planck equation as the Wasserstein gradient flow [Ottoet al. 90s-2000s]
δF∂t µ= ∇ · µ∇ [µ]δµ= ∇ · (µ(∇ log ρ − ∇ log π))= ∆ρ + ∇ · (ρ∇ log π)
I If π is the Lebesgue measure, we obtain the heat equation ∂t µ = ∆µ
δFI Note: The force field δFδµ [µ] and the “score” ∇ δµ [µ] are not accessible if µ isatomic. =⇒ “score-matching”. . .

Application: sampling and variational inference
Suppose π ∝ e −V (x ) , but with unknown normalizing constant, we want
inf DKL (µkπ).

µ∈P

Using the WGF, we have the Fokker-Planck equation∂t µ = ∇ · (µ(∇ log ρ(x ) − ∇V (x )))Suppose there is a single atom whose state is Xt (R.V.), it is pushed towards thevelocity field∇ log ρ(Xt ) − ∇V (Xt )We can construct gradient descent
Xt+1 = Xt + τ · ∇ log ρ(Xt ) − ∇V (Xt )Langevin Monte-Carlo forward-Euler discretization√
Xt+1 = Xt − τ · ∇V (Xt ) + 2τ Z , Z ∼ N(0, Id)

Application: (distributionally) robust learning with Otto’s WGF

We can use our WGF theory (invented 20yr ago; nothing new) to solve Wasserstein
DRO for robust learning (also adversarial robustness in [Sinha et al. 2017])
min sup Eµ l(θ, x ) − γ · W22 (µ, µ̂N )θµThe inner measure-update step is gradient ascent
Xt+1 = Xt + τ ∇l(θt , Xt )1
where τ = 2γ.Then the whole Wasserstein robust learning is simply gradientdescent-ascent (GDA).

Energy dissipation balance of WGFRecall the ODE case
11df (x (t)) = − kẋ k2 + k∇f (x )k2dt
22In (Prob(X̄ ), F , W2 ), Fenchel(-Young) yields the Energy dissipation balance(equality) [Ambrosio et al. 2007]
1d1F (µ(t)) = − |µ0 |W2 (t)2 − |∇− F |W2 (µ(t))2dt
22
1F (µ(t)) − F (µ(s)) = −2Z ts|µ0 |W2 (r )2 + |∇− F |W2 (µ(r ))2 dr

I metric speed with velocity vt : |µ0 |W2 (t) =I metric slope: |∇− F |W2 (µ(t)) =qRqR|vt |2 dµ2|∇ δF
δµ [µ] (x )| dµThe velocity field can be identified as vt = −∇ δF
δµ [µ]. EDB can then be used as thedefinition of gradient flows (curves of maximal slopes), even without GFE.
For (Boltzmann) entropy F (u) = ρ log ρ, EDB gives

Rd2
dt F (µ(t)) = − |∇ log ρ| ρ dx

Evolutionary variational inequality (EVI)λ : Wasserstein GF
Under a few technical assumptions and the so-called λ-geodesic-convexity of the energy
F , if along a geodesic curve γ,F (γ(s)) ≤ (1 − s)F (γ(0)) + sF (γ(1)) −λs(1 − s)W22 (γ(0), γ(1)), ∀s ∈ [0, 1].
2Then, there exists unique gradient flow solution satisfies (EVI)λ , for .
11 2W (µ(t), ν) ≤ e−λ(t−s) W22 (µ(s), ν) + Mλ (t − s)(F (ν) − F (µ(t))),2 22Zτ∀ν ∈ dom(F), Mλ (τ ) =e −λ(τ −s) ds.

0Set ν ∈ arginf µ F (µ), we have exponential convergence in-time and uniqueness ofgradient flow.

Thank you!
There are many other active research topics in GF for ML
I Gradient flow structure with kernel geometry [also some of my past / current works]
I Unbalanced transport and its gradient flow
I Applications: causal inference, mean-field NN, Nash equilibrium, offline RL, policy
optimization. . .

Reference
1. Mielke, A. An introduction to the analysis of gradients systems. Preprint at
http://arxiv.org/abs/2306.05026 (2023).
2. Otto, F. The geometry of dissipative evolution equations: the porous medium
equation. (2001).
3. Ambrosio, L. & Savaré, G. Gradient Flows of Probability Measures. in Handbook of
Differential Equations: Evolutionary Equations vol. 3 1–136 (Elsevier, 2007).
4. Santambrogio, F. Optimal transport for applied mathematicians. Birkäuser, NY 55,
94 (2015).
5. Peletier, M. A. Variational Modelling: Energies, gradient flows, and large deviations.
arXiv:1402.1990 [math-ph] (2014).