Source-linked AI summary
Learning Generative Models with Sinkhorn Divergences
Aude Genevay, Gabriel Peyré, Marco Cuturi
TL;DR
High-dimensional generative models can produce singular distributions, challenging likelihood-based comparison and making OT attractive but computationally difficult. The paper combines entropic smoothing with automatic differentiation through Sinkhorn iterations to train generative models, yielding a loss that interpolates between OT and MMD. Its scope remains bounded by unresolved theoretical questions, including Sinkhorn sample complexity and positivity.
Problem
Generative models can produce singular distributions, while OT comparison is computationally burdensome, unstable, and difficult to estimate robustly in high dimensions.
Method
The paper combines entropy-regularized OT, Sinkhorn iterations, and automatic differentiation to construct a tractable differentiable loss for generative models.
Results
The Sinkhorn loss interpolates between pure OT as ε approaches 0 and MMD as ε approaches infinity, while providing stable and accurate approximations for learning.
Takeaways & Limitations
Tuning ε allows the method to blend OT's geometry with MMD's favorable high-dimensional sample behavior and unbiased mini-batch gradient estimates.
Takeaways & Limitations
The positivity of the Sinkhorn loss and its sample-complexity properties remain unproved, and related dual-potential gradients can be difficult and unstable to compute.
Abstract
from arXiv · showhide
The ability to compare two degenerate probability distributions (i.e. two probability distributions supported on two distinct low-dimensional manifolds living in a much higher-dimensional space) is a crucial problem arising in the estimation of generative models for high-dimensional observations such as those arising in computer vision or natural language. It is known that optimal transport metrics can represent a cure for this problem, since they were specifically designed as an alternative to information divergences to handle such problematic scenarios. Unfortunately, training generative machines using OT raises formidable computational and statistical challenges, because of (i) the computational burden of evaluating OT losses, (ii) the instability and lack of smoothness of these losses, (iii) the difficulty to estimate robustly these losses and their gradients in high dimension. This paper presents the first tractable computational method to train large scale generative models using an optimal transport loss, and tackles these three issues by relying on two key ideas: (a) entropic smoothing, which turns the original OT loss into one that can be computed using Sinkhorn fixed point iterations; (b) algorithmic (automatic) differentiation of these iterations. These two approximations result in a robust and differentiable approximation of the OT loss with streamlined GPU execution. Entropic smoothing generates a family of losses interpolating between Wasserstein (OT) and Maximum Mean Discrepancy (MMD), thus allowing to find a sweet spot leveraging the geometry of OT and the favorable high-dimensional sample complexity of MMD which comes with unbiased gradient estimates. The resulting computational architecture complements nicely standard deep network generative models by a stack of extra layers implementing the loss function.
1 Introduction
Generative models can produce singular distributions that make likelihood-based fitting problematic, motivating geometric alternatives. The paper introduces Sinkhorn loss and an automatically differentiated computational scheme to make OT-based training tractable and stable.
- Motivation: Generative models may yield distributions supported on low-dimensional manifolds, making standard maximum-likelihood approaches flawed for such settings.The resulting density can be singular in the high-dimensional observation space.
- Prior approaches: GANs, VAEs, MMD, and OT provide likelihood-free or geometric alternatives for comparing generated and observed distributions.OT is especially suited to distributions with non-overlapping supports, but its original computation is expensive.
- Architecture: The computation is organized as a flow from fixed latent samples and input data to the Sinkhorn loss function.The figure depicts the loss computation for a fixed set of generated and observed samples.
- Contributions: The paper introduces Sinkhorn loss, an entropy-regularized OT loss whose smoothing parameter connects pure OT and MMD.The contributions combine a theoretical loss construction with a numerical learning scheme.
- Contributions: The proposed learning method uses automatic differentiation through Sinkhorn iterations to obtain a stable, differentiable approximation suitable for generative models.The architecture adds computational layers implementing the loss and supports stochastic gradient training on GPUs.
2 Minimum Kantorovich Estimation
The paper formulates generative-model fitting through OT and introduces entropy-regularized Sinkhorn divergences. These divergences retain OT's geometric behavior while approaching an MMD associated with the ground cost at large smoothing.
- Density fitting: The generative model distribution is defined as the push-forward µθ = gθ#ζ of a reference latent measure through a differentiable generator.The objective is to choose θ so generated samples resemble the observed dataset.
- Density fitting: Likelihood-based fitting can fail because low-dimensional generative models produce singular distributions without densities relative to a fixed reference measure.Weak metrics address this by incorporating spatial displacement and can compare singular measures.
- Optimal transport distances: OT compares probability measures through couplings with prescribed marginals and a ground cost, generalizing beyond Euclidean spaces and the 1-Wasserstein distance.The coupling formulation minimizes expected transport cost over joint measures whose marginals are the input distributions.
- Regularized OT: Entropy regularization makes OT efficiently solvable with Sinkhorn's algorithm and yields a differentiable loss.The regularized problem restricts admissible couplings to a scaling form.
- Sinkhorn loss: The Sinkhorn divergence subtracts self-comparison terms from regularized OT costs to define a loss for fitting µθ to the data distribution.The paper applies this divergence to the generative-model fitting objective.
- Limiting behavior: As ε approaches 0, Sinkhorn loss recovers pure OT; as ε approaches infinity, it approaches MMD with the transport cost as kernel.At infinite smoothing, the minimum-entropy coupling is the product of the marginals.
- OT versus MMD: Sinkhorn loss can combine OT geometry with MMD's O(1/√n) mini-batch sample-complexity rate and unbiased gradient estimates.The paper notes that Sinkhorn sample-complexity guarantees are not proved, although empirical behavior resembles MMD when ε is not too small.
- Caveats: The positivity of the Sinkhorn loss is not proved and is supported only by empirical evidence, while positive-definite kernels are required for meaningful MMD losses.The paper also distinguishes energy distance from a solution to high-dimensional OT estimation bias.
3 Sinkhorn AutoDiff Algorithm
The algorithm makes entropically regularized OT tractable and differentiable by approximating minibatch losses with finite Sinkhorn iterations and automatically differentiating those iterations. The resulting loss integrates with generative networks, supports learned costs, and offers GPU-efficient computation with accuracy improving as batch sizes and iteration counts grow.
- Algorithmic approximations: The method uses entropic regularization and two approximations: minibatch sampling and L-step Sinkhorn iterations, producing a loss amenable to automatic differentiation.The minibatch approximation enables stochastic gradient descent, while finite Sinkhorn iterations yield an algorithmic loss.
- Sinkhorn iterations: Sinkhorn regularization restricts couplings to a scaling form, making the optimal transport problem efficiently solvable and differentiable.The resulting iterations alternate matrix-vector multiplications and component-wise divisions.
- Sinkhorn iterations: The Gibbs kernel K encodes geometric information about generated samples’ proximity to the dataset and is repeatedly applied within the Sinkhorn layers.K depends on the cost matrix, which depends on the generator parameters θ.
- Accuracy and convergence: For L iterations, the proxy coupling is P_L = diag(a_L)K diag(b_L), and it approaches a solution of the regularized transport problem as ε → 0 and L → +∞.The convergence is linear, although its rate deteriorates as ε → 0.
- Accuracy and convergence: The proxy loss is smooth, fast, and stable to differentiate, while larger ε can improve stability and convergence enough to permit relatively few Sinkhorn iterations.This large-ε regime is motivated for large-scale, high-dimensional learning applications.
- Learned costs: The ground cost can be learned through a neural feature extractor f_ϕ that maps the data space X into a lower-dimensional space R^p, turning training into a min-max problem.This learned-cost variant is approximated in practice with minibatches and Sinkhorn iterations.
- Deep-network integration: The complete architecture combines generated and real samples into a pairwise cost matrix, then uses Sinkhorn steps for both forward loss evaluation and backward gradient computation.Samples are repeatedly drawn as push-forwards of latent samples while stochastic gradient descent updates the model.
- Deep-network integration: Automatic differentiation through L Sinkhorn steps has the same O(Lmn) computational complexity as Sinkhorn, with extra storage for the backward pass.Matrix-vector multiplications stream efficiently on GPUs, allowing the iterations to act as additional network layers.
4 Applications
The applications evaluate Sinkhorn-based training on ellipse fitting and image generation, showing stable fitting and sensitivity to regularization, batch size, and iteration count. On CIFAR10, a large regularization produced the best inception score, while closer-to-OT settings performed worse under the fixed iteration budget.
- Applications: The applications cover fitting labeled data with uniform distributions on ellipses and training image generators on MNIST and CIFAR10.The reported runs were stable when repeated.
- Data Fitting with Ellipses: Wasserstein-based fitting is illustrated for singular distributions supported on lower-dimensional manifolds, using K uniform ellipses to recover entire areas rather than Voronoi cells.Each ellipse is parameterized by a covariance-root matrix and center, with equal mass 1/K.
- Data Fitting with Ellipses: The ellipse experiment uses L = 20 Sinkhorn iterations and minibatches of m = 200, comparing regularized Sinkhorn loss with Gaussian-kernel MMD.The regularization parameter ε strongly influences the result.
- Tuning a Generative Neural Network: On MNIST, ε = 1 with L = 10 iterations produced a result similar to ε = 0.1 with L = 100 iterations, while m = 10 led to poor results.The generated digit manifold was examined across ε, m, and L.
- Tuning a Generative Neural Network: On CIFAR10, Sinkhorn and MMD models were compared using the same latent samples, and the Sinkhorn model with large regularization achieved the best inception score.Models closer to true OT scored poorly partly because their Sinkhorn iterations had not converged within the training-time budget.
- Tuning a Generative Neural Network: The results reflect a sample-complexity trade-off: performance can worsen nearer to OT metrics, making larger batch sizes potentially beneficial.This observation connects the empirical behavior to the regularization and minibatch choices.
Conclusion
The paper concludes that combining entropic smoothing with automatic differentiation yields a computational toolbox for training complex generative models with Sinkhorn divergence. It also establishes interpolation between OT and MMD while identifying sample complexity and positivity as future theoretical targets.
- Conclusion: The paper presents a computational toolbox for training large-scale generative models with Sinkhorn divergence.Its applicability is attributed to combining entropic smoothing with automatic differentiation.
- Conclusion: Sinkhorn divergence interpolates between classical OT and MMD, combining advantages attributed to both loss frameworks.The conclusion frames this interpolation as a central theoretical contribution.
- Conclusion: Future work should study the Sinkhorn divergence's sample complexity and positivity.
A.1 Sample Complexity
The experiments examine how entropic regularization affects Sinkhorn-divergence sample complexity and variance. Small regularization follows OT-like rates with lower variance, whereas large regularization approaches MMD-like rates.
- The study estimates the decay rate and standard deviation of empirical Sinkhorn divergences using independently sampled random measures.The experiments use uniform distributions on [0, 1]^d and numerically average over random draws.
- The experiments recover the known rates κ0,d = p/d for OT and κ∞,d = 1 for MMD.Here p = 3/2 for the OT rate.
- Small values of ε < 1 produce rates κε,d close to the OT rate κ0,d.
- Large values of ε > 1 produce rates κε,d that almost exactly match the MMD rate κ+∞,d = 1.
- The variance S2 ε,d(N) is significantly smaller for small ε values close to OT.
- Similar conclusions are obtained when testing other distributions, including a Gaussian.
A.2 Positivity
The paper numerically investigates whether the regularized Sinkhorn divergence is locally strictly positive for discrete measures. Across 10^4 additional realizations, the experiments support positivity under the tested perturbations.
- For ε = 0 and ε = +∞, the Sinkhorn divergence is zero if and only if the two measures are equal.
- For 0 < ε < +∞, whether zero divergence implies equal measures remains unknown and difficult to establish.
- The numerical test perturbs a discrete measure whose atom locations and weights are randomly generated.The locations are sampled in [0, 1]^2, weights are uniform on [1/2, 1], and the perturbation uses Gaussian variables.
- 10^4 additional realizations yielded the same outcome, with divergence positive for nonzero perturbations in the tested discrete cases.