Optimization and Dynamics: from Euclidean
Gradient Descent to Wasserstein Gradient Flow
Jia-Jie Zhu
Weierstrass Institute for Applied Analysis and Stochastics
Berlin, Germany

International Workshop of Intelligent Autonomous Learning Systems 2023
Kleinwalsertal, Austria. August 15th, 2023

Motivation & Introduction
Kernel Methods for Robust Learning under Distribution ShiftMotivation: Generative Adversarial Nets
Generative models.

Comparing two distributions
is a (dis-)similarity measure on thespace of probability distributions.
PGθ

Gθ?

(P, Q)}PZPX

Generative models

infGθ

(PGθ, PX),Here, PGθ is the distribution over the generated data Gθ(Z)𝒟𝒟

𝒟

where Z is sampled from a simple distribution such as N(0,I).

Figure credit: Z., W. Jitkrittum, internet meme

Inference as measure optimization
Given density up to a constant π(x) ∝ exp(−V(x))
Generate samples from π (or estimate π ψ(X) for some ψ)

infμ∈ℳ

KL(μ∥π).

Monte-Carlo Sampling via Langevin SDE

Xk+1 = Xk − ∇V(Xk) ⋅ τ +2τΔZk

where ΔZk ∼ N(0,1), τ is the step size. The state distribution

XT ∼ μT converges to π. N-particle approximation results in thenoisy SGD in optimization. This dynamics is equivalent to thePDE gradient ow in the Wasserstein space [Otto 96].

fl𝒟𝔼

Motivation: Langevin Monte-CarloMotivation: Structured Distribution Shift
Structured Distribution S
Structured Distr(Joint) DRO mostly robustifie
Causal Confounding

Learning Under Stru! Confounding
can leadto m(Joint) DRO
mostly!thoseConfounding canCausal confounding can lead to much stronger distribution shifts
thanX : Smoking, Y : Cancer, U: L
considered in (joint) distribution shift!

X : Smoking, Y : Ca

X: Smoking,Y: Cancer,U: Lifestyle
Y := gθ(X) + ϵU,[ϵU] = 0, but [ϵU | X] ≠ 0⟹ gθ(x) ≠ [Y | X = x]

UXTo estimate gθ robustly, we use conditional moment restriction
𝔼[ϵU | Z] = [Y − gθ(X) | Z] = 0 ℙZ-a.s.

Y :=g✓YU=)g✓

Y predTake into Xaccount genetic

Take intoUaccount g

Take into account genetic predisposition for nicotine addiction ZZg✓XY UHeiner Kremer

ZRoXg✓

Heiner Kremer
𝔼𝔼𝔼𝔼Learning Under Structured Distribu

Figure credit: H. Kremer

Motivation & Introduction
Kernel Methods for Robust Learning under Distribution ShiftDistributional robustness, but what kind?

Figure credit: The Princess Bride,a bedside story by your grandpa

From Statistical Learning to Distributionally Robust Learning
Empirical Risk MinimizationN•1
minl(θ, ξi), ξi ∼ P0θ N∑
i=1“Robust” under statistical uctuation
N11̂ ξ) ≤̂ ξ) + (l(θ,l(θ,)
P0iN∑
N
i=1

Distributionally Robust Optimization (DRO)
min supθQ∈ℳQl(θ, ξ)

Worst-case distribution Q within the ambiguity set ℳ[Delage & Ye 2010] in certain geometry.
Distrib. Shift• Not robust under data distribution shifts,
when Q ( ≠ P0)δξi

Why study new geometry?
New geometries leading to new elds of research andbreakthroughs:

flfi

𝒪

fl

fl

𝔼

𝔼

N1PN̂ :=δξi ⇒ P0 ≠ Q
N∑
i=1Information geometry [S. Amari et al.] e.g. descent inFisher-Rao geometry
Wasserstein Gradient ow [F. Otto et al.] e.g. FokkerPlanck equation as Wasserstein ow Figure credit: H. KremerOptimal Transportp-Wasserstein

Integral Prob. Metrics1-Wasserstein

WassersteinFisher-Rao

Kernel MMD

Total Variation

ϕ-divergences

Bregman divergences

Information Divergence

Figure credit: J. Zhu

Background: “Kernel Geometry”
De nition. Kernel Maximum-Mean Discrepancy (MMD)
associated with (PSD) kernel k (e.g., k(x, x′) := eMMD(P, Q) :=−|x−x′|2/σk(x, ⋅ )dP− k(x, ⋅ )dQ .
ℋ∫∫d)How does the MMD behave?
MMD with a broad kernel:
MMD P Q
sup f
1 EP f X

EQ f Y .

}MMD(P, Q)

MMD=1.8(Prob(ℝ ), MMD) is a (simple) metric space.
Dual formulation as an integral probability metric.

∫∥f∥ ≤1MMD(P, Q) = supf d(P − Q)

ℋ23/34fifi

￼𝒳￼

ℋ is the reproducing kernel Hilbert space ℋ (RKHS),
which satis es f(x) = ⟨f, ϕ(x)⟩ℋ, ∀f ∈ ℋ, x ∈ ,ϕ(x) := k(x, ⋅ ) is the canonical feature of ℋ.

Figure credit: W. Jitkrittum, J. Zhu

Previous work: Kernel DRO
Primal DRO (not solvable as it is)

Example. Robust least squares[El Ghaoui Lebret ’97]

(DRO) minθ

sup̂MMD(Q,P)≤ϵ2minimize l(θ, ξ) := ∥A(ξ) ⋅ θ − b∥2Ql(θ, ξ)Given historical samples ξ1, ξ2, …, ξN

Kernel DRO Theorem (simpli ed). [Z. et al. 2021]
DRO problem is equivalent to the dual kernel
machine learning problem, i.e., (DRO)=(K).
N1(K) minf(ξi) + ϵ∥f∥ℋ subject to l(θ, ⋅ ) ≤ f
θ,f∈ℋ N ∑i=1
Geometric intuition: dual kernel function f as robust
surrogate losses ( atten the curve)

ffl(θ, ⋅ )l(θ, ⋅ )(test distribution)

Robustifyingwith DRO(a) Robust least squares
lossf asplotwitness(test)
Figure 3: (a) Thisdepicts thetestfunction
loss of algorithms. Alran 10 independent trials. In each trial, we solved K-DRO to oof 500 samples. We then vary the perturbation from 0 to 4.
What if ff≡⇤ c ∈⇤ ℝ?
⇤+f.
(black)isthefunctionl(✓, ·). The pink bars depict a0bars the empirical distribution. We can observe that f0⇤ + f ⇤ tof
worst-case distribution P ⇤ (pink dots). Note f ⇤ (normalized)
cthe two distributions.

fi𝔼227fl(b) G4

Numerical studies

Our risk minimization scheme (11) can be straightforwardly used with stochastic gradient based
method for large-scale learning, e.g., with DNNs. We detail the training procedure in Algorithm 1.
Note that Step 3 of Algorithm 1 can be seen as a proximal algorithm, cf. [55]. ARKS (11) can also

Experiments of robustness under distribution shifts
Algorithm 1: Adversarially robust kernel smoothing (ARKS)
1: input: data sampler, initial iterate ✓0
2: for k = 0, 1, 2, . . . , T do
3:sample {⇠k } and find u⇤k by maximizing l(✓k , u)k(u, ⇠k ) w.r.t. u
4:update ✓ by stochastic gradient descent using estimate r✓ l(✓k , u⇤k )5: output: approximate solution ✓ ⇤ := ✓ T

Toy example. Uncertain least squares[El Ghaoui Lebret ’97, Z et al. 2021]

minimize l(θ, ξ) := ∥A(ξ) ⋅ θ − b∥22

Given historical samples ξ1, ξ2, …, ξN

Certi ed adversarially robust deep
Kernel stochastic model predictive
be interpreted as a form of adversarial training [38, 81, 27]: for each ⇠ , the inner maximization
learning
the upresenceoftheglasses
problemof (11) looks for(Classify
an adversarial example
that hurts the learnermost. In the case of control (MPC) with nonlinear
Gaussian RBF kernel, it is an exercise (we provide this proof in the appendix) to show that the innerusing
a 20-layer
DNN[Sinhaet al.
et al.of2022]
maximization
objectivein (11) has favorable
convexity
structuresfor 2017;
suitable Z
choices. The main constraints [Nemmour et al. & Z 2022]
i

intuition is that, by multiplying the loss l(u) by the kernel function k(u, x) which is strongly concave
near its peak, the resulting function is consequently locally concave too. We illustrate this intuition inFigure 3 (right). A thorough analysis for various other kernels choices is out of our current scope.
Next, we empirically demonstrate ARKS in Section 5 that ARKS can easily work with DNN models,
which is a limitation of typical existing DRO reformulation techniques.

642(test distribution)(a) Robust least squares loss

Numerical Experiments

5.1

Robust Learning under Adversarial Perturbations

x25°2In this section, we present empirical evaluations of our robust learning algorithm. We compare thefollowing algorithms when applicable: (A) ARKS Algorithm 1, (B) empirical risk minimization(ERM), (C) WRM [21], (D) (worst-case) robust optimization [7, 68] (in the appendix). Both ARKSand WRM have one hyper-parameter: bandwidth for ARKS and Lagrangian relaxation coefficienty for WRM. We do not test classical Wasserstein DRO algorithms, e.g., [48, 85, 64], since theycannot be applied to our test settings with general machine losses and DNN models. We further notethat ARKS (11) is a general robustification technique that can be applied to tasks beyond supervised
learning and adversarial training.

°4°50510x1

15

Wasserstein robust feedback
optimal control of nonlinear
dynamical systems [Zhong & Z 2023](b) Geometric interpretation

Robust reinforcement learning[Sinha et al. 2017]
re3: Robust(a) This plot depicts the test loss of algorithms. All error bars are in standard error. We
ar10trials. In each trial, we solved K-DRO to obtain ✓ ⇤ and tested it on a test dataset± 0.1independent
400.0 ± 0.000samples. We then vary the perturbation from 0 to 4. (b) (red) is the dual optimal solution
onments± 0.0 ⇤ 400.0 ± 0.0⇤+± 0.0f .400.0(black)isthefunctionl(✓, ·). The pink bars depict a worst-case distribution while the blue
± 0.0
± 0.0 400.0 ± 0.0
the empirical distribution. We can observe that f0⇤ + f ⇤ touches loss l(✓⇤ , ·) at the support of the⇤⇤onmentst-casedistribution
P(pinkdots).
Notef(normalized)
witnessfunctionof (right)
Figure 1: PGDcan
attackbe
withviewed
respect to k.kasonathe
Fashion-MNIST(left) and CIFAR-10± 4.7 334.0 ± 3.7
datasets. We show the classification error on perturbed test images versus the allowed magnitude of±wo
4.8 distributions.
400.0 ± 0.0
the adversarial perturbation . AKRS and WRM exhibit similar adversarial performance profiles;± 2.3 398.5 ± 0.3

00:8
0:6

400
350

0:4

250x2300

0:2

200
150
100fi

over 1000 trials

0:215000:00500

1000

1500

2000

Figure 8. Episode lengths during training.

AKRS becomes more robust as the kernel width increases, while WRM improves with a lower
Lagrangian penalty y. For all algorithms, we report the mean and standard deviation across 10

0:4
2:0

1:5

1:0

0:5
x1

0:0

0:5

Learning Under Structured Distribution Shifts

To relax the semi-in nite constraint in DRO reformulationsminθ

Experimental
Resultsl(θ,ξ)Qsup̂
MMD(Q,P)+λD
Learning Under Structured
DistributionShiftsϕ(Q∥ω)≤ϵNonlinear
Instrumental
Variable Regression

Dual reformulation. Adapted Experimental
from [Kremer etResultsal., Z. 2023]
Y := g (X ; ✓ ) + ⌫(U) + ✏θ,f∈ℋ {inf01−f + l Variable
X :=⌘(Z ) + µ(U) + ✏2
Nonlinear InstrumentalRegressionf+ϵ∥f∥+λϕ*̂ℋ
ωP()Z⇠P,✏⇠
N(0,)
Z
λ1/2}
Y := g (X ; ✓ ) + ⌫(U) + ✏01

hard cons.
relax.
MMD only
KL-MMD
Log-MMD

Z: geneticUZ

U: Lifestyle

Empirical likelihood formulation + similar techniques: ∀h ∈ ℋ,Q [(Y − gθ(X)) h(Z)] = 0T𝔼𝔼𝔼𝔼

12inf MMD (Q, P)̂ + λDϕ(Q∥ω) s.t.
Q 2||µ ° µ0 ||2X := ⌘(Z ) + µ(U) + ✏2
nicotine addiction(t)(1soft cons. ϕ*,log-barrier
(t)
=exp
ϕ*(t)
=−log−t)✓4 ✓3x ✓1KLglogwithg(x;✓)=✓+✓(x✓)+log
1+e✓31Z ⇠ PZ , ✏1/2 ⇠ N (0, ) 22ZXY
Example (Causality). Robustness against structured
CMR: E[Y g (X ; ✓)|Z ] = 0 PZ -a.s.
X: Smoking
Y: Cancer
distribution shifts instead of (joint-)DRO. [Kremer et al., Z. ✓4 ✓3with g (x; ✓) = ✓2 + ✓3 (x ✓1 ) + 2 log 1 + e x ✓1
DeepGMM2023]. Estimating gθ via conditional moment restriction
Functional GEL
1.0KMM
CMR: E[Y g (X ; ✓)|Z ] = 0 PZ -a.s.
[Y − gθ(X) | Z] = 0 ℙZ-a.s.

fi𝔼

Entropy-MMD & Conditional independence

0.5

0.0
103

sample size

104

X

Example. Nonlinear DR-chance constrained opt. [Nemmour et al., Z. 2022]
Tmin c x s . t .
x∈

inf

D(P,PN̂ )≤ϵαP( f(x, ξ) ≤ 0) ≥ 1 − αVaRCVaRf(x, ξ) nonlinear in uncertainty.
Idea: Rewrite

using the indicator function (or CVaR)f (x, ª)g(ª)
1(f (x, ª))

1.00
0.75

0.25

Apply the Kernel DRO Theorem and algorithm, informally

P f(x, ξ) ≤ ((0.50) ≥ 1 − α,N )

0.00
°0.25

1°2

Constraint°1

P = {P̂n }P bootstrap
P MMD rate

0.02
0.01

0ª

12

3

Objective

1.1
1.0
0.9
0.8

0.00

0.7

𝕀

𝒳𝒪

with large probability. The big-O term depends on the class of f.

°3cT xD(P,PN̂ )≤ϵP[ ( f(x, ξ) ≥ 0))] ≤ αP0 -CVaRsup𝔼𝔼DR Probablistic Robust Optimization

°0.01

102

sample size

103

102

sample size

103

More ML problems we can use dual functions
for measure optimization
Primal-dual optimization problems

inf F(μ) = sup ℰ( f )μ∈ℳf∈ℱ

Examples in ML [Dvurechensky, Z.]
Generative models

infGθZμ∈ℳ f∈ℱ { ∫(P, Gθ(Z)) = inf supf(x)dP(x) −θ∼μ∫f(gθ(z))dQ(z)}Distributionally robust optimization

infNsupθ MMD(μ,μ)≤ϵ̂1
inf sup μ(l − f ) +f(xi) + ϵ∥f∥ℋ.
μ[l(θ; x)] =N∑θ∈ℝd,f∈ℋ μ∈ℳi=1Wasserstein barycenter

cminαi [Wp(μ, νi)] = minαi supfi dμ + fidνi ,}∫μ∈ℳ ∑μ∈ℳ ∑ f ∈Ψ { ∫ici=1i=1𝔼n𝔼𝒟𝔼𝔼nSummary
This talk is mainly based on:Z., Jitkrittum, W., Diehl, M. & Schölkopf, B. Kernel DistributionallyRobust Optimization. AISTATS 2021
Kremer, H., Nemmour,Y. , Schölkopf, B. & Z. Estimation Beyond Data

• There are many important uses of the dual (kernel)

function for measure optimization: causal inference,barycenter problems, conditional moments, (robust)
control and RL• However, optimization over measures is a mathematically
non-trivial topic. We will now learn the optimization
perspective of gradient ow

Reweighting: Kernel Method of Moments. ICML 2023
Nemmour,Y. , Kremer, H., Schölkopf, B. & Z. MMD Distributionally
Robust Nonlinear Chance-Constrained Optimization with FiniteSample Guarantee. IEEE CDC 2022
Z., Kouridi, C., Nemmour,Y. & Schölkopf, B. Adversarially Robust
Kernel Smoothing. AISTATS 2022
P. Dvurechensky, Z., Kernel Mirror Prox and RKHS Gradient Flowfor Mixed Functional Nash Equilibrium. Preprint

flSlides will be available
Website: jj-zhu.github.io