Research map / Learning theory and algorithms / Optimal transport
Learning stochastic dynamics from snapshots through regularized unbalanced optimal transport
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes DeepRUOT, a deep learning solver for regularized unbalanced optimal transport that learns continuous unbalanced stochastic dynamics (growth/death included) from time-series snapshots via a Fisher-regularized ODE formulation.
Problem
Inferring continuous stochastic dynamics with unknown growth and death from sparsely time-resolved unnormalized snapshot distributions, where standard Schrödinger bridge and balanced OT methods fail to account for mass changes.
Approach
Reformulate regularized unbalanced OT with a Fisher information regularization that turns the stochastic (SDE) dynamics into a probability-flow ODE with a learned growth term, parameterized by neural networks for velocity, growth, and log-density, trained with energy, mass-aware unbalanced OT reconstruction, and Fokker-Planck losses.
Key results
- On the synthetic gene regulatory network at t=1, DeepRUOT reaches W1 0.044 and W2 0.058, versus 0.047/0.060 for Unbalanced OT and 0.174/0.303 for SF2M.
- On scRNA-seq hematopoiesis data at t=1, DeepRUOT reaches W1 0.145 and W2 0.173, versus 0.314/0.397 for Unbalanced OT and 0.167/0.213 for SF2M.
- Ablation on gene regulatory data shows removing growth term (w/o growth) degrades W2 at t=4 from ~0.168 to 0.877, matching the SF2M-like balanced behavior.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use DeepRUOT instead of Unbalanced OT | Gene regulatory network · Gene regulatory network, t=4 | W2: 0.124_{\pm 0.000} -> 0.121_{\pm 0.000} | Helped |
| Use DeepRUOT instead of SF2M | scRNA-seq · scRNA-seq, t=1 | W2: 0.213489_{\pm 0.006} -> 0.172878_{\pm 0.002} | Helped |
| Use DeepRUOT instead of Unbalanced SB | scRNA-seq · scRNA-seq, t=2 | W1: 0.128254_{\pm 0.003} -> 0.132411_{\pm 0.006} | Hurt |
| Add the growth term (ablation) | DeepRUOT · Gene regulatory network, t=4 | W2: 0.877_{\pm 0.011} -> 0.168_{\pm 0.007} | Helped |
| Add mass loss L_Mass (ablation) | DeepRUOT · Gene regulatory network, t=4 | W2: 0.816_{\pm 0.009} -> 0.168_{\pm 0.007} | Helped |
| Add Fokker-Planck constraint L_FP (ablation) | DeepRUOT · Gene regulatory network, t=4 | W2: 0.197_{\pm 0.011} -> 0.168_{\pm 0.007} | Hurt |
| Add the pre-training stage (ablation) | DeepRUOT · Gene regulatory network, t=4 | W2: 1.565_{\pm 0.028} -> 0.168_{\pm 0.007} | Helped |
| Add score matching during pre-training (ablation) | DeepRUOT · Gene regulatory network, t=4 | W2: 3.391_{\pm 0.007} -> 0.168_{\pm 0.007} | Helped |
| Add the training stage (ablation) | DeepRUOT · Gene regulatory network, t=4 | W2: 0.200_{\pm 0.015} -> 0.168_{\pm 0.007} | Helped |
Surprising findings
- On the hematopoiesis data at t=2, the Unbalanced SB baseline achieves a lower W1 (0.128) than DeepRUOT (0.132).
- Ablating the Fokker-Planck loss gives a lower W1 at early time points on the gene regulatory data than full DeepRUOT, though full DeepRUOT is better at later times.
Limitations
- Scalability to thousands of dimensions is not addressed; the authors note it would require simulation-free training methods.
- Relies on sparse snapshot time points and a hyperparameter scheduling strategy and two-stage training to stabilize optimization.
- Diffusion coefficient sigma had to be tuned differently across experiments (e.g., 0.05 for DeepRUOT vs 0.25 for SF2M on gene data).
Experimental setup
Datasets: synthetic gene regulatory network synthetic Gaussian mixtures mouse hematopoiesis scRNA-seq (lineage tracing) A549 EMT time-series scRNA-seq
Benchmarks: synthetic gene regulatory network synthetic Gaussian mixtures mouse hematopoiesis scRNA-seq A549 EMT scRNA-seq
Metrics: W1 distance W2 distance
Baselines: MIOFlow SF2M Unbalanced SB uAM Unbalanced OT UDSB
Builds on: Schrödinger bridge unbalanced dynamic optimal transport (WFR) conditional flow matching SF2M Neural ODE
Closest papers by meaning
- Variational Regularized Unbalanced Optimal Transport: Single Network, Least Action (NeurIPS 2025 · similarity 0.91)
- Modeling Cell Dynamics and Interactions with Unbalanced Mean Field Schrödinger Bridge (NeurIPS 2025 · similarity 0.88)
- Inferring stochastic dynamics with growth from cross-sectional data (NeurIPS 2025 · similarity 0.88)
- GENOT: Entropic (Gromov) Wasserstein Flow Matching with Applications to Single-Cell Genomics (NeurIPS 2024 · similarity 0.85)
- Joint Velocity-Growth Flow Matching for Single-Cell Dynamics Modeling (NeurIPS 2025 · similarity 0.85)
- Let's do the time-warp-attend: Learning topological invariants of dynamical systems (ICLR 2024 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.