Optimization over probability measures

I develop accelerated dynamics for optimization problems in the space of probability measures.

Hamiltonian flows and acceleration

Many problems in machine learning seek a probability distribution that minimizes an objective functional. Wasserstein gradient flow extends gradient descent to this setting, but incorporating momentum requires accounting for the geometry of the space of measures.

In our work on Hamiltonian flows for accelerated optimization, we introduce a momentum coordinate and evolve probability measures in phase space. This gives counterparts of heavy-ball, Nesterov, and variational acceleration methods. Using optimal transport and Lyapunov analysis, we establish accelerated continuous-time convergence rates under suitable convexity and regularity assumptions, supported by numerical experiments.

Four views: particles with position x, a measure over x, particles with position and momentum (x, v), and a joint measure over (x, v). Two particles sharing a position become distinguishable through momentum; the phase-space density has two lobes with the same position marginal. — Particle and measure views of position and phase space. Momentum distinguishes particles that share a position; taking the position marginal removes this information. Hamiltonian dynamics evolve the joint phase-space measure. Densities are shown schematically. Paper ↗

For papers, preprints, and related work, see Publications. Research code is available on GitHub ↗.