Outline

1. Optimal transport

2. Wasserstein gradient flows

3. Numerical method

4. Movies

1. Optimal transport

2. Wasserstein gradient flows

3. Numerical method

4. Movies

Setting

Two mass densities over $\Omega\subset\mathbb{R}^d$, with same total mass,

$$\int_\Omega \mu(x)dx=\int_\Omega\nu(y)dy$$

$\mu$

How to optimally transport $\mu$ to $\nu$?

Map $T\colon\Omega\to\Omega$, measure $\mu$

Pushforward: $(T_{\#}\mu)(A)=\mu(T^{-1}(A))$

D E F I N I T I O N

$\mu$

$T_{\#}\mu$

$T(x)$

Cost to move mass $\mu(dx)$ from $x$ to $T(x)$ is

$$|T(x)-x|^2 \,\mu(dx)$$

$T(x)$

$$\inf_{T}\int_\Omega |T(x)-x|^2 \,\mu(dx)=:W_2^2(\mu,\nu)$$

subject to: $T_{\#}\mu=\nu$

Features: 

1) A distance between $\mu$ and $\nu$

2) An optimal map $T$

3) Formally a Riemannian metric on the manifold of measures with geodesics $$\rho_t = ((1-t)x+tT(x))_{\#}\mu$$

Example 1

$2048\times 2048$ points

Example 2

Caffarelli's counterexample

$2048\times 2048$ points

Summary

A distance $W_2$ over measures, with “Riemannian” structure

1. Optimal transport

2. Wasserstein gradient flows

3. Numerical method

4. Movies

“$\dot\rho_t = -\mathrm{grad}_{W_2}U(\rho_t)$”

Functional $U(\rho)$, want to make sense of 

Look at ”implicit Euler”

(JKO)

The PDEs

$$\partial_t\rho+\mathrm{div}(\rho\, v)=0,$$

$$v =-\nabla\phi,$$

$$\phi=\delta U(\rho)$$

on domain $\Omega\subset\mathbb{R}^d$, no flux out and with initial condition $\rho(t=0)=\rho_0$.

When $\tau\to 0$,

$m\to 1: U(\rho) = \int_\Omega \rho\log\rho + V\rho$

Dual formulation

$$\sup_{\phi,\psi}\,\langle\psi,\mu\rangle - U^*(\phi)$$

over $(\phi,\psi)$ s.t.

$$\psi(x)-\phi(y)\le\frac{|x-y|^2}{2\tau}.$$

$$\inf_\rho U(\rho)+\frac{1}{2\tau}W_2^2(\mu,\rho).$$

Dual formulations

with $\phi^c(x)=\inf_y\phi(y)+\frac{|x-y|^2}{2\tau}$,

        $\psi^c(y)=\sup_x\psi(x)-\frac{|x-y|^2}{2\tau}$.

☇

$$\sup_\phi\langle\phi^c,\mu\rangle-U^*(\phi)=:J(\phi)$$

$$\sup_\psi\langle\psi,\mu\rangle-U^*(\psi^c)=:I(\psi)$$

Dual formulations

→ unconstrained concave maximization problems

→ Recover $\rho^*$ from $\phi^*$ by

$$\rho^*=\delta U^*(\phi^*)$$

$$\sup_\phi\langle\phi^c,\mu\rangle-U^*(\phi)=:J(\phi)$$

$$\sup_\psi\langle\psi,\mu\rangle-U^*(\psi^c)=:I(\psi)$$

The power of duality

$$U(\rho)=\int_\Omega u_\infty(\rho(x))+V(x)\rho(x)\,dx.$$

Summary

• Wasserstein gradient flow: $$\inf_\rho U(\rho)+\frac{1}{2\tau}W_2^2(\mu,\rho)$$

• Dual problem is better behaved: $$\sup_\phi J(\phi)$$ or $$\sup_\psi I(\psi)$$

1. Optimal transport

2. Wasserstein gradient flows

3. Numerical method

4. Movies

Back-and-forth algorithm

$H$ is the Sobolev space

$$\|h\|_H^2=\int_\Omega \alpha_1|\nabla h(x)|^2+\alpha_2|h(x)|^2\,dx$$

$$J(\phi)=\langle\phi^c,\mu\rangle-U^*(\phi)$$

$$I(\psi)=\langle\psi,\mu\rangle-U^*(\psi^c)$$

Recall

Jacobs, Lee, L. (’21)

A L G O R I T H M

Why $H$ ?

Assume that

$$0\le -\delta^2\!J(\phi)(h,h) \le  \lVert h\rVert_H^2$$

for any $\phi,h\in H$. Then the iterations

$$\phi_{k+1}=\phi_k+\nabla_{\!H} J(\phi_k)$$

converge to the supremum of $J$.

Fundamental lemma in optimization:

Why $H$ ?

Want Hessian bound

$$0\le -\delta^2\!J(\phi)(h,h)\le \|h\|_H^2$$

Why $H$ ?

$$U^*(\phi)=\int_\Omega u^*_\infty(\phi(x))\,dx$$

Back-and-forth is fast

Operations on a grid with $N$ points:

 - $c$-transform: $O(N)$

 - $T_{\phi_n} : O(N)$

 - $\Delta^{-1}:O(N\ln N)$

Summary

Novelty of the approach:

1) a careful analysis of the Hessian of the objective function 

2) the back-and-forth scheme which boosts the convergence

4. Movies!

Slow diffusion (porous medium eq)

$V(x)=-\sin(5\pi x_1)\sin(3\pi x_2)$

$512\times 512$ points

$m=2$

$m=4$

Slow diffusion

$m=4$

$V(x)=\|x-a\|^2$

$512\times 512$ points

Incompressible

$V(x)=\|x-a\|^2$

$1024\times 1024$ points

Aggregation-diffusion

$$U(\rho)=\int \rho(x)^3dx+\iint |x-y|^2\rho(x)\rho(y)\,dxdy$$

Thanks!

Paper: https://arxiv.org/pdf/2011.08151.pdf

Code: https://github.com/wonjunee/wgfBFMcodes

Doc: https://wasserstein-gradient-flows.netlify.app/