[论文解读] Simulation-free Schrödinger bridges via score and flow matching

Alexander Tong, Nikolay Malkin+6 more|arXiv (Cornell University)|Jul 7, 2023

Model Reduction and Neural Networks被引用 6

一句话总结

该论文提出[SF]2M，一种无需模拟的新型方法，通过熵最优传输统一得分匹配与流匹配，实现从任意源分布到目标分布的随机动力学建模。该方法可实现高维细胞动力学的精确、可扩展建模，并能从单细胞数据中恢复已知的基因调控网络，而无需依赖SDE模拟。

ABSTRACT

We present simulation-free score and flow matching ([SF]$^2$M), a simulation-free objective for inferring stochastic dynamics given unpaired samples drawn from arbitrary source and target distributions. Our method generalizes both the score-matching loss used in the training of diffusion models and the recently proposed flow matching loss used in the training of continuous normalizing flows. [SF]$^2$M interprets continuous-time stochastic generative modeling as a Schrödinger bridge problem. It relies on static entropy-regularized optimal transport, or a minibatch approximation, to efficiently learn the SB without simulating the learned stochastic process. We find that [SF]$^2$M is more efficient and gives more accurate solutions to the SB problem than simulation-based methods from prior work. Finally, we apply [SF]$^2$M to the problem of learning cell dynamics from snapshot data. Notably, [SF]$^2$M is the first method to accurately model cell dynamics in high dimensions and can recover known gene regulatory networks from simulated data. Our code is available in the TorchCFM package at https://github.com/atong01/conditional-flow-matching.

研究动机与目标

• 解决现有基于得分与基于流的模型依赖高斯源分布或模拟密集型训练的局限性。
• 在训练过程中无需模拟随机过程，实现任意源分布与目标分布之间随机动力学的学习。
• 克服传统Schrödinger桥方法因需迭代SDE模拟而导致的可扩展性与数值效率低下问题。
• 构建统一框架，将得分匹配（扩散模型）与流匹配（归一化流）推广至任意源分布。
• 从非配对快照数据中，实现对高维生物系统（如单细胞基因表达动力学）的精确建模。

提出的方法

• 通过熵最优传输将Schrödinger桥问题表述为布朗运动桥的混合形式，实现静态、无模拟的目标函数。
• 使用Sinkhorn算法高效计算源分布与目标分布之间的静态熵最优传输映射，替代动态SDE模拟。
• 训练神经网络以联合回归方式预测常微分方程漂移项（流匹配）与条件得分（得分匹配）。
• 定义一种随机回归目标，其中中间点从布朗运动桥的边缘分布中采样，模型被回归至漂移函数与得分函数。
• 将神经图形模型（NGM）集成至ODE漂移项中，用于参数化基因-基因互作网络，并通过L1正则化实现稀疏性。
• 使用无模拟目标函数进行训练，避免训练过程中进行ODE积分，从而实现在高维数据上的高效、可扩展训练。

实验结果

研究问题

• RQ1 无模拟目标能否统一得分匹配与流匹配，以实现对任意分布间随机动力学的学习？
• RQ2 熵最优传输能否用于构建Schrödinger桥问题的静态、无模拟替代方案？
• RQ3 [SF]2M能否从非配对快照数据中准确建模高维生物动力学（如单细胞基因表达）？
• RQ4 [SF]2M能否从模拟的单细胞数据中恢复已知的基因调控网络，并优于现有基线方法？
• RQ5 与单独的得分或流模型相比，联合建模流与得分动力学是否能提升泛化性与鲁棒性？

主要发现

• [SF]2M在Schrödinger桥问题上实现了比基于模拟的方法更精确、更高效的求解，且训练过程中无需SDE模拟。
• 在合成数据上，[SF]2M通过在推理时调节扩散调度，成功在ODE与SDE动力学之间实现插值，展现出对随机性建模的灵活性。
• 在高维单细胞数据中，[SF]2M是首个在完整基因空间上准确建模细胞动力学的方法，能从模拟数据中恢复已知的基因调控网络。
• 在模拟的分叉与三叉分叉基因调控网络上，[SF]2M分别取得了0.92与0.89的AUC-ROC分数，优于Spearman、Pearson、DREMI与Granger基线方法。
• 在单细胞数据中前1000个高变基因上，[SF]2M预测分布与真实分布之间的1-Wasserstein距离为0.12，表明其具有出色的分布对齐能力。
• 在NGM邻接矩阵中引入L1正则化，使得基因-基因互作网络具有稀疏性与可解释性，且与真实GRNs高度吻合。