# Fast and scalable Wasserstein-1 neural optimal transport solver for single-cell perturbation prediction

arXiv: https://arxiv.org/abs/2411.00614 · HTML: https://mlsanity.com/paper/2411.00614/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 15 of its 44 references that are in the map.
- First posted on arXiv 1 Nov 2024.

## Builds on (most cited first, 15 of 15)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Improved Training of Wasserstein GANs](https://mlsanity.com/paper/1704.00028/) (2017)
- [Computational Optimal Transport](https://mlsanity.com/paper/1803.00567/) (2018)
- [Input Convex Neural Networks](https://mlsanity.com/paper/1609.07152/) (2016)
- [Sorting out Lipschitz function approximation](https://mlsanity.com/paper/1811.05381/) (2018)
- [Optimal transport mapping via input convex neural networks](https://mlsanity.com/paper/1908.10962/) (2019)
- [Orthogonalizing Convolutional Layers with the Cayley Transform](https://mlsanity.com/paper/2104.07167/) (2021)
- [Do Neural Optimal Transport Solvers Work? A Continuous Wasserstein-2 Benchmark](https://mlsanity.com/paper/2106.01954/) (2021)
- [Supervised Training of Conditional Monge Maps](https://mlsanity.com/paper/2206.14262/) (2022)
- [Predicting Cellular Responses to Novel Drug Perturbations at a Single-Cell Resolution](https://mlsanity.com/paper/2204.13545/) (2022)
- [On amortizing convex conjugates for optimal transport](https://mlsanity.com/paper/2210.12153/) (2022)
- [Approximating Lipschitz continuous functions with GroupSort neural networks](https://mlsanity.com/paper/2006.05254/) (2020)
- [Kantorovich Strikes Back! Wasserstein GANs are not Optimal Transport?](https://mlsanity.com/paper/2206.07767/) (2022)
- [A new method for determining Wasserstein 1 optimal transport maps from Kantorovich potentials, with deep learning applications](https://mlsanity.com/paper/2211.00820/) (2022)

## Built on by (most cited first, 4 of 4)

- [Towards generalizable single-cell perturbation modeling via the Conditional Monge Gap](https://mlsanity.com/paper/2504.08328/) (2025)
- [Displacement-Sparse Neural Optimal Transport](https://mlsanity.com/paper/2502.01889/) (2025)
- [Modeling Cell-Cycle-Aware Single-Cell Drug Perturbation Responses](https://mlsanity.com/paper/2606.30695/) (2026)
- [CloneOT: predicting clone-specific perturbation responses with lineage-supervised neural inverse optimal transport](https://mlsanity.com/paper/2610.16093/) (2026)

## Abstract (arXiv)

Motivation: Predicting single-cell perturbation responses requires mapping between two unpaired single-cell data distributions. Optimal transport (OT) theory provides a principled framework for constructing such mappings by minimizing transport cost. Recently, Wasserstein-2 (W2) neural optimal transport solvers (textite.g., CellOT) have been employed for this prediction task. However, W2 OT relies on the general Kantorovich dual formulation, which involves optimizing over two conjugate functions, leading to a complex min-max optimization problem that converges slowly.  Results: To address these challenges, we propose a novel solver based on the Wasserstein-1 (W1) dual formulation. Unlike W2, the W1 dual simplifies the optimization to a maximization problem over a single 1-Lipschitz function, thus eliminating the need for time-consuming min-max optimization. While solving the W1 dual only reveals the transport direction and does not directly provide a unique optimal transport map, we incorporate an additional step using adversarial training to determine an appropriate transport step size, effectively recovering the transport map. Our experiments demonstrate that the proposed W1 neural optimal transport solver can mimic the W2 OT solvers in finding a unique and ``monotonic" map on 2D datasets. Moreover, the W1 OT solver achieves performance on par with or surpasses W2 OT solvers on real single-cell perturbation datasets. Furthermore, we show that W1 OT solver achieves 25 ~ 45× speedup, scales better on high dimensional transportation task, and can be directly applied on single-cell RNA-seq dataset with highly variable genes.  Availability and Implementation: Our implementation and experiments are open-sourced at https://github.com/poseidonchan/w1ot.