Source-linked AI summary
A Smoothed Dual Approach for Variational Wasserstein Problems
Marco Cuturi, Gabriel Peyré
TL;DR
Variational problems involving Wasserstein distances are computationally challenging because the distances themselves are hard to compute. The paper combines duality with entropic smoothing to obtain smooth optimization problems with closed-form objectives and derivatives, and applies the approach to barycenters and gradient flows. The numerical findings indicate stabilized barycenter computation and fast numerical schemes.
Problem
Variational problems involving Wasserstein distances are computationally challenging because they require minimizing quantities that are themselves difficult to compute.
Method
The paper combines a dual Wasserstein formulation with entropic smoothing to obtain smooth, differentiable, convex optimization problems with closed-form objectives and derivatives.
Results
The approach is applied to Wasserstein barycenters and gradient flows, with numerical findings that entropic smoothing stabilizes barycenter computation and yields fast numerical schemes.
Takeaways & Limitations
The framework provides a scalable and flexible way to minimize energies involving Wasserstein distances and spatial regularization terms.
Takeaways & Limitations
Direct computation of the smoothed Wasserstein distance and its gradient requires solving an optimization problem, such as with Sinkhorn fixed-point iteration.
Abstract
from arXiv · showhide
Variational problems that involve Wasserstein distances have been recently proposed to summarize and learn from probability measures. Despite being conceptually simple, such problems are computationally challenging because they involve minimizing over quantities (Wasserstein distances) that are themselves hard to compute. We show that the dual formulation of Wasserstein variational problems introduced recently by Carlier et al. (2014) can be regularized using an entropic smoothing, which leads to smooth, differentiable, convex optimization problems that are simpler to implement and numerically more stable. We illustrate the versatility of this approach by applying it to the computation of Wasserstein barycenters and gradient flows of spacial regularization functionals.
1. Introduction.
The paper develops a scalable, flexible framework for variational problems involving Wasserstein distances and broader regularization terms. It combines duality with entropic smoothing to support barycenters, spatially regularized energies, and gradient flows.
- Motivation: Variational learning tasks such as averaging and clustering histograms become challenging when they involve Wasserstein distances rather than Bregman divergences.Wasserstein-based formulations require optimizing quantities that are themselves computationally difficult.
- Contribution: The framework targets energies involving Wasserstein distances and more general regularization terms while remaining scalable and flexible.It exploits regularization, Legendre duality, and convex optimization tools.
- Applications: The paper positions the framework for applications including image segmentation, where derivatives with respect to two histograms are required.The corresponding formula is provided in the appendix.
- Contribution: The main contribution combines a dual Wasserstein formulation with entropic smoothing to obtain objectives and derivatives computable in closed form.The resulting optimization problem is smooth and is presented as a computational framework for variational Wasserstein problems.
- Applications: The approach is applied to Wasserstein barycenters, spatial regularization of barycenters, and gradient flows.The dual approach also supports non-separable energies such as image total variation.
2. Legendre Transforms of the Smoothed Wasserstein Distance.
The paper entropically regularizes Wasserstein optimal transport and studies its Fenchel-Legendre transform. For positive smoothing, the transform is smooth and can be evaluated with derivatives without solving a matrix-scaling problem, while the zero-smoothing limit selects a maximally entropic optimal coupling.
- Entropic smoothing: The entropic transport problem is strongly convex for γ > 0 and has a unique optimal coupling, whereas γ = 0 recovers the usual Wasserstein linear program.The cost matrix is assumed only to be symmetric and non-negative.
- Entropic smoothing: As γ approaches zero, the regularized solution converges to an optimal transport coupling with maximum entropy among optimal couplings.This selects a particular solution when the unregularized problem has multiple optimal couplings.
- Smoothed distance: For γ > 0, H_q is convex and continuously differentiable on positive histograms, with gradient given by the normalized dual optimizer.Computing H_q and its gradient directly requires solving the optimization problem defining the dual variables.
- Legendre transform: For γ > 0, nearest-neighbor assignments in the unregularized dual formulation are replaced by soft assignments.The transform is expressed using the Gibbs kernel K = e^(-M/γ) and α = e^(g/γ).
- Legendre transform: The Fenchel-Legendre transform and its derivatives can be computed without solving a matrix-scaling problem.This property is central to the paper’s computational framework.
3. Smooth Dual Algorithms For the Wasserstein Barycenter Problem.
Entropic smoothing and Fenchel-Legendre duality turn the Wasserstein barycenter problem into a smooth optimization problem with closed-form objectives and derivatives. The resulting algorithms include practical initialization strategies and can improve convergence and stability, although the unsmoothed problem may be ill-posed.
- 3.1. Smooth Dual Formulation of the WBP: For γ > 0, the smoothed barycenter problem is strictly convex, has a unique solution, and can be optimized using gradients obtained from Sinkhorn matrix-scaling problems.At γ = 0, the formulation reduces to the Wasserstein barycenter problem and becomes a linear program.
- 3.1. Smooth Dual Formulation of the WBP: Theorem 3.1 replaces direct minimization of regularized Wasserstein distances with minimization of a strictly convex function having closed-form objectives and gradients.The dual formulation is presented as a simpler computational route for the smoothed Wasserstein barycenter problem.
- 3.2. Optimization Algorithms: The smoothed dual formulation supports gradient descent, accelerated gradient methods, quasi-Newton methods, and truncated Newton methods.The dual objective has smooth variables, a linear equality constraint, and gradients and Hessians available in closed form; the resulting KKT system can be sparse.
- 3.3. Initialization Heuristic: The initialization differs by regularization: for γ > 0 it uses a normalized weighted geometric average of kernel columns, while for γ = 0 it selects a minimizer of Mq̄.The dual initialization is the same for smoothed and non-smoothed barycenter problems.
- 3.3. Initialization Heuristic: The proposed initialization gives exact primal and dual solutions when all target histograms are Dirac histograms, for every γ ≥ 0.For more general problems, the authors report that the dual initialization captures important features of the optimal solution more reliably than the primal initialization.
- 3.5. Performance on the Wasserstein Barycenter Problem: In experiments, smooth-dual L-BFGS converged faster under both smooth and non-smooth metrics, while naive subgradient optimization failed to converge in the tested examples.The authors also emphasize the importance of the proposed initialization for convergence.
- 3.5. Performance on the Wasserstein Barycenter Problem: The dual optimization framework extends beyond barycenters to variational problems involving more general functionals, including spatial regularization and gradient flows.The paper presents this versatility as a consequence of combining duality, smoothing, and convex optimization.
4. Regularized Problems.
The section develops regularized Wasserstein barycenters and solves them through dual and proximal optimization, then demonstrates total-variation effects on images, shapes, MEG data, and gradient flows.
- 4.1. Regularized Wasserstein Barycenters: The regularized barycenter admits a dual optimization formulation with primal-dual relationships, enabling the penalty to be handled through its conjugate.The construction removes one dual variable by exploiting the nonzero weight assumption and derives the dual through Fenchel-Rockafellar duality.
- 4.2. Resolution using First Order Proximal Splitting: Forward-Backward iterations converge when τ < 2/L, where L is the Lipschitz constant of the smooth term’s gradient; accelerated schemes such as FISTA can improve convergence.The method separates a smooth component from a proximal component and assumes the proximal operator of J* is computable.
- 4.3. Total Variation Regularization: Total variation uses a discrete gradient and regularization strength λ; β = 2 yields isotropic smoothing, while β = 1 favors horizontal and vertical edges.Isotropic variation tends to round corners and may merge clusters, whereas anisotropic variation penalizes horizontal and vertical derivatives independently.
- 4.4. Barycenters of Images: For 256 × 256 image grids, isotropic barycenters round input corners, whereas anisotropic barycenters favor horizontal and vertical edges.The experiments use forward finite differences with Neumann boundary conditions and the squared Euclidean metric.
- 4.5. Barycenters of MEG Data: Increasing total-variation regularization groups small clusters and helps significant activity emerge from noisy MEG histograms, while preserving sharp active-to-inactive transitions.The method compares unregularized, ℓ2, and increasingly regularized barycenters on recordings associated with left- or right-button responses.
- 4.6. Gradient Flow: Entropic smoothing enables larger gradient-flow problems and faster numerical schemes, at the cost of additional blurring that is acceptable for imaging denoising.The illustrated isotropic total-variation flow groups mass clusters and produces progressive percolation across the image.
- Conclusion: The paper concludes that closed-form dual gradients involving Gibbs-kernel multiplications stabilize barycenter computation and support regularized barycenters and gradient flows.It identifies entropic smoothing as crucial for stable computation and fast numerical schemes.
Appendix A. Legendre Transform with Respect to Two Histograms.
The appendix extends the Legendre transform of the entropically regularized Wasserstein distance to both histogram arguments. It establishes smoothness and expresses the gradients through a Gibbs-type optimizer.
- Joint Legendre transform: The map (p, q) ↦ Wγ(p, q) is convex in both histogram arguments, allowing a joint Legendre transform over (g, h).The appendix defines the transform for arbitrary dual vectors in R^n × R^n.
- Smoothness and closed form: W*γ is C∞ in (g, h), with α = e^(g/γ), β = e^(h/γ), and K = e^(−M/γ) entering the closed-form representation.The construction uses the Gibbs kernel and the transformed dual variables to express the optimizer.
- Gibbs optimizer: The maximization becomes a uniquely solvable maximum-entropy problem whose optimizer is a Gibbs distribution X⋆.Substituting X⋆ into the transform yields the stated expression and identifies the gradients with its marginal sums.
- Derivatives: The gradients with respect to g and h are given by the two marginals of X⋆, while the Hessian and gradient Lipschitz bound follow from this representation.The Hessian trace is bounded using the corresponding kernel expressions.