Learning non-equilibrium diffusions with Schrödinger bridges: from exactly solvable to simulation-free
Stephen Zhang, Michael Stumpf
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.