Transport with Support: Data-Conditional Diffusion Bridges

Ella Tamir, Martin Trapp, Arno Solin

Abstract:

The dynamic Schrödinger bridge problem provides an appealing setting for solving constrained time-series data generation tasks posed as optimal transport problems. It consists of learning non-linear diffusion processes using efficient iterative solvers.
Recent works have demonstrated state-of-the-art results (eg., in modelling single-cell embryo RNA sequences or sampling from complex posteriors) but are limited to learning bridges with only initial and terminal constraints. Our work extends this paradigm by proposing the Iterative Smoothing Bridge (ISB). We integrate Bayesian filtering and optimal control into learning the diffusion process, enabling the generation of constrained stochastic processes governed by sparse observations at intermediate stages and terminal constraints.
We assess the effectiveness of our method on synthetic and real-world data generation tasks and we show that the ISB generalises well to high-dimensional data, is computationally efficient, and provides accurate estimates of the marginals at intermediate and terminal times.

Submission Length: Regular submission (no more than 12 pages of main content)

Changes Since Last Submission: Addressed requests to make it clearer that our method is not an optimal transport solver but a generative method for dynamical systems

Code: https://github.com/AaltoML/iterative-smoothing-bridge

Assigned Action Editor: ~Alain_Durmus1

License: Creative Commons Attribution 4.0 International (CC BY 4.0)

Submission Number: 1309