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

要約

\ textBf {Motivation：}シングルセル摂動応答を予測するには、2つの対応のないシングルセルデータ分布間のマッピングが必要です。
Optimal Transport（OT）理論は、輸送コストを最小限に抑えることにより、このようなマッピングを構築するための原則的なフレームワークを提供します。
最近、Wasserstein-2（$ w_2 $）Neural Optimal Transport Solvers（\ textit {e.g。}、Cellot）がこの予測タスクに採用されています。
ただし、$ w_2 $ otは、2つのコンジュゲート関数を最適化する一般的なカントロビッチの二重製剤に依存しており、ゆっくりと収束する複雑なMIN-MAX最適化問題につながります。
\\ \ textbf {results：}これらの課題に対処するために、wasserstein-1（$ w_1 $）のデュアル定式化に基づいた新しいソルバーを提案します。
$ w_2 $とは異なり、$ w_1 $ dualは、単一の1-lipschitz関数にわたる最大化問題に対する最適化を簡素化するため、時間のかかるMIN-MAX最適化の必要性が排除されます。
$ w_1 $ $ dualを解くと、輸送方向のみが明らかになり、独自の最適な輸送マップを直接提供しませんが、敵対的なトレーニングを使用して適切な輸送ステップサイズを決定し、トランスポートマップを効果的に回復するための追加のステップを組み込みます。
私たちの実験は、提案されている$ W_1 $ NEURAL OPTIMAL TRASTRANCE SOLVERが、2Dデータセットで一意の「単調」マップを見つける際に$ W_2 $ OTソルバーを模倣できることを示しています。
さらに、$ w_1 $ o otソルバーは、実際のシングルセル摂動データセットで$ w_2 $ otソルバーと同等のパフォーマンスを実現します。
さらに、$ w_1 $ otソルバーは$ 25 \ sim 45 \ times $ speedupを達成し、高次元の輸送タスクでより良いスケーリングを行い、高可変遺伝子を持つ単一セルRNA-seqデータセットに直接適用できることを示します。
\\ \ textBf {可用性と実装：}実装と実験は、https：//github.com/poseidonchan/w1otでオープンソースを撮影しています。

要約(オリジナル)

\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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google