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

• URL: http://arxiv.org/abs/2411.00614v1
• Date: Fri, 01 Nov 2024 14:23:19 GMT
• Title: Fast and scalable Wasserstein-1 neural optimal transport solver for single-cell perturbation prediction
• Authors: Yanshuo Chen, Zhengmian Hu, Wei Chen, Heng Huang,
• Abstract summary: Optimal transport theory provides a principled framework for constructing such mappings. We propose a novel optimal transport solver based on Wasserstein-1. Our experiments demonstrate that the proposed solver can mimic the $W$ OT solvers in finding a unique and monotonic" map on 2D datasets.
• Score: 55.89763969583124
• License:
• Abstract: 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 ($W_2$) neural optimal transport solvers (\textit{e.g.}, CellOT) have been employed for this prediction task. However, $W_2$ 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. To address these challenges, we propose a novel solver based on the Wasserstein-1 ($W_1$) dual formulation. Unlike $W_2$, the $W_1$ 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 $W_1$ 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 $W_1$ neural optimal transport solver can mimic the $W_2$ OT solvers in finding a unique and ``monotonic" map on 2D datasets. Moreover, the $W_1$ OT solver achieves performance on par with or surpasses $W_2$ OT solvers on real single-cell perturbation datasets. Furthermore, we show that $W_1$ OT solver achieves $25 \sim 45\times$ speedup, scales better on high dimensional transportation task, and can be directly applied on single-cell RNA-seq dataset with highly variable genes. Our implementation and experiments are open-sourced at \url{https://github.com/poseidonchan/w1ot}.

Related papers

• Obtaining Lower Query Complexities through Lightweight Zeroth-Order Proximal Gradient Algorithms [65.42376001308064]
  We propose two variance reduced ZO estimators for complex gradient problems. We improve the state-of-the-art function complexities from $mathcalOleft(minfracdn1/2epsilon2, fracdepsilon3right)$ to $tildecalOleft(fracdepsilon2right)$.
  arXiv  Detail & Related papers  (2024-10-03T15:04:01Z)
• Relative-Translation Invariant Wasserstein Distance [82.6068808353647]
  We introduce a new family of distances, relative-translation invariant Wasserstein distances ($RW_p$) We show that $RW_p distances are also real distance metrics defined on the quotient set $mathcalP_p(mathbbRn)/sim$ invariant to distribution translations.
  arXiv  Detail & Related papers  (2024-09-04T03:41:44Z)
• Progressive Entropic Optimal Transport Solvers [33.821924561619895]
  We propose a new class of EOT solvers (ProgOT) that can estimate both plans and transport maps. We provide experimental evidence demonstrating that ProgOT is a faster and more robust alternative to standard solvers. We also prove statistical consistency of our approach for estimating optimal transport maps.
  arXiv  Detail & Related papers  (2024-06-07T16:33:08Z)
• Semi-Discrete Optimal Transport: Nearly Minimax Estimation With Stochastic Gradient Descent and Adaptive Entropic Regularization [38.67914746910537]
  We prove an $mathcalO(t-1)$ lower bound rate for the OT map, using the similarity between Laguerre cells estimation and density support estimation. To nearly achieve the desired fast rate, we design an entropic regularization scheme decreasing with the number of samples.
  arXiv  Detail & Related papers  (2024-05-23T11:46:03Z)
• Normalizing flows as approximations of optimal transport maps via linear-control neural ODEs [49.1574468325115]
  "Normalizing Flows" is related to the task of constructing invertible transport maps between probability measures by means of deep neural networks. We consider the problem of recovering the $Wamma$-optimal transport map $T$ between absolutely continuous measures $mu,nuinmathcalP(mathbbRn)$ as the flow of a linear-control neural ODE.
  arXiv  Detail & Related papers  (2023-11-02T17:17:03Z)
• A Specialized Semismooth Newton Method for Kernel-Based Optimal Transport [92.96250725599958]
  Kernel-based optimal transport (OT) estimators offer an alternative, functional estimation procedure to address OT problems from samples. We show that our SSN method achieves a global convergence rate of $O (1/sqrtk)$, and a local quadratic convergence rate under standard regularity conditions.
  arXiv  Detail & Related papers  (2023-10-21T18:48:45Z)
• Fast Computation of Optimal Transport via Entropy-Regularized Extragradient Methods [75.34939761152587]
  Efficient computation of the optimal transport distance between two distributions serves as an algorithm that empowers various applications. This paper develops a scalable first-order optimization-based method that computes optimal transport to within $varepsilon$ additive accuracy.
  arXiv  Detail & Related papers  (2023-01-30T15:46:39Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.