Implicit Bias of the JKO Scheme

Wasserstein gradient flow provides a general framework for minimizing an energy functional J over the space of probability measures on a Riemannian manifold (M,g). Its canonical time-discretization, the Jordan-Kinderlehrer-Otto (JKO) scheme, produces for any step size η>0 a sequence of probability distributions ρkη that approximate to first order in η Wasserstein gradient flow on J. But the JKO scheme also has many other remarkable properties not shared by other first order integrators, e.g. it preserves energy dissipation and exhibits unconditional stability for λ-geodesically convex functionals J. To better understand the JKO scheme we characterize its implicit bias at second order in η. We show that ρkη are approximated to order η2 by Wasserstein gradient flow on a modified energy [ Jη(ρ) = J(ρ) - fracη4intM BiglVert nablag fracδ Jδ ρ (ρ) BigrVert22 ,ρ(dx), ] obtained by subtracting from J the squared metric curvature of J times η/4. The JKO scheme therefore adds at second order in η a deceleration in directions where the metric curvature of J is rapidly changing. This corresponds to canonical implicit biases for common functionals: for entropy the implicit bias is the Fisher information, for KL-divergence it is the Fisher-Hyv"arinen divergence, and for Riemannian gradient descent it is the kinetic energy in the metric g. To understand the differences between minimizing J and Jη we study JKO-Flow, Wasserstein gradient flow on Jη, in several simple numerical examples. These include exactly solvable Langevin dynamics on the Bures-Wasserstein space and Langevin sampling from a quartic potential in 1D.