[论文解读] Proximal Optimal Transport Modeling of Population Dynamics

Charlotte Bunne, Laetitia Meng-Papaxanthos+2 more|arXiv (Cornell University)|Jun 11, 2021

Gaussian Processes and Bayesian Inference被引用 15

一句话总结

该论文提出JKOnet，一种可微分神经架构，通过约化最优传输（proximal optimal transport）利用Jordan-Kinderlehrer-Otto（JKO）流来建模群体动力学，从稀疏的时间分辨点云快照中学习能量函数。其在重建轨迹和捕捉单细胞RNA-seq数据中的生物异质性方面优于直接前向建模方法，对噪声具有鲁棒性，并能准确预测谱系分化。

ABSTRACT

We propose a new approach to model the collective dynamics of a population of particles evolving with time. As is often the case in challenging scientific applications, notably single-cell genomics, measuring features for these particles requires destroying them. As a result, the population can only be monitored with periodic snapshots, obtained by sampling a few particles that are sacrificed in exchange for measurements. Given only access to these snapshots, can we reconstruct likely individual trajectories for all other particles? We propose to model these trajectories as collective realizations of a causal Jordan-Kinderlehrer-Otto (JKO) flow of measures: The JKO scheme posits that the new configuration taken by a population at time $t+1$ is one that trades off an improvement, in the sense that it decreases an energy, while remaining close (in Wasserstein distance) to the previous configuration observed at $t$. In order to learn such an energy using only snapshots, we propose JKOnet, a neural architecture that computes (in end-to-end differentiable fashion) the JKO flow given a parametric energy and initial configuration of points. We demonstrate the good performance and robustness of the JKOnet fitting procedure, compared to a more direct forward method.

研究动机与目标

• 当仅有周期性、破坏性群体快照时，重建群体中单个粒子的轨迹。
• 将粒子构型的演化建模为一种因果性、基于能量的JKO流，平衡能量最小化与先前状态的Wasserstein接近度。
• 通过端到端可微分优化，从观测快照中学习参数化能量函数。
• 在噪声大、高维的设定下，相比直接前向建模方法，提升鲁棒性与预测准确性。
• 从时间分辨的单细胞数据中捕捉生物异质性，如胚胎干细胞分化过程中的谱系分支。

提出的方法

• JKOnet使用输入凸神经网络（ICNNs）计算JKO步骤，以参数化能量函数并确保凸性，从而实现可微分的最优传输流。
• 该模型将JKO步骤视为测度空间中的近端算子，最小化能量函数的同时保持与先前构型在Wasserstein距离上的接近。
• 采用双层优化问题通过最小化预测与观测点云快照之间的时间Wasserstein距离来训练能量参数。
• 该架构完全可微分，支持通过JKO步骤反向传播以更新能量函数参数。
• 该方法利用Brenier定理，确保流源自凸势函数的梯度，与最优传输原理一致。
• 该模型在合成数据和人类胚胎干细胞的真实单细胞RNA-seq数据上进行评估，并与直接前向建模基线方法进行比较。

实验结果

研究问题

• RQ1 可微分的近端最优传输框架能否仅从群体的周期性、破坏性快照中重建单个粒子的轨迹？
• RQ2 在存在噪声和数据有限的情况下，JKOnet在重建群体动力学方面与直接前向建模方法相比表现如何？
• RQ3 JKOnet在多大程度上能从时间分辨的单细胞数据中捕捉复杂的生物异质性，如胚胎干细胞分化过程中的谱系分支？
• RQ4 即使预测多步前进，基于JKO的方法是否能在长时间跨度内保持稳定的预测误差？
• RQ5 该模型能否学习到一个有意义且可解释的能量函数，以反映复杂生物系统背后的动态机制？

主要发现

• 在所有时间步长中，JKOnet在单步预测和全步预测的$W_{ε}$距离上均优于直接前向方法，性能稳定至第27天。
• 即使在预测整个轨迹（全步预测）时，模型仍保持较低的$W_{ε}$差异，表明对误差累积具有鲁棒性。
• JKOnet与前向方法均捕捉到了人类胚胎干细胞分化中的主要谱系分支（中胚层、内胚层、神经外胚层、神经嵴），但对代表性不足类别的预测准确性较低。
• 预测轨迹与单步预测相比，预测与真实谱系分布之间的$\ell_1$距离和Hellinger距离均显示性能下降。
• JKOnet更好地保持了生物异质性，PCA和PHATE可视化结果表明预测谱系分布与真实值对齐。
• 该模型展现出强大的泛化能力与稳定性，即使在噪声或数据稀疏的情况下，长时间跨度内预测误差也无显著增加。