Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Learning non-equilibrium diffusions with Schrödinger bridges: from exactly solvable to simulation-free

Stephen Y. Zhang, Michael P H Stumpf

TL;DR

This work extends Schrödinger bridges to non-equilibrium dynamics by using a multivariate Ornstein–Uhlenbeck reference process $d\mathbf{X}_t = \mathbf{A}(\mathbf{X}_t - \mathbf{m})\,dt + \boldsymbol{\sigma}\,d\mathbf{B}_t$, enabling non-conservative forces when $\mathbf{A}$ is asymmetric. It derives exact Gaussian SBP solutions for Gaussian endpoints (mvOU-GSB) and introduces a simulation-free learning approach mvOU-OTFM for general marginals by combining static entropic OT with score/flow matching. Empirically, mvOU-OTFM achieves higher accuracy and faster training than competing methods on synthetic and real single-cell data, and iterated mvOU-OTFM improves dynamic fidelity in oscillatory systems like the repressilator. The approach provides a practical, scalable framework for inferring non-equilibrium biological dynamics from snapshot data, with broad applicability to high-dimensional settings where traditional Brownian-based SBP struggles.

Abstract

We consider the Schrödinger bridge problem which, given ensemble measurements of the initial and final configurations of a stochastic dynamical system and some prior knowledge on the dynamics, aims to reconstruct the "most likely" evolution of the system compatible with the data. Most existing literature assume Brownian reference dynamics and are implicitly limited to potential-driven dynamics. We depart from this regime and consider reference processes described by a multivariate Ornstein-Uhlenbeck process with generic drift matrix $\mathbf{A} \in \mathbb{R}^{d \times d}$. When $\mathbf{A}$ is asymmetric, this corresponds to a non-equilibrium system with non-conservative forces at play: this is important for applications to biological systems, which are naturally exist out-of-equilibrium. In the case of Gaussian marginals, we derive explicit expressions that characterise the solution of both the static and dynamic Schrödinger bridge. For general marginals, we propose mvOU-OTFM, a simulation-free algorithm based on flow and score matching for learning the Schrödinger bridge. In application to a range of problems based on synthetic and real single cell data, we demonstrate that mvOU-OTFM achieves higher accuracy compared to competing methods, whilst being significantly faster to train.

Learning non-equilibrium diffusions with Schrödinger bridges: from exactly solvable to simulation-free

TL;DR

This work extends Schrödinger bridges to non-equilibrium dynamics by using a multivariate Ornstein–Uhlenbeck reference process , enabling non-conservative forces when is asymmetric. It derives exact Gaussian SBP solutions for Gaussian endpoints (mvOU-GSB) and introduces a simulation-free learning approach mvOU-OTFM for general marginals by combining static entropic OT with score/flow matching. Empirically, mvOU-OTFM achieves higher accuracy and faster training than competing methods on synthetic and real single-cell data, and iterated mvOU-OTFM improves dynamic fidelity in oscillatory systems like the repressilator. The approach provides a practical, scalable framework for inferring non-equilibrium biological dynamics from snapshot data, with broad applicability to high-dimensional settings where traditional Brownian-based SBP struggles.

Abstract

We consider the Schrödinger bridge problem which, given ensemble measurements of the initial and final configurations of a stochastic dynamical system and some prior knowledge on the dynamics, aims to reconstruct the "most likely" evolution of the system compatible with the data. Most existing literature assume Brownian reference dynamics and are implicitly limited to potential-driven dynamics. We depart from this regime and consider reference processes described by a multivariate Ornstein-Uhlenbeck process with generic drift matrix . When is asymmetric, this corresponds to a non-equilibrium system with non-conservative forces at play: this is important for applications to biological systems, which are naturally exist out-of-equilibrium. In the case of Gaussian marginals, we derive explicit expressions that characterise the solution of both the static and dynamic Schrödinger bridge. For general marginals, we propose mvOU-OTFM, a simulation-free algorithm based on flow and score matching for learning the Schrödinger bridge. In application to a range of problems based on synthetic and real single cell data, we demonstrate that mvOU-OTFM achieves higher accuracy compared to competing methods, whilst being significantly faster to train.

Paper Structure

This paper contains 45 sections, 6 theorems, 138 equations, 7 figures, 4 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Background and related work
  3. Schrödinger bridges
  4. Probability flows and (score, flow)-matching
  5. Probability flows
  6. Conditional flow matching
  7. Simulation-free Schrödinger bridges
  8. Related work
  9. Schrödinger bridges for non-conservative systems
  10. Multivariate Ornstein-Uhlenbeck bridges
  11. The Gaussian case
  12. The Non-Gaussian case -- score and flow matching
  13. Results
  14. Gaussian marginals: benchmarking accuracy
  15. Non-Gaussian marginals
  16. ...and 30 more sections

Key Result

Theorem 1

Consider the $d$-dimensional mvOU process eq:ref_sde. Conditioning on $(0, \bm{x}_0)$ and $(T, \bm{x}_T)$, the bridges of this process $\bm{Y}_t = \bm{X}_t | \{ \bm{X}_0 = \bm{x}_0, \bm{X}_T = \bm{x}_T \}, 0 \leq t \leq T$ are generated by the SDE where $\bm{\Lambda}_t = \int_0^{T-t} e^{-s \bm{A}} e^{-s \bm{A}^\top} \mathrm{d} s$ and $\bm{k}_t = e^{-(T - t) \bm{A}} (\bm{x}_T - \bm{m}) + \bm{m}$.

Figures (7)

  • Figure 1: The Schrödinger bridge problem with a multivariate Ornstein-Uhlenbeck reference process \ref{['eq:ref_sde']} can be solved via a generalised entropic transport problem and characterisation of the $\mathbb{Q}$-bridges. For non-Gaussian endpoints, score and flow matching provide a route to building neural approximations without simulation.
  • Figure 2: Gaussian Schrödinger Bridges. (a) (i) Marginals of the Gaussian Schrödinger bridge ($d = 10$) with mvOU reference (mvOU-GSB) (ii) Time dependent vector field generating the bridge. (b) Same as (a) but for Brownian reference.
  • Figure 3: Non-Gaussian marginals. (a)(i) Sampled trajectories from the mvOU-SB learned between non-Gaussian marginals in $d = 10$, shown in $(t, x_0)$ coordinates with score field shown in background. (a)(ii, iii) Sampled SDE (stochastic) and PF-ODE (deterministic) trajectories shown in $(x_0, x_1)$ coordinates, and reference drift shown in background. (b) Same as (a) but for Brownian reference.
  • Figure 4: Repressilator dynamics. (a) Repressilator population snapshots and trajectories. (b) (i) Ground truth vector field. (ii, iii) Inferred multi-marginal SB vector field with (fitted mvOU, Brownian) references. (c) Ground truth linearisation of system and drift $\bm{A}$ learned by mvOU-OTFM. (d) Leave-one-out error by iteration. (e) Illustration of leave-one-out interpolation between two example timepoints $p_{i-1}$ (blue), $p_{i+1}$ (green) with learned mvOU reference vs. Brownian reference.
  • Figure 5: Cell cycle scRNA-seq. (a) Streamlines of (i) transcriptomic vector field calculated from metabolic labelling data (ii) Fitted mvOU reference process, and Schrödinger bridge drift $\bm{v}_{\mathrm{SB}}$ learned by OTFM with (iii) mvOU reference and (iv) Brownian reference. (b) Marginal interpolation error between first and last snapshot as a function of the reference velocity scale parameter, $\gamma$.
  • ...and 2 more figures

Theorems & Definitions (7)

  • Theorem 1: SDE characterisation of multivariate Ornstein-Uhlenbeck bridge
  • Theorem 2: Score and flow for multivariate Ornstein-Uhlenbeck bridge
  • Theorem 3: Characterisation of mvOU-GSB
  • Proposition 1
  • Theorem 4: mvOU-OTFM solves \ref{['eq:SBP_dynamic']}
  • Lemma 1
  • proof