Difference of Convex Programming in the Wasserstein Space with Applications to MMD Optimization

2026-06-26

AI Summary

Optimizing non-convex functionals such as the Maximum Mean Discrepancy (MMD) and energy distance in Wasserstein space is highly challenging due to the lack of geodesic convexity, which hinders the theoretical analysis of standard first-order methods. This work introduces, for the first time, a difference-of-convex (DC) optimization framework into Wasserstein space by constructing effective DC decompositions for such functionals and extending the classical Convex–Concave Procedure (CCCP) to this setting. Under mild assumptions, we establish the local convergence of the proposed algorithm. Both theoretical analysis and empirical experiments demonstrate that the method converges faster and more stably than Wasserstein gradient descent, offering a provably convergent new pathway for optimizing non-convex probabilistic functionals.

Abstract

Optimizing functionals over the space of probability measures is now ubiquitous in machine learning. A widely used approach is to perform the optimization directly over the Wasserstein space, but many objective functionals of practical interest are non-convex along Wasserstein geodesics, making the analysis of standard first-order methods challenging. In this work, we study a class of objectives over the Wasserstein space that admit a difference-of-convex (DC) decomposition and we lift the classical convex-concave procedure (CCCP) to this setting. Under smoothness and strong convexity assumptions on the convex components of the decomposition, we prove almost stationarity along the iterates of the resulting algorithm. Our main focus is on the Maximum Mean Discrepancy (MMD) and the Energy Distance (ED) functionals, for which we develop explicit Wasserstein DC decompositions, and establish local convergence of the scheme under mild assumptions. Empirically, we show that well-chosen DC decompositions yield faster and more stable convergence than Wasserstein gradient descent on these MMD objectives.

Problem

• Wasserstein space
• non-convex optimization
• difference-of-convex
• Maximum Mean Discrepancy
• Energy Distance

Innovation

• Difference-of-Convex Programming
• Wasserstein Space
• Maximum Mean Discrepancy
• Convex-Concave Procedure
• Energy Distance