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.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Intermediate distribution restoration | Single-cell data (intermediate time points ti for i in {1, 2, 3}) | W1 Score0.793 | 28 | |
| Reconstruction of discrete branching dynamics | CITE 50D (t=3) | W111.65 | 24 | |
| Trajectory Interpolation | EB 50D (held-out time points) | Mean W18.646 | 24 | |
| Reconstruction of discrete branching dynamics | 5D EB t=1 | W10.556 | 24 | |
| Population Dynamics Interpolation | EB scRNA 5-dim PCA representation (leave-one-out) | W1 Distance0.793 | 21 | |
| Continuous-Time Dynamics Estimation | Synthetic Y-shaped first snapshot as initial state | L_DTW18.99 | 20 | |
| Continuous-Time Dynamics Estimation | Synthetic Arch first snapshot as initial state | L_DTW20.62 | 20 | |
| Trajectory Interpolation | Light V1 | W1 Error3.254 | 18 | |
| Trajectory Interpolation | Dendritic Stimulus | W1 Distance/Error4.333 | 18 | |
| Trajectory reconstruction | Gaussian Mixtures 1000D | W1 Distance3.543 | 18 |