From Gradient Flow Force-Balance toDistributionally Robust Learning
Jia-Jie Zhu
Weierstrass Institute for Applied Analysis and Stochastics
Berlin, Germany

European Conference on Computational Optimization (EUCCO)
Heidelberg University, Germany. September 26th, 2023

Distributional robustness, but what kind?

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

Motivation: From statistical learning to robust learning
Empirical Risk MinimizationN•1minl(θ, ξi), ξi ∼ P0
θ N∑
i=1
“Robust” under statistical uctuationN11̂ ξ) ≤̂ ξ) + (l(θ,l(θ,)
P0iN∑Ni=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:

fl

fl

fl

𝔼Figure credit: H. Kremer, J. Zhu

fi𝒪𝔼N1PN̂ :=δξi ⇒ P0 ≠ QN∑i=1

Information geometry [S. Amari et al.] e.g. descent inFisher-Rao geometryWasserstein Gradient ow [F. Otto et al.] e.g. FokkerPlanck equation as Wasserstein ow

Background: Kantorovich-Wasserstein Geometry
De nition. The p-Wasserstein distance between probability measures
dP, Q on ℝ (with p nite moments, p ≥ 1) is de ned through the following
Kantorovich problem
pWp (P, Q) := inf {∫p| x − y | dΠ(x, y)(1)(2)
π# Π = P, π# Π = Q}

(Dual Kantorovich problem)= sup { ψ1(x) dP(x) + ψ2(y) dQ(y) ψ1(x) + ψ2(y) ≤ | x − y | }∫∫pd2-Wasserstein space (Prob(ℝ ), W2) is a geodesic metric space.
Dynamic formulation: à la Benamou-Brenier∫0| vt | dμtdt μ0 = P, μ1 = Q, ∂t μt + div(vt μt) = 0}
2fifi

fi

2W2 (P, Q) = min {∫1xy

Background: MMD and interaction force
75

6.2 Bochner’s characterization

De nition. Kernel Maximum-Mean Discrepancy (MMD)
associated with (PSD) kernel k (e.g., k(x, x′) := eMMD(P, Q) :=2−|x−x′| /σ)k(x, ⋅ )dP− k(x, ⋅ )dQ .
ℋ∫∫

Fig. 6.2 The reconstruction of a smooth function using a Gaussian with scale parameter α = 10, α =100, α = 1000.

Proof Suppose that " is bounded and has a nonnegative and nonvanishing Fourier trans! ∈ L 1 (Rd ) is satisfied. Then the Fourier inversion formulaform. It suffices to show that "can be applied and Theorem 6.8 finishes the proof.
As in the proof of Bochner’s theorem we choose gm from Theorem 5.20 and get""(x)gm (x)d x(2π)d/2 "(0) = (2π)d/2 lim
m→∞ Rd"d/2! gm (ω)dω= (2π)
lim"(ω)!
MMD(P, Q)
m→∞ Rd"!
= kernel:"(ω)dω,
MMD with a broadd(Prob(ℝ ), MMD) is a (simple) metric space.
Dual formulation as an integral probability metric.

∫∥f∥ ≤1

MMD(P, Q) = supf d(P − Q)Rd

P Q ! issup
E f X
EQ f Y .
!MMD
since "
is nonnegative. Thus "in L 1f(Rd ). 1 P
If, conversely, " is positive definite then we know from Theorem 6.2 that " is bounded
and from Bochner’s theorem that " is the inverse Fourier transform of a nonnegative finite
Borel measure µ on Rd . Furthermore, using Theorem 5.16, Theorem 5.20 and Fubini’s
MMD=1.8
theorem we find that"!
!
"(ω)g"(x)= limm (ω − x)dωm→∞ Rd"T= lim"(ω)!
gm (ω)e−i x ω dωm→∞ Rd" "TT= lim (2π )−d/2
e−iω η dµ(η)!
gm (ω)e−i x ω dω
m→∞
Rd Rd
" "T!
gm (ω)e−iω (η+x) dωdµ(η)= lim (2π )−d/2m→∞
Rd Rd
"
= lim
gm (−η − x)dµ(η)ℋℋ is the reproducing kernel Hilbert space ℋ (RKHS),
which satis es f(x) = ⟨f, ϕ(x)⟩ℋ, ∀f ∈ ℋ, x ∈ ,ϕ(x) := k(x, ⋅ ) is the canonical feature of ℋ.
As an interaction energy for Wasserstein GF [Arbel et al.]

fifi

￼

￼

MMD (P, Q) =∫∫m→∞ Rd≥ 0.

k(x, y) d(P − Q)(x) d(P − Q)(y)𝒳2}

How does the MMD behave?

! is nonnegative. Now we can proceed as in the first part of the proof to show thatThus "! ∈ L 1 (Rd ) and &"&! L 1 (Rd ) = (2π)d/2 "(0). Hence "! cannot vanish identically.
"!
.
8CC B

1 5 6
8CC B
31 2 9 75

73

31 2 9 75 7 3
5 C5 B 8CC B

5 ,92 9 C85
9 7,0/1

1C

BD2:53C C C85 1 2 9 75

5 C523/34

B 6 DB5 1 19 12 5 1C

Figure credit: W. Jitkrittum, J. Zhu, H. Wendland

Gradient Flow Force-BalanceGradient ow facts
Otto's Gradient ow equation in the Wasserstein spaceδF∂t μ − ∇ ⋅ (μ ∇ [μ]) = 0δμe.g., di usion, Fokker Planck equation. It describes the “steepest”
dissipation of energy F in (Prob(X̄), W2). [Otto et al 2000s, Ambrosio 2005, …]
In a di erent avor, we can write it just like ODE gradient ow x· = − ∇f(x)in the primal rate-form

μ· = − Otto(μ) DF (DF is the (sub)di ., e.g., in the sense of Fréchet)Time-discretization yields the minimizing movement scheme (MMS)

fl

fl

ff

fl

fl

ff

ff

fl

𝕂𝒫

1 2
W2 (u, uk−1)“JKO Scheme” uk ∈ arg inf F(u) +u∈2τ

ODE ow: gradient descent

1x k ∈ arg min F(x) + ∥x − x k−1∥2.
2τx∈ℝdGradient ow force-balance
1 2W2 (u, uk−1)
Force-balance in Wasserstein MMS uk ∈ arg inf F(u) +u∈2τ
ϕDF + = const . , ϕ : “Kantorovich potential”
τForce:drive movement
e.g., entropy

Force-balance in ODE:
Dissipation Geometry:⊤⊤resist movement
xt − xt−1∇f(xt) += 0 ∈ X*e.g., viscosity

τIn practice, approximate ϕ (and hence −DF) based on data samples using
function approximators (force matching, score matching), NN/RKHS, e.g.,ϕ≈f=n∑
i=1αik(xi, ⋅ ) ∈ ℋ.

fl𝒫

We will now see two applications of this force-balance relation to robust learningRobust Learning under(Joint) Distribution ShiftKernel DRO under distribution shift
Primal DRO (not solvable as it is)(DRO) minθsup̂
MMD(Q,P)≤ϵ

MMD(Q, P)̂ + λDϕ(Q∥ω) ≤ ϵQl(θ, ξ)

Dual. Adapted from [Kremer et al., Z. 2023]

Kernel DRO Theorem (simpli ed). [Z. et al. 2021]
DRO problem is equivalent to the dual kernelmachine learning problem, i.e., (DRO)=(K).
N1(K) minf(ξi) + ϵ∥f∥ℋ subject to l(θ, ⋅ ) ≤ f
θ,f∈ℋ N ∑i=1Example. Robust least squares

min

Entropy regularization (“interior point method”)

2l(θ, ξ) := ∥A(ξ) ⋅ θ − b∥2−f + linf
P̂ f + ϵ∥f∥ℋ + λ ωϕ* (θ,f∈ℋ {λ )}(t)soft cons. ϕ*(t)=expKLlog-barrier ϕ* (t) = − log (1 − t)log

Geometric intuition: dual kernel function f as robust
surrogate losses ( atten the curve)ff𝔼(testfifl𝔼𝔼l(θ, ⋅ )l(θ, ⋅ )flForce-balance of Kernel DROPrimal DRO:

Lagrangian:minθsup̂MMD(Q,P)≤ϵmin supθ,γ≥0 μ∈Ql(θ, ξ)2̂ ) + γϵμ l(θ, x) − γ ⋅ MMD (μ, μN

2=: f ∈ ℋ

MMS in kernel-MMD

Dual kernel function f as robust
112Lk
2k
inf F(μ) + MMD (μ, μ ) ⟹ −D F =k(x, ⋅ )d(μ − μ )(x) + const .surrogate lossesτ∫μ∈2τ
atten the curve → force balance

Force-balance using function approximation RKHS functions, e.g.,−DF = f+f0, f =∑i=1αik(xi, ⋅ ) ∈ ℋ, f0 ∈ ℝD F = l(θ, ⋅ ) ⟹ force-balance relation: l(θ, ⋅ ) = f + f0 a.e.
𝔼(force matching, score matching)
𝔼𝒫𝒫L2nfl(θ, ⋅ )fl(θ, ⋅ )

Robust Learning under
Structured Distribution ShiftGeneralization errors

From statistical uctuation to structured distribution shift(Mild)

Introduction: Generalization in Machine Learning

chine Learning
Learning task(Strong)

Model mis-specificationf✓ 6= f 8✓ 2 ⇥What can go wrong?

Finite sample biasn⌧1

Distribution shifts
Ptest 6= Ptrain

y := f (x) + ✏n
en i.i.d. observations {(xi , yi )}i=1 of randomvariables(X , Y ) ⇠ Ptrain! Statistical learning!
Useflexible
modelsFigure credit: Heiner Kremer
ng valuesin X as⇥aY,learn apaper
modelf✓ 2020:(NN/non-parametric)
X ! Y and evaluate on Ptesttheory/regularization.
Publishedconferenceat ICLR

pirical risk minimization (ERM):
Common training examples
Test examples
ne a loss function ` : X ⇥y: waterbird
Y ⇥ ⇥ ! R+y: andminimize average
error
landbird
y: waterbird

Published asa conference paper at ICLR 2020
Publishedas a conference paper
at ICLR 2020
a: watera: land
a: land

Qtest

Waterbirds n̂
PtrainCelebAbackgroundbackgroundwaterbird+ landbackground

1X
Commontraining examples
Common training examples
Testexamples
min R(✓) :=`(xi , yi ; f✓y:)waterbird=y: E[`(X,Y;f)]✓y: blond hair
dark
hair
y:blond hairP̂
y:y:waterbird
y: landbird
y: landbird
waterbirdtrain✓2⇥
n a: female
a: malea: malea:a:water
a: landa: watera: landland
landbird
waterbird
i=1

Waterbirds

2f✓y:(x)).
contradiction

backgroundbackground

Heiner Kremery: blondy: hair
entailment
a: femalea: has negation
a: no negation[Heinze‐Deml& Meinshausen
2021]

Robust(P) Theeconomy ML via CMR(P) Read for Slate's take
CelebAflh e.g. `(x, y ; f✓ ) = (yHeiner Kremer! Robustness,Causality

Waterbirds

backgroundbackgroundbackground

y:y:blondblondhair
hair[Sagawaet al. 2020]
a: femaley: dark hair
a: male+ land Robust ML via
CMR+ watery: dark hair y: entailment
a: malea: has negation

14.12.2022(P)There was silence
CelebA

a: male4 / 26

Wasserstein/Kernel
DRO
Test examples
y: waterbird
notsuitablefor(strong)a: landbackgroundstructured
distribution5 / 26
14.12.2022
shifts!
y: blond hair
a: male(Joint) DRO mostl!
Confoundingca
Structured Distribution Shift — Causal Confounding
Causal confounding can lead to much stronger
distribution shifts than those considered in (joint)
distribution shift, e.g., DRO, adversarial robustness.

X : Smoking, Y : C

X: Smoking,Y: Cancer,U: LifestyleU

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

Regression minθY2in[∥Y − gθ(X)∥ ] or DRO does not work Take
into account𝔼𝔼𝔼𝔼

this case.

XU𝔼X

YKernel Method of Moment:Takeconditional
into account genetic pred
moment restriction for causal inference
URobustness against structured distribution shifts instead of(joint-)DRO. Estimating gθ via conditional moment restriction (CMR)[Y − gθ(X) | Z] = 0 ℙZ-a.s.

Generalized Empirical likelihood [Owen, 1988; Qin and Lawless, 1994] withCMR [Bierens, 1982]. Equivalently, generalized method of moment (GMM)

inf Dϕ(Q∥P)̂ s.t.

θ,Q∈Q [(Y − gθ(X))

Th(Z)] = 0, ∀h ∈ ℋZXg✓

YHeiner Kremer

Instrument: Genetic predisposition
for nicotine addiction ZKernel MoM [Kremer et al., Z. 2023] with CMR

𝔼𝔼𝒫

𝒫

1
2
inf
MMD (Q, P)̂ s.t.
θ,Q∈ 2Q [(Y − gθ(X))

Th(Z)] = 0

Rob

Lift the restriction that Q is anatomic distribution

Kernel MoM: duality and algorithm
θKMM= arg min R(θ)
θ12R(θ) := inf MMD (Q, P)̂ s.t.
Q∈ 2Q [(ψ(X; θ))T√(x)T hf (x) + ¥h(Z)] = 0Theorem. [Kremer et al., Z. 2023] The MMD pro le R(θ) has the strongly dual formn

11R(θ) =supf0 +f(xi, zi) − ∥f∥2ℱn∑
2f0 ∈ ℝ, f ∈ ℱ,
i=1
f0 + f(x, z) ≤ ψ(x; θ)T h(z)∀(x, z) ∈×.

Entropy regularization In nite constraint → soft-constraintT
ψ(X;θ)
h(Z)] = 0Q[𝒵𝒳

1infMMD2(Q, P)̂ + λDϕ(Q∥ω) s.t.
θ,Q∈ 2√(x)T h√(x)T h + ≤MMD onlyKL-reg
Log-regf (x) + ¥s.t.

h∈ℋxresults in an unconstrained dualfi𝔼𝔼fi𝔼𝔼𝒫𝒫12[f+f(X,Z)]−∥f∥Pn̂ 0ℱ−2T
φ*f+f(X,Z)−ψ(X;θ)h(Z))]ω[ ϵ ( 0x(t)soft cons. ϕ*(t)=exp
KLlog-barrier ϕ* (t) = − log (1 − t)log

Nonlinear Instrumental Variable Regression

Experimental ResultsKernel MoM: Nonlinear Instrumental
Variable
Y := g (X ; ✓0 ) +⌫(U) + ✏1 Regression
Nonlinear Instrumental Variable Regression
X := ⌘(Z ) + µ(U) + ✏2
Y := g (X ; ✓0 ) + ⌫(U) +Z✏1⇠ P , ✏ ⇠ N (0, )Z1/2
Y := g(X; θ0) + ν(U) + ϵ1
UZX := ⌘(Z ) + µ(U) + ✏2
X := η(Z) +Zμ(U)+ϵg,✓2✓✓⇠ PZ , ✏1/2 ⇠withN (0,)x✓43g (x; ✓) = ✓2 + ✓3 (x Z ✓1 ) + X2 log 1 +Ye 1

Z ∼ PZ, ϵ1/2 ∼ (0,σ)
CMR:
E[Yg(X;✓)|Z]=0P-a.s.
Z✓4 ✓3g(x; withθ) is gnonlinearx,θinboth.
x✓(x; ✓) = ✓2 + ✓3 (x ✓1 ) +log 1 + e 1
20.5

𝒩Takeaway. (Strong) structured distribution shifts(e.g., causal confounding) can be accounted forusing the Kernel MoM + CMR, but not (joint)
DRO, adversarial robustness, …

1.0||µ ° µ0 ||2CMR:
E[Yg(X;✓)|Z]=0P-a.s.
Z
Estimate θ using Kernel MoM with CMR

DeepGMM
Functional GELKMM

Heiner Kremer

0.0103

sample size

Robust Learning with Kernels

104Force-balance of Kernel MoM
Lagrangian:

1sup inf MMD2(Q, P)̂ + γ ⋅γ∈ℝ,h∈ℋ Q 2

Q [(Y − gθ(X))

Th(Z)]

12kMinimizing movement scheme (MMS) in MMD inf F(μ) +MMD (μ, μ )μ∈2γForce balance using function approximation, e.g., kernel functions
n1−DF = f + f0, f =αik([xi, yi, zi], ⋅ ) ∈ ℋ, f0 ∈ ℝτ∑
i=1Since DF = (Y − gθ(X))

Th(Z), the optimal force function approximates the moment function
f + f0 = (Y − gθ(X)) h(Z) a.e.

𝔼𝒫TSummary• We exploited explicitly parametrized dual force functions forrobust learning under joint and structured distribution
shifts. This is inspired by generalized force in gradient ows,
optimal transport, and mechanics.

• The gradient ow force-balance eqns give insights forconstructing robust learning algorithms.

• Kernel DRO: force gives the robusti ed surrogate loss

This talk is mainly based on:
1. Z., Jitkrittum, W., Diehl, M. & Schölkopf, B. Kernel DistributionallyRobust Optimization. AISTATS 2021
2. Kremer, H., Nemmour,Y. , Schölkopf, B. & Z. Estimation Beyond
Data Reweighting: Kernel Method of Moments. ICML 2023

slides & code available: jj-zhu.github.iofl(θ, ⋅ )• Kernel MoM: force gives the robusti ed moment function(Y − gθ(X)) h(Z )T

Postdoc position opening in Berlin: datadriven dynamics modeling for medicalimaging. Contact for info.

flfififlf