From Saddle Points Toward Global Minima: A Newton-Type Method on Wasserstein Space

TL;DR

This paper introduces a second-order Newton-type method for non-convex optimization in Wasserstein space, enabling efficient escape from saddle points and fast convergence to global minima.

Contribution

It proposes Wasserstein Saddle-Free Newton (WSFN), a novel second-order method that overcomes saddle points and improves convergence in Wasserstein space.

Findings

WSFN escapes saddle regions in polynomial time.

WSFN converges linearly to a global minimizer once in the neighborhood.

The method outperforms prior first-order approaches in convergence speed.

Abstract

We study the minimization of non-convex functionals over the Wasserstein space. While recent work has showed that perturbed Wasserstein gradient methods can avoid saddle points for benign landscapes, existing approaches remain essentially first-order and do not provide fast local convergence once the iterates enter a neighborhood of a global minimizer. We propose Wasserstein Saddle-Free Newton (WSFN), a second-order method that preconditions the Wasserstein gradient by a regularized square root of the squared Wasserstein Hessian. This construction preserves attraction toward directions of positive curvature while inducing repulsion along directions of negative curvature, thereby overcoming the tendency of standard Wasserstein Newton dynamics to be attracted to saddles. We also establish second-order sufficient optimality conditions on Wasserstein space for strict local minimality. Under regularity and benign landscape assumptions, we prove that WSFN escapes saddle regions and reaches an $\alpha$-neighborhood of a global minimizer in polynomial time, with improved dependence on saddle parameters compared with prior perturbed first-order methods. Once inside this neighborhood, we show that WSFN converges linearly in $L^2$-Wasserstein distance to a non-degenerate global minimizer. Finally, we present a particle-based implementation of the method.