Three-operator splitting with stale gradients for faster non-linear optimal transport

Jacob Lindbäck, David Alvarez-Melis, Mikael Johansson

TL;DR: We introduce an asynchronous adaptation of three-operator splitting for accelerating the computation of non-linear optimal transport.

Abstract: Scalable optimization for non-linear optimal transport (OT) poses unique challenges; it requires efficient memory management of large matrices, effective parallelization strategies suited for modern accelerators like GPUs, and theoretical guarantees that support practical implementation patterns. To address these challenges, we introduce a new algorithm based on three-operator splitting that reduces gradient computation costs by allowing gradient evaluations to run asynchronously and in parallel with other computations. Using monotone operator theory, we establish new convergence guarantees for this asynchronous adaptation and extend existing results to important non-convex problem classes, including Gromov–Wasserstein as a notable example. We validate our method through a series of experiments demonstrating improved accuracy and faster convergence for a broad range of problems

Code Dataset Promise: Yes

Code Dataset Url: https://github.com/lindbackjacob/atos-ot

Format Confirmation: I agree that I have read and followed the formatting instructions for the camera ready version.

Submission Number: 1145