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

Jacob Lindbäck, David Alvarez-Melis, Mikael Johansson
Proceedings of The 29th International Conference on Artificial Intelligence and Statistics, PMLR 300:2593-2601, 2026.

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

Cite this Paper

BibTeX
@InProceedings{pmlr-v300-lindback26a, title = { Three-operator splitting with stale gradients for faster non-linear optimal transport }, author = {Lindb{\"a}ck, Jacob and Alvarez-Melis, David and Johansson, Mikael}, booktitle = {Proceedings of The 29th International Conference on Artificial Intelligence and Statistics}, pages = {2593--2601}, year = {2026}, editor = {Khan, Emtiyaz and Li, Yingzhen and Solin, Arno and Ramdas, Aaditya}, volume = {300}, series = {Proceedings of Machine Learning Research}, month = {02--05 May}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v300/main/assets/lindback26a/lindback26a.pdf}, url = {https://proceedings.mlr.press/v300/lindback26a.html}, 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 } }
Endnote
%0 Conference Paper %T Three-operator splitting with stale gradients for faster non-linear optimal transport %A Jacob Lindbäck %A David Alvarez-Melis %A Mikael Johansson %B Proceedings of The 29th International Conference on Artificial Intelligence and Statistics %C Proceedings of Machine Learning Research %D 2026 %E Emtiyaz Khan %E Yingzhen Li %E Arno Solin %E Aaditya Ramdas %F pmlr-v300-lindback26a %I PMLR %P 2593--2601 %U https://proceedings.mlr.press/v300/lindback26a.html %V 300 %X 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
APA
Lindbäck, J., Alvarez-Melis, D. & Johansson, M.. (2026). Three-operator splitting with stale gradients for faster non-linear optimal transport . Proceedings of The 29th International Conference on Artificial Intelligence and Statistics, in Proceedings of Machine Learning Research 300:2593-2601 Available from https://proceedings.mlr.press/v300/lindback26a.html.

Related Material