SWAP: A Scalable, Warm-Startable, Anytime Permutation Solver for Optimal Transport

Abstract

Large-scale and high-dimensional discrete optimal transport often involves a tension between scalability and exact assignment feasibility: many scalable approaches optimize relaxed, factorized, or sparse transport plans rather than maintaining a one-to-one assignment throughout optimization. We introduce SWAP, an iterative solver that operates directly in permutation space for equally weighted point clouds. Starting from any feasible assignment, SWAP uses geometry-aware proposals to identify cost-decreasing local rearrangements. Consequently, every iterate is a valid permutation, so the one-to-one constraints are satisfied exactly throughout the optimization without rounding or feasibility correction. SWAP avoids constructing an $N\times N$ cost or coupling matrix and requires only $O(Nd)$ memory. We establish monotone descent, introduce a family of cycle-stability conditions, prove that SWAP reaches pairwise stability almost surely under generic nonparallel conditions, and derive sufficient conditions under which pairwise stability implies global optimality. Experiments on an ImageNet assignment with 640,500 points in 2,048 dimensions, synthetic benchmarks, and MERFISH spatial alignment demonstrate that SWAP achieves lower transport costs on challenging large-scale and high-dimensional problems, with substantially reduced runtime and memory compared with