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

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.

• The paper introduces the WCCCP algorithm that extends CCCP to Wasserstein space for optimizing MMD and ED through a difference-of-convex decomposition.
• It establishes theoretical convergence guarantees under smoothness and strong convexity while linking the approach to mirror descent and Bregman proximal descent.
• Empirical results show that WCCCP accelerates convergence and improves stability compared to traditional Wasserstein gradient descent and proximal gradient methods.

Difference-of-Convex Programming in Wasserstein Space: Theory and Applications to MMD Optimization

Introduction

This work introduces a novel optimization framework for functionals over the Wasserstein space that admit a difference-of-convex (DC) decomposition, specifically targeting objectives such as Maximum Mean Discrepancy (MMD) and Energy Distance (ED) (2606.27767). Standard first-order optimization methods in Wasserstein spaces rely on geodesic convexity of the objective, which does not hold for important functionals (e.g., MMD, ED). The authors propose an extension of the Convex-Concave Procedure (CCCP) to the Wasserstein space (WCCCP), derive explicit convergence results under smoothness and strong convexity assumptions, and demonstrate improved empirical performance over traditional Wasserstein Gradient Descent (WGD) and Proximal Gradient approaches.

Difference-of-Convex Optimization in Wasserstein Space

The optimization landscape in Wasserstein space is significantly complicated by geodesic non-convexity of many objective functionals. The DC programming paradigm enables efficient minimization by expressing such objectives $F$ as $F = F^+ - F^-$, with $F^+$ and $F^-$ totally convex. This work lifts CCCP and related DCA algorithms from Euclidean and Riemannian settings to the Wasserstein space, exploiting the geometry induced by optimal transport.

The Wasserstein CCCP (WCCCP) scheme operates by majorizing the concave component via linearization at each iteration, producing updates analogous to mirror descent and Bregman proximal descent in Wasserstein space. The theoretical analysis provided demonstrates non-increasing objective values along the iterates, with termination criteria directly linked to stationarity of the Wasserstein gradient.

Theoretical Convergence Properties

The authors establish sublinear convergence rates for the criterion reflecting stationarity and show that, under strong convexity of either $F^+$ or $F^-$, the iterates approach a stationary point of the objective. Notably, the WCCCP algorithm admits equivalences with mirror descent and Bregman proximal descent formulations. Under additional Lipschitz continuity of the Wasserstein gradient, the minimum gradient norm across iterates is bounded by the decrease in objective value and the convexity parameters.

Crucially, the scheme generalizes beyond previous literature that restricts DC decompositions to linear concave components, accommodating functionals with more flexible structure (e.g., sum of non-convex quadratic and linear terms, as in MMD).

DC Decomposition of MMD and Energy Distance Functionals

A main focus is explicit DC decomposition of the squared MMD with translation-invariant/radial kernels. For kernels $k(x,y) = \psi(x-y)$ admitting a DC split $\psi = \psi^+ - \psi^-$, the authors show both $F^+$ and $F^-$ inherit strong total convexity. Several constructions are provided for radial kernels:

• Gaussian kernel: Both algebraic (cosh/sinh) and Jordan decomposition schemes are explored, yielding different numerical stability profiles.
• Smoothed Riesz kernel and rational quadratic kernels: DC decompositions are derived analytically.

These decompositions ensure the applicability of WCCCP and guarantee local stationarity and convergence properties, with quantitative bounds on the Wasserstein gradient.

Numerical Experiments

Empirical results demonstrate that WCCCP outperforms WGD and Proximal Gradient on minimizing non-convex functionals like ED and MMD, with substantial acceleration and improved stability. The convergence of WCCCP is quantified on problems such as shape matching (spiral/cat/mnist/cifar distributions) and distribution fitting with the Energy Distance and MMD with Gaussian kernels.

Several numerical benchmarks indicate:

• Faster convergence rates: For equivalent computational budgets, WCCCP achieves lower objective values and more reliable convergence than WGD.
• Stability: Iterates of WCCCP are less prone to local minima and sample drift, especially with robust DC decompositions (cosh/sinh for Gaussian, Jordan decomposition for radial kernels).
• Complex targets: On high-dimensional image distributions (CIFAR10, MNIST), WCCCP demonstrates significant robustness and scalability relative to traditional schemes.

Empirical evidence underscores that performance of the scheme is highly dependent on the choice of DC decomposition for the kernel; choices affect both convergence speed and stability.

Practical and Theoretical Implications

The WCCCP framework expands the toolkit for optimization of probability measure functionals in Wasserstein space, accommodating important objectives lacking geodesic convexity. Practical implications include:

• Accelerated optimization in statistical inference, generative modeling, and domain adaptation, where MMD and ED are prevalent.
• Improved robustness for high-dimensional data applications, crucial for large-scale generative models and distribution-matching tasks.

Theoretically, the work bridges DC programming and optimal transport, showing how strong convexity and Lipschitz gradient conditions propagate from kernel decompositions to measure functionals. The framework highlights the importance of kernel algebraic properties and opens avenues for adaptive, automatic DC decompositions tailored to optimization geometry.

Limitations and Future Directions

The improved convergence of WCCCP is empirically connected to the chosen DC decomposition; current strategies are limited to simple decompositions without geometry-awareness. Further research should focus on:

• Automatic or adaptive DC decomposition selection leveraging algebraic and geometric insights.
• Extension to more general functionals and kernels beyond radial or translation-invariant forms.
• Refined theory for handling inexact solvers in inner optimization steps.

There is also an outstanding question regarding the full integration of Polyak-Łojasiewicz inequalities for DC functionals in Wasserstein spaces.

Conclusion

By lifting difference-of-convex optimization to the Wasserstein space and developing explicit DC decompositions for objectives such as MMD and ED, this work provides both theoretical guarantees and empirical improvements for non-convex measure functional optimization. The WCCCP algorithm achieves faster, more stable convergence than standard gradient-based approaches, contingent on the DC decomposition quality. This advances optimization in high-dimensional probability spaces, with immediate implications for statistical learning, generative modeling, and optimal transport methods.