Source-linked AI summary
Interpolating between Optimal Transport and MMD using Sinkhorn Divergences
Jean Feydy, Thibault Séjourné, François-Xavier Vialard, Shun-ichi Amari, Alain Trouvé, Gabriel Peyré
TL;DR
Comparing probability distributions requires losses that account for feature-space geometry rather than relying only on pointwise density comparisons. This paper studies Sinkhorn divergences, using geometric entropy and a numerical gradient scheme to establish theoretical guarantees and large-scale applicability. The divergences interpolate between MMD and OT while providing positivity, convexity, metrization of convergence in law, and scalable computation.
Problem
Simple dissimilarities such as total variation and relative entropy ignore feature-space distance, failing to metrize convergence in law and remaining unstable under support deformations.
Method
The paper develops Sinkhorn divergences through a Sinkhorn entropy and Bregman divergence, derives gradients using the fixed-point solver's converged equation, and provides an automatic-differentiation scheme.
Results
Sinkhorn divergences interpolate between MMD and OT and have positivity, convexity, metrization of convergence in law, and scalability to large datasets; gradients for 100,000 samples per measure take seconds.
Takeaways & Limitations
The resulting loss has theoretical foundations for geometric distribution comparison and supports machine-learning applications involving large datasets.
Takeaways & Limitations
MMD losses can exhibit vanishing gradients or similar artifacts near extreme support points, except for the trivial Dirac-mass case.
Abstract
from arXiv · showhide
Comparing probability distributions is a fundamental problem in data sciences. Simple norms and divergences such as the total variation and the relative entropy only compare densities in a point-wise manner and fail to capture the geometric nature of the problem. In sharp contrast, Maximum Mean Discrepancies (MMD) and Optimal Transport distances (OT) are two classes of distances between measures that take into account the geometry of the underlying space and metrize the convergence in law. This paper studies the Sinkhorn divergences, a family of geometric divergences that interpolates between MMD and OT. Relying on a new notion of geometric entropy, we provide theoretical guarantees for these divergences: positivity, convexity and metrization of the convergence in law. On the practical side, we detail a numerical scheme that enables the large scale application of these divergences for machine learning: on the GPU, gradients of the Sinkhorn loss can be computed for batches of a million samples.
1 Introduction
The paper introduces Sinkhorn divergences as geometry-aware losses that interpolate between Optimal Transport and MMD while addressing entropic bias and retaining useful theoretical and computational properties.
- Motivation: Point-wise divergences ignore feature-space geometry, can be unstable under support deformations, and do not metrize convergence in law.Ground distances can instead help gradient descent overcome spurious local minima.
- Existing geometric losses: Optimal Transport and MMD both account for geometry and metrize convergence in law, but OT is computationally challenging while MMD induces flatter geometry.MMD is cheaper and has lower sample complexity, whereas OT more faithfully lifts ground distances.
- Entropic regularization: Entropic regularization makes OT approximately computable on GPUs, but the resulting OTε cost is biased because its self-cost is generally nonzero.Minimizing OTε can shrink the fitted measure toward a support smaller than the target.
- Sinkhorn divergences: Sinkhorn divergences subtract autocorrelation terms from OTε, satisfy Sε(β, β) = 0, and interpolate between OT and MMD.This debiasing structure was introduced to correct entropic bias and resemble a squared kernel norm.
- Theory: Under compact-space and Lipschitz-cost assumptions, Sε is symmetric positive definite, smooth, convex in each input, and metrizes convergence in law.The result also covers bounded-support Euclidean measures with linear or quadratic ground costs, inducing Laplacian or Gaussian kernels.
- Practical implications: The debiased loss has the target distribution as the unique minimizer and costs roughly as much to compute as OTε.The paper also develops GPU methods for large-scale gradients, including computations with 100,000 samples per measure in seconds.
2 Proof of Theorem 1
The proof introduces Sinkhorn negentropy and uses its strict convexity, differentiability, and continuity to establish that Sinkhorn divergences are positive definite, convex, and metrize convergence in law.
- Proof strategy: The proof relies on a new Bregman divergence derived from a weak∗-continuous Sinkhorn entropy.This entropy-based construction is presented as independently interesting and supports the theoretical analysis.
- Regularized OT properties: OTε admits dual potentials that exist uniquely up to an additive constant and provide gradients of the regularized transport cost.The potentials satisfy Sinkhorn optimality conditions and can be canonically extended to continuous functions.
- Sinkhorn negentropy: The symmetric corrective term defines a Sinkhorn negentropy whose associated Bregman divergence is positive definite.The construction uses the Gibbs kernel and assumes it is positive and universal.
- Convexity and definiteness: Strict convexity of the Sinkhorn negentropy makes the Sinkhorn divergence convex in each input and ensures that zero divergence implies identical measures.The proof combines convexity of regularized OT with the corrective functional’s strict convexity.
- Convergence: Sinkhorn divergences metrize convergence in law on the stated measure space.The result is established as the final step of the theorem proof.
3 Computational scheme
The computational scheme represents sampled measures with weighted atoms, solves regularized transport through dual Sinkhorn iterations, and accelerates gradients and large-scale GPU computation with specialized updates and KeOps.
- Measure encoding: Sampled measures are encoded as normalized nonnegative weight vectors and coordinate tensors for their Dirac atoms.The weights have shapes [N] and [M], while coordinates have shapes [N, D] and [M, D].
- Sinkhorn iterations: The Sinkhorn algorithm alternately updates dual vectors until convergence, implementing block-coordinate ascent on the regularized OT dual problem.The loop starts from zero potentials and enforces the dual optimality conditions one block at a time.
- Convergence: Linear convergence makes 10 or 20 iterations sufficient when ε ⩾ .05 on the unit square with Earth Mover’s cost.The dual formulation also yields gradients robust to small perturbations in the potentials.
- Symmetric corrections: Symmetric Sinkhorn updates efficiently compute the autocorrelation corrections, with three iterations usually sufficient for accurate dual vectors.These corrections are used alongside the cross-correlation Sinkhorn iterations to implement Sε.
- Gradient computation: Implicit differentiation at convergence avoids backpropagating through the Sinkhorn loop and reduces computation time by a factor of 2-3.PyTorch graph surgery preserves automatic differentiation while using the explicit gradient formula.
- Large-scale computation: KeOps avoids storing [N, M] matrices by performing online log-sum-exp reductions on the GPU, enabling gradients with 100,000 samples per measure in seconds.The tensorized approach stops fitting GPU memory as sample counts exceed 10,000.
4 Numerical illustration
Gradient-flow experiments compare MMD, regularized OT, and Sinkhorn divergences in one and two dimensions, showing that debiasing removes entropic shrinkage while ε controls interpolation between MMD and OT geometry.
- Experimental setup: The experiments evolve a discrete model measure toward a target under Wasserstein-2 gradient flows of competing loss functionals.The support positions are integrated with an Euler scheme in one or two dimensions.
- One-dimensional flows: Figure 4 displays 1-D flows for N = M = 5000 using kernel density estimates and compares them with the exactly computable Earth Mover’s dynamics.The experiments vary ε between .10 and .01 for the Sinkhorn-based losses.
- Two-dimensional flows: Figure 5 displays 2-D measures as point clouds of N = M = 500 and includes the descent vector field −∇xiL(α, β).The visualization makes the evolution of the model measure’s support explicit.
- Debiasing: Regularized OT drives the model toward a shrunk target, whereas the de-biased Sinkhorn divergence removes this entropic bias.The Sinkhorn divergence is described as a reliable positive-definite loss for fitting distributions.
- Interpolation: The Sinkhorn flow interpolates between MMD dynamics at ε = +∞ and Optimal Transport dynamics at ε = 0.Intermediate ε values bridge energy-distance flows and Earth Mover’s motion according to the transport geometry.
5 Conclusion
The paper establishes Sinkhorn divergences as theoretically grounded geometric losses that interpolate between MMD and OT while supporting large-scale use.
- Sinkhorn divergences interpolate between MMD and OT and provide an additional parameter for geometric loss design.The parameter ε can be cross-validated in typical learning scenarios.
- The divergences satisfy positivity, convexity, metrization of convergence in law, and scalability to large datasets.
- The paper presents this as the first demonstrated firm theoretical foundation for a loss derived from entropic Optimal Transport.
A Standard results
This section introduces standard preliminary results on the Kullback-Leibler divergence and the SoftMin operator used later in the paper.
- The paper recalls established results for the Kullback-Leibler divergence and the SoftMin operator before presenting its proofs.The SoftMin operator is defined in equation (12).
A.1 The Kullback-Leibler divergence
The appendix develops foundational properties of the Kullback-Leibler divergence, including its dual representation, convexity, lower semicontinuity, nonnegativity, and SoftMin-related results.
- The Kullback-Leibler divergence is finite under absolute continuity and otherwise takes the value +∞.
- Dual formulation: The divergence can be expressed as an f-divergence and characterized through a dual concave optimization problem over bounded measurable functions.
- KL(α, β) is nonnegative and equals zero exactly when α = β.
- Dual formulation: The dual objective provides upper and lower bounds linked to KL(α, β), yielding the divergence representation.
- The KL divergence is jointly convex on M+(X) × M+(X).
- The divergence is weakly lower semi-continuous under weak convergence of both input measure sequences.
- SoftMin operator: The SoftMin operator interpolates between a minimum and a sum and is continuous under weak convergence of measures and uniform convergence of functions.
B.1 Dual Potentials
The paper analyzes entropic OT dual potentials on compact metric spaces, establishing their regularity, uniqueness up to constants, and continuous dependence on input measures.
- The analysis assumes a compact metric space with a ground cost that is κ-Lipschitz in both input variables.
- The Sinkhorn map T is used to prove existence of an optimal pair of dual potentials reaching the dual objective's maximum.
- Optimal potentials are extended to the whole feature space through an imposed condition related to the Sinkhorn map.
- Uniqueness: Optimal potential pairs are unique up to an additive constant.
- Regularity: The optimal dual potentials are κ-Lipschitz functions on the feature space.
- Continuity: When input measures converge weakly, normalized optimal potentials converge uniformly to the normalized potentials of the limiting measures.
B.2 Proof of Proposition 2
The proof establishes continuity and differentiability of the regularized OT cost by controlling optimal potentials under weak convergence and perturbations of the input measures.
- The proof bounds the variation ratio from below and above using optimal and suboptimal dual potentials.Continuity of the optimal potentials supplies the limiting control as t approaches 0.
- Weakly converging measures yield uniformly converging optimal dual potentials after fixing their value at an anchor point.The argument uses uniform convergence of potentials and a normalization at x_o.
- The continuous expression for OTε is independent of the optimal potentials' value at the anchor point.This removes the normalization choice from the continuity conclusion.
- Perturbations α + tδα and β + tδβ converge weakly to α and β as t approaches 0.The perturbation directions have zero total mass, as required for the variation argument.
B.3 Proof of Proposition 3
The proof constructs an optimal measure for the auxiliary entropy problem, then establishes uniqueness through weak compactness and strict convexity under the kernel assumptions.
- Symmetry of the concave problem provides an optimal potential pair on the diagonal with g = f.
- The optimization can be reformulated over positive measures while encoding absolute continuity through the KL term.The µ ≪ α constraint is removed using positivity of kε, while α ≪ µ is enforced by the log-density term.
- A minimizing sequence for Fε has uniformly bounded mass, enabling extraction of a weakly convergent subsequence.Positivity of the kernel on compact X × X supplies the mass bound, and Banach–Alaoglu provides weak compactness.
- Weak lower semicontinuity of Eε yields an optimal measure µα from the weak limit of the minimizing sequence.Continuity of kε is used in the lower-semicontinuity argument.
- Strict convexity of kε and the KL term makes µ ↦ Eε(α, µ) strictly convex, ensuring that µα is unique.The kernel is assumed positive universal, and strict convexity of the KL term is used in the contradiction argument.
B.5 Proof of the Metrization of the Convergence in Law
The proof combines weak continuity of the regularized OT cost with compactness of probability measures to show that Sinkhorn and Hausdorff divergences metrize convergence in law.
- Weak convergence αn ⇀ α implies that the regularized OT cost, Hε, and Sε converge continuously.Uniform convergence of dual potentials supports continuity of the regularized cost and the derived divergences.
- If Sε(αn, α) or Hε(αn, α) tends to zero, every weak subsequential limit of αn equals α by positive definiteness.
- Compactness of probability measures on compact X turns the unique subsequential adherence value into convergence of the full sequence.