[论文解读] Trajectory Inference via Mean-field Langevin in Path Space

Lénaïc Chizat, Stephen X. Zhang+2 more|arXiv (Cornell University)|May 14, 2022

Single-cell and spatial transcriptomics被引用 4

一句话总结

本文提出了一种无网格的、基于平均场的朗之万（MFL）动力学算法，用于轨迹推断，该算法利用薛定谔桥和噪声梯度下降，在路径空间中计算相对于维纳测度的最小熵估计器。该方法实现了对真实估计器的全局收敛，并具备端到端的理论保证，能够从稀疏的时间分辨数据中实现一致的推断，适用于单细胞RNA-seq和质量变化建模。

ABSTRACT

Trajectory inference aims at recovering the dynamics of a population from snapshots of its temporal marginals. To solve this task, a min-entropy estimator relative to the Wiener measure in path space was introduced by Lavenant et al. arXiv:2102.09204, and shown to consistently recover the dynamics of a large class of drift-diffusion processes from the solution of an infinite dimensional convex optimization problem. In this paper, we introduce a grid-free algorithm to compute this estimator. Our method consists in a family of point clouds (one per snapshot) coupled via Schrödinger bridges which evolve with noisy gradient descent. We study the mean-field limit of the dynamics and prove its global convergence to the desired estimator. Overall, this leads to an inference method with end-to-end theoretical guarantees that solves an interpretable model for trajectory inference. We also present how to adapt the method to deal with mass variations, a useful extension when dealing with single cell RNA-sequencing data where cells can branch and die.

研究动机与目标

• 开发一种无网格、可扩展的算法，用于计算路径空间中的最小熵估计器，该问题在传统上因无限维凸优化而计算不可行。
• 通过证明所提出的MFL动力学全局收敛至真实最小熵估计器，为轨迹推断提供端到端的理论保证。
• 将该方法扩展以处理数据中的质量变化，如细胞分支和死亡，这些在单细胞RNA测序应用中十分常见。
• 为缺乏严格收敛保证的非凸深度学习基轨迹推断方法，提供一种理论基础坚实且可解释的替代方案。

提出的方法

• 该方法采用一组点云（每个时间快照对应一组）通过熵正则化最优传输或薛定谔桥进行耦合，以在时间上建模动力学。
• 将推断问题表述为路径空间中的平均场朗之万动力学，其中粒子通过在数据拟合与熵正则化组合目标上的噪声梯度下降演化。
• 动力学由一个随机微分方程控制，通过相对于维纳测度的相对熵，在对观测边缘分布的保真度与平滑性之间取得平衡。
• 通过调整步长 $\eta$、噪声水平 $\sigma$ 和温度 $\tau$ 的退火策略，确保稳定收敛。
• 引入动力学的改进版本以处理质量变化，通过引入已知的增长率 $g(t,x)$，实现对细胞增殖和死亡的建模。
• 该方法在基于CPU的框架中实现，无需显式边界约束，依赖于粒子在观测数据支撑集内的稳定性。

实验结果

研究问题

• RQ1 能否设计一种无网格的随机优化方法，以理论收敛保证计算路径空间中的最小熵估计器？
• RQ2 平均场朗之万动力学如何全局收敛至真实轨迹推断解？其收敛速率如何？
• RQ3 该方法能否扩展以处理单细胞数据中的质量变化（如细胞分支和死亡），而这些在标准传输模型中无法捕捉？
• RQ4 与现有方法（如Global Waddington-OT）相比，所提出的MFL方法在重建精度和对稀疏采样鲁棒性方面表现如何？

主要发现

• MFL动力学以指数速率全局收敛至最小熵估计器，为轨迹推断提供了强有力的理论保证。
• 该方法在合成和真实单细胞数据上均达到最先进性能，能量距离最小化时，MFL的 $\lambda = 0.025$，gWOT的 $\lambda = 0.01$。
• 即使每个时间点仅有50个粒子，该算法仍能成功重建复杂的动力学，如分叉过程和时变势场。
• 引入已知的增长率 $g(t,x)$ 能够准确建模质量变化（如细胞增殖），在生物上更现实的情境中显著提升重建保真度。
• 该方法在整个优化过程中保持稳定，无需显式边界条件，因为粒子始终位于观测数据的支撑集中。
• 该方法对超参数选择具有鲁棒性，在多个数据集和参数设置下均表现出一致性能，经能量距离和可视化重建质量验证。