Trajectory Inference with Smooth Schrödinger Bridges

Wanli Hong, Yuliang Shi, Jonathan Niles-Weed

TL;DR: We extend the multi-marginal Schrödinger Bridge problem by generalizing its reference process to a smooth Gaussian process, retaining computational tractability with polynomial-time complexity.

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 multi-marginal 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.

Lay Summary: Imagine trying to track the motion of identical ants in a colony using snapshots taken over time. Since they all look the same, it’s hard to tell which ant is which between photos. Due to computational constraints, prior methods only compared adjacent photos to guess the ants' paths, but this often failed because movements depend on longer patterns (like an ant speeding up or changing direction over time). Our new method finds these longer patterns. Instead of only looking one step ahead, it additionally infers information of velocities and accelerations of each ant, ensuring their paths make sense as smooth, natural motions (no sudden jumps). Our approximation scheme makes this algorithm scalable when the dataset expands and it outperforms existing methods on numerous datasets.

Link To Code: https://github.com/WanliHongC/Smooth_SB

Primary Area: Probabilistic Methods->Gaussian Processes

Keywords: Trajectory Inference, Schrodinger Bridge, Message Passing, Probabilisitic Graphical Models

Submission Number: 7187