Research map / Learning theory and algorithms / Optimal transport

Log-Averaged Mirror Prox for Fast, Large-Scale Optimal Transport in Linear Space

NeurIPS 2026 · Learning theory and algorithms › Optimal transport · approach: Sinkhorn and entropic OT algorithms · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes LAMP, a linear-space (O(n+m) memory) dual-only implementation of primal-dual mirror prox for OT/EOT that is equivalent to PDMP, with a last-iterate certificate for direct unregularized OT and GPU-optimized log-domain reductions.

Problem

Large-scale (entropic) optimal transport requires O(nm) storage for primal-dual first-order methods with good last-iterate complexity, which is prohibitive for large n, m and hard to map to GPUs.

Approach

Shows that primal mirror descent steps on the row-stochastic reparameterized OT penalty equal dual log-averaging with a temperature-annealing scalar sequence, so primal iterates can be recovered from O(n+m) dual variables, keeping PDMP's guarantees while using dense GPU reductions.

Key results

• LAMP (eta=0) outperforms Sinkhorn-based baselines on DOTmark with l_inf and l_1 ground costs at n=m=4096, converging extremely rapidly; Sinkhorn overtakes it for l_2^2 costs with tuned regularization.
• A fused+warp-tiled LogSumExp kernel runs 252.5 ms vs PyKeOps 332 ms at n=65536 on an RTX 4090.
• LAMP runs 512x512 color transfer on a single L40s using about 38 MB, where explicit primal-dual methods would need about 512 GB; scales to n=m=2^18.

Surprising findings

• Unregularized LAMP with eta=0 and fixed step sizes performs competitively without problem-specific tuning, outperforming well-tuned Sinkhorn variants on several cost types.

Limitations

• Last-iterate convergence (infeasibility) for the performant eta=0 regime is not proven; the paper only gives a bound dependent on infeasibility.
• LAMP converges slowly for l_2^2 costs where ||C||_inf = O(n), and Sinkhorn can overtake it.
• Theoretical complexity guarantees hold only for the conservative eta>0 parameter set.

Experimental setup

Datasets: DOTmark cell omics datasets

Benchmarks: DOTmark cell omics similarity color transfer 512x512

Metrics: objective gap plus infeasibility primal-dual gap wall-clock time infeasibility

Baselines: Sinkhorn APDAMD Accelerated Sinkhorn Bregman HPD Annealed Sinkhorn Dual Extrapolation PyKeOps

Builds on: PDMP mirror prox Sinkhorn Round (Altschuler et al.) Dual extrapolation

Closest papers by meaning

• A fast and accurate splitting method for optimal transport: analysis and implementation (ICLR 2022 · similarity 0.88)
• A Truncated Newton Method for Optimal Transport (ICLR 2025 · similarity 0.88)
• A Provably Convergent and Practical Algorithm for Gromov–Wasserstein Optimal Transport (NeurIPS 2026 · similarity 0.84)
• Progressive Entropic Optimal Transport Solvers (NeurIPS 2024 · similarity 0.84)
• Safe and Sparse Newton Method for Entropic-Regularized Optimal Transport (NeurIPS 2024 · similarity 0.84)
• On Partial Optimal Transport: Revising the Infeasibility of Sinkhorn and Efficient Gradient Methods (AAAI 2024 · similarity 0.83)