Duality from Distributionally Robust Learning to
Gradient Flow Force-Balance
Jia-Jie Zhu
Weierstrass Institute for Applied Analysis and Stochastics
Berlin, Germany

DP4ML Workshop, International Conference on Machine Learning (ICML)
Honolulu Hawaii, USA, July 29th, 2023,

Duality in this talk:Primal Measures vs Dual Functions

Duality in this talk:Primal Measures vs Dual Functions
Primal-dual optimization problems

inf F(μ) = sup ℰ( f )

μ∈ℳ

f∈ℱ

Duality in this talk:Primal Measures vs Dual Functions
Primal-dual optimization problems

inf F(μ) = sup ℰ( f )

μ∈ℳ

f∈ℱ

Examples in ML
Generative models

inf
Gθ

Z

μ∈ℳ f∈ℱ { ∫(P, Gθ(Z)) = inf sup

f(x)dP(x) −θ∼μ∫f(gθ(z))dQ(z)}Distributionally robust optimization

inf

N

sup

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

Wasserstein barycentercminαi [Wp(μ, νi)] = minαi supfi dμ + fidνi ,
}
∫
μ∈ℳ ∑
μ∈ℳ ∑ f ∈Ψ { ∫
i
c
i=1
i=1
𝔼

n

𝔼

𝒟

𝔼

𝔼

n

Optimal Transportp-Wasserstein

Integral Prob. Metrics

1-Wasserstein

WassersteinFisher-Rao

Kernel MMD

Total Variation

ϕ-divergences

Bregman divergences

Information Divergence

Figure credit: J. Zhu

Optimal Transportp-Wasserstein

Integral Prob. Metrics

1-Wasserstein

WassersteinFisher-Rao

Kernel MMD

Total Variation

ϕ-divergences

Bregman divergences

Information Divergence

Figure credit: J. Zhu

Static: Duality ofDistributionally Robust Learning

Distributional robustness, but what kind?

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

From Statistical Learning to Distributionally Robust Learning

From Statistical Learning to Distributionally Robust Learning
Empirical Risk MinimizationN1minl(θ, ξi), ξi ∼ P0
θ N∑
i=1

From Statistical Learning to Distributionally Robust Learning
Empirical Risk MinimizationN1minl(θ, ξi), ξi ∼ P0
θ N∑
i=1

𝒪
fl

𝔼

• “Robust” under statistical uctuationN11̂ ξ) ≤̂ ξ) + (
l(
θ,
l(
θ,
)
P0
i
N∑
N
i=1

From Statistical Learning to Distributionally Robust Learning
Empirical Risk MinimizationN1minl(θ, ξi), ξi ∼ P0
θ N∑
i=1

• “Robust” under statistical uctuationN11̂ ξ) ≤̂ ξ) + (
l(
θ,
l(
θ,
)
P0
i
N∑
N
i=1

• Not robust under data distribution shifts,
when Q ( ≠ P0)
δξi

𝒪

fl

𝔼

N
1
PN̂ :=
δξi ⇒ P0 ≠ Q
N∑
i=1From Statistical Learning to Distributionally Robust Learning
Empirical Risk MinimizationN1minl(θ, ξi), ξi ∼ P0
θ N∑
i=1

• “Robust” under statistical uctuationN11̂ ξ) ≤̂ ξ) + (
l(
θ,
l(
θ,
)
P0
i
N∑
N
i=1

• Not robust under data distribution shifts,
when Q ( ≠ P0)
δξi

𝒪

fl

𝔼

𝔼

N
1
PN̂ :=
δξi ⇒ P0 ≠ Q
N∑
i=1Distributionally Robust Optimization (DRO)
min supθQ∈ℳ

Ql(θ, ξ)

From Statistical Learning to Distributionally Robust Learning
Empirical Risk Minimization
N

•

1
min
l(θ, ξi), ξi ∼ P0
θ N∑
i=1
“Robust” under statistical uctuationN11̂ ξ) ≤̂ ξ) + (
l(
θ,
l(
θ,
)
P0
i
N∑
N
i=1

• Not robust under data distribution shifts,
when Q ( ≠ P0)
δξi

𝒪

fl

𝔼

𝔼

N
1
PN̂ :=
δξi ⇒ P0 ≠ Q
N∑
i=1Distributionally Robust Optimization (DRO)
min supθQ∈ℳ

Ql(θ, ξ)Worst-case distribution Q within the ambiguity set ℳ[Delage & Ye 2010] in certain geometry.

From Statistical Learning to Distributionally Robust Learning
Empirical Risk Minimization
N

•

1
min
l(θ, ξi), ξi ∼ P0
θ N∑
i=1
“Robust” under statistical uctuationN11̂ ξ) ≤̂ ξ) + (
l(
θ,
l(
θ,
)
P0
i
N∑
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

𝒪

𝔼

𝔼

N
1
PN̂ :=
δξi ⇒ P0 ≠ Q
N∑
i=1
fl

Figure credit: H. Kremer

From Statistical Learning to Distributionally Robust Learning
Empirical Risk Minimization
N

•

1
min
l(θ, ξi), ξi ∼ P0
θ N∑
i=1
“Robust” under statistical uctuationN11̂ ξ) ≤̂ ξ) + (
l(
θ,
l(
θ,
)
P0
i
N∑
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?

𝒪

𝔼

𝔼

N
1
PN̂ :=
δξi ⇒ P0 ≠ Q
N∑
i=1
fl

Figure credit: H. Kremer

From Statistical Learning to Distributionally Robust Learning
Empirical Risk Minimization
N

•

1
min
l(θ, ξi), ξi ∼ P0
θ N∑
i=1
“Robust” under statistical uctuationN11̂ ξ) ≤̂ ξ) + (
l(
θ,
l(
θ,
)
P0
i
N∑
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:

𝒪

𝔼

𝔼

N
1
PN̂ :=
δξi ⇒ P0 ≠ Q
N∑
i=1
fi

fl

Figure credit: H. Kremer

From Statistical Learning to Distributionally Robust Learning
Empirical Risk Minimization
N

•

1
min
l(θ, ξi), ξi ∼ P0
θ N∑
i=1
“Robust” under statistical uctuationN11̂ ξ) ≤̂ ξ) + (
l(
θ,
l(
θ,
)
P0
i
N∑
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:

𝒪

𝔼

𝔼

N
1
PN̂ :=
δξi ⇒ P0 ≠ Q
N∑
i=1Information geometry [S. Amari et al.] e.g. descent in
Fisher-Rao geometryfiflFigure credit: H. Kremer

From Statistical Learning to Distributionally Robust Learning
Empirical Risk Minimization
N

•

1
min
l(θ, ξi), ξi ∼ P0
θ N∑
i=1
“Robust” under statistical uctuationN11̂ ξ) ≤̂ ξ) + (
l(
θ,
l(
θ,
)
P0
i
N∑
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:

fl

fi

𝒪

fl

fl

𝔼

𝔼

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

Background: Wasserstein Geometry

Background: Wasserstein Geometry
De nition. The p-Wasserstein distance between probability measuresdP, Q on ℝ (with p nite moments, p ≥ 1) is de ned through the following
Kantorovich problem

fi

fi

∫

p

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

fi

fi

∫

p

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

fi

fi

∫

p

| x − y | dΠ(x, y)(1)(2)
π# Π = P, π# Π = Q}fip
Wp (P, Q) := inf {xyBackground: Wasserstein Geometry
De nition. The p-Wasserstein distance between probability measuresdP, Q on ℝ (with p nite moments, p ≥ 1) is de ned through the following
Kantorovich problemp
Wp (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 | }
∫
∫
p

d

fi

fi

fi

2-Wasserstein space (Prob(ℝ ), W2) is a geodesic metric space.

xyBackground: Wasserstein Geometry
De nition. The p-Wasserstein distance between probability measuresdP, Q on ℝ (with p nite moments, p ≥ 1) is de ned through the following
Kantorovich problemp
Wp (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 | }
∫
∫
p

d

2-Wasserstein space (Prob(ℝ ), W2) is a geodesic metric space.
Dynamic formulation: Benamou–Brenierd| vt | dμtdt μ0 = P, μ1 = Q, μt + div(vt μt) = 0}
∫
dt
0
2

fi

fi

fi

2
W2 (P, Q) = min {
∫

1

x

y

Background: Wasserstein Geometry
De nition. The p-Wasserstein distance between probability measuresdP, Q on ℝ (with p nite moments, p ≥ 1) is de ned through the following
Kantorovich problemp
Wp (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 | }
∫
∫
p

d

2-Wasserstein space (Prob(ℝ ), W2) is a geodesic metric space.
Dynamic formulation: Benamou–Brenier
2
W2 (P, Q) = min {∫1d| vt | dμtdt μ0 = P, μ1 = Q, μt + div(vt μt) = 0}
∫
dt
0
2

Example. Entropy-OT [Cuturi 2013] Duality leads to faster computation

fi

fi

fi

inf cdΠ + λDϕ(Π∥P ⊗ Q)∫xy

Background: Wasserstein Geometry
De nition. The p-Wasserstein distance between probability measuresdP, Q on ℝ (with p nite moments, p ≥ 1) is de ned through the following
Kantorovich problemp
Wp (P, Q) := inf {∫p| x − y | dΠ(x, y)

x

y

(1)
(2)
π# Π = P, π# Π = Q}

(Dual Kantorovich problem)= sup { ψ1(x) dP(x) + ψ2(y) dQ(y) ψ1(x) + ψ2(y) ≤ | x − y | }
∫
∫
p

d

2-Wasserstein space (Prob(ℝ ), W2) is a geodesic metric space.
Dynamic formulation: Benamou–Brenier
2
W2 (P, Q) = min {∫1d| vt | dμtdt μ0 = P, μ1 = Q, μt + div(vt μt) = 0}
∫
dt
0
2

Example. Entropy-OT [Cuturi 2013] Duality leads to faster computation

fi

fi

fi

inf cdΠ + λDϕ(Π∥P ⊗ Q)∫

Figure credit: Wiki., M. Cuturi, A. Genevay

Background: “Kernel Geometry”

Background: “Kernel Geometry”
De nition. Kernel Maximum-Mean Discrepancy (MMD)
associated with (PSD) kernel k (e.g., k(x, x′) := eMMD(P, Q) :=
d

￼

￼
fi

−|x−x′| /σk(x, ⋅ )dP− k(x, ⋅ )dQ .
ℋ∫∫(Prob(ℝ ), MMD) is a (simple) metric space.

2

)

}

MMD(P, Q)

Background: “Kernel Geometry”
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 .
ℋ
∫
∫

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∥ ≤1

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

ℋ

23/34

fi

fi

￼

𝒳

￼

ℋ 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, H. Kremer

Background: “Kernel Geometry”
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 .
ℋ
∫
∫

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∥ ≤1

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

ℋ is the reproducing kernel Hilbert space ℋ (RKHS),
which satis es f(x) = ⟨f, ϕ(x)⟩ℋ, ∀f ∈ ℋ, x ∈ ,ϕ(x) := k(x, ⋅ ) is the canonical feature of ℋ.
Example. Entropy-MMD [Kremer et al. 2023]

MMD(Q, P)̂ + λDϕ(Q∥ω)

23/34

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

fi

fi

￼

𝒳

￼

Duality leads to “interior point method” for prob. distributions
Figure credit: W. Jitkrittum, J. Zhu, H. Kremer

Previous work: Kernel DRO

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

sup

̂
MMD(Q,P)≤ϵ𝔼(DRO) minQl(θ, ξ)Previous work: Kernel DRO
Primal DRO (not solvable as it is)(DRO) minθsup

̂MMD(Q,P)≤ϵQl(θ, ξ)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

fi

𝔼

1
(K) minf(ξi) + ϵ∥f∥ℋ subject to l(θ, ⋅ ) ≤ f
θ,f∈ℋ N ∑i=1

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

̂MMD(Q,P)≤ϵQl(θ, ξ)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=1
Geometric intuition: dual kernel function f as robust
surrogate losses ( atten the curve)ffl(θ, ⋅ )fifl𝔼l(θ, ⋅ )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∥2

Ql(θ, ξ)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).
N

1(K) minf(ξi) + ϵ∥f∥ℋ subject to l(θ, ⋅ ) ≤ f
θ,f∈ℋ N ∑i=1

(test distribution)(a) Robust least squares loss

Figure 3: (a) This plot depicts the test loss of algorithms. Al
ran 10 independent trials. In each trial, we solved K-DRO to oof 500 samples. We then vary the perturbation from 0 to 4.
f0⇤ + f ⇤ . (black) is the function l(✓⇤ , ·). The pink bars depict abars the empirical distribution. We can observe that f0⇤ + f ⇤ toworst-case distribution P ⇤ (pink dots). Note f ⇤ (normalized) c
the two distributions.

Geometric intuition: dual kernel function f as robust
surrogate losses ( atten the curve)ffl(θ, ⋅ )l(θ, ⋅ )
fi

𝔼

227
fl

(b) G

4

Numerical studies

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∥2

Ql(θ, ξ)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).
N

1(K) minf(ξi) + ϵ∥f∥ℋ subject to l(θ, ⋅ ) ≤ f
θ,f∈ℋ N ∑i=1

(test distribution)

Robustifying
with DRO(a) Robust least squares
loss

f asplot
witness(test)
Figure 3: (a) This
depicts thetestfunction
loss of algorithms. Al
ran 10 independent trials. In each trial, we solved K-DRO to oof 500 samples. We then vary the perturbation from 0 to 4.
f0⇤ + f ⇤ . (black) is the function l(✓⇤ , ·). The pink bars depict abars the empirical distribution. We can observe that f0⇤ + f ⇤ tofworst-case distribution P ⇤ (pink dots). Note f ⇤ (normalized)cthe two distributions.

Geometric intuition: dual kernel function f as robust
surrogate losses ( atten the curve)ffl(θ, ⋅ )l(θ, ⋅ )
fi

𝔼

227
fl

(b) G

4

Numerical studies

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∥2

Ql(θ, ξ)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).
N

1(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)

Robustifying
with DRO(a) Robust least squares
loss

f asplot
witness(test)
Figure 3: (a) This
depicts thetestfunction
loss of algorithms. Al
ran 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 a
0bars the empirical distribution. We can observe that f0⇤ + f ⇤ tofworst-case distribution P ⇤ (pink dots). Note f ⇤ (normalized)cthe two distributions.

fi

𝔼

227
fl

(b) G

4

Numerical studies

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∥2

Ql(θ, ξ)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).
N

1(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)

Robustifying
with DRO(a) Robust least squares
loss

f asplot
witness(test)
Figure 3: (a) This
depicts thetestfunction
loss of algorithms. Al
ran 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 a
0bars the empirical distribution. We can observe that f0⇤ + f ⇤ tofworst-case distribution P ⇤ (pink dots). Note f ⇤ (normalized)cthe two distributions.

fi

𝔼

227
fl

(b) G

4

Numerical studies

Duality perspective

Duality
perspective
2-Wasserstein

Kernel DRO [Z. et al. 2021]min

Primal:

𝔼

𝔼

θ

sup

̂
W2(P,P)≤ϵ

l(θ,
ξ)
P

Primal:

min
θ

sup

̂
MMD(P,P)≤ϵ

l(θ,
ξ)
P

Duality
perspective
2-Wasserstein

Kernel DRO [Z. et al. 2021]min

Primal:θ

sup̂
W2(P,P)≤ϵ

l(θ,
ξ)
P

N

1
2
λ∥⋅∥
2
min
lθ (ξi) + λϵ
θ,λ>0 N ∑
i=1

Dual:

where Moreau envelope
2

𝔼

𝔼

λ∥⋅∥
2
lθ (x) := sup l(θ, u) − λ∥u − x∥uPrimal:

min
θ

sup

̂
MMD(P,P)≤ϵ

l(θ,
ξ)
P

Duality
perspective
2-Wasserstein

Kernel DRO [Z. et al. 2021]min

Primal:θ

sup̂
W2(P,P)≤ϵ

l(θ,
ξ)
P

Primal:

N

1
2
λ∥⋅∥
2
min
lθ (ξi) + λϵ
θ,λ>0 N ∑
i=1

Dual:

where Moreau envelope
2

λ∥⋅∥
2
lθ (x) := sup l(θ, u) − λ∥u − x∥
u

min
θ

sup

̂
MMD(P,P)≤ϵ

l(θ,
ξ)
P

N

Dual:

1
min
f(ξi) + ϵ∥f∥ℋ
θ,f∈ℋ N ∑i=1

s . t . l(θ, ξ) ≤ f(ξ), ∀ξ a.e.
f𝔼𝔼l(θ, ⋅ )Duality
perspective
2-Wasserstein

Kernel DRO [Z. et al. 2021]min

Primal:θ

sup̂
W2(P,P)≤ϵ

l(θ,
ξ)
P

Primal:

N

1
2
λ∥⋅∥
2
min
lθ (ξi) + λϵ
θ,λ>0 N ∑
i=1

Dual:

where Moreau envelope
2

λ∥⋅∥
2
lθ (x) := sup l(θ, u) − λ∥u − x∥u𝔼𝔼Considerations from WGF theory- l is nonconvex (e.g., DNN, g-non-cvx)- Nonlinear (in measure) energies

min
θ

sup

̂
MMD(P,P)≤ϵ

l(θ,
ξ)
P

N

Dual:

1
min
f(ξi) + ϵ∥f∥ℋ
θ,f∈ℋ N ∑i=1

s . t . l(θ, ξ) ≤ f(ξ), ∀ξ a.e.
fl(θ, ⋅ )Duality
perspective
2-Wasserstein

Kernel DRO [Z. et al. 2021]min

Primal:θ

sup̂
W2(P,P)≤ϵ

l(θ,
ξ)
P

Primal:

N

1
2
λ∥⋅∥
2
min
lθ (ξi) + λϵ
θ,λ>0 N ∑
i=1

Dual:

where Moreau envelope
2

λ∥⋅∥
2
lθ (x) := sup l(θ, u) − λ∥u − x∥
u

fl

𝔼

𝔼

Considerations from WGF theory- l is nonconvex (e.g., DNN, g-non-cvx)- Nonlinear (in measure) energies

min
θ

sup

̂
MMD(P,P)≤ϵ

l(θ,
ξ)
P

N

Dual:

1
min
f(ξi) + ϵ∥f∥ℋ
θ,f∈ℋ N ∑i=1

s . t . l(θ, ξ) ≤ f(ξ), ∀ξ a.e.
fl(θ, ⋅ )Nonlinear kernel approx. as robust
surrogate losses ( atten the curve)

Dynamic: Duality ofGradient Flow

From static DRO to JKO scheme for gradient flows
DRO’s Wasserstein measure optimization is not new.
min
◊

sup

W2 (P,P̂)Æ‘

EP l(◊, ›)2
min sup EP l(◊, ›) ≠ “ · W2 (P, P̂)◊P

From static DRO to JKO scheme for gradient flows
DRO’s Wasserstein measure optimization is not new.
min
◊

sup

W2 (P,P̂)Æ‘

EP l(◊, ›)2
min sup EP l(◊, ›) ≠ “ · W2 (P, P̂)◊
P

Wasserstein gradient flow [Otto et al. 90s-2000s]. The Fokker-Planck equation”Fˆt µ + Ò · (µÒ [µ]) = 0”µ
is the gradient-flow equation of energy F in (Prob(X̄ ), W2 ).

From static DRO to JKO scheme for gradient flows
DRO’s Wasserstein measure optimization is not new.
min
◊

sup

W2 (P,P̂)Æ‘

EP l(◊, ›)2
min sup EP l(◊, ›) ≠ “ · W2 (P, P̂)◊
P

Wasserstein gradient flow [Otto et al. 90s-2000s]. The Fokker-Planck equation”Fˆt µ + Ò · (µÒ [µ]) = 0”µ
is the gradient-flow equation of energy F in (Prob(X̄ ), W2 ).
Jordan-Kinderlehrer-Otto (JKO) scheme or Minimizing Movement Scheme (MMS):
k+1

µ

1 2
k
œ inf F (µ) + W2 (µ, µ )µœP
2·generalizes the DRO dual reformulation of DRO to nonlinear-in-measure F .

Duality in gradient flow dynamics: nonlinear ODE
ẋ (t) = ≠Òf (x (t))

Duality in gradient flow dynamics: nonlinear ODE
ẋ (t) = ≠Òf (x (t))

ẋ (t) œ X provides the rate (or velocity) (we can see)

Duality in gradient flow dynamics: nonlinear ODE
ẋ (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)Duality in gradient flow dynamics: nonlinear ODE
ẋ (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 ,ú

If X

IR : X æ X is the Riesz isomorphism.
úúúúX : u̇ œ ˆR (µ, ≠DF ) µ Tu M (rate) vs 0 œ DF + ˆR(µ, µ̇) µ Tu M (force)Duality in gradient flow dynamics: nonlinear ODE
ẋ (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 ,ú

If X

IR : X æ X is the Riesz isomorphism.
úúúúX : u̇ œ ˆR (µ, ≠DF ) µ Tu M (rate) vs 0 œ DF + ˆR(µ, µ̇) µ Tu M (force)

Energy dissipation balance (equality) via Fenchel(-Young) duality and optimalitydf (x (t)) =X ú ÈÒf (x (t)), ẋ ÍXdt

Duality in gradient flow dynamics: nonlinear ODE
ẋ (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 ,ú

If X

IR : X æ X is the Riesz isomorphism.
úúúúX : u̇ œ ˆR (µ, ≠DF ) µ Tu M (rate) vs 0 œ DF + ˆR(µ, µ̇) µ Tu M (force)

Energy dissipation balance (equality) via Fenchel(-Young) duality and optimalityd2f (x (t)) =X ú ÈÒf (x (t)), ẋ ÍX = ≠ÎÒf (x (t))Î
dt

Duality in gradient flow dynamics: nonlinear ODE
ẋ (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 ,ú

If X

IR : X æ X is the Riesz isomorphism.
úúúúX : u̇ œ ˆR (µ, ≠DF ) µ Tu M (rate) vs 0 œ DF + ˆR(µ, µ̇) µ Tu M (force)

Energy dissipation balance (equality) via Fenchel(-Young) duality and optimality
!
"
d
1
1
2
2
2
f (x (t)) =X ú ÈÒf (x (t)), ẋ ÍX = ≠ÎÒf (x (t))Î = ≠ Îẋ Î + ÎÒf (x )Î
dt
22

Duality in gradient flow dynamics: nonlinear ODE
ẋ (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 ,ú

If X

IR : X æ X is the Riesz isomorphism.
úúúúX : u̇ œ ˆR (µ, ≠DF ) µ Tu M (rate) vs 0 œ DF + ˆR(µ, µ̇) µ Tu M (force)

Energy dissipation balance (equality) via Fenchel(-Young) duality and optimality
!
"
d
1
1
2
2
2
f (x (t)) =X ú ÈÒf (x (t)), ẋ ÍX = ≠ÎÒf (x (t))Î = ≠ Îẋ Î + ÎÒf (x )Î
dt
2
2
Energy does not necessarily decrease along non-solutions, i.e., only inequality
!1
"
d
1
2
2
f (z(t)) Ø ≠ ÎżÎ + ÎÒf (z(t))Î .
dt
22

Duality in the Wasserstein gradient flow

Wasserstein gradient flow in the rate form (primal; vs. force-balance)
ˆt µ = ≠KOtto (µ) DFDuality in the Wasserstein gradient flow

Wasserstein gradient flow in the rate form (primal; vs. force-balance)
ˆt µ = ≠KOtto (µ) DF = Ò · (µÒDF ).

Duality in the Wasserstein gradient flow

Wasserstein gradient flow in the rate form (primal; vs. force-balance)
ˆt µ = ≠KOtto (µ) DF = Ò · (µÒDF ).

In (Prob(X̄ ), F , W2 ), Fenchel(-Young) duality yields the Energy dissipation balance(equality) [Ambrosio et al. 2007]
d
1 Õ
1 ≠
2
2
F (µ(t)) = ≠ |µ |W2 (t) ≠ |Ò F |W2 (µ(t))
dt22Duality in the Wasserstein gradient flow

Wasserstein gradient flow in the rate form (primal; vs. force-balance)
ˆt µ = ≠KOtto (µ) DF = Ò · (µÒDF ).

In (Prob(X̄ ), F , W2 ), Fenchel(-Young) duality yields the Energy dissipation balance(equality) [Ambrosio et al. 2007]
d
1 Õ
1 ≠
2
2
F (µ(t)) = ≠ |µ |W2 (t) ≠ |Ò F |W2 (µ(t))
dt22For (Boltzmann) entropy F (u) = ﬂ log ﬂ, the metric slope is2|Ò F |W2 (µ(t)) =≠⁄2|Ò log ﬂ| ﬂ dx

Duality in the Wasserstein gradient flow

Wasserstein gradient flow in the rate form (primal; vs. force-balance)
ˆt µ = ≠KOtto (µ) DF = Ò · (µÒDF ).

In (Prob(X̄ ), F , W2 ), Fenchel(-Young) duality yields the Energy dissipation balance(equality) [Ambrosio et al. 2007]
d
1 Õ
1 ≠
2
2
F (µ(t)) = ≠ |µ |W2 (t) ≠ |Ò F |W2 (µ(t))
dt22For (Boltzmann) entropy F (u) = ﬂ log ﬂ, the metric slope is2|Ò F |W2 (µ(t)) =≠⁄2|Ò log ﬂ| ﬂ dx

However, for some nonlinear (in measure) energy (e.g., in variational inference)”F
F (µ) = DKL (µÎﬁ),[µ] = log ﬂ ≠ log ﬁ,”µ
density ﬂ :=

dµ
”F
and
force
field
[µ]
are
not
accessibleifµ
isatomic.
dL”µ

Kernel gradient flow as dual space force-balance
Motivated by the “Kernel DRO-type” derivation in [Zhu et al.’21, Kremer et al.’23],

Kernel gradient flow as dual space force-balance
Motivated by the “Kernel DRO-type” derivation in [Zhu et al.’21, Kremer et al.’23],
Proposition(informal). The gradient flow equation for (P(X̄ ), F , MMD) is given by thedual space (force-balance) kernel gradient flow”Fk ú µ̇ = ≠g œ H, where Òg = Ò[µ] µ-a.e.
”µ
where convolution k ú µ :=sk(x , ·)µ( dx ).

Kernel gradient flow as dual space force-balance
Motivated by the “Kernel DRO-type” derivation in [Zhu et al.’21, Kremer et al.’23],
Proposition(informal). The gradient flow equation for (P(X̄ ), F , MMD) is given by thedual space (force-balance) kernel gradient flow”Fk ú µ̇ = ≠g œ H, where Òg = Ò[µ] µ-a.e.
”µ
where convolution k ú µ :=sk(x , ·)µ( dx ). If F is entropy, Òg “matches the score” .

Kernel gradient flow as dual space force-balance
Motivated by the “Kernel DRO-type” derivation in [Zhu et al.’21, Kremer et al.’23],
Proposition(informal). The gradient flow equation for (P(X̄ ), F , MMD) is given by thedual space (force-balance) kernel gradient flow”Fk ú µ̇ = ≠g œ H, where Òg = Ò[µ] µ-a.e.
”µ
where convolution k ú µ :=sk(x , ·)µ( dx ). If F is entropy, Òg “matches the score” .

Compared with the Wasserstein GF of entropy, our kernel geometry approximates the(unavailable) “score function” Òg = Ò log ﬂ in a principled geometry.
This gives the interpretation of the dual kernel function in dynamics
g is the approximate (thermodynamic) force field.

Back to (kernel) robust learning
Motivated by our insight so far, we have a “dynamic formulation” of the dual DRO
problem [Zhu et al. 2021]min
sup
EP l(◊, ›),◊MMD(P,P̂)Æ‘Back to (kernel) robust learning
Motivated by our insight so far, we have a “dynamic formulation” of the dual DRO
problem [Zhu et al. 2021]min
sup
EP l(◊, ›),◊MMD(P,P̂)Æ‘the distribution shift (a.k.a. adversarial attack) is modeled by the dynamical system of
the dual force-balance kernel gradient flowk ú µ̇ = ≠g,µ(0) = P̂, µ(T ) = P.

where Òg(x ) approximates the gradient Òl(◊, ›). (see also an alternative using kernel
mirror prox. [Dvurechensky & Zhu])

SummarySummary• Two types of duality: static (DRO) to dynamic (GF):Summary• Two types of duality: static (DRO) to dynamic (GF):• duality in dynamics: force vs. rateSummary• Two types of duality: static (DRO) to dynamic (GF):• duality in dynamics: force vs. rate
• Role of the dual kernel function in this talk

Summary• Two types of duality: static (DRO) to dynamic (GF):• duality in dynamics: force vs. rate
• Role of the dual kernel function in this talkfl

• robust surrogate loss ( atten the curve)Summary• Two types of duality: static (DRO) to dynamic (GF):• duality in dynamics: force vs. rate
• Role of the dual kernel function in this talk• robust surrogate loss ( atten the curve)

fl• optimal test fcn. for two-sample test

Summary• Two types of duality: static (DRO) to dynamic (GF):• duality in dynamics: force vs. rate
• Role of the dual kernel function in this talk• robust surrogate loss ( atten the curve)• optimal test fcn. for two-sample test

flfi

• approximate force

eldSummary• Two types of duality: static (DRO) to dynamic (GF):• duality in dynamics: force vs. rate
• Role of the dual kernel function in this talk• robust surrogate loss ( atten the curve)• optimal test fcn. for two-sample test
• approximate force

eld

fl

fi

• Our dual kernel scheme (DRO, GF) is designed to treatSummary• Two types of duality: static (DRO) to dynamic (GF):• duality in dynamics: force vs. rate
• Role of the dual kernel function in this talk• robust surrogate loss ( atten the curve)• optimal test fcn. for two-sample test
• approximate force

eld• Our dual kernel scheme (DRO, GF) is designed to treat• Energy that’s the integral of nonlinear functions or
nonlinear in measures

F(μ) = V dμ,∫

F(μ) = ϕ(ρ)∫(μ = ρ ⋅ ℒ)

flfi

which are challenging for computation using the WGF(complication due to W-geodesics).

Summary• Two types of duality: static (DRO) to dynamic (GF):• duality in dynamics: force vs. rate
• Role of the dual kernel function in this talk• robust surrogate loss ( atten the curve)• optimal test fcn. for two-sample test
• approximate force

eld• Our dual kernel scheme (DRO, GF) is designed to treat• Energy that’s the integral of nonlinear functions or
nonlinear in measures

F(μ) = V dμ,∫

F(μ) = ϕ(ρ)∫(μ = ρ ⋅ ℒ)

which are challenging for computation using the WGF(complication due to W-geodesics).

• Other important uses of the dual kernel function: Causal
flfiinference, conditional moments, (robust) control and RL

Summary• Two types of duality: static (DRO) to dynamic (GF):• duality in dynamics: force vs. rateZ., Jitkrittum, W., Diehl, M. & Schölkopf, B. Kernel Distributionally• Role of the dual kernel function in this talk• robust surrogate loss ( atten the curve)

Robust Optimization. AISTATS 2021
Kremer, H., Nemmour,Y. , Schölkopf, B. & Z. Estimation Beyond Data
Reweighting: Kernel Method of Moments. ICML 2023
P. Dvurechensky, Z., Kernel Mirror Prox and RKHS Gradient Flow

• optimal test fcn. for two-sample test
• approximate force

References:for Mixed Functional Nash Equilibrium. Preprint

eld• Our dual kernel scheme (DRO, GF) is designed to treat

Download the slides:• Energy that’s the integral of nonlinear functions or
nonlinear in measures

F(μ) = V dμ,∫

F(μ) = ϕ(ρ)∫(μ = ρ ⋅ ℒ)

which are challenging for computation using the WGF(complication due to W-geodesics).

• Other important uses of the dual kernel function: Causal
flfiinference, conditional moments, (robust) control and RL

Website: jj-zhu.github.io