Multi-marginal temporal Schrödinger Bridge Matching from unpaired data
Abstract
Many natural dynamic processes –such as in vivo cellular differentiation or disease progression– can only be observed through the lens of static sample snapshots. While challenging, reconstructing their temporal evolution to decipher underlying dynamic properties is of major interest to scientific research. Existing approaches enable data transport along a temporal axis but are poorly scalable in high dimension and require restrictive assumptions to be met. To address these issues, we propose Multi-Marginal temporal Schrödinger Bridge Matching (MMtSBM) from unpaired data, extending the theoretical guarantees and empirical efficiency of Diffusion Schrödinger Bridge Matching (Shi et al., 2023) by deriving the Iterative Markovian Fitting algorithm to multiple marginals in a novel factorized fashion. Experiments show that MMtSBM retains theoretical properties on toy examples, achieves state-of-the-art performance on real-world datasets such as transcriptomic trajectory inference in 100 dimensions, and, for the first time, recovers couplings and dynamics in very high-dimensional image settings. Our work establishes multi-marginal Schrödinger bridges as a practical and principled approach for recovering hidden dynamics from static data.
Lay Summary
Many natural continuous processes cannot be directly observed over time. Instead we often only have access to static snapshots of data. Typically, following the evolution of the genetic expression of a single cell through time is mostly impossible currently, because measuring it actually requires destroying the cell. The same can be said of many imaging techniques of cells or tissues. Many methods have been developed to recover continuous processes from such static snapshots, mostly for data like genetic expression. However they are hard to apply to images because of the heavy computational cost they incur, and they use restrictive assumptions on the trajectories they can generate. To address these issues we developed a method, MMtSBM, that makes minimal assumptions and works well on images. It extends a recent successful algorithm that already performs well on images but does not build full trajectories through time. MMtSBM also finds the most optimal process, in the sense of the least varying one relatively to an abstract reference. We demonstrate that not only does MMtSBM perform very well on genetic expression data, but also works well on diverse types of images.