Trajectory Inference with Smooth Schrödinger Bridges

TL;DR

This work introduces Smooth Schrödinger Bridges (SSB) by replacing the Brownian reference with a smooth Gaussian process prior, notably Matérn kernels, to obtain regular trajectory estimates for trajectory inference. The authors lift the SB problem to phase space, where GAPs yield a Gauss–Markov representation that enables belief propagation to solve the lifted entropic OT efficiently, with polynomial-time guidance in $K$. They develop an approximate belief propagation scheme based on finite orthonormal expansions, achieving practical runtimes and strong empirical performance on synthetic tracking tasks and real single-cell RNA-seq–inspired data, substantially outperforming or matching state-of-the-art methods. The approach offers a principled, flexible framework for smooth trajectory inference with broad applicability, while noting dimensionality as a key future challenge and pointing to noisy-observation extensions as avenues for future work.

Abstract

Motivated by applications in trajectory inference and particle tracking, we introduce Smooth Schrödinger Bridges. Our proposal generalizes prior work by allowing the reference process in the Schrödinger Bridge problem to be a smooth Gaussian process, leading to more regular and interpretable trajectories in applications. Though naïvely smoothing the reference process leads to a computationally intractable problem, we identify a class of processes (including the Matérn processes) for which the resulting Smooth Schrödinger Bridge problem can be lifted to a simpler problem on phase space, which can be solved in polynomial time. We develop a practical approximation of this algorithm that outperforms existing methods on numerous simulated and real single-cell RNAseq datasets. The code can be found at https://github.com/WanliHongC/Smooth_SB

Figures (12)

• Figure 1: Comparison between the classical SB and smooth SB with Matérn prior on a trajectory inference task for the petal dataset HugMagTon22. Colored points are observed data and paths are inferred trajectories. Our smooth SB approach generates much smoother and more concentrated trajectories.
• Figure 2: Comparison between the classical SB and smooth SB. The observations come from 20 independent trajectories of a Matérn Gaussian process with $\nu = 3.5$. The two pictures depict the posterior identification of particles obtained by solving the SB problem with two different priors. Our smooth SB approach recovers trajectories much more accurately.
• Figure 3: A portion of the graphical model corresponding to the joint law of $\boldsymbol{\omega}$ and $\boldsymbol{\eta}$.
• Figure 4: Visualization of orbits in 3D space. The colors of trajectories depend on the starting point. The second row is the visualization of the XY-space projection of the corresponding above plot.
• Figure 5: Visualization of trajectory inference on various datasets by lifted SB, for the $5$D Dyngen Tree data, we visualize the 2D projection of the first three dimensions in the second row.

Theorems & Definitions (11)

• Lemma 3.1
• Definition 3.2: RasWil06
• Theorem 3.3: Saa12
• Proposition 3.4: HarSAr10
• Lemma 4.1
• Theorem 5.1
• Lemma 6.1
• Theorem 7.1
• proof
• Lemma B.1