快速且可扩展的 Wasserstein-1 神经最优传输求解器用于单细胞扰动预测
机器学习
2025-04-22 v2 基因组学
摘要
动机：预测单细胞扰动响应需要在两个未配对的单细胞数据分布之间建立映射。最优传输（OT）理论通过最小化传输成本，为构建此类映射提供了原则性框架。近期，Wasserstein-2（W₂）神经最优传输求解器（例如 CellOT）已被用于此预测任务。然而，W₂ OT 依赖于通用的 Kantorovich 对偶公式，涉及对两个共轭函数的优化，导致一个复杂的极小-极大优化问题，收敛缓慢。结果：为解决这些挑战，我们提出了一种基于 Wasserstein-1（W₁）对偶公式的新求解器。与 W₂ 不同，W₁ 对偶将优化简化为对单个 1-Lipschitz 函数的极大化问题，从而消除了耗时的极小-极大优化的需要。虽然求解 W₁ 对偶仅揭示传输方向而不直接提供唯一的最优传输映射，但我们结合对抗训练来确定合适的传输步长，有效地恢复了传输映射。我们的实验表明，所提出的 W₁ 神经最优传输求解器能够在二维数据集上模仿 W₂ OT 求解器找到唯一且单调的映射。此外，W₁ OT 求解器在真实单细胞扰动数据集上的性能与 W₂ OT 求解器持平或超越。此外，我们证明 W₁ OT 求解器实现了 25∼45 倍的速度提升，在高维传输任务上具有更好的可扩展性，并且可以直接应用于具有高可变基因的单细胞 RNA-seq 数据集。可用性与实现：我们的实现和实验已开源，地址为 https://github.com/poseidonchan/w1ot。
引用
@article{arxiv.2411.00614,
title = {Fast and scalable Wasserstein-1 neural optimal transport solver for single-cell perturbation prediction},
author = {Yanshuo Chen and Zhengmian Hu and Wei Chen and Heng Huang},
journal= {arXiv preprint arXiv:2411.00614},
year = {2025}
}
备注
ISMB/ECCB 2025