Share your thoughts, 1 month free Claude Pro on usSee more
WorkDL logo mark

Simulation-free Schr\"odinger bridges via score and flow matching

About

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\"odinger 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.

Alexander Tong, Nikolay Malkin, Kilian Fatras, Lazar Atanackovic, Yanlei Zhang, Guillaume Huguet, Guy Wolf, Yoshua Bengio• 2023

Related benchmarks

TaskDatasetResultRank
Intermediate distribution restorationSingle-cell data (intermediate time points ti for i in {1, 2, 3})
W1 Score0.793
28
Reconstruction of discrete branching dynamicsCITE 50D (t=3)
W111.65
24
Trajectory InterpolationEB 50D (held-out time points)
Mean W18.646
24
Reconstruction of discrete branching dynamics5D EB t=1
W10.556
24
Population Dynamics InterpolationEB scRNA 5-dim PCA representation (leave-one-out)
W1 Distance0.793
21
Continuous-Time Dynamics EstimationSynthetic Y-shaped first snapshot as initial state
L_DTW18.99
20
Continuous-Time Dynamics EstimationSynthetic Arch first snapshot as initial state
L_DTW20.62
20
Trajectory InterpolationLight V1
W1 Error3.254
18
Trajectory InterpolationDendritic Stimulus
W1 Distance/Error4.333
18
Trajectory reconstructionGaussian Mixtures 1000D
W1 Distance3.543
18
Showing 10 of 68 rows

Other info

Follow for update