Transport with Support: Data-Conditional Diffusion Bridges

Published: 09 Dec 2023, Last Modified: 17 Sept 2024Accepted by TMLREveryoneRevisionsBibTeXCC BY 4.0
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
Supplementary Material: zip
Assigned Action Editor: ~Alain_Durmus1
License: Creative Commons Attribution 4.0 International (CC BY 4.0)
Submission Number: 1309
Loading