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

作者：Yanshuo Chen, Zhengmian Hu, Wei Chen, Heng Huang · 发表于：Bioinformatics · 年份：2025 · DOI：10.1093/bioinformatics/btaf253 · 被引用次数：5 · 研究领域：Single-cell and spatial transcriptomics、Stochastic Gradient Optimization Techniques、Genomics and Chromatin Dynamics

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 (e.g. CellOT) have been used 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...