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

Basque Center for Applied Mathematics: BCAM
Bilbao, Basque Country, Spain. October 31st, 2023

Big picture:measure optimization

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 π.

This stochastic dynamics is equivalent to the deterministic
PDE gradient ow in the Wasserstein space [Otto 96].

fl𝒟

𝔼

Motivation: Langevin Monte-CarloMotivation: From statistical learning to distrib. robust learning
Empirical Risk MinimizationN•1
minl(θ, ξi), ξi ∼ P0
θ N∑
i=1

Distributionally Robust Optimization (DRO)
min supθ“Robust” under statistical uctuation
N11̂ ξ) ≤̂ ξ) + (l(θ,l(θ,)
P0iN∑
Ni=1Q∈ℳ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 principled elds ofresearch and breakthroughs for computation

fi

fl

fl

𝔼

Figure credit: H. Kremer, J. Zhu

fl𝒪

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

Information 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Optimal Transportp-Wasserstein

Integral Prob. Metrics

1-Wasserstein

WassersteinFisher-Rao

Kernel MMD

Total Variation

ϕ-divergences

Bregman divergences

Information Divergence

Figure credit: J. Zhu

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}
2fi

fi

fi

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

Kernel maximum-mean discrepancy
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 68CC 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, heat conduction, Fokker Planck equation“steepest” dissipation of energy. [Otto et al 2000s, Ambrosio 2005, …]

The Wasserstein gradient system that generates the WGF is (Prob(X̄), F, W2)

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)

flfl

ff

fl

fl

ff

ff

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

ODE ow: (ℝd, F,∥∥2)
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

Distributional robustness, but what kind?

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

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 kernel
machine learning problem, i.e., (DRO)=(K).
N

1(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)=exp
KLlog-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 DRO
Primal DRO:

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

2=: f ∈ ℋ

MMS in kernel-MMD

11L2k
2kinf F(μ) + MMD (μ, μ ) ⟹ −D F =k(x, ⋅ )d(μ − μ )(x) + const . Dual kernel function f as robust
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)
𝔼𝒫𝒫L2

nfl(θ, ⋅ )fl(θ, ⋅ )

Robust Learning under
Structured Distribution Shift(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✓

Mean regression minθ2[∥Y − gθ(X)∥ ] andY

Take into account𝔼𝔼𝔼

𝔼(distributionally) robust optimization does not work in thiscase.

XU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] with
CMR [Bierens, 1982]. Equivalently, generalized method of moment (GMM)

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

θ,Q∈Q[Y − gθ(X) | Z] = 0 ℙZ-a.s.

ZXg✓

YHeiner Kremer𝔼𝒫𝒫12
infMMD (Q, P)̂ s.t.
θ,Q∈ 2Q [(Y − gθ(X))

TRobInstrument: Genetic predisposition
for nicotine addiction ZKernel MoM [Kremer et al., Z. 2023] with CMR𝔼𝔼Xh(Z)] = 0, ∀h ∈ ℋ
Lift the restriction that Q is anatomic distribution

Kernel MoM: duality and algorithm
θKMM= arg min R(θ)
θ1
2R(θ) := inf MMD (Q, P)̂ s.t.
Q∈ 2Q [(ψ(X; θ))T√(x)T hf (x) + ¥h(Z)] = 0, ∀h ∈ ℋ

Theorem. [Kremer et al., Z. 2023] The MMD pro le R(θ) has the strongly dual form
n1
1R(θ) =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 only
KL-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 Results

Kernel 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) + ϵ1UZ
X := ⌘(Z ) + µ(U) + ✏2X := η(Z) +Zμ(U)+ϵg,✓2✓✓⇠ PZ , ✏1/2 ⇠withN (0,)
x✓43g (x; ✓) = ✓2 + ✓3 (x Z ✓1 ) + X2 log 1 +Ye 1Z ∼ 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 GEL
KMM

Heiner Kremer

0.0103

sample size

Robust Learning with Kernels

104Force-balance of Kernel MoM
Lagrangian:1
sup inf MMD2(Q, P)̂ + γ ⋅γ∈ℝ,h∈ℋ Q 2

Q [(Y − gθ(X))

Th(Z)]

12k
Minimizing 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

for robust learning under joint and structured distribution
shifts.

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

• Kernel DRO: force gives the robusti ed surrogate loss

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

Website for slides, code Beginner tutorial
https://jj-zhu.github.io/ on gradient ows

fl(θ, ⋅ )• Kernel MoM: force gives the robusti ed moment function(Y − gθ(X)) h(Z )TWorkshop on Optimal Transport - OPT & ML
Berlin, March 2024
fififlflf