Online Sinkhorn Algorithms
- Online Sinkhorn is a scalable iterative method that computes entropic-regularized optimal transport distances by updating dual potentials directly from data streams.
- It employs stochastic approximation and mini-batch updates to achieve robust convergence and reduced memory usage compared to batch methods.
- Extensions such as Mirror and Compressed Online Sinkhorn enhance efficiency by controlling computational complexity while preserving convergence rates.
Online Sinkhorn algorithms constitute a family of scalable iterative methods for computing entropic-regularized optimal transport (OT) distances and more general convex objectives over transport polytopes, leveraging sample streams and stochastic approximation. Unlike classical batch Sinkhorn methods that operate on static, finite measures, online Sinkhorn approaches update nonparametric representations of dual potentials or transport plans directly from incoming sample streams. These methods are robust to noise, require only one pass through data, and have favorable convergence guarantees, making them suitable for large-scale, streaming, or online machine learning applications.
1. Foundational Problem Formulations
The core objective is to compute the entropic-regularized OT cost between two probability measures and on a compact space with cost function ,
where , and denotes the set of couplings. The dual formulation is smooth and strictly convex,
For discrete measures , , optimization is over the transport polytope
for general convex , with special cases including linear OT and entropic-regularized OT (Mensch et al., 2020, Wang et al., 2023, Ballu et al., 2022).
2. Stochastic and Online Sinkhorn Algorithms
The classic Sinkhorn algorithm alternates soft -transforms to update dual potentials on discretized measures, requiring full batch access. Mensch and Peyré introduced a stochastic framework for streaming data:
- Dual potentials are represented by growing kernel expansions:
and similarly for in terms of past samples .
- Update steps use mini-batches :
with vanishing, and unbiased Monte Carlo approximations of soft -transforms.
- Per-iteration cost is for batch size and kernel centers (Mensch et al., 2020, Wang et al., 2023).
For the discrete convex setup, the Mirror Sinkhorn algorithm executes a single-loop of mirror descent with a negative-entropy mirror map, alternating Sinkhorn-style marginal projections:
- Mirror step: ,
- Normalizations alternate row- or column-scaling via Diag operators,
- Averaging: Optional running average in stochastic settings,
- Output: Last or averaged iterate, depending on the regime (Ballu et al., 2022).
3. Convergence Analysis and Theoretical Guarantees
Online Sinkhorn algorithms admit strong convergence properties:
- For streaming settings with growing batch sizes , decreasing step sizes , and , the expected variation-norm error for the duals decays as , where is the total sample count (Wang et al., 2023).
- Previous results only established rates; refined analyses lead to sharper algebraic exponents tied to (Mensch et al., 2020, Wang et al., 2023).
- In the discrete finite case, Mirror Sinkhorn yields, for -Lipschitz convex , online regret bounds and constraint violation sums (Ballu et al., 2022).
- When is strongly convex and smooth in KL divergence, the method obtains aggregate error up to logarithmic factors.
- The Bregman decomposition underpins the proofs, leveraging contraction and Pinsker’s inequalities.
4. Algorithmic Extensions: Compression and Mirror Descent
To counteract the unbounded memory and kernel-evaluation cost, Compressed Online Sinkhorn periodically reduces the support of the dual potentials using moment-preserving compression:
- Compression by moment-matching: Replace the kernel expansion with a reduced atomic measure matching selected functional moments (e.g., polynomials in Gaussian Quadrature for , Fourier moments for any ). For Fourier-based compression, error bounds are for .
- Theoretical result: If (compressed) satisfies , and , the compressed algorithm preserves the convergence rate for an appropriate compression scheme (Wang et al., 2023).
Mirror Sinkhorn is further extensible to non-linear, stochastic, or multi-marginal objectives and to dual-space variants.
5. Empirical Performance and Implementation Considerations
Empirical studies validate that:
- Standard Online Sinkhorn matches batch Sinkhorn objectives with substantially less memory and enables online warm-starting of batch algorithms, yielding speed-ups (e.g., for 3D shape matching).
- Compressed Online Sinkhorn reduces wall-clock times by (e.g., in 5D GMMs), maintaining matching error decay envelopes with reduced evaluation and storage costs.
- In convex discrete settings, Mirror Sinkhorn achieves zero-bias convergence and outperforms entropic-regularized Sinkhorn in accuracy, especially under stochastic-noise settings (Ballu et al., 2022).
- Typical choices for regularization (), step-sizes (), batching (), and compression levels are empirically justified (Wang et al., 2023).
6. Comparisons, Advantages, and Limitations
| Algorithm | Per-iteration Cost | Memory (w.r.t. N) | Constraint Consistency | Rate |
|---|---|---|---|---|
| Batch Sinkhorn | Finite sample bias | (dual error) | ||
| Online Sinkhorn | Consistent | , growing | ||
| Mirror Sinkhorn (discrete) | Marginal error | , (strong convex.) | ||
| Compressed Online Sinkhorn | Preserved rate | Same as Online, lower asymp. cost |
- Online Sinkhorn is single-pass and streaming-friendly, in contrast to two-phase batch approaches (Mensch et al., 2020, Wang et al., 2023).
- Mirror Sinkhorn avoids nested loops and the need to know a target precision in advance (Ballu et al., 2022).
- Compression schemes further control memory and arithmetic complexity, yielding asymptotic computational exponents as low as $2 + 1/a$ (Wang et al., 2023).
- A limitation is that marginal constraint satisfaction in stochastic methods is typically , slower than exponentially fast convergence in highly entropic batch Sinkhorn. No rate is guaranteed in the non-strongly convex case for Mirror Sinkhorn (Ballu et al., 2022).
7. Applications and Extensions
Online Sinkhorn methods are applicable across OT-based tasks in machine learning, including domain adaptation, WGAN training, barycenter computation, and matching in high-dimensional or streaming-data settings:
- These algorithms extend naturally to unbalanced OT, multi-marginal OT, and semi-discrete OT, and to robust optimization under gradient noise.
- Possible future directions include adaptive annealed scaling of projections, more efficient dual-space stochastic mirror-descent variants, and applications to nonconvex or measure-valued settings (Ballu et al., 2022, Wang et al., 2023, Mensch et al., 2020).