arXiv:2411.00614cs.LGq-bio.GN2024-11被引 10

用W1最优传输加速单细胞扰动预测,速度快45倍且效果不差。

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

  • 改用W1对偶形式,将复杂双函数优化简化为单函数最大化,计算更高效。
  • 在真实单细胞数据上性能媲美甚至超过传统W2方法,且支持高维基因数据。
  • 适合需要快速建模细胞状态变化的生物学家和计算研究人员使用。

预测单细胞扰动响应需映射两个未配对的单细胞数据分布。最优传输(OT)理论为此提供了原则性框架,通过最小化传输成本构建映射。近期,基于Wasserstein-2(W2)的神经最优传输求解器(如CellOT)被用于该任务。然而,W2 OT依赖于广义Kantorovich对偶形式,需优化两个共轭函数,导致复杂的极小极大优化问题,收敛缓慢。为此,我们提出一种基于Wasserstein-1(W1)对偶形式的新求解器。与W2不同,W1对偶将优化简化为对单一1-Lipschitz函数的最大化,无需耗时的极小极大优化。虽然仅通过解W1对偶无法直接获得唯一最优传输映射,我们引入对抗训练以确定合适的传输步长,有效恢复映射。实验表明,所提W1神经最优传输求解器可在二维数据上模拟出与W2相似的唯一且“单调”的映射。在真实单细胞扰动数据集上,其性能达到或超越现有W2求解器。此外,该方法实现25~45倍的速度提升,高维任务下扩展性更好,并可直接应用于含高度可变基因的单细胞RNA-seq数据。

原文摘要 · Abstract (English)

\textbf{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 ($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. \\ \textbf{Results:} 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. \\ \textbf{Availability and Implementation:} Our implementation and experiments are open-sourced at https://github.com/poseidonchan/w1ot.

单细胞最优传输速度优化生物信息

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。