Efficient displacement convex optimization with particle gradient descent

TL;DR

This paper analyzes particle gradient descent for optimizing displacement convex functions, providing theoretical guarantees on convergence and complexity, with applications to neural network function approximation.

Contribution

It establishes the first theoretical guarantees for particle gradient descent on displacement convex functions with finite particles.

Findings

O(1/ε²) particles suffice for ε-optimal solutions

Complexity bounds are improved for smooth displacement convex functions

Applications demonstrated in neural network function approximation

Abstract

Particle gradient descent, which uses particles to represent a probability measure and performs gradient descent on particles in parallel, is widely used to optimize functions of probability measures. This paper considers particle gradient descent with a finite number of particles and establishes its theoretical guarantees to optimize functions that are \emph{displacement convex} in measures. Concretely, for Lipschitz displacement convex functions defined on probability over $\mathbb{R}^d$, we prove that $O(1/\epsilon^2)$ particles and $O(d/\epsilon^4)$ computations are sufficient to find the $\epsilon$-optimal solutions. We further provide improved complexity bounds for optimizing smooth displacement convex functions. We demonstrate the application of our results for function approximation with specific neural architectures with two-dimensional inputs.