Learning stochastic dynamics from snapshots through regularized unbalanced optimal transport

## Abstract

Abstract Reconstructing dynamics using samples from sparsely time-resolved snapshots is an important problem in both natural sciences and machine learning. Here, we introduce a new deep learning approach for solving regularized unbalanced optimal transport (RUOT) and inferring continuous unbalanced stochastic dynamics from observed snapshots. Based on the RUOT form, our method models these dynamics without requiring prior knowledge of growth and death processes or additional information, allowing them to be learnt directly from data. Theoretically, we explore the connections between the RUOT and Schrödinger bridge problem and discuss the key challenges and potential solutions. The effectiveness of our method is demonstrated with a synthetic gene regulatory network. Compared with other methods, our approach accurately identifies growth and transition patterns, eliminates false transitions, and constructs the Waddington developmental landscape.

###### keywords:

## Contents

- 1 Introduction 1
- 2 Preliminaries and Backgrounds 2
2.1 Schrödinger Bridge Problem 2.1
2.2 Optimal Transport 2.2
- 3 Regularized Unbalanced Optimal Transport 3
3.1 RUOT with Fisher Regularization 3.1
3.2 RUOT in the General Case 3.2
- 4 Connections with Schrödinger Bridge 4
4.1 Connections with Branching Schrödinger Bridge 4.1
4.2 Connections with Schrödinger Bridge on Weighted Path 4.2
- 5 Learning RUOT through Deep Learning 5
5.1 Energy Function in Theorem 3.2 5.1
5.2 Reconstruction Loss 5.3
5.3 Learning Probability Flow ODE and Score Function 5.3
5.4 Training 5.4
- 6 Numerical Results 6
6.1 Synthetic Dynamics 6.1
6.2 Learning Waddington’s Developmental Landscapes 6.2
- 7 Conclusion 7

## 1 Introduction

In machine learning and natural science, a fundamental problem is the coupling of distributions, often high-dimensional, based on the observed samples. For instance, central to generative models as Variational Autoencoders (VAEs) (Kingma and Welling, ) is to learn a nonlinear mapping from a complex data distribution to a simpler latent space, where the underlying structure of the data can be more easily captured and generated. In multi-modal data analysis, integration of distribution from disparate data types is essential to provide a more fused and comprehensive understanding of the system (Lahat et al., ). Particularly in biology, this has been applied in multi-omics data integration, where distributions from genomics, transcriptomics, proteomics as well as spatial information are aligned into the unified representation of cellular states (Demetci et al., ; Cao et al., ; Gao et al., ; Cang and Zhao, ).

Recently, there has been growing interest in understanding the dynamics of how distributions are coupled over time. In generative machine learning, this is exemplified by diffusion models (Ho et al., ; Sohl-Dickstein et al., ), which improves the VAE models by the gradual transformation of a simple distribution into a more complicated data distribution. The diffusion models could be formulated in the framework of stochastic differential equations (SDEs) (Song et al., ), and useful for interpolating paths between arbitrary distributions and learning to model the resulting dynamics (Tong et al., ; Neklyudov et al., ; De Bortoli et al., ; Wang et al., ).

The analysis of time series data in single-cell RNA sequencing (scRNA-seq) provides an important application scenario of the high dimensional probability distribution coupling and dynamical inference techniques. Due to the destructive nature of current mainstream sequencing technologies, the continuous trajectories of gene expression dynamics within individual cells are not present in the data. Instead, we can only observe changes in their distributions. Consequently, reconstructing the trajectories of individual cells from samples collected at sparse temporal points in the experiment is crucial for understanding development and other biological processes (Schiebinger et al., ; Sha et al., ; Peng et al., ; Lavenant et al., ; Jiang et al., ; Jiang and Wan, ; Bunne et al., , ; Tong et al., ).

The trajectory inference algorithms in scRNA-seq data have been extensively studied (Saelens et al., ), and optimal transport (OT)-based methods have emerged as a central tool for datasets with temporal resolution (Bunne et al., ). Static OT methods, for instance, have been applied to infer the ancestors or descendants of given cells across time points (Schiebinger et al., ; Klein et al., ). Meanwhile, in many cases, there is a need to learn the continuous dynamics of particles over time and fit the mechanistic model that transforms the source distribution into the target distribution. This has led to the development of various dynamical OT methods, as well as flow-based generative methods, including Benamou-Brenier formulation (Benamou and Brenier, ) based approaches (Tong et al., ; Huguet et al., ; Zhang et al., ; Ruthotto et al., ; Liu et al., ; Cheng et al., ; Wan et al., ; Pooladian et al., ), unbalanced dynamic OT based approaches (Peng et al., ; Sha et al., ; Tong et al., ), Gromov-Wasserstein (GW) OT approaches (Klein et al., ), continuous normalizing flows (CNF) and conditional flow matching methods (CFM) (Lipman et al., ; Tong et al., ; Albergo and Vanden-Eijnden, ; Liu et al., ; Jiao et al., ; Gao et al., ; Liu et al., ; Wu et al., ; Jin et al., ; Chow et al., ; Cheng et al., ). Despite the advances, these methods have not fully taken the stochastic dynamical effects into account, especially the intrinsic noise in gene expression and cell differentiation (Zhou et al., , ), which is prevalent in biological processes on single-cell level (Elowitz et al., ).

Pertinent to stochastic dynamics, the Schrödinger bridge (SB) problem aims to identify the most likely stochastic transition path between two arbitrary distributions relative to a reference stochastic process (Léonard, ). The SB problem or its variants has been applied in a wide range of context, including
scRNA-seq analysis and generative modeling. These approaches, for instance, include static methods (Lavenant et al., ; Pariset et al., ; Ventre et al., ; Chizat et al., ; Shi et al., ; De Bortoli et al., ; Pooladian and Niles-Weed, ; Liu et al., ), dynamic methods (Koshizuka and Sato, ; Neklyudov et al., , ; Zhang et al., ; Bunne et al., ; Chen et al., ; Albergo et al., ; Wang et al., ; Jiao et al., ; Zhou et al., ; Liu et al., ; Zhou et al., ), and flow-matching methods (Tong et al., ). However, these methods either fail to account for unnormalized distributions resulting from cell growth and death (Sha et al., ), or require prior knowledge of these processes (e.g. growth/death rate) (Schiebinger et al., ; Lavenant et al., ; Pariset et al., ; Chizat et al., ) or additional information (e.g. cell lineage) (Ventre et al., ).

Recently, regularized unbalanced optimal transport (RUOT) offers a promising approach for modeling such stochastic unbalanced continuous dynamics (Baradat and Lavenant, ; Chen et al., ; Buze and Duong, ; Janati et al., ). It can be viewed as an unbalanced relaxation of the dynamic formulation of the Schrödinger bridge problem. In (Baradat and Lavenant, ), the authors have provided insights on the connection between certain RUOT and branching Schrödinger bridges. However, computational methods for learning unbalanced stochastic dynamics from samples or such dynamic formulations of certain unbalanced Schrödinger bridges without prior knowledge or additional information are still lacking.

Here we develop a new deep learning method (DeepRUOT) for learning general RUOT and inferring continuous unbalanced stochastic dynamics from samples based on the derived Fisher regularization form without requiring prior knowledge. We demonstrate the effectiveness of our method using a synthetic gene regulatory data. Compared to the common (balanced) Schrödinger bridge method, our RUOT-based approach accurately identifies growth and transition patterns, eliminates false transitions, and constructs the Waddington developmental landscape. On the theoretical side, in this paper we provide a straightforward derivation for the dynamical form of the Schrödinger bridge and then summarize and explore the relationship between RUOT and the Schrödinger bridge problem.

The rest of the paper is organized as follows. In Section 2, we provide preliminaries and background on the Schrödinger bridge and optimal transport, along with a straightforward proof of the dynamical formulation of Schrödinger bridge. Section 3 derives the RUOT problem and its Fisher regularization form. In Section 4, we summarize and explore the relationship between the RUOT problem and the SB problem. Section 5 introduces our deep learning method for learning the RUOT problem from data. In Section 6, we present a synthetic gene regulatory data to validate our framework and demonstrate its effectiveness. We summarize our contributions and discuss possible future directions in Section 7.

## 2 Preliminaries and Backgrounds

In this section, we will first review the formulation of the dynamic diffusion Schrödinger Bridge problem. Subsequently, we will provide new proof of the dynamical form of the Schrödinger Bridge. Finally, we will revisit the definitions of the dynamical formulations of optimal transport, regularized optimal transport, and unbalanced optimal transport respectively.

### 2.1 Schrödinger Bridge Problem

The Schrödinger Bridge problem aims to identify the most likely evolution between a given initial distribution $\nu_{0}$ and a terminal distribution $\nu_{1}$ (assumed to have density in this paper), relative to a given reference stochastic process. Formally, this problem can be formulated as the minimization of the Kullback-Leibler (KL) divergence in the optimal control perspective (Dai Pra, ) as below:

$$ $\min_{\mu^{\boldsymbol{X}}_{0}=\nu_{0},~{}\mu^{\boldsymbol{X}}_{1}=\nu_{1}} \mathcal{D}_{\mathrm{KL}}\left(\mu^{\boldsymbol{X}}_{[0,1]}\|\mu^{\boldsymbol{ Y}}_{[0,1]}\right),$ (2.1) $$

where $\mu^{\boldsymbol{X}}_{[0,1]}$ denotes the probability measure induced by stochastic process ${\boldsymbol{X}}_{t}\ \left(0\leq t\leq 1\right)$ defined on the space of all continuous paths $C([0,1],\mathbb{R}^{d})$, with the distribution of ${\boldsymbol{X}}_{t}$ at given time $t$ characterized by the measure $\mu^{\boldsymbol{X}}_{t}$ with density function $p(\boldsymbol{x},t)$. In this work, we consider ${\boldsymbol{X}}_{t}$ the stochastic process characterized by the following stochastic differential equation (SDE):

$$ $\mathrm{d}{\boldsymbol{X}}_{t}=\boldsymbol{b}\left({\boldsymbol{X}}_{t},t \right)\mathrm{d}t+\boldsymbol{\sigma}\left({\boldsymbol{X}}_{t},t\right) \mathrm{d}{\boldsymbol{W}}_{t},$ (2.2) $$

and the reference measure
$\mu^{\boldsymbol{Y}}_{[0,1]}$ is chosen as the probability measure induced by the process $\mathrm{d}\boldsymbol{Y}_{t}=\boldsymbol{\sigma}(\boldsymbol{Y}_{t},t)\mathrm{
d}{\boldsymbol{W}}_{t}$, where ${\boldsymbol{W}}_{t}\in\mathbb{R}^{d}$ is the standard multidimensional Brownian motion defined on a probability space $(\Omega,\mathcal{F},\mathbb{P})$ with $\mathbb{P}$ the Wiener measure for the coordinate process ${\boldsymbol{W}}_{t}(\omega)=\omega(t)$, and $\boldsymbol{\sigma}(\boldsymbol{x},t)\in\mathbb{R}^{d\times d}$ is the diffusion matrix which is typically assumed as bounded, coercive and invertible. Under this formulation, the solution to Eq. (2.1) is also referred to as the *diffusion Schrödinger bridge*. The diffusion Schrödinger bridge also possesses a dynamic formulation, which can be formally stated as the following theorem.

###### Theorem 2.1 .

Consider the diffusion Schrödinger bridge problem

$$ $\min_{\mu^{\boldsymbol{X}}_{0}=\nu_{0},~{}\mu^{\boldsymbol{X}}_{1}=\nu_{1}} \mathcal{D}_{\mathrm{KL}}\left(\mu^{\boldsymbol{X}}_{[0,1]}\|\mu^{\boldsymbol{ Y}}_{[0,1]}\right),$ (2.3) $$

where $\mu^{\boldsymbol{Y}}_{[0,1]}$ the reference measure induced by $\mathrm{d}\boldsymbol{Y}_{t}=\boldsymbol{\sigma}(\boldsymbol{Y}_{t},t)\mathrm{
d}{\boldsymbol{W}}_{t}$. Then Eq. (2.3) is equivalent to

$$ $\inf_{\left({p},\boldsymbol{b}\right)}\int_{0}^{1}\int_{\mathbb{R}^{d}}\left[ \frac{1}{2}\boldsymbol{b^{T}}(\boldsymbol{x},t)\boldsymbol{a}^{-1}(\boldsymbol {x},t)\boldsymbol{b}(\boldsymbol{x},t)\right]{p}(\boldsymbol{x},t)\mathrm{d} \boldsymbol{x}\mathrm{d}t,$ (2.4) $$

where the infimum is taken over all function pairs $\left({p},\boldsymbol{b}\right)$ such that ${p}(\cdot,0)=$ ${\nu_{0}}$, ${p}(\cdot,1)={\nu_{1}}$, ${p}(\boldsymbol{x},t)$ absolutely continuous, and

$$ $\partial_{t}{p}(\boldsymbol{x},t)=-\nabla_{\boldsymbol{x}}\cdot\left({p}( \boldsymbol{x},t)\boldsymbol{b}(\boldsymbol{x},t)\right)+\frac{1}{2}\nabla_{ \boldsymbol{x}}^{2}:\left(\boldsymbol{a}(\boldsymbol{x},t){p}(\boldsymbol{x},t )\right),$ (2.5) $$

where $\nabla_{\boldsymbol{x}}^{2}:\left(\boldsymbol{a}(\boldsymbol{x},t){p}(
\boldsymbol{x},t)\right)=\sum_{ij}\partial_{ij}(a_{ij}p)$ and $\boldsymbol{a}(\boldsymbol{x},t)=\boldsymbol{\sigma}(\boldsymbol{x},t)
\boldsymbol{\sigma}^{T}(\boldsymbol{x},t)$,
with vanishing boundary condition: $\displaystyle\lim_{|\boldsymbol{x}|\rightarrow\infty}p(\boldsymbol{x},t)=0$.

This theorem or its variants have been stated and proven in various forms, such as in (Gentil et al., ; Chen et al., ; Dai Pra, ). Here, we provide a simple direct proof for illustration.

###### Proof 2.2 .

For Schrödinger bridge problems, the optimal solution can be found within the class of SDEs:

$$ ${\boldsymbol{X}}_{t}\sim\mu^{\boldsymbol{X}}_{t}:\quad\mathrm{d}{\boldsymbol{X }}_{t}=\boldsymbol{b}\left({\boldsymbol{X}}_{t},t\right)\mathrm{d}t+ \boldsymbol{\sigma}\left({\boldsymbol{X}}_{t},t\right)\mathrm{d}{\boldsymbol{W }}_{t},$ $$

The Fokker-Planck equation for the SDE is

$$ $\partial_{t}{p}(\boldsymbol{x},t)=-\nabla_{\boldsymbol{x}}\cdot\left({p}( \boldsymbol{x},t)\boldsymbol{b}(\boldsymbol{x},t)\right)+\frac{1}{2}\nabla_{ \boldsymbol{x}}^{2}:\left(\boldsymbol{a}(\boldsymbol{x},t){p}(\boldsymbol{x},t )\right),$ $$

where $\boldsymbol{a}(\boldsymbol{x},t)=\boldsymbol{\sigma}(\boldsymbol{x},t)
\boldsymbol{\sigma}^{T}(\boldsymbol{x},t)$.
Then we need to compute the KL divergence of these two stochastic processes. Define the auxiliary variable

$$ $\gamma({\boldsymbol{x}},t)=\boldsymbol{\sigma}({\boldsymbol{x}},t)^{-1} \boldsymbol{b}({\boldsymbol{x}},t)$ $$

and apply the Girsanov’s theorem, we get

$$ $\displaystyle\mathcal{D}_{\text{KL}}\left(\mu^{\boldsymbol{X}}_{[0,1]}||\mu^{ \boldsymbol{Y}}_{[0,1]}\right)=\mathbb{E}_{\mathbb{P}}\left[\mathcal{E}_{1} \left(\int_{0}^{1}\gamma({\boldsymbol{Y}}_{t},t)\mathrm{d}{\boldsymbol{W}}_{t} -\frac{1}{2}\int_{0}^{1}\left\|\gamma({\boldsymbol{Y}_{t}},t)\right\|_{2}^{2} \mathrm{d}t\right)\right],$ (2.6) $$

where $\{\mathcal{E}_{t}\}_{t\in[0,1]}$ is an exponential martingale with respect to the Brownian filtrations $\{\mathcal{F}_{t}\}$ defined as

$$ $\mathcal{E}_{t}=\exp\left(\int_{0}^{t}\gamma({\boldsymbol{Y}}_{t},t)\mathrm{d} {\boldsymbol{W}}_{t}-\frac{1}{2}\int_{0}^{t}\left\|\gamma({\boldsymbol{Y}_{t}} ,t)\right\|_{2}^{2}\mathrm{d}t\right).$ $$

It is well-known that the exponential martingale $\mathcal{E}_{t}$ satisfies the SDE

$$ $\mathrm{d}\mathcal{E}_{t}=\mathcal{E}_{t}\gamma({\boldsymbol{Y}}_{t},t)\mathrm {d}{\boldsymbol{W}}_{t},\quad\mathcal{E}_{0}=1.$ $$

So we obtain

$$ $\mathcal{E}_{1}=\int_{0}^{1}\mathcal{E}_{s}\gamma({\boldsymbol{Y}}_{s},s) \mathrm{d}{\boldsymbol{W}}_{s}+1.$ $$

Since $\mathbb{E}\int_{0}^{1}\gamma({\boldsymbol{Y}}_{t},t)\mathrm{d}{\boldsymbol{W}}
_{t}=0$, we get

$$ $\displaystyle\mathbb{E}\left[\mathcal{E}_{1}\int_{0}^{1}\gamma({\boldsymbol{Y} }_{t},t)\mathrm{d}{\boldsymbol{W}}_{t}\right]$ $\displaystyle=\mathbb{E}\left[\int_{0}^{1}\mathcal{E}_{t}\gamma(\boldsymbol{Y} _{t},t)\mathrm{d}{\boldsymbol{W}}_{t}\cdot\int_{0}^{1}\gamma(\boldsymbol{Y}_{t },t)\mathrm{d}{\boldsymbol{W}}_{t}\right],$ $\displaystyle=\mathbb{E}\int_{0}^{1}\mathcal{E}_{t}\left\|\gamma(\boldsymbol{Y }_{t},t)\right\|_{2}^{2}\mathrm{d}t,$ $\displaystyle=\mathbb{E}\int_{0}^{1}\mathbb{E}\left(\left\|\gamma(\boldsymbol{ Y}_{t},t)\right\|_{2}^{2}\mathcal{E}_{1}\mid\mathcal{F}_{t}\right)\mathrm{d}t,$ $\displaystyle=\int_{0}^{1}\mathbb{E}\left(\mathcal{E}_{1}\left\|\gamma( \boldsymbol{Y}_{t},t)\right\|_{2}^{2}\right)\mathrm{d}t,$ $\displaystyle=\mathbb{E}\left[\mathcal{E}_{1}\int_{0}^{1}\left\|\gamma( \boldsymbol{Y}_{t},t)\right\|_{2}^{2}\mathrm{d}t\right].$ $$

Combine with Eq. (2.6), we obtain that

$$ $\displaystyle\mathcal{D}_{\text{KL}}(\mu^{\boldsymbol{X}}_{[0,1]}||\mu^{ \boldsymbol{Y}}_{[0,1]})$ $\displaystyle=\mathbb{E}\left[\mathcal{E}_{1}\left(\int_{0}^{1}\frac{1}{2} \left\|\gamma(\boldsymbol{Y}_{t},t)\right\|_{2}^{2}\mathrm{d}t\right)\right]$ $\displaystyle=\mathbb{E}\int_{0}^{1}\frac{1}{2}\left\|\gamma(\boldsymbol{X}_{t },t)\right\|_{2}^{2}\mathrm{d}t$ $\displaystyle=\int_{0}^{1}\int_{\mathbb{R}^{d}}\left[\frac{1}{2}\boldsymbol{b^ {T}}(\boldsymbol{x},t)\boldsymbol{a}(\boldsymbol{x},t)^{-1}\boldsymbol{b}( \boldsymbol{x},t)\right]{p}(\boldsymbol{x},t)\mathrm{d}\boldsymbol{x}\mathrm{d }t.$ $$

The proof is done.

We refer to Eq. (2.4) and Eq. (2.5) as the *dynamic diffusion Schrödinger bridge*.

###### Theorem 2.1 .

###### Proof 2.2 .

### 2.2 Optimal Transport

The optimal transport problem has been extensively studied in various fields. Given two probability distributions $\alpha\in\mathbb{R}_{+}^{n}$ and $\beta\in\mathbb{R}_{+}^{m}$, its primary goal is to find the optimal coupling $\pi\in\mathbb{R}_{+}^{n\times m}$ to transport a given distribution of mass (or resources) from one location to another while minimizing the cost associated with the transportation. The static optimal transport problem can be mathematically formulated as:

$$ $\displaystyle\min_{\pi\in\boldsymbol{\Pi}(\alpha,\beta)}\langle\pi,c\rangle, \quad\boldsymbol{\Pi}(\alpha,\beta)=\left\{\pi\in\mathbb{R}_{+}^{n\times m}: \pi\mathbf{1}_{m}=\alpha,\pi^{T}\mathbf{1}_{n}=\beta,\pi\geq 0\right\}.$ $$

The cost matrix $c\in\mathbb{R}^{n\times m}$ defines the transportation cost between each pair of points, where $c_{ij}:=c(x_{i},y_{j})$ represents the cost of moving a unit mass from point $x_{i}$ to point $y_{j}$.
We refer to (Peyré et al., ) for more details. Next we briefly stated some well-known results on the dynamical formulation of the OT.

#### Dynamical Optimal Transport

The formulation is also known as the Benamou-Brenier formulation (Benamou and Brenier, ), which can be stated as follows.

$$ $\begin{gathered}\mathcal{W}\left(\nu_{0},\nu_{1}\right)=\inf_{\left({p}( \boldsymbol{x},t),\boldsymbol{b}(\boldsymbol{x},t)\right)}\int_{0}^{1}\int_{ \mathbb{R}^{d}}\frac{1}{2}\|\boldsymbol{b}(\boldsymbol{x},t)\|_{2}^{2}{p}( \boldsymbol{x},t)\mathrm{d}\boldsymbol{x}\mathrm{d}t,\\ \mbox{s.t. }\partial_{t}p+\nabla\cdot\left(\boldsymbol{b}(\boldsymbol{x},t)p \right)=0,\ p|_{t=0}=\nu_{0},p|_{t=1}=\nu_{1}.\end{gathered}$ $$

Compared to SB, in OT the distributions are connected by a deterministic transport equation instead of diffusion. It can be shown that this dynamical formulation corresponds to a static Kantorovich’s optimal transport problem with cost function $c(x,y)=\|x-y\|_{2}^{2}$.

#### Regularized Optimal Transport

The regularized optimal transport is defined by a dynamical form of the general Schrödinger bridge by taking the diffusion rate as constant and scaled by $\sigma^{2}$ in Eq. (2.4)).

$$ $\begin{gathered}\mathcal{W}\left(\nu_{0},\nu_{1}\right)=\inf_{\left({p}( \boldsymbol{x},t),\boldsymbol{b}(\boldsymbol{x},t)\right)}\int_{0}^{1}\int_{ \mathbb{R}^{d}}\frac{1}{2}\|\boldsymbol{b}(\boldsymbol{x},t)\|_{2}^{2}{p}( \boldsymbol{x},t)\mathrm{d}\boldsymbol{x}\mathrm{d}t,\\ \mbox{s.t. }\partial_{t}p+\nabla\cdot\left(\boldsymbol{b}(\boldsymbol{x},t)p \right)-\frac{\sigma^{2}}{2}\Delta p=0,\ p|_{t=0}=\nu_{0},p|_{t=1}=\nu_{1}, \end{gathered}$ $$

It can be demonstrated that as $\sigma^{2}$ approaches zero, the solution to this problem converges to that of the Benamou-Brenier problem (Mikami and Thieullen, ).

#### Unbalanced Optimal Transport

To account for unnormalized marginal distributions and effects such as growth and death, an unbalanced optimal transport problem with Wasserstein–Fisher–Rao (WFR) metric has been proposed (Chizat et al., , ) or its extensions (Gangbo et al., ). Here we adopt the WFR unbalanced optimal transport.

$$ $\begin{gathered}\mathcal{W}\left(\nu_{0},\nu_{1}\right)=\inf_{\left({p}( \boldsymbol{x},t),\boldsymbol{b}(\boldsymbol{x},t),{g}(\boldsymbol{x},t)\right )}\int_{0}^{1}\int_{\mathbb{R}^{d}}\left(\frac{1}{2}\|\boldsymbol{b}( \boldsymbol{x},t)\|_{2}^{2}+\alpha|{g}(\boldsymbol{x},t)|_{2}^{2}\right){p}( \boldsymbol{x},t)\mathrm{d}\boldsymbol{x}dt,\\ \mbox{s.t. }\partial_{t}p+\nabla\cdot\left(\boldsymbol{b}(\boldsymbol{x},t)p \right)=g(\boldsymbol{x},t)p,\ p|_{t=0}=\nu_{0},p|_{t=1}=\nu_{1}.\end{gathered}$ $$

Here $g(x,t)$ is a scalar function that denotes the growth or death rate of particles at the state $x$ and time $t$, and is also optimized in the total energy term. $\alpha$ is the hyperparameter of weight. We should also note that in this case, $\nu_{0}$ and $\nu_{1}$ are not necessarily the normalized probability densities, but are generally densities of masses.

## 3 Regularized Unbalanced Optimal Transport

In this section, we first introduce the RUOT problem, then derive the Fisher regularization form of RUOT, and discuss its connections with existing literature.

### 3.1 RUOT with Fisher Regularization

To simplify the notation and illustrate a commonly used setting, we take a special case considering $\boldsymbol{a}(\boldsymbol{x},t)=\sigma^{2}(t)\boldsymbol{I}$, and the general case is left to Section 3.2 for further discussion. Inspired by unbalanced dynamic optimal transport, the dynamical formulation Theorem 2.1 suggests a natural approach to relaxing the mass conservation constraint by introducing a growth/death term $gp$ in Eq. (2.5) as

$$ $\partial_{t}p=-\nabla_{\boldsymbol{x}}\cdot\left(p\boldsymbol{b}\right)+\frac{ 1}{2}\nabla_{\boldsymbol{x}}^{2}:\left(\sigma^{2}(t)\boldsymbol{I}p\right)+gp.$ $$

Meanwhile, we also define a loss functional in Eq. (3.7) which incorporates the growth penalization and Wasserstein metric considered in Eq. (2.4) (with a rescaling by $\sigma^{2}(t)$). We refer to this formulation as the
*regularized unbalanced optimal transport*.

###### Definition 3.1 (Regularized unbalanced optimal transport) .

Consider

$$ $\inf_{\left(p,\boldsymbol{b},g\right)}\int_{0}^{1}\int_{\mathbb{R}^{d}}\frac{1 }{2}\left\|\boldsymbol{b}(\boldsymbol{x},t)\right\|_{2}^{2}{p}(\boldsymbol{x}, t)\mathrm{d}\boldsymbol{x}\mathrm{d}t+\int_{0}^{1}\int_{\mathbb{R}^{d}}\alpha \Psi\left(g(\boldsymbol{x},t)\right){p}(\boldsymbol{x},t)\mathrm{d}\boldsymbol {x}\mathrm{d}t,$ (3.7) $$

where the infimum is taken all pairs $\left(p,\boldsymbol{b}\right)$ such that $p(\cdot,0)=$ $\nu_{0},p(\cdot,T)=\nu_{1},p(\boldsymbol{x},t)$ absolutely continuous, and

$$ $\partial_{t}p=-\nabla_{\boldsymbol{x}}\cdot\left(p\boldsymbol{b}\right)+\frac{ 1}{2}\nabla_{\boldsymbol{x}}^{2}:\left(\sigma^{2}(t)\boldsymbol{I}p\right)+gp.$ (3.8) $$

with vanishing boundary condition: $\displaystyle\lim_{|\boldsymbol{x}|\rightarrow\infty}p(\boldsymbol{x},t)=0$.

Note that in the definition if $g(\boldsymbol{x},t)=0$, $\Psi(g)=+\infty$ unless $g=0$ and $\Psi(0)=0$ it degenerates to the *regularized optimal transport* problem. If $\sigma(t)\rightarrow 0$ and $\Psi\left(g(\boldsymbol{x},t)\right)=|g(\boldsymbol{x},t)|^{2}$ this degenerates to the *unbalanced dynamic optimal transport* with Wasserstein-Fisher-Rao (WFR) metrics. Meanwhile when $\sigma(t)$ is constant, it coincides with the definition of RUOT provided in (Baradat and Lavenant, ). Then, we can reformulate Definition 3.1 with the following Fisher information regularization.

###### Theorem 3.2 .

The regularized unbalanced optimal transport problem (3.7) is equivalent to

$$ $\inf_{\left(p,\boldsymbol{v},g\right)}\int_{0}^{1}\int_{\mathbb{R}^{d}}\left[ \frac{1}{2}\left\|\boldsymbol{v}(\boldsymbol{x},t)\right\|_{2}^{2}+\frac{ \sigma^{4}(t)}{8}\left\|\nabla_{\boldsymbol{x}}\log p\right\|_{2}^{2}-\frac{ \sigma^{2}(t)}{2}\left(1+\log p\right)g+\alpha\Psi\left(g\right)\right]{p}( \boldsymbol{x},t)\mathrm{d}\boldsymbol{x}\mathrm{d}t,$ (3.9) $$

where the infimum is taken all pairs $\left(p,\boldsymbol{v},g\right)$ such that $p(\cdot,0)=$ $\nu_{0},p(\cdot,T)=\nu_{1},p(\boldsymbol{x},t)$ absolutely continuous, and

$$ $\partial_{t}p=-\nabla_{\boldsymbol{x}}\cdot\left(p\boldsymbol{v}(\boldsymbol{x },t)\right)+g(\boldsymbol{x},t)p.$ (3.10) $$

with vanishing boundary condition: $\displaystyle\lim_{|\boldsymbol{x}|\rightarrow\infty}p(\boldsymbol{x},t)=0$.

###### Proof 3.3 .

From (3.8) we obtain

$$ $\displaystyle\partial_{t}p$ $\displaystyle=-\nabla_{\boldsymbol{x}}\cdot\left(p\boldsymbol{b}\right)+\frac{ 1}{2}\nabla_{\boldsymbol{x}}^{2}:\left(\sigma^{2}(t)\boldsymbol{I}p\right)+g( \boldsymbol{x},t)p,$ $\displaystyle=-\nabla_{\boldsymbol{x}}\cdot\left(\left(\boldsymbol{b}-\frac{1} {2}\sigma^{2}(t)\nabla_{\boldsymbol{x}}\log p\right)p\right)+{g}(\boldsymbol{x },t)p.$ $$

Using the change of variable $\boldsymbol{v}(\boldsymbol{x},t)=\boldsymbol{b}(\boldsymbol{x},t)-\frac{1}{2}
\sigma^{2}(t)\nabla_{\boldsymbol{x}}\log p$, we see that it is equivalent to

$$ $\partial_{t}p=-\nabla_{\boldsymbol{x}}\cdot\left(p\boldsymbol{v}(\boldsymbol{x },t)\right)+g(\boldsymbol{x},t)p.$ $$

Correspondingly, the integrand in the objective functional becomes

$$ $\int_{0}^{1}\int_{\mathbb{R}^{d}}\left[\frac{1}{2}\left\|\boldsymbol{v}( \boldsymbol{x},t)\right\|_{2}^{2}+\frac{\sigma^{4}(t)}{8}\left\|\nabla_{ \boldsymbol{x}}\log p\right\|_{2}^{2}+\frac{1}{2}\left\langle\boldsymbol{v}( \boldsymbol{x},t),\sigma^{2}(t)\nabla_{\boldsymbol{x}}\log p\right\rangle+ \alpha\Psi\left(g\right)\right]{p}(\boldsymbol{x},t)\mathrm{d}\boldsymbol{x} \mathrm{d}t.$ $$

Letting $H\left(p\right):=\int p\log p$ be the entropy, we have

$$ $\displaystyle H\left(p_{1}\right)-H\left(p_{0}\right)$ $\displaystyle=\int_{0}^{1}\partial_{t}H\left({p}(\boldsymbol{x},t)\right) \mathrm{d}t$ $\displaystyle=\int_{0}^{1}\int_{\mathbb{R}^{d}}\left(1+\log{p}(\boldsymbol{x}, t)\right)\partial_{t}{p}(\boldsymbol{x},t)\mathrm{d}\boldsymbol{x}\mathrm{d}t$ $\displaystyle=\int_{0}^{1}\int_{\mathbb{R}^{d}}\left(1+\log{p}(\boldsymbol{x}, t)\right)\cdot\left(-\nabla_{\boldsymbol{x}}\cdot\left({p}(\boldsymbol{x},t) \boldsymbol{v}(\boldsymbol{x},t)\right)+g{p}(\boldsymbol{x},t)\right)\mathrm{d }\boldsymbol{x}\mathrm{d}t,$ $\displaystyle=\int_{0}^{1}\int_{\mathbb{R}^{d}}{p}(\boldsymbol{x},t)\left \langle\nabla_{\boldsymbol{x}}\log{p}(\boldsymbol{x},t),\boldsymbol{v}( \boldsymbol{x},t)\right\rangle+\left(1+\log{p}(\boldsymbol{x},t)\right)gp \mathrm{d}\boldsymbol{x}\mathrm{d}t.$ $$

Therefore,

$$ $\displaystyle\int_{0}^{1}\int_{\mathbb{R}^{d}}\left[\left\langle\sigma^{2}(t) \nabla_{\boldsymbol{x}}\log p,\boldsymbol{v}(\boldsymbol{x},t)\right\rangle \right]{p}(\boldsymbol{x},t)\mathrm{d}\boldsymbol{x}\mathrm{d}t$ $\displaystyle=H\left(p_{1}\right)-H\left(p_{0}\right)$ $\displaystyle-\int_{0}^{1}\int_{\mathbb{R}^{d}}\sigma^{2}(t)\left(1+\log{p}( \boldsymbol{x},t)\right)gp\mathrm{d}\boldsymbol{x}\mathrm{d}t.$ $$

###### Remark 3.4 .

When $g=0$ and $\Psi(0)=0$, then Eq. (3.9) is same with the dynamic entropy-regularized optimal transport form as discussed in (Bunne et al., ; Li et al., ; Gentil et al., ; Pooladian and Niles-Weed, ; Neklyudov et al., ; Léger and Li, )

###### Remark 3.5 .

An additional term in Eq. (3.9)

$$ $\mathcal{I}(p)=\int_{\mathbb{R}^{d}}\|\nabla_{\boldsymbol{x}}\log{p}( \boldsymbol{x},t)\|_{2}^{2}p(\boldsymbol{x},t)\mathrm{d}\boldsymbol{x}$ $$

is referred to as the *Fisher information*. Notably, when considering growth/death factors, Eq. (3.9) includes not only the Fisher-Rao metric but also an additional cross-term
$\int_{\mathbb{R}^{d}}-\frac{1}{2}\sigma^{2}(t)\left(1+\log{p}(\boldsymbol{x},t
)\right)g(\boldsymbol{x},t)p(\boldsymbol{x},t)\mathrm{d}\boldsymbol{x}.$

###### Remark 3.6 .

In (Buze and Duong, ), the authors defined a RUOT problem with nonlinear Fokker-Planck equation constraints. In Appendix A, we show that the proposed form is indeed consistent with the RUOT defined here, when $\sigma(t)$ is constant and $\Psi\left({g}(\boldsymbol{x},t)\right)$ take the quadratic form i.e. $\Psi\left({g}(\boldsymbol{x},t)\right)=|{g}(\boldsymbol{x},t)|_{2}^{2}$ . It has been conjectured that the proposed form is the higher order approximation scheme to the unbalanced optimal transport problem as $\sigma\rightarrow 0$.

From the Fokker-Plank equation and the proof of Theorem 3.2, the original SDE
$\mathrm{d}{\boldsymbol{X}}_{t}=\left(\boldsymbol{b}\left({\boldsymbol{X}}_{t},
t\right)\right)\mathrm{d}t+{\sigma}\left(t\right)\boldsymbol{I}\mathrm{d}{
\boldsymbol{W}}_{t}$
can be transformed
into the probability flow ODE

$$ $\mathrm{d}{\boldsymbol{X}}_{t}=\underbrace{\left(\boldsymbol{b}\left({ \boldsymbol{X}}_{t},t\right)-\frac{1}{2}\sigma^{2}(t)\nabla_{\boldsymbol{x}} \log p({\boldsymbol{X}}_{t},t)\right)}_{\boldsymbol{v}\left({\boldsymbol{X}}_{ t},t\right)}\mathrm{d}t.$ $$

Conversely, if the probability flow ODE’s drift $\boldsymbol{v}(\boldsymbol{x},t)$ the diffusion rate $\sigma(t)$ and the *score function* $\nabla_{\boldsymbol{x}}\log p(\boldsymbol{x},t)$ are known, then the
the SDE’s drift term $\boldsymbol{b}(\boldsymbol{x},t)$ can be determined by

$$ $\boldsymbol{b}(\boldsymbol{x},t)=\boldsymbol{v}(\boldsymbol{x},t)+\frac{1}{2} \sigma^{2}(t)\nabla_{\boldsymbol{x}}\log p(\boldsymbol{x},t).$ $$

So to specify an SDE is equivalent to specify the probability flow ODE and its score function.

###### Definition 3.1 (Regularized unbalanced optimal transport) .

###### Theorem 3.2 .

###### Proof 3.3 .

###### Remark 3.4 .

###### Remark 3.5 .

###### Remark 3.6 .

### 3.2 RUOT in the General Case

For the general case, we consider the following equation

$$ $\partial_{t}p=-\nabla_{\boldsymbol{x}}\cdot\left(p\boldsymbol{b}\right)+\frac{ 1}{2}\nabla_{\boldsymbol{x}}^{2}:\left(\boldsymbol{a}(\boldsymbol{x},t)p\right )+gp,$ (3.11) $$

with the corresponding metric tensor

$$ $\int_{0}^{1}\int_{\mathbb{R}^{d}}\frac{1}{2}\left\|\boldsymbol{b}(\boldsymbol{ x},t)\right\|_{2}^{2}p(\boldsymbol{x},t)\mathrm{d}\boldsymbol{x}\mathrm{d}t+ \int_{0}^{1}\int_{\mathbb{R}^{d}}\alpha\Psi\left(g(\boldsymbol{x},t)\right)p( \boldsymbol{x},t)\mathrm{d}\boldsymbol{x}\mathrm{d}t,$ (3.12) $$

where $\alpha$ and $\Psi:\mathbb{R}\rightarrow[0,+\infty]$ is the growth penalization.
Then we can define the
*regularized unbalanced optimal transport* problem in the same way.

###### Definition 3.7 (Regularized unbalanced optimal transport II) .

Consider

$$ $\inf_{\left(p,\boldsymbol{b},g\right)}\int_{0}^{1}\int_{\mathbb{R}^{d}}\frac{1 }{2}\left\|\boldsymbol{b}(\boldsymbol{x},t)\right\|_{2}^{2}{p}(\boldsymbol{x}, t)\mathrm{d}\boldsymbol{x}\mathrm{d}t+\int_{0}^{1}\int_{\mathbb{R}^{d}}\alpha \Psi\left(g(\boldsymbol{x},t)\right){p}(\boldsymbol{x},t)\mathrm{d}\boldsymbol {x}\mathrm{d}t,$ (3.13) $$

where the infimum is taken all pairs $\left(p,\boldsymbol{b}\right)$ such that $p(\cdot,0)=$ $\nu_{0},p(\cdot,T)=\nu_{1},p(\boldsymbol{x},t)$ absolutely continuous, and

$$ $\partial_{t}p=-\nabla_{\boldsymbol{x}}\cdot\left(p\boldsymbol{b}\right)+\frac{ 1}{2}\nabla_{\boldsymbol{x}}^{2}:\left(\boldsymbol{a}(\boldsymbol{x},t)p\right )+gp.$ (3.14) $$

with vanishing boundary condition: $\displaystyle\lim_{|\boldsymbol{x}|\rightarrow\infty}p(\boldsymbol{x},t)=0$.

We similarly can
reformulate Definition 3.7 with the following Fisher information regularization.

###### Theorem 3.8 .

Consider the regularized unbalanced optimal transport problem Definition 3.7, it is equivalent to

$$ $\displaystyle\inf_{\left(p,\boldsymbol{v},g\right)}\int_{0}^{1}\int_{\mathbb{R }^{d}}\frac{1}{2}\left\|\boldsymbol{v}(\boldsymbol{x},t)\right\|_{2}^{2}{p}( \boldsymbol{x},t)\mathrm{d}\boldsymbol{x}\mathrm{d}t+\frac{1}{8}\left\| \boldsymbol{\sigma}(\boldsymbol{x},t)\nabla_{\boldsymbol{x}}\cdot\log\left( \boldsymbol{a}(\boldsymbol{x},t)p\right)\right\|_{2}^{2}{p}(\boldsymbol{x},t) \mathrm{d}\boldsymbol{x}\mathrm{d}t$ (3.15) $\displaystyle+\frac{1}{2}\left\langle\boldsymbol{v}(\boldsymbol{x},t), \boldsymbol{a}(\boldsymbol{x},t)\nabla_{\boldsymbol{x}}\cdot\log\left( \boldsymbol{a}(\boldsymbol{x},t)p\right)\right\rangle{p}(\boldsymbol{x},t) \mathrm{d}\boldsymbol{x}\mathrm{d}t+\alpha\Psi\left(g(\boldsymbol{x},t)\right) {p}(\boldsymbol{x},t)\mathrm{d}\boldsymbol{x}\mathrm{d}t,$ $$

where the infimum is taken all pairs $\left(p,\boldsymbol{v},g\right)$ such that $p(\cdot,0)=$ $\nu_{0},p(\cdot,T)=\nu_{1},p(\boldsymbol{x},t)$ absolutely continuous, and

$$ $\partial_{t}p=-\nabla_{\boldsymbol{x}}\cdot\left(p\boldsymbol{v}(\boldsymbol{x },t)\right)+g(\boldsymbol{x},t)p.$ (3.16) $$

with vanishing boundary condition: $\displaystyle\lim_{|\boldsymbol{x}|\rightarrow\infty}p(\boldsymbol{x},t)=0$.

The proof is left to Appendix B. From the proof of Theorem 3.8, the original SDE
$\mathrm{d}{\boldsymbol{X}}_{t}=\left(\boldsymbol{b}\left({\boldsymbol{X}}_{t},
t\right)\right)\mathrm{d}t+\boldsymbol{\sigma}\left(\boldsymbol{X}_{t},t\right
)\mathrm{d}{\boldsymbol{W}}_{t}$
can be transformed
into the probability flow ODE

$$ $\mathrm{d}{\boldsymbol{X}}_{t}=\underbrace{\left(\boldsymbol{b}\left({ \boldsymbol{X}}_{t},t\right)-\frac{1}{2}\boldsymbol{a}({\boldsymbol{X}}_{t},t) \nabla_{\boldsymbol{x}}\cdot\log\left(\boldsymbol{a}({\boldsymbol{X}}_{t},t) \boldsymbol{p}\left({\boldsymbol{X}}_{t},t\right)\right)\right)}_{\boldsymbol{ v}\left({\boldsymbol{X}}_{t},t\right)}\mathrm{d}t.$ $$

Conversely, if the probability flow ODE’s drift $\boldsymbol{v}(\boldsymbol{x},t)$ the diffusion rate $\boldsymbol{a}(\boldsymbol{x},t)$ and the function $\nabla_{\boldsymbol{x}}\cdot\log\left(\boldsymbol{a}(\boldsymbol{x},t){p}(
\boldsymbol{x},t)\right)$ are known, then the
the SDE’s drift term $b(\boldsymbol{x},t)$ can be determined by

$$ $b(\boldsymbol{x},t)=\boldsymbol{v}(\boldsymbol{x},t)+\frac{1}{2}\boldsymbol{a} (\boldsymbol{x},t)\nabla_{\boldsymbol{x}}\cdot\log\left(\boldsymbol{a}( \boldsymbol{x},t){p}(\boldsymbol{x},t)\right).$ $$

So to specify an SDE is equivalent to specify the probability flow ODE and its score function.

###### Definition 3.7 (Regularized unbalanced optimal transport II) .

###### Theorem 3.8 .

## 4 Connections with Schrödinger Bridge

Although the RUOT problem can be derived from the dynamical formulation of a Schrödinger bridge relaxation, the relationship between the RUOT problem and the original static Schrödinger bridge problem, the choice of the growth penalty form in RUOT as well as its microscopic interpretation, remains unsolved. In this section, we will discuss two potential microscopic interpretations with the choice of the growth penalty form and identify the problems and challenges associated with each.

### 4.1 Connections with Branching Schrödinger Bridge

In (Baradat and Lavenant, ), the authors explore the relationship between the RUOT and the branching Schrödinger bridge, with the corresponding reference process identified as branching Brownian motion. We will begin by reviewing their principal findings and subsequently highlight the associated issues and challenges.

A Branching Brownian motion (BBM) can be characterized by its diffusivity $\varepsilon$, branching rate $\lambda$ (where each particle is associated with an independent exponential distribution with parameter $\lambda$ ), and the law of offspring denoted by $p=(p_{k})_{k\in\mathbb{N}}\in\mathcal{P}(\mathbb{N})$ (The probability measures on $\mathbb{N}$ with nonnegative sequences summing to one), where $p_{k}$ is the probability of producing $k$ descendants at a branching event. We define $q_{k}=\lambda p_{k}$ and introduce the generating function of $\boldsymbol{q}$ as $\Phi_{\boldsymbol{q}}$.

$$ $\Phi_{\boldsymbol{q}}:z\mapsto\sum_{k\in\mathbb{N}}q_{k}z^{k}=\lambda\sum_{k \in\mathbb{N}}p_{k}z^{k}.$ $$

The branching mechanism is entirely determined by $\boldsymbol{q}$ , since $\lambda$ and $\boldsymbol{p}$ can be recovered from $\boldsymbol{q}$ by $\lambda=\sum_{k\in\mathbb{N}}q_{k}$ and $p_{k}=\lambda^{-1}q_{k}$ for $k\in\mathbb{N}$ . Therefore, these parameters are sufficient to fully characterize the BBM with diffusivity $\nu$ , branching mechanism $\boldsymbol{q}=\lambda\boldsymbol{p}$ , and initial law $R\sim\text{BBM}(\varepsilon,\boldsymbol{q},\nu_{0})$ . Their core result indicates that consider Definition 3.1 where $\sigma(t)=\sigma$ is constant and $\varepsilon=\sigma^{2}$ then if the growth penalty term is defined as follows, it can correspond to a branching Schrödinger process.

$$ $\Psi_{\varepsilon,\boldsymbol{q}}^{*}(s)=\varepsilon\left(\Phi_{\boldsymbol{q} }\left(e^{s/\varepsilon}\right)e^{-s/\varepsilon}-\Phi_{\boldsymbol{q}}(1) \right)=\varepsilon\sum_{k=0}^{+\infty}q_{k}\left\{\exp\left((k-1)\frac{s}{ \varepsilon}\right)-1\right\},$ (4.17) $$

and we define $\Psi_{\varepsilon,\boldsymbol{q}}(r)=\sup_{s\in\mathbb{R}}rs-\Psi_{\varepsilon
,\boldsymbol{q}}^{*}(s)$. So by following their result, we can present some examples.

#### The case of only loss mass

This example has been included in their work. We consider $\varepsilon=\lambda=1$, $p_{0}=1$ and $\sum_{k\geq 2}p_{k}=0$, which means the particle can only die at the branching event. Substituting it into Eq. (4.17), we have

$$ $\Psi_{1,\boldsymbol{q}}^{*}(s)=\exp\left(-s\right)-1.$ $$

Then by computing its Legendre transform we obtain

$$ $\Psi_{1,\boldsymbol{q}}(r)=1+r-r\log(-r),$ $$

where $r<0$. This form is consistent with the formula discussed in (Chen et al., ).

#### The case of only gain mass

Similarly, if we consider $\varepsilon=\lambda=1$, $p_{2}=1$, which means the particle can only divide into two parts at the branching event. Then by a similar calculation, we obtain

$$ $\Psi_{1,\boldsymbol{q}}(r)=1-r+r\log(r),$ $$

which $r>0$.

#### The case of both gain mass and loss mass

We consider $\varepsilon=\lambda=1$, $p_{0}=p_{2}=\frac{1}{2}$, which means the particle can divide into two parts or die at the same probability at the branching event. This is also presented in their work, which we have

$$ $\Psi_{1,\boldsymbol{q}}^{*}(s)=\cosh(s)-1.$ $$

However in this case the Legendre transform would be more complicated but it has some asymptotic properties discussed in their work.

#### General case

Although we have illustrated several examples, challenges persist, particularly in the context of more general scenarios. Specifically, when the growth penalty term cannot be expressed as $\Psi_{\varepsilon,\boldsymbol{q}}$, the connection to the corresponding stochastic process remains unclear. For instance, this ambiguity arises in cases where in the RUOT the $\Psi(g)$ takes the quadratic form i.e., the WFR metric, or when $\Psi(g)$ is in a linear form, such as $|g|$. These issues still require specific treatment and to our knowledge, remain an open problem.

### 4.2 Connections with Schrödinger Bridge on Weighted Path

From an alternative perspective on the microscopic interpretation of RUOT, we aim to consider particles on the weighted path space, namely:

$$ $\mathrm{d}\mu^{\boldsymbol{G}}_{[0,1]}=\exp\left(\int_{0}^{1}g(\boldsymbol{X}_ {s},s)\mathrm{d}s\right)\mathrm{d}\mu^{\boldsymbol{X}}_{[0,1]},$ $$

where $\mu^{\boldsymbol{X}}_{[0,1]}$ represents the probability measure induced by the SDE (Eq. (2.2)).
Here for example we consider a simple case when $g(\boldsymbol{x},t)=C$. Then following the procedure of Theorem 2.1.

$$ $\displaystyle\mathcal{D}_{\text{KL}}(\mu^{\boldsymbol{G}}_{[0,1]}||\mu^{ \boldsymbol{Y}}_{[0,1]})$ $\displaystyle=\mathbb{E}_{\mu^{\boldsymbol{G}}_{[0,1]}}\left(\int_{0}^{1} \gamma({\boldsymbol{Y}_{t}},t)\mathrm{d}{\boldsymbol{W}}_{t}-\frac{1}{2}\left \|\gamma({\boldsymbol{Y}_{t}},t)\right\|_{2}^{2}\mathrm{d}t+C\right),$ (4.18) $\displaystyle=\mathbb{E}\left(\mathcal{E}_{1}\exp(C)\right)\left(\int_{0}^{1} \frac{1}{2}\left\|\gamma({\boldsymbol{Y}_{t}},t)\right\|_{2}^{2}\mathrm{d}t \right)+C\exp(C),$ $\displaystyle=\mathbb{E}\exp(C)\int_{0}^{1}\frac{1}{2}\left\|\gamma({ \boldsymbol{X}_{t}},t)\right\|_{2}^{2}\mathrm{d}t+C\exp(C).$ $$

Here $\mu^{\boldsymbol{Y}}_{[0,1]}$ is the distribution of stochastic process induced by $\mathrm{d}\boldsymbol{Y}_{t}=\boldsymbol{\sigma}(\boldsymbol{Y}_{t},t)\mathrm{
d}{\boldsymbol{W}}_{t}$.
Our derivation reveals that, from a path weighting perspective, even when the growth term is a constant, it still does not correspond to the form of RUOT. This suggests that the microscopic interpretation from this angle may need further investigation to determine whether a RUOT problem could be recovered from this starting point.

## 5 Learning RUOT through Deep Learning

Given unnormalized distributions at $T$ discrete time points, $\boldsymbol{X}_{i}\sim\mu_{i}$ for fixed timepoints $i\in\{0,\ldots,T-1\}$, we aim to learn a continuous stochastic dynamic process satisfying the RUOT problem from the data. Similarly, to simplify the exposition and illustrate a commonly used case, we consider Definition 3.1 with the stochastic dynamics

$$ $\mathrm{d}{\boldsymbol{X}}_{t}=\left(\boldsymbol{b}\left({\boldsymbol{X}}_{t}, t\right)\right)\mathrm{d}t+{\sigma}\left(t\right)\boldsymbol{I}\mathrm{d}{ \boldsymbol{W}}_{t}.$ $$

As previously discussed in Theorem 3.2, we approach this by transforming the problem into learning the drift $\boldsymbol{v}(\boldsymbol{x},t)$ of the probability ODE and its score function $\frac{1}{2}\sigma^{2}(t)\nabla_{\boldsymbol{x}}\log p(\boldsymbol{x},t)$. We parameterize $\boldsymbol{v}(\boldsymbol{x},t)$, ${g}(\boldsymbol{x},t)$, and $\frac{1}{2}\sigma^{2}(t)\log{p}(\boldsymbol{x},t)$ using neural networks $\boldsymbol{v}_{\theta},~{}g_{\theta}$ and $s_{\theta}$ respectively.

Figure: Figure 1: Overview of DeepRUOT.
Refer to caption: x1.png

### 5.1 Energy Function in Theorem 3.2

To compute equation Eq. (3.9), direct calculation is infeasible due to the high dimensionality. Thus, we need to transform it into an equivalent form that can be evaluated using Monte Carlo methods. Adapting from the approach in (Sha et al., ), we have the following theorem.

###### Theorem 5.1 .

If smooth density $p(\boldsymbol{x},t):\mathbb{R}^{d}\times[0,T]\rightarrow\mathbb{R}^{+}$, velocity field $\boldsymbol{v}(\boldsymbol{x},t):\mathbb{R}^{d}\times[0,T]\rightarrow\mathbb{R
}^{d}$ and growth rate $g(x,t):\mathbb{R}^{d}\times[0,T]\rightarrow\mathbb{R}$ satisfy

$$ $\left\{\begin{array}[]{c}\partial_{t}p(\boldsymbol{x},t)+\nabla\cdot( \boldsymbol{v}(\boldsymbol{x},t)p(\boldsymbol{x},t))={g}(\boldsymbol{x},t)p( \boldsymbol{x},t),\\ p(\boldsymbol{x},0)=p_{0}(\boldsymbol{x}),\end{array}\right.$ $$

for all $0\leq t\leq T$ with $\frac{\mathrm{d}\boldsymbol{x}(t)}{\mathrm{d}t}=\boldsymbol{v}(\boldsymbol{x},t)$ and $\boldsymbol{x}(0)=\boldsymbol{x}_{0}$, then for any measurable function $f(\boldsymbol{x},t):\mathbb{R}^{d}\times[0,T]\rightarrow\mathbb{R}^{d}$, we have

$$ $\int_{0}^{1}\int_{\mathbb{R}^{d}}f(\boldsymbol{x},t)p(\boldsymbol{x},t)\mathrm {d}\boldsymbol{x}\mathrm{d}t=\mathbb{E}_{\boldsymbol{x_{0}}\sim p_{0}}\int_{0} ^{1}f(\boldsymbol{x},t)e^{\int_{0}^{t}g(\boldsymbol{x},s)\mathrm{d}s}\mathrm{d }t.$ $$

According to this theorem, Eq. (3.9) is equivalent to the following form

$$ $\displaystyle W_{0,T}=\int_{0}^{T}\int_{\mathbb{R}^{d}}\left[\frac{1}{2}\left \|\boldsymbol{v}(\boldsymbol{x},t)\right\|_{2}^{2}+\frac{\sigma^{4}(t)}{8} \left\|\nabla_{\boldsymbol{x}}\log p\right\|_{2}^{2}-\frac{\sigma^{2}(t)}{2} \left(1+\log p\right)g+\alpha\Psi\left(g\right)\right]p\mathrm{d}\boldsymbol{x }\mathrm{d}t,$ (5.19) $\displaystyle=\mathbb{E}_{\boldsymbol{x_{0}}\sim p_{0}}\int_{0}^{1}\left[\frac {1}{2}\left\|\boldsymbol{v}(\boldsymbol{x},t)\right\|_{2}^{2}+\frac{\sigma^{4} (t)}{8}\left\|\nabla_{\boldsymbol{x}}\log p\right\|_{2}^{2}-\frac{\sigma^{2}(t )}{2}\left(1+\log p\right)g+\alpha\Psi\left(g\right)\right]e^{\int_{0}^{t}g( \boldsymbol{x},s)\mathrm{d}s}\mathrm{d}t,$ $$

which we compute through Monte Carlo sampling and Neural ODE solver.

###### Theorem 5.1 .

### 5.2 Reconstruction Loss

The reconstruction loss aims to match the final distribution in Theorem 3.2, which is composed of two parts: the mass matching loss and the optimal transport loss in matching the normalized distributions. Specifically, we first address the mass matching loss to align the total number of cells. Subsequently, we normalize the distributions according to the matched masses, and then use these weights to perform optimal transport matching. This process is similar to using unbalanced optimal transport to evaluate the distance between two distributions.

For the first part, we propose a local mass-matching strategy.
We denote $A_{i}$ ($i=0,1,\cdots T-1$) as the dataset observed at different time points, and define $\phi_{\theta}:\mathbb{R}^{d}\times\mathcal{T}\rightarrow\mathbb{R}^{d|\mathcal
{T}|}$ as the function representing the Neural ODE, with $\mathcal{T}$ being a set of time indices. Starting with an initial set $A_{0}$, the Neural ODE function $\phi_{\theta}$ predicts the subsequent sets of data points over the time indices in $\mathcal{T}$. Specifically, the predicted sets $\widehat{A}_{1},\ldots,\widehat{A}_{T-1}$ are obtained by applying the Neural ODE function to $A_{0}$ and the full range of time indices from $1$ to $T-1$, i.e.,
$\widehat{A}_{1},\ldots,\widehat{A}_{T-1}=\phi_{\theta}\left(A_{0},\{1,\ldots,T
-1\}\right).$
Assume the sample size of $N_{0}$ at the initial time point, with the relative total masses at subsequent time points denoted as $n_{t_{k}}=N_{k}/N_{0}$ for $k=0,1,\cdots,T-1$. Along the trajectory, each sampled particle has a weight $w_{i}(t)$ and it follows that

$$ $\frac{\mathrm{d}\log w_{i}(t)}{\mathrm{d}t}=g(\boldsymbol{x}_{i}(t),t),\ \ w_{ i}(0)=1/N_{0}.$ $$

We then establish a mapping $h_{k}$ from points at time $t_{k}$ to the sampled particles (total number $N_{0}$) predicted at time $t_{k}$, which links the real data points in $A_{k}$ to its closest point in sampled particles at $\widehat{A}_{k}$, i.e.,

$$ $h_{k}:A_{k}\rightarrow\widehat{A}_{k},\ h_{k}(\boldsymbol{{x}}_{t_{k}})=\text{ argmin}_{\boldsymbol{y}\in\widehat{A}_{k}}\|\boldsymbol{{x}}_{t_{k}}- \boldsymbol{y}\|_{2}^{2}.$ $$

The mass matching error at time $t_{k}$ is then defined as

$$ $M_{k}=\sum_{i=1}^{N_{0}}\Big{\|}w_{i}(t_{k})-\text{card}\left(h_{k}^{-1}( \boldsymbol{x}_{i}(t_{k}))\right)\frac{1}{N_{0}}\Big{\|}_{2}^{2}.$ $$

Here $\text{card}(A)$ denotes the cardinality (i.e., the number of elements for finite set) of set $A$. This error metric ensures the matching of local masses. Furthermore, when $M_{t}=0$, we have

$$ $\sum_{i=1}^{N_{0}}w_{i}(t_{k})=\sum_{i=1}^{N_{0}}\text{card}\left(h_{t}^{-1}( \boldsymbol{x}_{i}(t_{k}))\right)/N_{0}=N_{k}/N_{0}=n_{t_{k}}.$ $$

So the local matching loss also encourages the matching of total mass, since the left hand side is the unbiased estimator of $\int_{\mathbb{R}^{d}}p(\boldsymbol{y},t_{k})d\boldsymbol{y}$.

Once we have determined the weights $w_{i}(t)$, we utilize these weights to perform optimal transport matching of the distributions

$$ $D_{0,T}:=\sum_{k=1}^{T-1}\mathcal{W}_{2}\big{(}\hat{\boldsymbol{w}}^{k}, \boldsymbol{w}(t_{k})\big{)},$ $$

where $\hat{\boldsymbol{w}}^{k}=(1/N_{k},1/N_{k},...,1/N_{k})$ is the uniform distribution of $A_{k}$ at time $t$, $\boldsymbol{w}(t_{k})=(w_{1}(t_{k}),w_{2}(t_{k}),...,w_{N_{0}}(t_{k}))/\sum_{i
=1}^{N_{0}}w_{i}(t_{k})$ is the predicted weight distribution of sample particles at time $t_{k}$ and $\mathcal{W}_{2}$ represents the Wasserstein distance between $A_{k}$ and $\widehat{A}_{k}$ with normalized distribution defined by $\hat{\boldsymbol{w}}^{k}$ and $\boldsymbol{w}(t_{k})$. Thus the reconstruction loss is

$$ $R_{d}=\lambda_{m}\sum_{k=1}^{T-1}M_{k}+\lambda_{d}D_{0,T},$ (5.20) $$

where $\lambda_{m}$ and $\lambda_{d}$ are hyperparameters.

### 5.3 Learning Probability Flow ODE and Score Function

We now integrate all components to derive the deep learning-based solver for learning the probability flow ODE. We utilize fully connected neural networks $\boldsymbol{v}_{\theta},g_{\theta},s_{\theta}$ to estimate $\boldsymbol{v}(\boldsymbol{x},t)$, $g(\boldsymbol{x},t)$ and $\frac{\sigma^{2}}{2}\log{p}(\boldsymbol{x},t)$ respectively. The energy function is computed using Eq. (5.19) and the reconstruction error is computed using Eq. (5.20). Additionally, it is necessary to incorporate a physics-informed loss (PINN-loss) (Raissi et al., ) to constrain the relationships among the three neural networks, i.e. Eq. 3.10. Also we utilize a Gaussian mixture model to estimate the initial distribution, ensuring that it satisfies the initial boundary conditions.

$$ $W_{\text{PINN}}=\left\|\partial_{t}p+\nabla_{\boldsymbol{x}}\cdot\left(p \boldsymbol{v}(\boldsymbol{x},t)\right)-g(\boldsymbol{x},t)p\right\|_{2}+ \lambda_{w}\left\|p(x,0)-p_{0}\right\|_{2}.$ $$

So overall, the loss function is defined as the weighted sum of the cost, the reconstruction error and the PINN loss:

$$ $L=R_{d}+\lambda_{m}W_{0,T}+\lambda_{p}W_{\text{PINN}},$ $$

with the loss weights hyper-parameters $\lambda_{m}$ and $\lambda_{p}$. Here, $\boldsymbol{x}=\boldsymbol{x}(t)$ is defined by the following equation.

$$ $\begin{cases}\frac{\mathrm{d}\boldsymbol{x}(t)}{\mathrm{d}t}=\boldsymbol{v}( \boldsymbol{x}(t),t),\\ \boldsymbol{x}(t_{i})=\boldsymbol{x}_{i}.\end{cases}$ $$

This allows us to develop a neural network algorithm for solving the RUOT problem (Algorithm 1).

Figure: Algorithm 1 Learning Regularized Unbalanced Optimal Transport

### 5.4 Training

Here we adopt a multi-stage training approach to deal with multiple loss terms and stabilize the training process. Initially, we used reconstruction error to train $\boldsymbol{v}$ and $g$, ensuring a required matching. Subsequently, we fixed $\boldsymbol{v}$ and $g$ and employed a flow-matching method to learn an initial, well-optimized log density function $s_{\theta}(\boldsymbol{x},t)$. Following this, we refined the log density function based on the PINN loss. Finally, using $\boldsymbol{v}$, $g$, and the optimized log density function as the starting point, we use Algorithm 1 to obtain the final result.

Our approach to learning the initial log density function is conducted by conditional flow matching (Lipman et al., ; Tong et al., ). Previously, we have learnt a matching. Then we perform the dynamics building on this to generate samples, which is akin to conducting an unbalanced optimal transport. We perform conditional flow matching (CFM) within such samples same as (Tong et al., ). Specifically, we choose pair samples $(\boldsymbol{x}_{0},\boldsymbol{x}_{1})$ from the optimal transport plan $q(\boldsymbol{x}_{0},\boldsymbol{x}_{1})$ and then construct Brownian bridges between these pair samples. We similar first consider $\sigma(t)=\sigma$ is constant. The log density function is then matched with these Brownian bridges, i.e., $p(\boldsymbol{x},t\mid(\boldsymbol{x}_{0},\boldsymbol{x}_{1}))=\mathcal{N}
\left(\boldsymbol{x};t\boldsymbol{x}_{1}+(1-t)\boldsymbol{x}_{0},\sigma^{2}t(1
-t)\right)$ and $\nabla_{\boldsymbol{x}}\log p(\boldsymbol{x},t\mid(\boldsymbol{x}_{0},
\boldsymbol{x}_{1}))=\frac{t\boldsymbol{x}_{1}+(1-t)\boldsymbol{x}_{0}-
\boldsymbol{x}}{\sigma^{2}t(1-t)},~{}t\in[0,1].$ We then utilize the fully connected neural networks $s_{\theta}(\boldsymbol{x},t)$ to approximate the log density function $\frac{1}{2}\sigma^{2}\log{p}(\boldsymbol{x},t)$ with a weighting function $\lambda_{s}$. Then we have

$$ $L_{us}=\lambda_{s}^{2}\|\nabla_{\boldsymbol{x}}s_{\theta}-\frac{1}{2}\sigma^{2 }\nabla_{\boldsymbol{x}}\log{p}(\boldsymbol{x},t)\|_{2}^{2}.$ $$

And the correspongding CFM loss is

$$ $L_{s}=\mathbb{E}_{Q^{\prime}}\lambda_{s}^{2}\|\nabla_{\boldsymbol{x}}s_{\theta }-\frac{1}{2}\sigma^{2}\nabla_{\boldsymbol{x}}\log p(\boldsymbol{x},t\mid( \boldsymbol{x}_{0},\boldsymbol{x}_{1}))\|_{2}^{2},$ $$

where $Q^{\prime}=(t\sim\mathcal{U}(0,1))\otimes q(\boldsymbol{x}_{0},\boldsymbol{x}_
{1})\otimes p(\boldsymbol{x},t\mid(\boldsymbol{x}_{0},\boldsymbol{x}_{1}))$.
We take the weighting function as

$$ $\lambda_{s}(t)=\frac{2}{\sigma^{2}}\sigma\sqrt{t(1-t)}=\frac{2\sqrt{t(1-t)}}{ \sigma}.$ $$

Then we have

$$ $\displaystyle L_{S}=$ $\displaystyle\lambda_{s}(t)^{2}\left\|\nabla_{\boldsymbol{x}}s_{\theta}( \boldsymbol{x},t)-\frac{\sigma^{2}}{2}\nabla_{\boldsymbol{x}}\log p\left( \boldsymbol{x}\mid\boldsymbol{x_{0}},\boldsymbol{x_{1}}\right)\right\|^{2},$ $\displaystyle=\left\|\lambda_{s}(t)\nabla_{\boldsymbol{x}}s_{\theta}( \boldsymbol{x},t)-\lambda_{s}(t)\frac{\sigma^{2}}{2}\nabla_{\boldsymbol{x}} \log p\left(\boldsymbol{x}\mid\boldsymbol{x_{0}},\boldsymbol{x_{1}}\right) \right\|^{2},$ $\displaystyle=\left\|\lambda_{s}(t)\nabla_{\boldsymbol{x}}s_{\theta}( \boldsymbol{x},t)+\epsilon_{1}\right\|^{2},$ $$

where $\epsilon_{1}\sim\mathcal{N}(0,1)$.
This is also numerically stable. For the case where $\sigma(t)$ is not constant, a similar approach can be applied. One may refer to their paper for further details.

In summary, by integrating all the discussed methodologies, we derive the Algorithm 2 for training the regularized unbalanced optimal transport in practice.

Figure: Algorithm 2 Training Regularized Unbalanced Optimal Transport

In practical implementation, we employ Algorithm 2. The temporal integral and ODEs were numerically solved using the Nerual ODE solver (Chen et al., ). The gradients of the loss function with respect to the parameters in the neural networks for $\boldsymbol{v}(\boldsymbol{x},t)$, $g(\boldsymbol{x},t)$ and $s(\boldsymbol{x},t)$ were computed using neural ODEs with a memory-efficient implementation. The Adam optimizer was employed to update the gradients.
For computing the Wasserstein distance between discrete distributions, we utilize the implementation provided by the Python Optimal Transport library (POT) (Flamary et al., ).

## 6 Numerical Results

In this section, we evaluate the effectiveness of the unbalanced dynamic diffusion Schrödinger bridge solver on synthetic dynamics from two perspectives: (1) Compared to the balanced diffusion Schrödinger bridge, our approach accurately recovers the correct growth and transition processes, eliminating false transitions caused by neglecting growth. (2) Compared to unbalanced optimal transport, the results obtained from the stochastic dynamics are more realistic, and we successfully learn the Waddington’s developmental landscape. We take $\Psi\left({g}(\boldsymbol{x},t)\right)=|{g}(\boldsymbol{x},t)|_{2}^{2}$ in the following discussions.

### 6.1 Synthetic Dynamics

Inspired by (Sha et al., ), we adopt a same three-gene simulation model to explore the stochastic dynamics of gene regulation, as illustrated in Fig. 2(a). The system dynamics are governed by the following set of stochastic ordinary differential equations (ODEs):

$$ $\displaystyle\frac{\mathrm{d}X_{1}}{\mathrm{~{}d}t}=\frac{\alpha_{1}X_{1}^{2}+ \beta}{1+\alpha_{1}X_{1}^{2}+\gamma_{2}X_{2}^{2}+\gamma_{3}X_{3}^{2}+\beta}- \delta_{1}X_{1}+\eta_{1}\xi_{t},$ $\displaystyle\frac{\mathrm{d}X_{2}}{\mathrm{~{}d}t}=\frac{\alpha_{2}X_{2}^{2}+ \beta}{1+\gamma_{1}X_{1}^{2}+\alpha_{2}X_{2}^{2}+\gamma_{3}X_{3}^{2}+\beta}- \delta_{2}X_{2}+\eta_{2}\xi_{t},$ $\displaystyle\frac{\mathrm{d}X_{3}}{\mathrm{~{}d}t}=\frac{\alpha_{3}X_{3}^{2}} {1+\alpha_{3}X_{3}^{2}}-\delta_{3}X_{3}+\eta_{3}\xi_{t}.$ $$

Genes $X_{1}$ and $X_{2}$ mutually inhibit each other and self-activate, forming a toggle switch. An external signal $\beta$ uniformly activates both $X_{1}$ and $X_{2}$ independently of gene expression levels. Gene $X_{3}$ inhibits the expression of both $X_{1}$ and $X_{2}$. Here, $X_{i}(t)$ represents the concentration of gene $i$ at time $t$, with $\alpha_{i}$ and $\gamma_{i}$ denoting the strengths of self-activation and inhibition, respectively. The parameters $\delta_{i}$ represent gene degradation rates, while $\eta_{i}\xi_{t}$ accounts for the stochastic effects with additive white noise. The probability of cell division correlates positively with the expression of $X_{2}$, calculated as $g=\frac{X_{2}^{2}}{1+X_{2}^{2}}$. Upon division, cells inherit the gene expression states $(X_{1}(t),X_{2}(t),X_{3}(t))$ of the parent cell, subject to independent perturbations $\eta_{d}\mathcal{N}(0,1)$ for each gene, and transition independently thereafter.

The resulting gene regulatory and cellular dynamics are illustrated in Fig. 2, where a quiescent region and an area exhibiting both transition and growth can be observed. We focus on the projection of these dynamics onto the two-dimensional space of $(X_{1},X_{2})$ since on $X_{3}$ it remains quiescent.

Figure: (a) Illustration of gene regulatory network.
Refer to caption: x2.png

Figure: (a) Balanced Schrödinger bridge (SF2M).
Refer to caption: x4.png

In Fig. 3, we compare our proposed method with the balanced diffusion Schrödinger bridge method as described in (Tong et al., ). Our findings indicate that neglecting the growth factor in the balanced diffusion Schrödinger bridge method leads to a false transition, incorrectly attracting cells in the quiescent state to the transition and growth region.
The underlying reason for this discrepancy is that the growth factor causes a continuous increase in the number of cells. If the growth factor is ignored, the regions with increasing cell numbers will require more cells to transition into them to maintain balance, which ultimately leads to false transitions.

Our approach, which explicitly incorporates growth dynamics, successfully eliminates false transitions and achieves results consistent with the ground truth dynamics (Fig. 3(b)). Furthermore, the growth rates inferred by our method closely match the ground truth, demonstrating the robustness and accuracy of our unbalanced dynamic diffusion Schrödinger
bridge model after normalization(Fig. 4).

Figure: (a) Ground truth growth rates.
Refer to caption: x6.png

### 6.2 Learning Waddington’s Developmental Landscapes

In both biology and physics, the Waddington’s landscape metaphor is a well-known model for representing cell fate decision process. This metaphor posits that meta-stable cell states correspond to wells in potential landscape, and state-transitions could be modelled as hopping between potential wells. Constructing such a potential landscape has been widely studied (Zhou and Li, ; Shi et al., ; Zhao et al., ; Li and Wang, ; Wang et al., ; Li and Wang, ; Bian et al., , ; Zhou et al., , ), however it still remains a challenging problem using the single-cell omics data. The energy landscape is defined by

$$ $U=-\sigma^{2}\log p_{\text{ss}}/2$ $$

where $p_{\text{ss}}$ is the steady-state PDF satisfying the Fokker-Planck equation

$$ $-\nabla\cdot(p_{ss}\boldsymbol{b})+\frac{\sigma^{2}}{2}\Delta p_{\text{ss}}=gp _{\text{ss}}.$ $$

In this work, leveraging our model, we can naturally infer the time-varying potential energy landscape through the learned log density function, and use $-\sigma^{2}(t)\log p(\boldsymbol{x},t)/2$ represent the landscape at time $t$. The lower the energy function indicates more stable cell fates in landscape.

Figure: Figure 5: (a)The Waddington developmental landscape learnt at $t=1$ ($\sigma=0.25$). (b) The constructed landscape at $t=4$ ($\sigma=0.25$).
Refer to caption: x8.png

In Fig. 5, we can observe the temporal evolution of potential landscape in the synthetic gene regulation system. The quiescent cells always occupy a potential well, indicating its stability. Interestingly, the location of potential well belonging to transitional cell population are moving in the consistent direction with ground truth cell-state transition dynamics. Overall, the results suggest the accuracy and usefulness of DeepRUOT algorithm.

## 7 Conclusion

We have introduced a new deep learning method for learning regularized unbalanced optimal transport (RUOT) and continuous unbalanced stochastic dynamics from time-series snapshot data. By leveraging Fisher regularization, our method transforms a stochastic differential equation (SDE) problem into an ordinary differential equation (ODE) constraint. Through the use of neural network modeling for growth and death, our framework models dynamics without requiring prior knowledge of these processes. Additionally, it provides a new approach for learning RUOT problems.

Furthermore, we have summarized and explored the connections between RUOT and the Schrödinger bridge, specifically the branching Schrödinger bridge and the Schrödinger bridge on weighted paths, highlighting the associated challenges. We have demonstrated the effectiveness of our method on synthetic gene regulatory data, showing its ability to eliminate false transitions caused by neglecting growth processes.

Future directions involve extending the methodology to learn the cell latent embedding space jointly with the dynamics, and developing more computationally efficient algorithms. We expect future work to address these aspects and improve trajectory inference problems in various scenarios, including multi-omics data (e.g. single-cell ATAC-seq, spatial transcriptomics) and beyond biology in machine learning.

## Acknowledgments

We thank Dr. Hugo Lavenant (Bocconi) for helpful discussions on Section 4.1 and Qiangwei Peng (PKU) for helpful suggestions on algorithm. We also thank Prof. Chunhe Li (Fudan), Juntan Liu (Fudan) and Prof. Qing Nie (UCI) for insightful discussions. We acknowledge the support from the High-performance Computing Platform of Peking University for computation.

## Appendix A The equivalent formulation of RUOT with nonlinear Fokker-Planck equation constraints

In RUOT problem, when $\Psi\left({g}(\boldsymbol{x},t)\right)$ take the quadratic form i.e. $\Psi\left({g}(\boldsymbol{x},t)\right)=|{g}(\boldsymbol{x},t)|_{2}^{2}$, one can define a formulation with nonlinear Fokker-Planck equation constraints as follows.

###### Definition A.1 .

We introduce

$$ $\operatorname{UOT}^{\mathrm{D}}\left(\nu_{0},\nu_{1}\right):=\inf_{(p, \boldsymbol{b},g)\in\mathcal{C}\left(\nu_{0},\nu_{1}\right)}\int_{0}^{1}\int_{ \mathbb{R}^{d}}\frac{1}{2}\left\|\boldsymbol{b}(\boldsymbol{x},t)\right\|_{2}^ {2}p\mathrm{d}\boldsymbol{x}\mathrm{d}t+\int_{0}^{1}\int_{\mathbb{R}^{d}} \alpha\left|g(\boldsymbol{x},t)\right|_{2}^{2}p\mathrm{d}\boldsymbol{x}\mathrm {d}t,$ $$

where, for $(p,\boldsymbol{b},g)$ belonging to appropriate functions spaces,

$$ $\mathcal{C}\left(\nu_{0},\nu_{1}\right):=\left\{p,\boldsymbol{b},g\left| \partial_{t}p=-\nabla_{\boldsymbol{x}}\cdot\left(p\boldsymbol{b}\right)+\frac{ 1}{2}\nabla_{\boldsymbol{x}}^{2}:\left(\sigma^{2}(t)\boldsymbol{I}p\right)+gp- \frac{\sigma^{2}(t)}{4\alpha}(1+\log p)p\right|\right\},$ $$

where $p_{0}=\nu_{0}$ and $p_{1}=\nu_{1}$.

To show its equivalence with RUOT, we study the property of its Fisher information regularisation form as below.

###### Theorem A.2 .

When $\Psi\left({g}(\boldsymbol{x},t)\right)=|{g}(\boldsymbol{x},t)|_{2}^{2}$, the regularized unbalanced optimal transport problem (3.9) is equivalent to

$$ $\inf_{\left(p,\boldsymbol{v},\widetilde{g}\right)}\int_{0}^{1}\int_{\mathbb{R} ^{d}}\left[\frac{1}{2}\left\|\boldsymbol{v}(\boldsymbol{x},t)\right\|_{2}^{2}+ \frac{\sigma^{4}(t)}{8}\left\|\nabla_{\boldsymbol{x}}\log p\right\|_{2}^{2}+ \alpha|\widetilde{g}(\boldsymbol{x},t)|^{2}-\frac{\sigma^{4}(t)}{16\alpha}(1+ \log p)^{2}\right]p\mathrm{d}\boldsymbol{x}\mathrm{d}t$ (1.21) $$

where the infimum is taken all pairs $\left(p,\boldsymbol{v},\widetilde{g}\right)$ such that $p(\cdot,0)=$ $\nu_{0},p(\cdot,T)=\nu_{1},p(\boldsymbol{x},t)$ absolutely continuous, and

$$ $\partial_{t}p=-\nabla_{\boldsymbol{x}}\cdot\left(p\boldsymbol{v}(\boldsymbol{x },t)\right)+\widetilde{g}(\boldsymbol{x},t)p+\frac{\sigma^{2}(t)}{4\alpha}(1+ \log p)p.$ (1.22) $$

with vanishing boundary condition: $\displaystyle\lim_{|\boldsymbol{x}|\rightarrow\infty}p(\boldsymbol{x},t)=0$.

###### Proof A.3 .

Based on the formulation of Theorem 3.2, we proceed by introducing the following change of variables:

$$ $\widetilde{g}=g-\frac{\sigma^{2}(t)}{4\alpha}(1+\log p).$ (1.23) $$

Substituting this variable into Eq. (3.9), then it follows immediately.

From Theorem A.2, we can find that if we introduce a new term in the continuity equation Eq. (3.8) as:

$$ $\partial_{t}p=-\nabla_{\boldsymbol{x}}\cdot\left(p\boldsymbol{b}\right)+\frac{ 1}{2}\nabla_{\boldsymbol{x}}^{2}:\left(\sigma^{2}(t)\boldsymbol{I}p\right)+gp- \frac{\sigma^{2}(t)}{4\alpha}(1+\log p)p.$ $$

Then Eq. (1.22) now yields the standard continuity equation:

$$ $\partial_{t}p=-\nabla_{\boldsymbol{x}}\cdot\left(p\boldsymbol{v}(\boldsymbol{x },t)\right)+\widetilde{g}(\boldsymbol{x},t)p.$ $$

###### Definition A.1 .

###### Theorem A.2 .

###### Proof A.3 .

## Appendix B Proof for Theorem 3.8

###### Proof B.1 .

From Definition 3.7 we know that
the Fokker-Planck equation for the SDE is

$$ $\displaystyle\partial_{t}p$ $\displaystyle=-\nabla_{\boldsymbol{x}}\cdot\left(p\boldsymbol{b}\right)+\frac{ 1}{2}\nabla_{\boldsymbol{x}}^{2}:\left(\boldsymbol{a}(\boldsymbol{x},t)p\right )+g(\boldsymbol{x},t)p,$ $\displaystyle=-\nabla_{\boldsymbol{x}}\cdot\left(\left(\boldsymbol{b}-\frac{1} {2}\boldsymbol{a}(\boldsymbol{x},t)\nabla_{\boldsymbol{x}}\cdot\log\left( \boldsymbol{a}(\boldsymbol{x},t)p\right)\right)p\right)+{g}(\boldsymbol{x},t)p.$ $$

with minimization problem

$$ $\inf_{\left(p,\boldsymbol{b}\right)}\int_{0}^{1}\int_{\mathbb{R}^{d}}\frac{1}{ 2}\left\|\boldsymbol{b}(\boldsymbol{x},t)\right\|_{2}^{2}{p}(\boldsymbol{x},t) \mathrm{d}\boldsymbol{x}\mathrm{d}t+\int_{0}^{1}\int_{\mathbb{R}^{d}}\alpha \Psi\left(g(\boldsymbol{x},t)\right){p}(\boldsymbol{x},t)\mathrm{d}\boldsymbol {x}\mathrm{d}t.$ $$

Using the change of variable $\boldsymbol{v}(\boldsymbol{x},t)=\boldsymbol{b}(\boldsymbol{x},t)-\frac{1}{2}
\boldsymbol{a}(\boldsymbol{x},t)\nabla_{\boldsymbol{x}}\cdot\log\left(
\boldsymbol{a}(\boldsymbol{x},t)p\right)$, we see that it is equivalent to

$$ $\partial_{t}p=-\nabla_{\boldsymbol{x}}\cdot\left(p\boldsymbol{v}(\boldsymbol{x },t)\right)+g(\boldsymbol{x},t)p.$ $$

On the other hand, since $\left\|\boldsymbol{b}(\boldsymbol{x},t)\right\|_{2}^{2}=\left\|\boldsymbol{v}(
\boldsymbol{x},t)\right\|_{2}^{2}+\frac{1}{4}\left\|\boldsymbol{a}(\boldsymbol
{x},t)\nabla_{\boldsymbol{x}}\cdot\log\left(\boldsymbol{a}(\boldsymbol{x},t)p
\right)\right\|_{2}^{2}+2\left\langle\boldsymbol{v}(\boldsymbol{x},t),\frac{1}
{2}\boldsymbol{a}(\boldsymbol{x},t)\nabla_{\boldsymbol{x}}\cdot\log\left(
\boldsymbol{a}(\boldsymbol{x},t)p\right)\right\rangle$, the integrand in the objective then becomes

$$ $\displaystyle\int_{0}^{1}\int_{\mathbb{R}^{d}}\frac{1}{2}\left\|\boldsymbol{v} (\boldsymbol{x},t)\right\|_{2}^{2}{p}(\boldsymbol{x},t)\mathrm{d}\boldsymbol{x }\mathrm{d}t+\frac{1}{8}\left\|\boldsymbol{\sigma}(\boldsymbol{x},t)\nabla_{ \boldsymbol{x}}\cdot\log\left(\boldsymbol{a}(\boldsymbol{x},t)p\right)\right\| _{2}^{2}{p}(\boldsymbol{x},t)\mathrm{d}\boldsymbol{x}\mathrm{d}t$ $\displaystyle+\frac{1}{2}\left\langle\boldsymbol{v}(\boldsymbol{x},t), \boldsymbol{a}(\boldsymbol{x},t)\nabla_{\boldsymbol{x}}\cdot\log\left( \boldsymbol{a}(\boldsymbol{x},t)p\right)\right\rangle{p}(\boldsymbol{x},t) \mathrm{d}\boldsymbol{x}\mathrm{d}t+\alpha\Psi\left(g(\boldsymbol{x},t)\right) {p}(\boldsymbol{x},t)\mathrm{d}\boldsymbol{x}\mathrm{d}t,$ $$

###### Proof B.1 .

## Peer Reviews

confidence: 2
contribution: 4
initial_rating: 8
presentation: 3
questions: - is not knowing the death/growth rates (what you call prior knowledge) such a big advantage?
- Why did you project the data onto a 2d manifold for the "Real Single-Cell Population Dynamics" experiment? Isn't it possible to be done on a higher dimensional space, e.g. gene expression of PCA space?
- The method involves multiple neural networks (v, g, s) and complex loss functions, which could require substantial computational resources as well as necessary tuning efforts, potentially limiting its practical application in real-time analysis. Would be great to get a comment on this.
soundness: 3
strengths: - The experimental results seem very convincing
- The paper provides thorough theoretical foundations with formal theorems and proofs connecting RUOT to the Schrödinger bridge problem.
summary: The paper introduces DeepRUOT, a new deep-learning approach for solving regularized unbalanced optimal transport (RUOT) and inferring continuous unbalanced stochastic dynamics from observed snapshots. The method models dynamics without requiring prior knowledge of growth and death processes, allowing these to be learned directly from data through a Fisher regularization formulation. The authors explore theoretical connections between RUOT and the Schrödinger bridge problem while reformulating RUOT to transform stochastic differential equations into more computationally tractable ordinary differential equations (a constraint). When tested on both synthetic gene regulatory networks and real single-cell RNA sequencing data, DeepRUOT outperformed existing methods in constructing meaningful developmental landscapes.
weaknesses: - it remains unclear how the loss function weights (cost, reconstruction error and the PINN loss) are set
- Figure 1 and its caption are not that helpful

Typos:
- Theorem 3.1: "probelm" -> "problem"
- page 8: "a increase" -> "an increase"

---

confidence: 3
contribution: 3
initial_rating: 3
presentation: 3
questions: - Regarding the connection with SB, could you either make the impact of the connection more explicit, or considering removing that part from the paper for clarity ?

- Could you please restructure Section 5 to help the readers follow the reasoning until the final loss and algorithm ? Having a more detailed pseudo-code description would already help significantly.

- Can you clarify Figure 2 - the predictions dots in the ground truth panel ? 

- Could you please add the results of an ablation experiment where you disable the learnable growth rate in your model ?
soundness: 3
strengths: - Population dynamics from snapshot data is an important problem, especially in biological application such as single cell trajectories inference.
- The proposed method appears to surpass the performance of previous methods.
- The related works section is extensive.
summary: This paper proposes a new approach for learning dynamics from the observation of distributions at different time steps. The paper proposes to use an unbalanced optimal transport approach that allows modeling particle growth or death over time, which is an important aspect of single cell dynamics modeling. The authors compare their approach on datasets from syntethic gene regulatory networks and single cell trajectories dynamics. The paper shows favorable performance results compared to previous methods.
weaknesses: - One of the most important weaknesses of this paper at this stage is its clarity. In its current form, the paper is poorly written, and the ideas are poorly motivated. The main limitation from prior works identified is the unabalanced aspect e.g. "However, computational methods for learning high-dimensional unbalanced stochastic dynamics from snapshots without prior knowledge of growth and death or additional information are still lacking.". Yet, the paper only indirectly addresses that issue, which confuses the reader. Despite the favorable experimental results, the current presentation of this paper prevents it from being appreciated by the community.

- Section 5, which is the core of your method, is extremly confusing. For instance, the reconstruction loss is not introduced previously such that it's extremely challenging to follow the development until the final loss in Equation 10. The algorithm 1 does not provide any help either in much needed clarification. The algorithm is too laconic and does not refer to corresponding parts of the text.

- It's not clear to me why the connection with the Schrodinger bridge is necessary. As far as my reading goes, it seems that you only need to leverage the dynamic version of regularized unbalanced OT. As such, it does not seem that the connection with Schrodinger bridges helps the reader understand the method, and you don't leverage that connection later on in your method.

- Theorem 4.1 uses $v$ which is not introduced previously. It's later undertstood that it corrresponds to the vector field of the corresponding ODE but this hinders the clarity of the paper.

- The presentation of the results in Figure 2 raises questions. Why does panel (b), with the true dynamics has ground truth, predicted and trajectories, if it is the ground truth dynamics ? What is the "predicted" in this case ? 

- Because the most important contribution of this paper relies on modeling the growth rate, an ablation study is required where the growth rate part of the method is disabled (e.g. g = 0). This would allow to understand the impact of it on the final performance. 

- Despite the extensive related works section, the experiments only compare with two baselines. Baselines such as Bunne et al. or Yutong et al. could also be considered.

Sha, Yutong, et al. "Reconstructing growth and dynamic trajectories from single-cell transcriptomics data."

---

confidence: 5
contribution: 2
initial_rating: 3
presentation: 2
questions: It is unclear to me exactly what the computational / numerical cost for allowing varying growth rates is in this framework. From my understanding both the deterministic and stochastic setting are extremely efficient, but it has thus far been numerically challenging to incorporate varying growth rates. How scalable is DeepRUOT? Can it be applied to the higher dimensional settings (20,50,100,1000D) in previous works?

I don’t understand the motivation behind the reconstruction loss R_d. What is its purpose and why does it correspond to a reconstruction loss vs. say an OT loss?

Comments:

\Psi(g) is not explicitly defined (although is fairly clear from context).

> So to specify an SDE is equivalent to specify the probability flow
ODE and its score function.

Probably missing an article here. 

### Overall

Overall this work has great potential, however I believe further empirical study would greatly improve its value to the community. I would be eager to reconsider my evaluation given further empirical study relative to existing unbalanced dynamic OT methods and a better empirical understanding of the various components of training and loss weighting.
soundness: 3
strengths: * Provided code provides additional clarity on experimental setting.
* Tackles an important problem in trajectory inference from population level data.
summary: This paper introduces DeepRUOT, a method for reconstructing dynamics from snapshots using regularized unbalanced optimal transport. This work tackles both the unbalanced and stochastic settings simultaneously, which few works have attempted thus far. An algorithm is proposed to learn drift, score, and growth functions from snapshot data including both synthetic and single-cell transcriptomics data.
weaknesses: * I found it extremely difficult to understand the multistage training procedure for DeepRUOT. There are some stages described in the main text, but then initial log density training also in the appendix? While there are a number of losses and initializations proposed to stabilize training, it is unclear to me what effects these have on final performance and how difficult RUOT is to train in practice. The paper would be much stronger if it was made clear which parts of the training procedure are important practically potentially through an ablation study on different components of the training procedure. 
* For a largely empirical work proposing a new algorithm to tackle this problem, the empirical studies are limited. The DeepRUOT is only benchmarked against balanced transport methods. It would greatly strengthen this paper if it was also compared to other unbalanced transport methods such as Action matching [Neklyudov et al. 2023] and unbalanced diffusion Schrödinger bridge [Pariset et al. 2023]. It would also be very useful to know how the many different weighting parameters affect performance, and how sensitive training is to them.
* It is difficult to tell what parts of the theory are novel and which are minor adaptations from prior work. Theorem 4.1 seems like a subclass of Baradat & Lavenant for the Fischer information case. I don’t believe varying sigma(t) over time changes anything theoretically and as far as I understand it is not used in practice.

Meta-review: This paper uses deep learning to learn stochastic dynamics from discretely observed data. They connect the RUOT problem to Schrödinger bridge, which reviewers found novel and in particular comes alongside a careful treatment of unbalanced effects. The methodology is novel and is also motivated by a compelling application to single-cell RNA data. This is a strong submission with several novel contributions

Decision: Accept (Oral)