[论文解读] A Computational Framework for Solving Wasserstein Lagrangian Flows

Kirill Neklyudov, Rob Brekelmans+4 more|arXiv (Cornell University)|Oct 16, 2023

Model Reduction and Neural Networks被引用 4

一句话总结

本文提出了一种统一的深度学习框架，通过利用拉格朗日作用泛函的对偶公式，求解Wasserstein拉格朗日流，实现无需轨迹反向传播或访问最优耦合的高效优化。该方法通过可定制的动能和势能项灵活整合多样先验知识，在单细胞轨迹推断任务中优于先前方法。

ABSTRACT

The dynamical formulation of the optimal transport can be extended through various choices of the underlying geometry (kinetic energy), and the regularization of density paths (potential energy). These combinations yield different variational problems (Lagrangians), encompassing many variations of the optimal transport problem such as the Schrödinger bridge, unbalanced optimal transport, and optimal transport with physical constraints, among others. In general, the optimal density path is unknown, and solving these variational problems can be computationally challenging. We propose a novel deep learning based framework approaching all of these problems from a unified perspective. Leveraging the dual formulation of the Lagrangians, our method does not require simulating or backpropagating through the trajectories of the learned dynamics, and does not need access to optimal couplings. We showcase the versatility of the proposed framework by outperforming previous approaches for the single-cell trajectory inference, where incorporating prior knowledge into the dynamics is crucial for correct predictions.

研究动机与目标

• 解决单细胞RNA-seq数据中从稀疏、边缘观测中推断轨迹的不适定问题。
• 将多种最优传输变体（如Schrödinger桥、非平衡OT、物理约束OT）统一于单一计算框架之下。
• 通过拉格朗日公式中可定制的动能和势能项，灵活地将先验知识融入动力学建模。
• 开发一种可微且可扩展的方法，避免对单个粒子轨迹进行模拟或反向传播。
• 通过拉格朗日的对偶优化，展示在单细胞轨迹推断和图像生成中优于现有方法的性能。

提出的方法

• 该框架将轨迹推断表述为最小化拉格朗日作用泛函，其由动能（Wasserstein-2或Wasserstein-Fisher-Rao几何）和势能（如熵或空间势能）组成。
• 采用拉格朗日的对偶公式推导出可处理的优化目标，避免对单个粒子轨迹进行模拟或微分。
• 该方法使用深度神经网络参数化向量场 $ s_t(x, \theta) $ 和密度路径 $ \rho_t(x, \eta) $，通过双目标训练以匹配边缘约束。
• 优化过程采用双目标形式，实现端到端训练，无需显式轨迹展开或耦合估计。
• 在图像生成中，该方法通过哈密顿-雅可比条件实现单步采样，其中零加速度意味着 $ X_1 = X_0 + \nabla s_0(X_0) $。
• 通过组合不同能量项，该框架支持多种配置，为特定生物或物理约束提供归纳偏置。

实验结果

研究问题

• RQ1 一个统一的深度学习框架能否在单一公式下求解多种最优传输问题，包括非平衡OT、Schrödinger桥和物理约束OT？
• RQ2 如何在变分框架中灵活编码关于细胞动力学的先验知识以用于轨迹推断？
• RQ3 拉格朗日的对偶优化能否避免轨迹反向传播的计算负担，同时保持精度？
• RQ4 与现有方法相比，该框架在单细胞RNA-seq数据中的轨迹推断性能提升程度如何？
• RQ5 该框架能否通过最优传输映射实现无需迭代ODE求解的单步图像生成？

主要发现

• 所提出的框架在单细胞轨迹推断任务中优于先前方法，尤其在通过定制拉格朗日量灵活整合生物先验知识时表现更优。
• 该方法在低维和高维单细胞数据集（包括Embryoid Body、CITE-seq和Multiome数据）上均达到最先进性能。
• 在图像生成方面，该框架通过哈密顿-雅可比条件实现单步采样，仅需一次函数评估即可达到与迭代ODE求解器相当的质量。
• 对偶优化方法避免了轨迹反向传播和耦合估计，显著降低了计算成本。
• 通过调整拉格朗日分量，该框架成功泛化至多种最优传输变体，包括非平衡OT和Schrödinger桥。
• 在MNIST和CIFAR-10上的实验表明，通过学习的向量场实现的单步生成可产生与迭代扩散模型相当的高保真度样本。