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

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

Problem Statement
Optimizing functionals over the space of probability measures is central to machine learning tasks ranging from variational inference to generative modeling. A standard approach involves performing optimization directly over the Wasserstein space $(P_2(\mathbb{R}^d), W_2)$ by discretizing Wasserstein gradient flows. While effective for convex functionals (e.g., KL divergence, f-divergences), many practically relevant objectives are non-convex along Wasserstein geodesics. Notable examples include the squared Wasserstein distance, the Sliced-Wasserstein distance, and the Maximum Mean Discrepancy (MMD). Standard first-order methods like Wasserstein Gradient Descent (WGD) struggle in these non-convex settings, often requiring delicate noise injection or optimization in alternative spaces (e.g., Wasserstein-MMD or Wasserstein-Fisher-Rao spaces) that alter the distribution weights or mix implicit steps.

Methodology
The authors propose a new particle-based algorithm tailored for non-convex functionals that admit a Difference-of-Convex (DC) decomposition within the Wasserstein space. The core methodology involves lifting the classical Convex-Concave Procedure (CCCP) from Euclidean spaces to the Wasserstein space, resulting in the Wasserstein Convex-Concave Procedure (WCCCP).

1. DC Decomposition in Wasserstein Space: The objective functional $F$ is decomposed as $F(\mu) = F_+(\mu) - F_-(\mu)$, where both $F_+$ and $F_-$ are totally convex functionals on $P_2(\mathbb{R}^d)$. Unlike prior work that restricted the concave part to potential energies (linear functionals), this approach allows $F_-$ to be any differentiable functional.
2. Algorithmic Scheme: At each iteration $k$, given a measure $\mu_k$, the algorithm solves a majorization-minimization problem. It linearizes the concave part $F_-$ around $\mu_k$ to construct an upper bound and minimizes the resulting convex surrogate over the space of transport maps $T \in L^2(\mu_k)$:
   $$T_{k+1} = \arg\min_{T \in L^2(\mu_k)} \left( F_+(T\#\mu_k) - \langle \nabla_{W_2}F_-(\mu_k), T - \text{Id} \rangle_{L^2(\mu_k)} \right)$$
   The next iterate is then defined as $\mu_{k+1} = (T_{k+1})\#\mu_k$.
3. Implementation: The inner minimization is performed via gradient descent on $L^2(\mu_k)$. For empirical distributions (particles), this translates to updating particle positions. The authors also establish connections between WCCCP and Mirror Descent/Bregman Proximal Descent in the Wasserstein space.
4. MMD Specifics: For the Maximum Mean Discrepancy, the authors propose explicit DC decompositions by splitting the kernel $k(x,y) = \psi(x-y)$ itself into convex components $\psi = \psi_+ - \psi_-$. They provide specific decompositions for translation-invariant kernels, including Gaussian and Riesz kernels, using Jordan decomposition or algebraic splitting (e.g., $\cosh/\sinh$).

Key Contributions

• WCCCP Algorithm: Introduction of the Wasserstein Convex-Concave Procedure, extending DC optimization to the Wasserstein space without restricting the concave component to linear potentials.
• Theoretical Convergence:
  • Proof that the objective value $F(\mu_k)$ is non-increasing along the iterates.
  • Establishment of almost stationarity results. Under assumptions of smoothness and strong convexity (total convexity) of the components, the authors prove sublinear convergence rates for the squared distance between iterates and the norm of the Wasserstein gradient. Specifically, $\min_{k} \|\nabla_{W_2}F(\mu_k)\|^2 = O(1/K)$.
  • Demonstration that the scheme is equivalent to Bregman Proximal Descent and Mirror Descent under specific conditions.
• Explicit DC Decompositions for MMD: The paper provides a systematic way to decompose MMD functionals with translation-invariant kernels into DC forms, satisfying the necessary convexity and smoothness conditions for convergence guarantees.
• Empirical Validation: Experiments on Energy Distance and MMD with Gaussian kernels demonstrate that well-chosen DC decompositions yield faster and more stable convergence compared to standard WGD and the Wasserstein Proximal Gradient (Forward-Backward) method.

Results

• Energy Distance: On synthetic data (spiral and cat shapes) and CIFAR10 images, WCCCP converges significantly faster than WGD and the Forward-Backward (FB) scheme. The algorithm successfully minimizes the Energy Distance to match complex target distributions.
• MMD with Gaussian Kernel: Minimizing MMD with a Gaussian kernel is notoriously sensitive to bandwidth and initialization. The authors show that WGD and the FB scheme (using a standard decomposition) often get stuck in local minima where particles drift away from the target. In contrast, WCCCP using specific kernel decompositions (Jordan or $\cosh/\sinh$) achieves better convergence. The Jordan decomposition is noted for being smoother and less prone to numerical instability, while the $\cosh/\sinh$ decomposition ensures all particles eventually converge, albeit at a slower rate.
• Comparison: The results highlight that the performance of WCCCP is highly dependent on the choice of the DC decomposition. A "well-chosen" decomposition is critical for outperforming standard gradient descent.

Significance and Claims
The paper claims to provide a robust framework for optimizing non-convex functionals in the Wasserstein space, addressing a gap where standard gradient methods fail due to non-convexity. By lifting the CCCP to this setting, the authors offer a theoretically grounded alternative that avoids the weight-mixing issues of previous MMD-specific approaches. The work emphasizes that while the algorithm is general, its practical success relies heavily on the specific DC decomposition of the kernel. The authors modestly note that future work is required to understand the impact of different decompositions more deeply and to design automatic, adaptive strategies for selecting them. Additionally, the theoretical analysis currently assumes exact solutions to the inner optimization problems, and accounting for inexact solvers is identified as a necessary direction for future research.