UAI 2023
Conference paper

Aligned Diffusion Schrödinger Bridges

Download paper


Diffusion Schrödinger Bridges ( DSBs) have re- cently emerged as a powerful framework for recov- ering stochastic dynamics via their marginal ob- servations at different time points. Despite numer- ous successful applications, existing algorithms for solving DSBs have so far failed to utilize the structure of aligned data, which naturally arises in many biological phenomena. In this paper, we propose a novel algorithmic framework that, for the first time, solves DSBs while respecting the data alignment. Our approach hinges on a com- bination of two decades-old ideas: The classical Schrödinger bridge theory and Doob’s h-transform. Compared to prior methods, our approach leads to a simpler training procedure with lower vari- ance, which we further augment with principled regularization schemes. This ultimately leads to sizeable improvements across experiments on syn- thetic and real data, including the tasks of predicting conformational changes in proteins and temporal evolution of cellular differentiation processes