[论文解读] Manifold Interpolating Optimal-Transport Flows for Trajectory Inference

Guillaume Huguet, D S Magruder+5 more|arXiv (Cornell University)|Jun 29, 2022

Traffic Prediction and Management TechniquesEngineering被引用 19

一句话总结

MIOFlow 提出了一种基于神经ODE的框架，通过结合流形学习与最优传输，从静态、时间稀疏的快照中学习随机的、连续的种群动态。它使用测地线自编码器将数据嵌入潜在空间，其中距离反映一种新颖的多尺度扩散测地线距离，从而在复杂流形上实现精确的插值与轨迹推断，优于归一化流和薛定谔桥在建模分叉和生物动态方面的表现。

ABSTRACT

We present a method called Manifold Interpolating Optimal-Transport Flow (MIOFlow) that learns stochastic, continuous population dynamics from static snapshot samples taken at sporadic timepoints. MIOFlow combines dynamic models, manifold learning, and optimal transport by training neural ordinary differential equations (Neural ODE) to interpolate between static population snapshots as penalized by optimal transport with manifold ground distance. Further, we ensure that the flow follows the geometry by operating in the latent space of an autoencoder that we call a geodesic autoencoder (GAE). In GAE the latent space distance between points is regularized to match a novel multiscale geodesic distance on the data manifold that we define. We show that this method is superior to normalizing flows, Schrödinger bridges and other generative models that are designed to flow from noise to data in terms of interpolating between populations. Theoretically, we link these trajectories with dynamic optimal transport. We evaluate our method on simulated data with bifurcations and merges, as well as scRNA-seq data from embryoid body differentiation, and acute myeloid leukemia treatment.

研究动机与目标

• 解决从高维数据的静态、时间有序快照中推断连续、随机种群动态的挑战。
• 在无法进行直接纵向追踪的系统中建模轨迹推断，例如单细胞RNA测序。
• 确保推断的轨迹尊重底层数据流形的内在几何结构，而非高维环境空间。
• 开发一种计算高效的算法，避免使用现有生成模型中常见的高斯先验和确定性流的限制。
• 实现在复杂生物系统中对未观测时间点和单个细胞轨迹的精确插值与重建。

提出的方法

• MIOFlow 使用神经常微分方程（Neural ODEs）来建模不同时间点的静态种群快照之间的连续流。
• 它采用测地线自编码器（GAE），通过在潜在空间中正则化以保持数据流形上的新型多尺度扩散测地线距离。
• 扩散测地线距离通过基于核的图拉普拉斯定义，并在样本无限增多的极限下收敛到流形上的真实测地线距离。
• 通过使用流形基距离的最优传输对流进行惩罚，确保时间点之间的传输成本与数据的内在几何结构一致。
• 该方法在GAE的潜在空间中运行，通过在低维流形上学习动态来降低计算成本并提高泛化能力。
• 通过潜在空间中点的初始采样自然引入随机性，避免了如确定性流中所需的显式噪声注入。

实验结果

研究问题

• RQ1 我们能否在不假设高斯先验的情况下，从未知时间稀疏的静态快照中学习连续、随机的种群动态？
• RQ2 如何确保推断的轨迹尊重底层数据流形的内在几何结构？
• RQ3 我们能否在潜在空间中构建一个数据驱动的、多尺度的距离度量，使其收敛到流形上的真实测地线距离？
• RQ4 MIOFlow 在建模复杂动态（如分叉和合并）方面，与归一化流、薛定谔桥和TrajectoryNet等现有方法相比表现如何？
• RQ5 MIOFlow 能否准确重建单个细胞轨迹，并捕捉分化和治疗反应等生物过程中基因表达的动态变化？

主要发现

• 在具有分叉和合并的合成数据上，MIOFlow 在插值种群之间时优于归一化流和薛定谔桥，生成的轨迹更平滑且测地线准确性更高。
• 在模拟数据中，MIOFlow 成功建模了发散轨迹，其流准确捕捉了底层流形的分支结构。
• 在类胚体分化的单细胞RNA测序数据中，MIOFlow 重建了与已知发育标志物一致的生物上合理的谱系轨迹。
• 在急性髓系白血病（AML）治疗数据中，MIOFlow 揭示了存活细胞向白血病干细胞（LSC）特征转变，后期时间点中SLPI、MYB和B2M的表达增加。
• 结合扩散测地线距离的测地线自编码器实现了准确的流形重建，且潜在空间中的流保留了可解释的生物距离，可解码为具有生物学意义的基因表达模式。
• 由于避免了雅可比行列式迹的计算，MIOFlow 的每次函数评估推理速度优于归一化流，其复杂度为O(k)而非O(k^2)，其中k为环境维度。