Learning non-equilibrium diffusions with Schr\"odinger bridges: from exactly solvable to simulation-free

TL;DR

This paper extends Schr"odinger bridge methods to non-equilibrium systems driven by Ornstein-Uhlenbeck processes with non-gradient forces, providing explicit solutions for Gaussian cases and a fast, simulation-free algorithm for general marginals, with applications to biological data.

Contribution

It introduces a novel approach for non-equilibrium Schr"odinger bridges using Ornstein-Uhlenbeck processes and develops a simulation-free algorithm for practical applications.

Findings

Explicit solutions for Gaussian marginals in non-equilibrium Schr"odinger bridges.

The mvOU-OTFM algorithm is faster and more accurate than existing methods.

Successful application to single-cell data demonstrating practical utility.

Abstract

We consider the Schr\"odinger 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 modelling systems driven by the gradient of a potential energy. 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 in which non-gradient forces are at play: this is important for applications to biological systems, which naturally exist out-of-equilibrium. In the case of Gaussian marginals, we derive explicit expressions that characterise exactly the solution of both the static and dynamic Schr\"odinger bridge. For general marginals, we propose mvOU-OTFM, a simulation-free algorithm based on flow and score matching for learning an approximation to the Schr\"odinger 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.