TRANSPORT WITH SUPPORT: DATA-CONDITIONAL DIFFUSION BRIDGES

Abstract

The dynamic Schrödinger bridge problem provides an appealing setting for posing optimal transport problems as learning non-linear diffusion processes and enables efficient iterative solvers. Recent works have demonstrated state-of-the-art results (e.g., in modelling single-cell embryo RNA sequences or sampling from complex posteriors) but are typically limited to learning bridges with only initial and terminal constraints. Our work extends this paradigm by proposing the Iterative Smoothing Bridge (ISB). We combine learning diffusion models with Bayesian filtering and optimal control, allowing for 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 and 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.

1. INTRODUCTION

Generative diffusion models have gained increasing popularity and achieved impressive results in a variety of challenging application domains, such as computer vision (e.g., Ho et al., 2020; Song et al., 2021a; Dhariwal & Nichol, 2021 ), reinforcement learning (e.g., Janner et al., 2022) , and time series modelling (e.g., Rasul et al., 2021; Vargas et al., 2021; Tashiro et al., 2021; Park et al., 2022) . Recent works have explored connections between denoising diffusion models and the dynamic Schrödinger bridge problem (SBP, e.g., Vargas et al., 2021; De Bortoli et al., 2021; Shi et al., 2022) to adopt iterative schemes for solving the dynamic optimal transport problem more efficiently. The solution of the SBP that correspond to denoising diffusion models is then given by the finite-time process, which is the closest in Kullback-Leibler (KL) divergence to the forward noising process of the diffusion model under marginal constraints. Data is then generated by time-reversing the process. In many applications, the interest is not purely in modelling transport between an initial and terminal state distribution In naturally occurring generative processes, we typically observe snapshots of realizations along intermediate stages of individual sample trajectories (see Fig. 1 ). Such problems arise in medical diagnosis (e.g., tissue changes and cell growth), demographic modelling, environmental dynamics, and animal movement modelling-see Fig. 4 for modelling bird migration and wintering patterns. Recently, constrained optimal control problems have been explored by adding additional fixed path constraints (Maoutsa et al., 2020; Maoutsa & Opper, 2021) or modifying the prior processes (Fernandes et al., 2021) . However, defining meaningful fixed path constraints or prior processes for the optimal control problems can be challenging, while sparse observational data are accessible in many real-world applications. In this work, we propose the Iterative Smoothing Bridge (ISB), an iterative method for solving control problems under sparse observational data constraints and constraints on the initial and terminal distribution. We perform the conditioning by leveraging the iterative pass idea from the Iterative Proportional Fitting procedure (IPFP) (Kullback, 1968; De Bortoli et al., 2021) procedure and applying differentiable particle filtering (Reich, 2013; Corenflos et al., 2021) within the outer loop. Integrating sequential Monte Carlo methods (e.g., Doucet et al., 2001; Chopin & Papaspiliopoulos, 2020) into the IPFP framework in such a way is non-trivial and can be understood as a novel iterative version of the algorithm by Maoutsa & Opper (2021) but with more general terminal constraints and path constraints defined by data.

