Unbalanced Diffusion Schrödinger Bridge

Matteo Pariset,Ya-Ping Hsieh,Charlotte Bunne,Andreas Krause,Valentin De Bortoli

2023-06-15

Abstract:Schrödinger bridges (SBs) provide an elegant framework for modeling the temporal evolution of populations in physical, chemical, or biological systems. Such natural processes are commonly subject to changes in population size over time due to the emergence of new species or birth and death events. However, existing neural parameterizations of SBs such as diffusion Schrödinger bridges (DSBs) are restricted to settings in which the endpoints of the stochastic process are both probability measures and assume conservation of mass constraints. To address this limitation, we introduce unbalanced DSBs which model the temporal evolution of marginals with arbitrary finite mass. This is achieved by deriving the time reversal of stochastic differential equations with killing and birth terms. We present two novel algorithmic schemes that comprise a scalable objective function for training unbalanced DSBs and provide a theoretical analysis alongside challenging applications on predicting heterogeneous molecular single-cell responses to various cancer drugs and simulating the emergence and spread of new viral variants.

Machine Learning

What problem does this paper attempt to address?

The paper attempts to address the issue of how to handle population size changes when modeling the temporal evolution of populations in physical, chemical, or biological systems. Traditional Diffusion Schrödinger Bridge (DSB) methods assume that both the initial and final states are probability measures and that mass is conserved. However, in practical applications, population size may change due to the emergence of new species or birth and death events. Therefore, this paper introduces an unbalanced DSB method to model the temporal evolution of marginal distributions with arbitrary finite mass.

Specifically, the paper addresses the following issues:

1. **Limitations of traditional DSB methods**: Existing neural parameterized DSB methods are only applicable when both the initial and final states are probability measures and mass is conserved. This is unrealistic in many practical application scenarios because population size changes over time.

2. **Modeling population size changes**: To overcome this limitation, the paper proposes an unbalanced DSB method by deriving time-reversal formulas for stochastic differential equations that include birth and death terms, to model the temporal evolution of marginal distributions with arbitrary finite mass.

3. **Algorithm implementation**: The paper proposes two new algorithmic frameworks for training unbalanced DSB and provides theoretical analysis and applications in challenging tasks such as predicting the response of heterogeneous molecular single cells to cancer drugs and simulating the emergence and spread of new virus variants.

4. **Theoretical contributions**: The paper derives time-reversal formulas for stochastic differential equations that include birth and death mechanisms and demonstrates how these formulas can be applied to solve the unbalanced DSB problem.

In summary, this paper aims to improve the modeling capability of population size changes by introducing the unbalanced DSB method, thereby better applying it to practical problems in fields such as biology, chemistry, and physics.