Transport with Support: Data-Conditional Diffusion Bridges
Authors: Ella Tamir
Martin Trapp
Arno Solin
Venue: Accepted by TMLR ()

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.

Published in Transactions on Machine Learning Research (11/2023)
Transport with Support: Data-Conditional Diffusion Bridges
Ella Tamir
Gila.tamir@oolllo
Deportment of Computer Science
Aolto University
Martin Trapp
martin.trapp@ualto,
Department of Computer Science
Aalto Universiby
Arno Solin
arno.solin@aalto..
Deportment of Computer Science
Aolto University
Reviewed OpenReview: https: // otwprr/veeenrveeeeeemmmem net/ forum? id= 1309
Abstract
The dynamic Schrödinger bridge problem provides appealing setting for solving constrained time-series data generation tasks posed optimal transport problems. It consists of learning non-linear diffusion processes using efficient iterative solvers. Recent works have demonstrated state-offthe-arl results (e.g., 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 Itcrative 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.
1 Introduction
Generative diffusion models have gained increasing popularity and achieved impressive results in a variely 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 scries 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 then acts as a denoising diffusion model in finite time and is the closest. Kullback-Leibler (KT.) divergence the forward noising process of the noising model under marginal constraints. Data may then be generated by time reversal of the process, í.e., through the denoising process.
In many applications, the interest is not purely in modelling transport between initial and terminal state distribution. For example, in naturally occurring generative processos, we typically observe snapshots of realizations along intermediale slages 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-ste soe Fig. 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, 2022) modifying the prior processes (Fernandes et al., 2021). However,

---

Published in Transactions on Machine Learning Research (11/2023)
Time, E
Sparse observations
N((00,0) TT.))
N((10,0)
mo ~ (0,
N(O, I
(a) Unconstrained transport (Schrödinger bridge)
(b) Constrained transport
(Iterative smoothing bridge)
Figure 1: Illustrative example transport between an initial unit Gaussian and a shifted unit Gaussian at the terminal time T. Unconstrained transport on the left and the solution constrained by sparse observations (w) on the right. Colour coding of the initial points is only for distinguishing the paths.
defining meaningful fixed path constraints or prior processes for the optimal control problems can be challenging, while sparse observational data are accessible many real-world applications.
In this work, we propose the Ilerative Smoothing Bridge (ISB), an iterative method for learning a dynamical system for time-series data generation under constraints on both the initial and terminal distribution and sparse observational data constraints. The sparse observational constraints act way to encourage the paths sampled from the transport process to lic close to the observed data points. We perform the conditioning by leveraging the iterative pass idea from the Tterative Proportional Fitting procedure (IPFP, see Kullback, 1968; De Bortoli et al., 2021) and applying differentiable particle filtering (Reich, 2013; Corenflos al., 2021) within the outer loop. Integrating scquential Monte Carlo methods (e.g., Doucet et al., 2001; Chopin & Papaspiliopoulos, 2020) into the IPFP framework in such way is non-trivial and be understood general marginal constraints and additional path constraints defined by data.
novel iterative version of the algorithm by Maoutsa & Opper (2022) but with
We summarize the contributions as follows. (i) We propose a novel method for learning a dynamical model where the terminal constraints match a bridge problem and additional constraints are placed in form of sparse observations, inspired by optimal transport approaches. (ii) Thereof, we utilize the strong connections between the constrained bridging problem and particle filtering in sequential Monte Carlo, extending those links from pure inference learning. Additionally, (iii) we demonstrate practical efficiency and show that the iterative smoothing bridge approach scales to high-dimensional data.
1.1 Related Work
Schrödinger bridges The problem of learning a stochastic process moving samples from one distribution to another can be posed as a type of transport problem known as a dynamic Schrödinger bridge problem (SBP, e.g., Schrödinger, 1932; Léonard, 2014), where the marginal densities of the stochastic process are desired to resemble given reference measure. In machine learning literature, the problem has been studied through learning the drift function of the dynamical system (De Bortoli et al., 2021; Wang et al., 2021; Vargas et al., 2021; Bunne et al., 2022; Shi et al., 2023). When SDE system also defines the reference measure, the bridge problem bocomes constrained optimal control problem (c.g., Caluya & Halder, 2022; 2021; Chen et al., 2021; Liu et al., 2022), which has been leveraged in learning Schrödinger bridges by Tianrong Chen (2022) through forward- backward SDEs. Moreover, neural stochastic control has been studied Zhang et al. (2022). An optimal control problem with both initial and terminal distribution constraints and a fixed path constraint has been studied in Maoutsa et al. (2020) and Maoutsa & Opper (2022), where particle filtering is applied continuous path constraints but the boundary constraints are defined by single point. Maoutsa
2

---

Published in Transactions on Machine Learning Research (11/2023)
(2023) studies problem setting where stochastic dynamics are inferred based on sparse observation, but with geometric constraints rather than in the form of a Schrödinger bridge problem.
Diffusion models in machine learning The recent advances in diffusion models in machine learning literature have been focused on generating samples from complex distributions defined by data through transforming samples from an easy-to-sample distribution by a dynamical system (e.g., Ho et al., 2020; Song et al., 2021b;a; Nichol & Dhariwal, 2021). The concept of reversing SDE trajectories via score-based learning (Hyvärinen & Dayan, 2005; Vincent, 2011) has allowed for models scalable enough to be applied to digh-dimensional data sets directly in the data space. In earlier work, score-baseds diffusson models have been applied to problems where the dynamical system itself is of interest, for example, for the problem of time series amputation in Tashiro et al. (2021), inverse problems in imaging in Song et al. (2022) and for importance sampling Doucet et al. (2022). Interpreting the diffusion modelling problem as optimal control has recently been studied in Berner et al. (2022). Other dynamical models parametrized by neural networks have been applicd to modelling latent time-scrics based on observed snapshots of dynamics (Rubanova et al., 2019; Li et al., 2020), but without further constraints on the initial or terminal distributions.
State-space models In their general form, state-space models combine a latent space dynamical system with an observation (likelihood) model. Evaluating the latent state distribution based on observational data can be performed by applying particle filtering and smoothing (Doucet et al., 2000) or by approximations of the underlying state distribution of a non-linear state-space model by a specific model family, for instance, a Gaussian (see Särkkä, 2013, for an overview). Speeding up parameter inference and learning in state-space models has been widely studied (e.g., Schön et al., 2011; Svensson & Schön, 2017: Kokkala et al., 2014). Particle smoothing can be connected to Schrödinger bridges via the two-filter smoother (e-g., Bresler, 1986; Briers et al., 2009; Hostettler, 2015), where the smoothing distribution is estimated by performing filtering both forward from the initial constraint and backwards from the terminal constraint. We refer to Mitter (1996) and Todorov (2008) for a more detailed discussion on the connection of stochastic control and filtering and to Chopin & Papaspiliopoulos (2020) for an introduction to particle filters.
Data Assimilation methods Data assimilation (DA) methods leverage techniques from state-space literature to 'asssimiatte' observations into a mechanistic model in order to inform the model dynamics based mcasurements (e.g., Asch et al., 2016; Wang et al., 2000). Approaches based on DA have found widc-spread use in scientific applications and have been extended to incorporate sparse observationa data, for example, in numerical weather predictions (Whitaker et al., 2009), modelling cell state evolution in transitions (Mendez et al., 2020), or in oceanographic scenarios (Beiser et al., 2023). A crucial difference to our work is that DA relies on a precise mechanistic model while our approach is data-driven, providing additional flexibility modelling scenarios where formulating precise model dynamics not possible.
2 Background
Let C= T[ Rd denote the space of continuous functions from [0, to R2 and let B(C) denote the Borel o-algebra on C. Let (To, denote the space of probability measures on (C, B(C) such that the marginals at 0,T coincide with probability densitics mo and тт. respectively. The KI. divergence from measure Q to mcasure P is written as DKL IQ || IP], where we assume that Q < P. For modelling the time dynamics, we assume a (continuous-time) state-space model consisting of a non-linear latent Itó SDE (see, e-g., Øksendal, 2003; Särkkä & Solin, 2019) in [o, 7 x Rd with drift function fa(-) and diffusion function g(-), and a Gaussian observation model, i.c.,
Xo To, dx, = So () di + g()) dB.,
(1)
and Yı N(ya Xx, o² La) li=e, where the drift function fa 1R x 10, T + IRd is mapping modelled by a neural network (NN) parameterized by e θ, diffusion g [0, - R and 3r denotes standard d-dimensional Brown- ian motion. x, denotes the latent stochastic process and y. denotes the observation-space process. In practice, we consider the contimnous-discret time setting, where the process is observed at discrete time instances lg such that observational data can be given in terms of collection of onput-output, pairs {(1,) 1;))
3

---

Published in Transactions on Machine Learning Research (11/2023)
2.1 Schrödinger Bridges and Optimal Control
The Schrödinger bridge problem (SBP, Schrödinger, 1932; Léonard, 2014) is entropy-regullrized optimal transport problem where the optimality is measured through the KL divergence from a reference measure IP to the posterior Q, with fixed initial and final donsitics To and Ty, Le.,
min DKL [Q] (Fa.T))
(2)
In this work, consider only the case where the measures IP and Q are constructed as the marginals of an SDE, i.e., Qr is the probability measure of the marginal of the SDE in Eq. (1) at time t, whereas P, corresponds to the probability measure of the marginal of a reference SDE dx; = f(x,, t) de + g(t) dB,, at time 4, where we call ʃ the reference drift. Under the optimal control formulation of the SBP (Caluya & Halder, 2021) the KL divergence in Eq. (2) reduces to
E [V."
1 2g(1)2 (x),) - f(x,, O)| de
(3)
where the expectation is over paths from Eq. (1). Rüschendorf & Thomsen (1993) and Ruschendorf (1995) showed that solution to the SBP can be obtained by iteratively solving two half-bridge problems using the Iterative Proportional Fitting procedure (IPFP) for I = 0, l,...,,,. steps,
Q21+1 arg min DKL Q Qz] and
Q212 arg min DKL [QI Q2411, (E..(TO)))
(4)
(+57))
where Qo is set as the reference measure, and (TT..) .) and A(, TT) denote the sets of probability measures with only either the marginal at time or time T coinciding with xo or wy, respectively. Recently, the IPFP to solving Schrödinger bridges has been adapted as machine learning problem (Bernton el. al., 2019; Vargas et al., 2021; De Bortoli et al., 2021). In practice, the interval 7) is discretized and the forward drift fo and the backward drift by of the corresponding reverse-time process (Haussmann & Pardoux, 1986; Fölleer, 1988) are modelled by NNs. Under the Gaussian transition approximations, each step in the discrete-time diffusion model can be reversed by applying an objective based on mean-matching.
3 Methods
Given an initial and terminal distribution To and TT, we are interested in learning a data-conditional bridge between Ty and Ty. Let D = {(t,, be a set of M sparsely observed values, i.e., only a few or no observations are made at each point in time and let the state-space model of interest be given by Eq. (1). Note that we deliberately use (t), (instead of (fx. yx)) to highlight that we allow for multiple observations at the same time point tg. Our aim is to find parametcrization of the drift function fo such that evolving N particles xL, with xi ~ no (with = 1, ...... N), according to Eq. (1) will result in samples x from the terminal distribution тy. Inspired by the IPFP by De Bortoli et al. (2021), which decomposes the SBP into finding two half-bridges, we propose to iteratively solve two modified half-bridge problems where the additional sparse observations are accounted for simultancously. For this, let
dx; = S1,0(xe; 1) di + g(() dz, = bi (2c. 1) dt + g(t) dß.,
X To:
(5)
Zg NT,
(6)
denote the forward and backward SDE at iteration / - 1,2,.... L, where B, is the reverse-time Brownian motion. For simplicity, we denote 3, = B. when the direction of the SDE is clear.
To learn the Iterative Smoothing Bridge dynamics, itcratively employ the following steps: 1 evolve forward particle trajectories according to Eq. (5) with drift fi-1 and filter w.r.t. the observations ((4+)s)) J.M. 2 learn the drift function bi for the reverso-time SDE, 3 evolve backward particle trajectorics according to Eq. (6) with the drift by learned in step and filter w.r.t. the observations ()) and 4 learn the drift. function fio for the forward SDE based the backward particles. Fig. 2 illustrates the forward and
4

---

Published in Transactions on Machine Learning Research (11/2023)
mo
Forward
S
S
11
SS 8
Backward 3
Forward
S
S
11
Sparse observations
S
S
Backward
time f
t=0
Figure 2: Sketch of a diffusion bridge between 2D data distribution (ma) and an isotropic Gaussian (TT) constrained by sparse observations (x). The forward diffusion at the first iteration (ISB 1) learns to account for the sparse observations but does not converge to the correct terminal distribution (t = T), and the backward diffusion vice versa. After iterating (ISB 6), the forward and backward diffusions converge to the
correct targets and are able to account for the sparse observational data.
backward process of our iterative scheme for data-conditioned denoising diffusion bridge. Next, we will go through steps 4 in detail and introduce the Iterative Smoothing Bridge method for data-conditional diffusion bridges.
3.1 The Iterative Smoothing Bridge
The Iterative Smoothing Bridge (ISB) mothod iteratively generates particle filtering trajectories (steps and 3 in Fig. 2) and learns the parameterizations of the forward and backward drift functions fi,0 and b (steps and ) by minimizing modified version of the mean-matching objective presented by De Bortoli et al. (2021). Note that steps and 4 are dependent applying differential resampling in the particle filtering steps 1 and 3 for reversing the generated trajectorics. We will now describe the forward trajectory generating step and the backward drift learning step 2 in detail. Steps 3 and are given by application of and on their reverse-time counterparts.
Step (and 3 ): Given a fixed discrctization of the time interval [0, T] denoted as [te) with t, and tk- T, denote the time step lengths as A tk+1 t.. By truncating the Itô itô-Taylor series of the SDE, consider Euler-Maruyama (e.g., Ch. 8 in Särkkä & Solin, 2019) type of discretization for the continuous-time problem. We give the time-update of the ith particle at time b evolved according to Eq. (5),
i.e.,
X = x + ft-1, Ax + g(tx-1)) A El
(7)
where E ~ N(O, I). Notice that have not yet conditioned the observational data. In step 3, the particles 2 of the backward SDE Eq. (6) are similarly obtained. The SDE dynamics sampled in steps I and 3 apply the learned drift functions fi-1.0 and by from the previous step and do not require sampling from the underlying SDE model. For times l at which observations available, we set x² = x (and
5

---

Published in Transactions on Machine Learning Research (11/2023)
zi 1 respectively) and otherwise compute the particle filtering weights w/, based on the observations ((,,,,) E 114 = 4x) for resampling. See Sec. 3.2 for details on the particle filtering.
For resampling, we employ a differentiable resampling procedure, where the particles and weights (8i, wi,) are transported to uniformly weighted particles (xi,, x)) by solving an intropy-regularizes optimal transport problem (Cuturi, 2013; Peyré & Cuturi, 2019; Corenflos et al., 2021) (see App. D). Through application of the s-regularized optimal transport map (e) E RN*N (see Corenflos et al., 2021) the particles are resampled via the map to x = x (()), where X1, E R denotes the stacked particles {xf, Ji at time ty before resampling.
K,N Step 2 (and 4 4) Given the particles { k=1,==1 we now aim to learn the drift function for the respective reverse-time process. The purpose of this step is to find mean-matching reversal of the trajectories, in other words we aim to find fi.e such that it best explains the change each trajectory i = 1, 1.2...., N and particle h = 1, I......,N. .N. optimization step, where the loss outside observation times will match the approach in De Bortoli et al. (2021) and the loss at observation times is motivated by a smoothing of trajectories (see App. A.2 for a discussion).
observe from {xi,} } to (x J for We will review the loss functions used for this
In case no observation is available at time tk. we apply the mean-matching loss based on Gaussian transition approximation proposed in De Bortoli et al. (2021):
no = |6 (x) tx+1) - Xint1 - fi-1.0(x 4s) + xi, + fi-1,(x ,t + Aall
(8)
In case observation is available at time 4 the particle values X will be coupled through the optimal transport map. Therefore, the transition density is a sum of Gaussian variables (see App. A for details and a derivation), and the mean-man meaamiig loss is therefore given by:
4+1,obs =
- X - -f-110 4.) + EW Teeaa (x%, + f (x) 112. (9)
The derivation of Eq. (9) relies on mean-matching the trajectories as in De Bortoli et al. (2021) combined with applying the differentiable resampling optimal transport map T = xt, on all the particles to obtain the transition density Px,, Ix|,_ (Xx.) at observation times, resulting in Gaussian distribution
dependent on all the particles. In addition, we apply the property that the reverse drift should satisfy
(x 1x+1) 10-1,00X11 tx) g(tx+1) Pt241:
(10)
where P is the particle filtering density after differential resampling at time tx+1. Thus the impact of observations to the loss function is two-fold, the observations define the value of the transport matrix T(e) and the marginal score V Pt The use of the reverse drift Eq. (10) is further motivated by the smoothing reverse presented in App. A.2, where we discuss how Eq. (10) matches a backwards controlled drift.
The overall objective function is a combination of both loss functions, with the respective mean- mean-matching loss depending on whether l is an observation time. The final loss function is written as:
N
K
(())
Σ
#0 + E2 no ( C
y., =0
(11)
where Icoud. denotes indicator function that returns '1' iff the condition is true, and '0' otherwisc. Consequently, the parameters é of bi, are learned by minimizing Eq. (11) through gradient descent. In practice, a cache of trajectorics (xi, 1 is maintained through training of the drift functions, and refreshed at fixed number of inner loop iterations, as in De Bortoli et al. (2021), avoiding differentiation over the SDE generation computational graph. The calculations for step 4 follow similarly. We present high-level description of the ISB steps in Alg. 1.
K.N
The learned backward drift bis can be interpreted as an analogy of the backward drift in Maoutsa & Opper (2022), connecting our approach to solving optimal control problems through Hamilton-Jacole equations, see App. A.2 for an analysis of the backwards SDE and the control objective. While we are generally considering problem settings where the number of observations is low, we propose that letting M + oo yields the underlying marginal distribution, see Prop. 2 in App. A.3.
6

---

Published in Transactions on Machine Learning Research (11/2023)
Algorithm The Iterative Smnothhng Bridge
Input: Marginal constraints (To, kr), observations {1) initial drift function fo.o, iterations L, discretization steps K, number of particles N, observation noise schedule k(I)
Output: Learned forward and backward drift bo) for - to L do
Initialize forward particles {xb)i=1 wo
for k: == to do Generate { using (xi-1h=1 J if Observations at is then {x DiffResample
Backward process Forward process
Eq. (5)
end if end for
Optimize the forward loss function w.r.t. Φ Initialize backward particles {2k J Wr
b Eq. (11)
for h b= K to 1 do 2 using { if Observations at & then (xk-1) W. DiffResample ( (())
Generate 1
Eq. (6)
end if
end for
Optimize the backwards loss function w.r.t.
> Eq. (15)
end for
3.2 Computational Considerations
The ISB algorithm is a generic approach to learn data-conditional diffusion bridges under various choices of, e.g., the particle filter proposal density or the reference drift. Next, we cover practical considerations for the implementation of the method and highlight the model choices in the experiments.
Multiple observations per time step Naturally, we can make more than one observation in time fk, denoted as Dr, y(t,,y,) E Dit, ) To compute particle weights wh, for the ith particle we consider only the H-nearest neighbours of x in Dr, instead of all observations in Dr.- By restricting to the H-nearest neighbours, denoted as D we introduce an additional locality to the proposal density computation, which can be helpful in the case of multimodality. On the other hand, letting H > 1 results in weights which take into account the local density of the observations, not only the distance to the nearest, neighbour. In experiments with few observations, we set H = 1, the choice of H is discussed when we have set the value higher.
single point
Particle filtering proposal The proposal density chosen for the ISB is the bootstrap filter, where the proposal matches the Gaussian transition density p(X², x(2_1). Assuming a Gaussian noise model N(0, o²I), the unnormalized log-weights for the 7th particle at time le given by log wt, 11/20² Σyyepi ||x]_ 11
While restrict our approach in practice to the bootstrap filter, applying more sophisticated proposals such as in the auxiliary particle filter (Pitt & Shephard, 1999) could improve the results in some cases, although restricting the problem to a linear observation model.
Observational noise schedule In practice, using a constant observation noise variance x² can result an itcrative scheme which does not have stationary point as L-+ oc. Even the Iearned drift function fL, was optimal, the filtering steps 1 and 3 would alter the trajectories unless all particles would have uniform weights. Thus, introduce noise schedule k(l) which ensures that the observation noise increases in the number of ISB iterations, causing ISB to converge to the IPFP (De Bortoli et al., 2021) as L - oc. We found that letting the observation noise first decrease and then increase (in the spirit of simulated annealing) often performmd strictly increasing observation noise schedule. The noise schedule is studied in App. C, where we derive the property that letting L + xo yields IPFP.
Drift initialization Depending on the application, one may choose to incorporate additional information by selecting appropriate initial drift. A possible choice includes pre-trained neural network drift learned to
7

---

Published in Transactions on Machine Learning Research (11/2023)
1=0
1/44
t t-1/2
t-4
t=T
Figure 3: 2D toy experiments from scikit-learn with both cases starting from a Gaussian: The TWO CIRCLES (top) and TWO MOONS (bottom) data scts, with observations (red markers) constraining the problem. For the circles, the 10 circular observations at t 0.5 first force the method to create circle that then splits into two; in the lower plot the observations at t E [0.25, 0.5, 0.75) split the data into clusters before joining them into two moons. See Fig. 6 in the Appendix for the IPFP result.
transport To to Æт without accounting for observations. However, starting from drift for the unconstrained SBP can be problematic in cases where the observations are far away from the unconstrained bridge. To encourage exploration, one may choose fo = 0 for the initial drift. In various problem settings, we found a zero drift and starting from the SBP to be successful in the experiments. See App. C for discussion.
4 Experiments
To assess the properties and performance of the ISB, we present range of experiments that demonstrato how the iterative learning procedure can incorporate both observational data and terminal constraints. We start with simple examples that build intuition (cf. Fig. 1 and Fig. 2) and show standard ML benchmark tasks. For quantitative assessment, we design experiment with non-linear SDE for which the marginal distributions are available in closed-form. Finally, we demonstrate our model both in a highly multimodal bird migration task, conditioned image generation, and in a single-cell embryo RNA modelling problem. Ablation studies are found in App.
Experiment setup In all experiments, the forward and backward drift functions fo and by are parametrized neural networks. For low-dimensional experiments, we apply the MLP block design in De Bortoli et al. (2021), and for the image experiment an U-Net as in Nichol & Dhariwal (2021). The latent state SDE was simulated by Euler- Maruyama with a fixed lime-step of 0.01 over 100 steps and 1000 particles if not otherwise stated. All cow-dimensional (at most d = 5) experiments were run on a MacBook Pro laptop CPU, whereas the image experiments used a single NVIDIA A100 GPU and ran for 5 10 min. Notice that since ISB only performs particle filtering outside the stochastic gradient training loop, the training runtime in the same order as the earlier Schrödinger bridge image generation experiments of De Bortoli et al. (2021). Thus we omit any wall-clock timings. Full details for all the experiments are included in App. B.
All experiment settings include mumber of hyperparameter choices, some typical to all diffusion problems and some specific to particle filtering and smoothing. The diffusion g(t) is pre-determined function not optimized during training. We divide the experiments into two main subsets: problems of 'sharpening to achieve a data distribution' and 'optimal transport problems'. In the former, the initial distribution has a support overlapping with the terminal distribution and the process noise level g(t) goes from high to low as time progresscs. Conversely, in the latter setting, the particles sampled from the initial distribution must travel reach the support of the terminal distribution, and we chose to use a constant process noise level. Perhaps the most significant choice of hyperparameter is the observational noise level, as it imposes preference on how closely should the observational points be followed, see App. C.1 for details.
2D toy examples We show illustrative results for the TWO MOONS and CIRCLES from scikit-learn. We add ar- tificial observation data to bias the processes. For the circles, the observational data consists of 10 points, spaced evenly on the cirele. The points are all observed simultancously, at halfway through the process, forcing the marginal density of the generating SDE collapse to the small circle, and then to expand. For the TWO MOONS, the observational data is collected from 10 trajectories of diffusion model, which generates the two moons from
8

---

Published in Transactions on Machine Learning Research (11/2023)
Summer
Example sightings during migration
Winter
Bird observations
ISB rosult
Figure 4: Bird migration example. The top row describes nesting and wintering areas and example sightings during migration. The bottom shows the marginal densities of the ISB model from the initial to terminal distribution, matching bird sightings along the migration.
noise, and these trajectorics are then observed at throe points time. Results are visualizod in Fig. 3 (sce videos in supplement). For reference, have included plots of the IPFP dynamics in the supplement, see Fig. 6.
Quantitative comparison the Beneš SDE In order to quantify how observing process in between its initial and terminal states steers the ISB model to areas with higher likelihood, test its performance Beneš SDE model (see, e.g. Särkkä & Solin, 2019). The Beneš SDE is a non-linear one-dimensional SDE of form dz; = tanh(z;) di + d3; with 3o = O, but its marginal density is available in closed-form, allowing for negative ligliikl evaluation. We simulate trajectories from the Beneš SDE and from the reverse drift and stack the reversed trajectorics. The terminal distribution is shifted and scaled so that the Beneš SDE itself does not solve the transport problem from To to æT, see App. B.2 for details and visualizations of the procossos.
We fit Schrödinger bridge model with no observational data a baseline, using the Beneš SDE drift as the reference model. The ISB model is initialized with a zero-drift model (not with the Beneš as reference). thus making learning more challenging. We compare the models in terms of negative log prodictive density Table 1, where we see that the ISB model captures the process well on average (over the entire time-horizon) and selected marginal times.
regular scasonal transport problem, where birds move (typically North-Southe along a flyway, between breeding and wintering grounds. We take this motivating example of constrained optimal transport, where the geographical and constraints and preferred routes are accounted for by bird sighting data (see Fig. 4 top). By adapting data from Ambrosini et al. (2014) and Pellegrino et al. (2015), we propose a simplified data set for geese migration in Europe (OIBMD: ornithologically implausible bird migration data; available the supplement). We applied the ISB for iterations, with a lincar observation noise schedule from 0.2, and constant diffusion noise 0.05. The drift function initialized as pero-function, and thus the method did not rely separately fit model optimized for generating the wintering distribution based on the breeding distribution. For comparison, we include the Schrödinger bridge results in App. B.3.
Bird migration Bird migration be
Constraining an image generation process We demonsttrle that the ISB approach scales well to high-dimensional inputs by studying a proof-of-concept image generation task. We modify the diffusion generative process of the MNIST (LeCun et al., 1998) digit 8 by artificial obscrvations steering the dynamical system the middle of the generation process. While the concept of observations in case of image generation is somewhat unnatural, it showcases the scalability of the method to high-dimensional data spaces. Here, the
9

---

Published in Transactions on Machine Learning Research (11/2023)
Table 1:
Results for the Benes experiment. We report the negative log predictive density (NLPD, lower better) of the Benes marginal likelihood over generated particles at the initial and terminal distri- butions and at the middle of the transport process.
Table 2: Results for the single-cell embryo RNA exper-
iment. We compare ISB to TrajectoryNet, IPML, and our implementation of IPFP. Unlike the other meth- ods, our model is able to utilize the intermediate data distributions while training.
NLPD
Earth mover's distance
МETROD
AVERAGE
MIDDEE
END
METHOD
t=
t=1
1=2
1=3
[=T
Schrödinger B
4.787
3.565
0.1919
TrajectoryNet
0.62
1.15
1.49
1.26
0.99
Iterative smoothing B
3.557
2.985
0.1567
IPML
0.34
1.13
1.35
1.01
0.49
IPFP (no obs)
0.57
1.53
1.86
1.32
0.85
ISB (single-coll obs)
0.57
1.04
1.24
0.94
0.83
drift is initialized using a pre-trained neural network obtained by first running a Schrödinger bridge model for image generation. The process is then given an observation in the form of a bottom-half of MNIST digit 8 in the middle of the dynamical process. As the learned model uses information from the observation both before and after the observation time, the lower half of the image is sharper than the upper half. We provide further details on this experiment and sampled trajectories in App. B.4, and an ablation of multi-modal MNIST image generation in App. B.5
Single-cell embryo RNA-seq Lastly, evaluated our approach an Embryoid body scRNA-seq time course (Tong et al., 2020). The data consists of RNA measurements collected over five time ranges from developing human embryo system. No trajectory information is available, instcad we only have access to snapshots of RNA data. This leads to a data set over 5 time ranges, the first from days 0-3 and the last from days 15-18. In the experiment, we followed the protocol by Tong et al. (2020), reduced the data dimensionality to d 5 using PCA, and used the first and last time ranges as the initial and terminal constraints. All other time ranges are considered observational data, Contrary to the other experiments, intermediate data are imprecise (only time range of multiple days is known) but abundant.
We learned the ISB using a zero drift and compared it against an unconditional bridge obtained through the IPFP (De Bortoli et al., 2021)-see Fig. 5. The ISB learns to generate trajectories with marginals closer to the observed data while performing comparably to the IPFP at the initial and terminal stages. This improvement is also verified numerically in Table 2, showing that the ISB obtains a lower Earth mover's distance between the generated marginals and the observational data than IPFP. Additionally, Table lists the performance of previous works that do not use the intermediate data during training (Tong et al., 2020) only use it to construct an informative reference drift (Vargas et al., 2021), see App. B.6 for details. In both cases, ISB outperforms the other approaches w.r.t. the intermediate marginal distributions (t (t=1 1, 1,2,3), while IPML (Vargas et al., 2021) outperforms ISB at the initial and terminal stages due to its data-driven reference drift. Notice that while we reduced the dimensionality via PCA to 5 for fair comparisons to Vargas et al. (2021), the ISB model would also allow modelling the full state-space model, with observations in the high-dimensional gene space and a latent SDE.
5 Discussion and Conclusion
The dynamic Schrödinger bridge problem provides an appealing setting for posing optimal transport problems as learning non-linear diffusson processes and enables efficient iterative solvers. However, while recent works have wrate-of-the-art performance in many complex application domains, they are typically limited to learning bridges with only initial and terminal constraints dependent on observed data. In this work, we have extended this paradigm and introduced the Iterative Smoothing Bridge (ISB), an itcrative algorithm for generating data-condiiional smoothing bridges. For this, we leveraged the strong connections between the constrained bridging problem and particle filtering in sequential Monte Carlo, extending them from pure inference to learning. We thoroughly assessed the applicability and flexibility of our approach in various experimental settings, including synthetic data sets and complex real-world scenarios (e.g., bird migration, conditional image generation, and modelling single-cell RNA RNA-sequencing time-series). Our experiments showed that ISB
10

---

Published in Transactions on Machine Learning Research (11/2023)
Psincipal axte #1
Den
PA #2
(a) Schrödinger bridge (via IPFP)
(b) Iterative Smoothing Bridge
Figure 5: Illustration of the trajectorics of the single-cell experiment for the Schrödinger bridge (a) and the ISB (b), projected onto the first two principal components. The first five trajectories are highlighted in colour, and intermediate observation densities visualized as slices.
generalizes well to high-dimensional data, computationally efficient, and provides accurate estimates of the marginals at initial, terminal, and intermediate times.
Accurately modelling the dynamics of complex systems under both path constraints induced by sparse observations and initial and terminal constraints is key challenge in many application domains. These include biomedical applications, demographic modelling, and environmental dynamics, but also machine learning specific applications such as reinforcement learning, planning, and time-series modelling. All these applications have in common that the dynamic nature of the problem is driven by the progression of time, and not only the progression a generative process as often is the case it, e.g., generalive image models. Thus, constraints over intermediate stages have a natural role and interpretation in this wider set of dynamic diffusion modelling applications. We believe the proposed ISB algorithm opens up new avenues for diffusion models in relevant real-world modelling tasks and will be stimulating for future work. Recent work suggest close connections between sequential DA, commonly applied in many real-world scientific domains, and the Schrödinger bridge problem (Reich, 2019) further emphasising the potential for future scientific applications of our work by exploiting these links. Morcover, in the future more sophisticated observational models, alternative strategies to account for multiple observations, and different noise schedules could be explored. Lastly, the proposed approach could naturally be extended to other types of optimal transport problems, such as the Wasserstein barycenter, frequently employed case of the multi-marginal optimal transport problem.
A reference implementation of the ISB model can be found at https://github. wwt.hub..cc//aaattM..... sherativessmoohhngg bridge.
Acknowledgements and Disclosure of Funding
Authors acknowledge funding from the Academy of Finland (grants 339730, 324345, and 347279). We also acknowledge the computational resources provided by the Aalto Science-IT project and CSC - IT Center for Science, Finland. We wish to thank Adrien Corenflos for sharing an implementation of differentiable resampling in Py Torch, and Prakhar Verma for comments on the manuscript.
References
Roberto Ambrosini, Riccardo Borgoni, Diego Rubolini, Beatrice Sicurella, Wolfgang Fiedler, Franz Bairlein, Stephen R. Baillie, Robert A. Robinson, Jacquie A. Clark, Fernando Spina, and Nicola Saino. Modelling the progression of bird migration with conditional autoregressive models applied to ringing data. PLoS ONE, 9(7):1-10, 07 2014.
Mark Asch, Mare Bocquet, and Maëlle Nodet. Data Assimilation: Methods, Algorithms, and Applications. STAM, 2016.
11

---

Published in Transactions on Machine Learning Research (11/2023)
Florian Beiser, Håvard Heitlo Holm, and Jo Eidsvik. Comparison of ensemble- data assimilation methods for sparse oceanographic data. orViv preprint or 0wXi:2330.071197, 2023.
Julius Berner, Lorenz Richter, and Karen Ulrrich. An optimal control perspective on diflusion-based generative modeling. In NewrIPS 2022 Workshop on Score-Based Methods, 2022.
Espen Bernton, Jeremy Heng, Arnaud Doucet, and Pierre E. Jacob. Schrödinger bridge samplers. arXiv preprint Xiv: 1917,70,,, 2019.
Yoram Bresler. Two-filter formulae for discrete-time non-linear Bayesian smoothing. International Journal of Control, 43(2))629 641, 1986.
Mark Bricrs, Arnaud Doucet, and Simon Maskell. Smoothing algorithms for state-space models. Annals of the Institute of Statislical Malhemalics, 62:61 89, 02 2009.
Charlotte Bunne, Ya-Ping Hsieh, Marco Cuturi, and Andreas Krause. Recovering stochastic dynamics via Gaussian Schrödinger bridges. arXiv preprint arXiv:2202 05722, 2022.
Kenneth F. Caluya and Abhishek Halder. Reflected Schrödinger bridge: Density control with path constraints. In American Control Conference, ACC 2021, pp. 1137-1142. IEEE, 2021.
Kenneth F. Caluya and Abhishek Halder. Wasserstein proximal algorithms for the Schrödinger bridge problem: Density control with nonlinear drift. IEEE Transactions on Automatic Control, 67(3):1163 1178, 2022.
Yongxin Chen, Tryphon T. Georgiou, and Michele Pavon. Stochastic control liaisons: Richard Sinkhorn meets Gaspard Monge on a Schrödinger Bridge. SIAM Review, 63(2):249 313, 2021.
Nicolas Chopin and Omiros Papaspiliopoulos. An Introduction to Sequential Monte Carlo. Springer, 2020.
Adrien Corenflos, James Thornton, George Deligiannidis, and Arnaud Doucet. Differentiable particle filtering via entropy-regularized optimal transport. In Proceedings of the 38th International Conference on Machine Learning, ICML, 2021, volume 139 of Proceedings of Machine Learning Research, pp. 2100 2111. PMLR, 2021.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In Advances in Newral Information Processing Systems 26 (NIPS), pp. 2292 2300. Curran Associates, Inc., 2013.
Valentin De Bortoli, James Thoruton, Jeremy Heng, and Arnaud Doucet. Diffusion Schrödinger bridge with applications to score-based generative modeling. In Advances in Neural Information Processing Systems 34 (NeurIPS), pp.
Curran Associates, Inc., 2021.
Prafulla Dhariwal and Alexander Quinn Nichol. Diffusion models beat GANs on image synthesis. In Advances in Neural Information Processing Systems 35 (NeurIPS), pp. 8780 8794. Curran Associates, Inc., 2021.
Arnaud Doucet, Simon Godsill, and Christophe Andrieu. On sequential Monte Carlo sampling methods for Bayesian filtering. Statistics and Computing, 10(3):197 208, 2000.
Arnaud Doucet, Nando De Freitas, and Neil James Gordon. Sequential Monte Carlo methods in practice. Statisties for Engineering and Information Science. Springer, 2001.
Arnaud Doucet, Will Sussman Grathwohl, Alexander G. D. G. Matthews, and Heiko Strathmann. Score- based diffusion meets annealed importance sampling. In Advances in Neural Information Processing Systems 35, 2022.
David Lopes Fernandes, Francisco Vargas, Carl Henrik Ek, and Neill DF Campbell. Shooting Schrödinger's cat. In Proceedings of the Fourth Symposium on Advances in Approzimate Bayesion Inference (AABI), 2021.
Hans Föllmer. Random fields and diffusion processes. In École d' d'Ééé de Probabilités de Caint-Floun XV-XVII, 1985-81 87, pp. 101-203 Springer, 1988.
12

---

Published in Transactions on Machine Learning Research (11/2023)
Uhrich G Haussmann and Etienne Pardoux. Time reversal of diffusions. The Annals of Probability, pp- 1188-1205, 1986.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advonces in Neural Information Processing Systems 33 (NeurIPS); pp. 0840-6851 Curran Associates, Inc., 2020.
Roland Hostettler. A two filter particle smoother for Wiener state-space systems. In 2015 IEFE Conference on Control Applications (CCA), pp. 412 417, 2015.
Aapo Hyvärinen and Peter Dayan. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(4):695-709, 2005.
Michael Janner, Yilun Du, Joshua B. Tenenbaum, and Sergey Levine. Planming with diffusion for flexible behavior synthesis. In Kamalika Chaudhuri, Stefanic Jegelka, Le Song, Csaba Szepesvári, Gang Niu, and Sivan Sabato (eds.), International Conference Machine Learning (TCML), volume 162, pp. 9902-99029000 -915. PMLR, 2022.
Juho Kokkala, Arno Solin, and Simo Särkkä. Expectation maximization based parameter estimation by sigma-point and particle smoothing. In Proceedings of the 1 17th International Conference Information Fusion (FUSION), pp. 1-8, 2014.
Solomon Kullback. Probability densities with given marginals. The Annals of Matheraatiaal Statistics, 39(4): 1236-1243) 1968.
Yann LeCun, Corinna Cortes, and Christopher J.C. Burges. The MNIST database of handwritten digits, 1998. URL http://yann... http:///aannneeun....mm www./ooo//ee/eeeeeee //////ddbmmmistt
Xuechen Li, Ting-Kam Leonard Wong, Ricky T. Q. Chen, and David Duvenaud. Scalable gradients for stochastic differential equations. In Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statistics (AISTATS), volume 108 of Proceedings of Machine Learning Research, pp. 3870 3882. PMLR, 2020.
Guan Guan-Horng Liu, Tianrong Chen, Oswin So, and Evangelos Theodorou. Deep generalized schrödinger bridge.
In Advances in Neural Information Processing Systems 35, 2022.
Christian Léonard. survey of the Schrödinger problem and some of its connections with optimal transport. Discrete & Continuous Dynomical Systems, 34(4):1533 1574, 2014.
Dimitra Maoutsa. Geometric constraints improve inference of sparsely observed stochastic dynamics. arXiv preprint wrXi:2230.00422 2023.
Dimitra Maoutsa and Manfred Opper. Deterministic particle flows for constraining stochastic nonlinear systems. Physical Review Research, 4(4), 2022.
Dimitra Maoutsa, Scbastian Reich, and Manfred Opper. Interacting particle solutions of Fokker-Planck equations Uhrough gradient log density estimation. Entropy, 22:802, 07 2020.
Mario J Mendez, Matthew J Hoffman, Elizabeth M Cherry, Christopher A Lemmon, and Scth H Weinberg. assimiaaiiiiation approach to predict epithelial mesenchymll transition. Biophysical
Cell fate forecasting: Journal, 118(7):1749-1768 2020.
S.K. Mitter. Filtering and stochastic control: a historical perspective. IEEE Control Systems Mogazine, 16 (3):67-76, 1996.
Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. Tn Proceedings of the 38th International Conference Machine Learning (ICML), volume 139 of Proceedings of Machine Learning Research, pp. 8162-8161 8171. PMLR, 2021.
Bernt Oksendal. Stochastic Differential Equations: An Introduction with Applications. Springer, New York,
NY, sixth edition, 2003.
13

---

Published in Transactions on Machine Learning Research (11/2023)
Sung Woo Park, Kyungjae Lee, and Junseok Kwon. Neural markov controlled SDE: for continuous-time data. In International Conference on Learning Represenlations (TCLR), 2022.
Stochastic optimization
Irene Pellegrino, Marco Cucco, Arne Follestad, and Mathieu Boos. Lack of genetic structure in greylag goose (Anser anser) populations along the European Atlantic flyway. PeerJ, 3:e1161, 2015.
Gabriel Peyré and Marco Cuturi. Computational optimal transport. Foundations and Trends in Machine Learning, 11(5-)):355 607, 2019.
Michael K. Pitt and Neil Shephard. Filtering via simulation: Auxiliary particle filters. Journal of the American Statistical Association, 94(446):590-55) 1999.
Kashif Rasul, Abdul-Saboor Sheikh, Ingmar Schuster, Urs M. Bergmann, and Roland Vollgraf. Multivariate probabilistic time scries forecasting via conditioncd normalizing flows. In International Conference on Learning Represenlations (ICLR), 2021.
Sebastian Reich. A nonparametric ensemble transform method for Bayesian inference. SIAM Journal Scentific Compuling, 35, 2013.
Sebastian Reich. Data assimilation: The Schrödinger perspective. Acta Numerica, 28:635-711, 2019.
Yulia Rubanova, Tian Qi Chen, and David K Duvenaud. Latent ordinary differential equations for irregularly- sampled time series. In Advances in Neural Information Processing Systems 32 (NeurIPS), pp. 5321-5331. Curran Associates, Inc., 2019.
Ludger Ruschendorf. Convergence of the iterative proportional fitting procedure. The Annals of Statistics, 23 (4):1160-1174, 1995.
L. Rüschendorf and W. Thomsen. Note the Schrödinger equation and I-projections. Statistics & Probability Letters, 1695)3369-375 1993.
Simo Särkkä. Bayesian Filtering and Smoothing. Cambridge University Press, Cambridge, UK, 2013.
Simo Särkkä and Arno Solin. Applied Stochastic Differential Equations. Cambridge University Press, Cambridge, UK, 2019.
E. Schrödinger. Sur la théorie relativiste de l'électron et I'interprétation de la mécanique quantique. Annales de l'institut Henri Poincaré, 2(4):269 310, 1932.
Thomas B. Schön, Adrian Wills, and Brett Ninness. System identification of nonlinear state-space models. Automatica, 47(1):39 49, 2011.
Yuyang Shi, Valentin De Bortoli, George Deligiannidis, and Arnaud Doucet. Conditional simulation using diffussin Schrödinger bridges. In 38th Conference on Uncertainly in Arlificial Intelligence. UAI, 2022.
Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusion schrödinger bridge matching. Xixiv preprint arXim:2303. 16852, 2023.
Yang Song, Conor Durkan, lain Murray, and Stefano Ermon. Maximum likelihood training of score-basod diffusion models. In Advances in Neural Information Processing Systems 35 (NeurIPS), pp. 1415-1428. 1428. Curran Associates, Inc., 2021a.
Yang Song, Jascha Sohl-Dicsstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference Learning Representations (ICLR), 2021b.
Yang Song, Liyue Shen, Lei Xing, and Stefano Ermon. score-based generative models. In International Conference on Learning Representations (ICLR), 2022.
Solving inverse problems in medical imaging with
Andreas Svensson and Thomas B. Schön. A flexible state-space model for learning nonlinear dynamical systems. Automaticn, 80:189-198 199, 2017.
14

---

Published in Transactions on Machine Learning Research (11/2023)
Yusuke Tashiro, Jiaming Song, Yang Song, and Stefano Ermon. CSDI: Conditional score-based diffusion models for probabilistic time series imputation. In Advonces in Neural Information Processing Systems 35 (NeurIPS), pp. 24804 24816. Curran Associates, Inc., 2021.
Evangelos A. Theodorou Tianrong Chen, Guan-Horng Liu. Likelihood training of schrödinger bridge using forward-backward SDEs theory. In International Conference Learning Representations (ICLR), 2022.
Emanuel Todorov. General duality between optimal control and estimation. In Proceedings of the 47th IEEE Conference on Decision and Control, pp. 4286-4292, 2008.
Alexander Tong, Jessie Huang, Guy Wolf, David van Dijk, and Smita Krishnaswamy. Trajectorynet: A dynamic optimal transport network for modeling cellular dynamics. In ICML, volume 119 of Proceedings of Machine Learning Research, pp. 9526 9536. PMLR, 2020.
Francisco Vargas, Pierre Thodoroff, Austen Lamacraft, and Neil Lawrence. Solving Schrödinger bridges via maximum likelihood. Entropy, 23(9):1134, 2021.
Pascal Vincent. A connection between score matching and denoising autoencoders. Neurol Computation, 23 (7):1661-167,, 2011.
Bin Wang, Xiaolei Zou, and Jiang Zhu, Data assimilation and its applications. Proceedings of the National Academy of Sciences, 97(21):11144 11144, 2000.
Gefei Wang, Yuling Jiao, Qian Xu, Yang Wang, and Can Yang. Deep generative learning via Schrödinger bridge. In Proceedings of the 38th International Conference on Machine Learning (ICML), volume 139 of PMLR, 2021.
Proceedings of Machine Learning Research, pp.
Jeffrey S Whitaker, Gilbert P Compo, and Jean-Noël Thépaut. A comparison of variational and ensemble- based data assimilation systems for reanalysis of sparse observations. Monthly Weather Review, 137(6): 1991 1999, 2009.
Jingdong Zhang, Qunxi Zhu, and Wei Lin. Neural stochastic control. In Advances in Neural Information Processing Systems 85, 2022.
15

---

Published in Transactions on Machine Learning Research (11/2023)
A
Method Details
We present the details of the objective function derivation in App. A.1 and explain the connection of the backward drift function to Hamilton Jacobi equations in App. A.2. In App. A.3, we discuss the behaviour of our model at the limit M - oc, that is, when the observations fully represent the marginal densities of the stochastic process.
A.1 Deriving the Mean-matching Loss at Observation Times
Recall that the forward loss is written as
N
K
(
Σ
Σ
Rohs
( + naba (o) (o) -0
(12)
where the loss at observations & k,o (Ø) and loss elsewhere nobe (Ø) are
1k -101,0(x ( A+t1)A Ximm 11 Me + xix + St (x
(13)
Lobss 1b tx+1) xhats f1-1,(xx ttt)a
1
+ Em=I (x%, (x-10(x),, (x) 11?, (14)
C..+
For convenience, we state the backward loss functions which follow similarly to their forward versions. The backward loss defined as
100)
(0) 21 + 2 (O)o
(15)
where the loss at observations T kobo and loss elsewhere & kk (0) are
T T1 o I xtan b1,0 tx) + +x², (x, :x tx?
(16)
1 = 11 (++1) x(u+1 bi bL 4)4x
+ Cei Em Tecli (x%, + bx, (x) 4x) 112. (17)
Proposition 1. Define the forward SDE as
dx; = fi,e(xe, g(t) dBr,
X0 no.
(18)
and a backward SDE drift as
( (x+1) f1-1,0 (x ( g(tx+1) Peates
(19)
where Peuti is the particle filtering density after differential resampling at time 1441. Then bi, (X 1x+1) minimizes the loss function
C001, = 11br, / lx+1) Ar X:.1: - St 1 1111"
(x 1x
DN Tersa (x), + fx- II?, (20)
+ Ce,
where we denote Ce. 0(tx11)=a; Var ENN (n and {Kt JN are the particles before resam- pling.
16

---

Published in Transactions on Machine Learning Research (11/2023)
Proof sketch. Our objective is to find a backward drift function (x, 1e+1) as in Eq. (19). Notice that at observation times la, this is not equivalent to finding the reverse drift of the SDE forward transition and differential resampling combined, since the drift function fi-18 alone does not map the particles {xl, Ji to the particles fxt 1.V. We will derive a loss function for learning the backward drift as in Eq. (19) below, leaving the discussion on why it a meaningful choice of backward drift to App. A.2. Our derivation closely follows the proof of Proposition sake of completeness.
De Bortoli et al. (2021); but we provide the details here for the
(xx) and apply it to derive the observation time loss kkobb The derivation for the loss f, obs is skipped since it is as in the proof of Proposition 3 in De Bortoli et al. (2021). Suppose that at lx, there are observations. By definition, the particles before resampling (xt M are generated by the Gaussian transition density
First, give the transition density Px., Ixi,
pl pl
x(x)) = N( xt, 6,Ji(x), (4) (A
(21)
Recall that the resampled particles are defined weighted average of all the particles, x = X7, Te)
Thus, the transition density from {xt J to the particles (x) /1 JM is also Gaussian,
N
p(Xit.n xxx,) = N(R Σ
Ter + An.fi-1 g(tx+1)
(22)
We will derive the loss function Eq. (9) by modifying the mean matching proof in De Bortoli et al. (2021)
by the transition mean Eq. (22) and the backward drift dcfinition Eq. (19). Using the partiele filtering
approximation, the marginal density can be decomposed as Ptuta (xx=)) = Σ. Pta (xt) xx (xx+1). By substituting the transition density Eq. (22) it follows that
(x)) + fi- (x-1,((x0x Il (E N
Xiats )
1 Σ Pte (x({ exp
Phuta (xte1)
(23)
7
2g(te+1)&Cr,A4
where Z is the normalization constant of Eq. (22). As in the proof of Proposition 3 of De Bortoli et al. (2021), derive expression for the score function. Since V In Ptats (Xta+1)= Valees Prezi (N ) first
(
manipulate V Vxi 11 Pez+i (X4241):
( [E EN Tessaa + 2g(tx+1)
Xix 1l2
1
V Vxtg| 12 Pla+s (X1e+1)
Σ
V
p(xt
Z
1
Σ p(x),)
Σ ((tt+))2ABCCa. (Tie)an (x), + fh-1,0(30x tx)) xeen))
Z
EnN'n T(essn (xt, + f1-1, (xxx: bx))
x )
X 112
)
exp
(24)
2g(1x+1)
Peati (Kee44 Pxxun 1x] (2E
Substituting Ptù (af)
Pat (ate 1 (xc) to the equation above gives
N
(Tes)a + f1 (x1,, (x),,))) -Xatt)
Ple+s (Xie41) Exyya Plzta (Xten1)
Px xx (xx)
(25)
i=1
and dividing by Ptsti(x yields
N
N (Tie (x Xlete)
In Pi (X(e)) =
Σ
Pxg |xesss (xx:)
(26)
22
17

---

Published in Transactions on Machine Learning Research (11/2023)
Substituting Eq. (26) to the definition of the optimal backward drift Eq. (19) gives
(x (x+1) fi th) g(tx+1) V In Plz+s (x++1)
-1-1,0 (x (
(To),nn (x),, + fi 11-11,0 1,0 (xxxxtx))) x(mi)) (27)
- -((+++1)
Σ Psy
xi (Xies1)
g((x+1) A
where taking (Xte4 tx) inside the sum yields
(x (k+1)= =
Σ px 1x (X(e+1)
N Eis + ((-,)((
Σ
Xte11 - /Ax). (28)
Ce.i
Cs.i
n==l
Multiplying the equation above by Au gives
Xix12 - tt).
A tx+1)
Σ
(X + ( taa))
(29)
Ce.
Thus we may the objective for finding the optimal backward drift bi
Xlets 1, C2,s
In (x411)
(x))
tx)
1 En Teean (x% + (x), 11 (30) Cr,p
+
Notice that if the weights before resampling are uniform, then T(e) = IN, and for all i 1,2,...... N. it holds that Ce. since all but one of the terms in the sum g(fe+1) Var (En Tech., vanish. Similarly, for one-hot weights Ce.; = 1. In practice, we set the constant Ce. = 1 as in Eq. (9) and observe good empirical performance with the simplified loss function.
A.2 Connection to Hamilton-Jacobi Equations
We connect the backward drift function bı, (X 11) = Hamilton-Jacobi equations for stochastic control through following the setting of Maoutsa & Opper (2022), which applies the drift (x,, L) v In (x.)) for a backwards SDE initialized at my-
+x) g(xx+1)2 In Ptata (Xte4a) ) to the
Consider stochastic control problem with path constraint U (xx, f), optimizing the following loss function,
Σ J = N t
20(t) ()((),)) - (12 di -n (xx+),
(31)
with the paths, x sampled as trajectories from the SDE
Xy ~ Nus
dx, = fx t) dt + g()) dp,
(32)
and the loss In ((x))) measures distance from the distribution
Since we set the path constraint via observational data, our method resembles setting U'x),t) = 0 when is not an observation time, and (xx)) = - log p(y xx), where p(y |x() the observation model.
Let (x) denote the marginal density of the controlled (drift fa) SDE at time In Maoutsa & Opper (2022), the marginal density is decomposed
qe (x)) (x)
(33)
18

---

Published in Transactions on Machine Learning Research (11/2023)
where (x( is solution to backwards Siteer-Planc-Kolomogeee (FPK) partial differential equation starting from (pT(x) = nr, and the density evolves as in
dip((x) ())p((x) + U Uxx ())x(xx)
(34)
dt
where c) is the adjoint FPK operator to the uncontrolled system. The density pt(x) corresponds to the forward filtering problem, initialized with mo,
dxr(x) Lx(p((x)) (xx)))())
(35)
dt
where Ey is the FPK operator of the uncontrolled SDE (with drift f)- The particle filtering trajectories {x,,)" ) generated in our method are samples from the density defined by Eq. (35). In the context of our method, the path constraint matches the log-weights of particle filtering at observation times and is zero elsewhere.
In Maoutsa & Opper (2022), a backward evolution for de is applied, using the backwards time ğr--(x) = dr (x), yielding a backwards SDE starting from (x) = (x reweighted according to Rт. The backward samples from q are generated following the SDE dynamics
dx; = (f(xt, - + Vl)pr + g(t) dB+-
(36)
We have thus selected the backward drift dr, to match the drift of qe (z), the backward controlled density. Intuitively, our choice of bi is drift which generates the smoothed particles when initialized at { the terminal state of the forward SDE. The discrepancy between कт and the distribution induced by {x {iT then motivates the use of an iterative scheme after learning to simulate from (z).
A.3 Observing the Full Marginal Density
Suppose that at time te, let the number of observations grow unbounded. We analyse the behaviour of our model at the resampling step, at the limit M + oc for the number of observations and - 0 for the observation noise. When applying the bootstrap proposal, recall that we combined the multiple observations to compute the log-weights as
log wt
1
Σ MYEDH
Xx - y,
(37)
202
which works well in practice for the sparse-data settings we have considered. Below we analyse the behaviour of an alternative way to combine the weights and show that given an infinite number of observations, it creates samples from the true underlying distribution.
Proposition 2. Let (x,, be set of particles and { the observations at time tv. Assume that the observations have been sampled from density Ptz and that for all it holds that xi, E supp(pt.). Define the particle weights
1M
log wh i log
Σ exp(-1Ix/, -
(38)
HMM
MEEDM
where Z is the normalization conslant of the observation model Gaussian p(y x()) Then for each particle x its weight satisfies
lim lim wh Ptz (Iu).
(39)
- M -000
Proof sketch. We drop the a and (M) from the weight notation for simplicity of notation, but remark that the particle filtering weights are dependent on both quantitics. Consider the number of particles N fixed, and denote the d-dimensional sphere centered at x as B(xi, r). Since each particle xi, lies in the support of the true underlying marginal density Peuy then for any radius 0 such that B(x supp(pr.). ) and
19

---

Published in Transactions on Machine Learning Research (11/2023)
H >> 0, may choose M high enough that the points yj E D satisfy yy E B(x),,.). r). It follows from
Eq. (38) that
w/.
Σ
exp(-I1xi, - -2,12/200)
(40)
i,fe
For any >> 0 and with observation noise o = cr, we may set c, (AM) so that the sum above approximates the integral
Write
[ r)] Dise ply (ely)dd
(41)
By applying the Lebesque differentiation theorem, we obtain that for almost every x we have lim,- w = Pte (x since as a - 0, the density ply x(.) collapses to the Dirac delta of x)_ 0
Prop. 2 can be interpreted as the infinite limit of a kernel density estimate of the true underlying distribution. Resampling accurately reweights the particles so that the probability of resampling particle x is proportional to the density Pt. compared the other particles. Notice that the result does not guarantee that the particles will cover the support of Pta. since we did not assume that the drift initialization generates marginal density at time & covering its support.
B Experimental Details
B.1 2D Toy Data Sets
For the constrained transport problem for two-dimensional scikit-learn, the observational data we chose to use was different for each of the three data sets presented; TWo MOONS, TWO CIRCLES and the S-shape. All three experiments had the same discretization (t E [0, 0.99), A 0.01), learning rate 0.001, and differentiable resampling regularization parameter = 0.01. The process noise g(t)² follows a linear schedule from 0.001 to 1, with low noise at time t = 0 and high noise at t = 0.99, and each iteration of the ISB method trains the forward and backward drift networks each for 5000 iterations, with batch size 256. When rumning on a Macbook Pro CPU, it took approximately 6 mimutes to complete for the TWO CIRCLES experiment for instance, while the exact runtime varies based on factors such as the number of observations and the number of ISB iterations required. Other hyperparameters are explained below.
Two moons The observational data consists of 10 points selected from the Schrödinger bridge trajectories, all observed at E [0.25, 0.5, 0.75] with an exponential observation noise schedule <tl) = 1.25' The ISB was run for 6 epochs and initialized with a drift from the pre-trained Schrödinger bridge model from the unconstrained problem.
Two circles The observational data consists of 10 points which lic evenly distributed on a circle, obscrvod at = 0.5 with an exponential observational noise schedule x(I) 0.5 1.25'-1. The ISB was run for 6 epochs and initialized with a drift. from the pre-trained Schrödinger bridge model from the unconstrained problem.
S-shape The observational data consists of 6 points, with pairs being observed at times E [0.4, 0.6). We used a bilinear observational noise schedule with linear decay for the first half of the iterations from x())2 = 4 to x(L.2)2² = and a linear ascend for the second half of the iterations from s( <(L.2))² = 1 to x(L)² = 4. The ISB ran for epochs, with a zero drift initialization.
B.2 The Beneš SDE
In the Beneš SDE experiment, we obtain the sparse observational data from sampled Beneš SDE trajectories while the terminal state is a shifted and scaled (3 + 52T) version of a Beneš marginal density. As the Beneš trajectorics were first generated by simulating the SDE until t = 6 and then in reverse from t = 6 to t = 0, we set T = 11.97. We apply the analytical expression for the Bencš marginal density for computing log Pe (x),
1 cosh(x)
1 2 (-21
1 (x- - x-x))² )
px(x)= = /2 cx))(xx)
exp
exp
(42)
28
20

---

Published in Transactions on Machine Learning Research (11/2023)
0000
O
(a) 0.00
(b) 0.25
(c) 1 0.50
(d) 0.75
(e) 1 0.99
Figure 6: The unconditioned Schrödinger bridges for the scikit-learn 2D experiment, corresponding to the Fig. 3. The observations (red markers) were not used in any way during training but are included in the figure for reference, show that the unconstrained dynamics greatly differ from the learned ISB model dynamics.
See the Beneš SDE trajectories in Fig. 8a. As expected, the transport model with no observations performs well in the generative task, but its trajectories cover also some low-likelihood space around t = 6 (in the middle part in Fig. 8b). The obscrvations for the ISB model were sampled from the generated trajectorics, 10 observations at 10 random time-instances (see Fig. 8c)
Both the unconstrained Schrödinger bridge model and the ISB model were run for 3 iterations, using learning rate of 0.001 for the neural networks. Likely due to the fact that the problem was only one-dimensional, the convergence of the Schrödinger bridge process which matches the desired terminal state was fast, and we chose not to run the model for higher number of ISB iterations, see Fig. 7 for a comparison of the trained model marginal densities and the true terminal distribution тy. We set the observation noise schedule to the constant 0.7, and at each iteration of the ISB or the unconstrained Schrödinger bridge the drift neural networks were trained for 5000 iterations each with the batch size 256, and the trajectories were refreshed every 500 itcrations with a cache size of 1000 particles. The number of nearest neighbours to compare was H = 10.
True
ISB
Schrödinger
Figure A kernel density estimate of the Beneš SDE terminal state. We compare कe to the Schrödinger bridge and ISB terminal states. Both unconstrained Schrödinger bridge and TSB terminal states succeed in representing कт well, with the Schrödinger bridge terminal state more closely matching कт near its mean.
21

---

Published in Transactions on Machine Learning Research (11/2023)
20
20
2
8
12
12
8
12
Time,
Time,
Time,
(a) Trajectories of the Beneš SDE
(b) SBP trajectories
(c) ISB trajectories (ours)
Figure 8: Comparison of the solution for the SBP (with Beneš SDE reference drift) and the ISB (with zero initial drift) on the Beneš SDE under sparse observations (o). The target distribution vr slightly shifted and scaled from the Beneš SDE. Even if the SBP has the true model as reference drift, its trajectories
degenerate into unimodal distribution, while the ISB manages to cover both modes even if only sparse obscrvations are available.
B.3 The Bird Migration Data Set
The ISB model learned bird migration trajectories which transport the particles from the Northern Europe summer habitats to the southern winter habitats, see Fig. 11 for comparison of a Schrödinger bridge and ISB. Since the problem lies on a sphere, Schrödinger bridge methods adjusted for learning on Riemannian manifolds could have been applied here. For simplicity, we mapped the problem to a two-dimensional plane using a Mercator projection and solved the problem on a [0. 5] x [0, 5] square. The SDE had the discretization E [0, 0.99), A& = 0.01 and a constant process noise g(t)² = 0.05. The model was trained for 12 iterations, and initialized with a zero drift, while the observational data was chosen by the authors to promote learning trajectories clearly different from the unconstrained transport trajectories. The observation noise schedule was piecewise linear (starting at 2, going to 0.1 at iteration 6, then rising linearly to reach 2 at iteration 12). At cach ISB itcration, the neural networks were trained for 5000 iterations cach, and the trajectories were refreshed every 1000 ilerations. We used a batch size of 256 and a learning rate of 0.001.
B.4 The MNIST Generation Task
Applying state-space model approaches such as particle filtering and smoothing to generative diffusion models directly in the observation space (that is, not in a lower-dimensional latent space) has to our knowledge not been explored before. Some experimental design choices had great impact the training objectives sensibility, as the observational data is completely artificial and its timing during the process modifies the filtering distribution significantly. As the MNIST conditional generative model was trained to display the scalability of our method beyond low-dimensional toy examples, we did not further explore optimizing the hyperparameters or the observation model. To avoid the background noise in MNIST images in the middle of the generative process impacting the particle filtering weights excessively, the observation model is a Gaussian with masked inputs equal to zero in pixels where the observation image is black, see Fig. 9 for sampled trajectories. The figure shows the progression of seven samples, where the lower half of the eight resemble the observation target.
The SDE was run for time t E [0. 0.5]. with the digit cight observed at t = 0.38. The ISB method was applied for 10 iterations, with a discretization / E [0, 0.495), A = 0.005, and the process noise g(t)² followed a linear schedule from 0.0001 to 1. At each iteration of the method, the forward and backward drift neural networks were trained for 5000 iterations with a batch size of 256, and the trajectory cache regenerated every 1000 iterations. The observational data consisted of a single sample of a lower half of the digit eight, observed at time L 0.38. The observation noise schedule was constant x(I) 0.3.
22

---

Published in Transactions on Machine Learning Research (11/2023)
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
6
8
8
8
8
8
8
8
8
8
8
8
8
8
6
8
8
6
8
8
8
8
8
8
8
8
8
8
8
8
8 8
8
0
8
8
8
8
8
8
8
0 noize
T 0.5
Observation
blasing the lower half of Images
Figure 9: Model trajectories for MNIST digit '8' conditioned on lower-loop of a single '8' at { = 0.38 to bias the lower half of the digits to look alike, with the effect still visible at terminal time T.
7
9
9
7
9
9
9
9
9
8
8
8
8
8
8
8
8
7
9
9
9
9
9
9
9
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
TO noise
0.5
Obeervation
bissing the upper half of images
Model trajectories for MNIST digits '8' and '9' conditioned on a upper-loop of a single '8' at = 0.38 to bias the upper half of the digits to look alike, with the effect still visible at terminal time T, while observing both modes at the end of the trajectory.
Figure 10:
B.5 MNIST for Multi-modal Data Generation
In addition to the MNIST experiment explained in App. B.4, we evaluated the performance of our model a multi-modal generation task. The reference model was Schrödinger bridge trained via IPFP from normal distribution noise to MNIST digits eight and nine, and the observation was a single upper loop of a figure cight. Our goal was to generate trajectorios that match the obscrvation but still generate both digits eight and nine by the end of the trajectory. Based on the results in Fig. 10, ISB succeeds in the lask. Most hyperparameters were kept the same as explained in App. B.4, but the cache size if IPFP training was increased to 5000 and the number of trajectories in the particle filtering step of ISB was increased to 1000, to encourage sufficiently versatile sample of both modes.
B.6 Single-Cei Cell Data Set
We directly use the preprocessed data from the TrajectoryNet (Tong et al., 2020) repository. A major difference between our implementation and Vargas et al. (2021) is the reference drift. We set the reference drift to zero, which means that we utilize the intermediate data only as observations the state-space model. On the contrary, Vargas et al. (2021) fits mixture model of 15 Gaussians the combined data set (across
23

---

Published in Transactions on Machine Learning Research (11/2023)
all measurement times) and sets the reference drift to the gradient of the log-likelihood of the mixture model. Effectively, such reference drift aids in keeping the SDE trajectories within the support of the combined data set. We remark that if the intermediate observed marginals had clearly disjoint support, combining all the data would cause the mixture model to have 'gaps' and could cause an unstable reference model drift. Thus we consider our approach of seiting the reference drift. to zero as more generally applicable.
As in Vargas et al. (2021); we set the process noise to g(t) = 1 and model the SDE between time & [0, 4]. The learning rate is set to 0.001 with batch size of 256 and the mumber of neural network training iterations equal to 5000. We apply the ISB for 6 iterations. We perform filtering using 1000 points from the intermediate data sets, but compute the Earth mover's distance by comparing it to all available data. As the observational data at T= 1, 2, 3 consists of a high number of data points, the parameters H (number of nearest neighbours) and a (observation noise) need to be carefully set. We set = 10 to only include the close neighbourhood of each particle and set the observation noise schedule as constant 0.7.
C Computational Considerations
In Sec. 3.2, we raised a number of important computational considerations for the constrained transport problem. Below we discuss them in detail, analyzing the limit - oc from the perspective of setting the observation noise schedule in App. C.I, and presenting ablation results modifying the initial drift in the bird migration experiment in App. C.2. Finally, study the impact of the observation noise schedule effective sample size in the filtering step in App. C.3.
C.1 Discussion on Observation Noise
We briefly mentioned in Sec. 3.2 that when letting L - oo, the choice of observation noise should be carefully planned in order for the ISB procedure to have a stationary point. Here we explain why an unbounded observation noise schedule x(l) implies convergence to the IPF method for uncontrolled Schrödinger bridges (De Bortoli et al., 2021), when using a nearest neighbour boolstrap filter as the proposal density.
Proposition 3. Let Ω E Rd be bounded domain where both the obsercations and SDE trajectories lie, and let the particle filtering weights {wi IN at ISB iteration 1 be
log wf.. 2x(1)²
Σ | - 111
(43)
YJED!!
If the schedule k(l) is unbounded with respect to I, then for any 8 there exists 1' such that for the normalized weights it holds
1
1
N 5.
(44)
Proof sketch. Since x(I) is unbounded, for any S 0 3'' such that s(l') 2 S. We choose the value of S so that the following derivation yields Eq. (44).
Let S (0.5R-1DDf] di)m and apply the property that ||x|, -,, < diam( to Eq. (43),
log wh.s
Σ
|x
ysll²
25²
EJED!! $t
(45)
ΣyYeDH ||x] - 11,112
Eyseny diam (12)
2
R- RD! diam2 2 R.
R-1ID diam
The bound above for the unnormalizod weights, and the normalized log-weights are defined as
log 2h.1e log -logwh log
Σ exp(log w/ w..
(46)
24

---

Published in Transactions on Machine Learning Research (11/2023)
where for the normalizing constant it holds that
N
N = log(N),
log
exp(log w an) = log
Σ
(47)
since u is the value of a probability density and thus always w/ . S Combining Eq. (46), Eq. (45) and
Eq. (47), follows that.
log uf (- log(N)) 2 - R
(48)
where taking exponentials on both sides gives
1
1P.4
2 -1-1-eexpp
-1x-exx(-1BB)) exp(-=1)) N
(49)
N
Since the weights are normalized, even the largest particle weight wt can differ from N as much as every smaller weight in total lies under
N
tr - N + (N (N-1 exp(-x) ()) h
(50)
implying that for any weight o6 it holds that
1 + s (N - 11) (1-((-ppp(-B))) N
< - exp(-R),
(51)
N
and selecting R = - log(1 - 5) is sufficient for < 0
Effectively, the above derivation implies that for an unbounded observation noise schedule x(I), the particle weights will converge to uniform weights. Since performing differentiable resampling on uniform weights implies that (e) IN. the ISB method trajectory generation step and the objective in training the backward drift converge to those of the IPF method for solving unconstrained Schrödinger bridges. Intuitively, this means that at the limit - o0, our method will focus on reversing the trajectories and matching the terminal distribution while not further utilizing information from the observations.
C.2 Ablation on Initial Drift
We conducted an ablation study on drift initialization for the bird migration problem. As the distributions To and Ty (as pictured in Fig. 11) are complex, we consider the problem setting to be interesting for setting fo as the unconstrained transport problem drift. To this end, we trained a Schrödinger bridge model for 10 epochs, and trained an ISB model with the same hyperparameter selections as explained in App. B.3, using the Schrödinger bridge as the initialization. Compare the two bottom rows of Fig. 11 to see selection of marginal densities of the two processes. Based on visual analysis of the densities, it seems that the zero drift and pre-trained diffusion model initializations produce similar results around the observations, although the Schrödinger bridge initialization gave slightly sharper results at the terminal time.
C.3 Effective Sample Size
We studied the Effective Sample Size (ESS) of the particles generated during the filtering steps 1 and 3 Ideally, our method would steer the particle trajectories near the observations quickly over the iterations, to allow for efficient training. Furthermore, as we additionally constrain the system to match the initial and terminal distributions, it is desired that the ESS will eventually rise to a high number, indicating that the particle filtering step no longer greatly adjusts the trajectories. In order to assess these properties, computed the ESS in the scikit-learn TWO CIRCLES experiment, see App. B.1 for other experimental details. In Table 3, we have reported three scenarios: (i) an increasing noise schedule without lcarning of the drift function, (ii) an ISB model with constant observation noise, (iii) the proposed ISB model with the increasing noise schedule. Note that the first scenario only serves as a control to better isolate the impact of increasing noise ISB. Table indicates that our approach succoods in directing the particle trajectories and consistently obtains higher ESS even when comparing the constant noise level ISB to an increasing noise level but no learning.
25

---

Published in Transactions on Machine Learning Research (11/2023)
Summer
Intermediate observations during migration
Winter
Bird observation
Unconditional
Conditional (transport drift) Conditional (scro drift)
Figure 11: Top row: The first map image the left describes the initial position of the birds, and the final one on the right depicts their position after migration. The observational data in the middle are bird obscrvations during migration, at given timestamps. Second row: Marginal densities Schrödinger bridge model from the initial terminal distribution, without using the observations. Third row: Marginal densities of our model, using both initial and terminal distributions and observational data and zero drift initialization. Bottom row: Same as the third row, but with the second-row dynamics initialization.
Table 3: Average and standard deviation of effective sample size (ESS) results over five seeds for the scikit-learn Two CIRCLES experiment, developing over six ISB iterations with 1000 particles. We observe a rapid increase in ESS when the trajectories are steered towards the observations in the optimization step the ISB compared only increasing the observing noise, and diminishing impact of the filtering step as ESS nears 1000.
ESS ITERATION
NOISE SCHEDULE
I=2
1=3 3
1=
1
I=
No learning and noise schedule 0.5 1.25'-1
6
9 2
14
27
+ 5
57 +
121 + 11
ISB with constant noise schedule 0.5
6 2
+ 21
165 +
300 +
24
456 37
564 46
ISB with noise schedule 0.5 1.25/-1
612
91
250 +
442
1
39
593 1 64
700 39
26

---

Published in Transactions on Machine Learning Research (11/2023)
D Differentiable Resampling
In the ISB model steps 1 and 3 presented in Sec. 3.1, we applicd differentiable resampling (see Corenflos et al., 2021). Resampling itself is a basic block of particle filtering. A differentiable resampling step transports the particles and weights (xt, w(,)) to a uniform distribution over a set of particles through applying the differentiable ensemble transport map ((e), that is
(x) (f, + (X,, (()), 1/N) (xl, 1/N),
(52)
where X,, E Nxd denotes the stacked particles {xl, LN at time tx before resampling and x denotes the particles post resampling. Here we give the definition of the map (e) and review the regularized optimal transport problem which has to be solved to compute it. We partly follow the presentation Sections 2 and of Corenflos et al. (2021); but directly apply the notation explaining the transport problem rather than the algorithm used to solve it.
for particles and weights and focus
The standard particle filtering resampling step consists of sampling N particles from the categorical distribution defined by the weights {wi, 1 resulting in the particles with large weights being most likely to be repeated multiple times. A result from Reich (2013) gives the property that the random resampling step can be approximated by a deterministic ensemble transform T. In heuristic terms, the ensemble transform map will be selected so that the particles (xi, J.W will be transported with minimal cost, while allowing all the weights to be uniform.
Let j and be atomic measures, / = Ei i= "i and v -= EE N where & is the Dirac delta at æ. Then μ is the particle filtering distribution before resampling. Define the elements a cost matrix C N as C.. = ||x x22 and the 2- Wasserrstti distance between two atomic measures
N min PES(a.)
(2,
CrpPas. 1=1
(53)
Above the optimal matrix P to be found within S(u, v), which is a space consisting of mixtures of particles to N particles such that the marginals coincide with the weights of μ and v, formally
S(UDD) PE 1
Σ 2=1
W
Σ Pas
(54)
The entropy-regularizee Wasserstein distance with regularization parameter is then
N ΣΣ Pis CCs (Cis Eau++elog log wi,
Pr.
W = min PES(um) =1
(55)
The unique minimizing transport map of the above Wasserstein distance is denoted by POPT and the ensemble transport map is then set as T(e) = NPOPT, This means that we can find the matrix T(e) via minimizing the regularized Wasserstein distance, which is done by applying the iterative Sinkhorn algorithm for regulie-re optimal transport (Cuturi, 2013).
27


Reviews:
[{"review_id": "RJL4LdJ0x6", "reviewer": "Reviewer_wuTy", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "- Strengths: \n\n  - I found the paper to be an interesting read and well written with clear motivations. The problem tackled is interesting, as demonstrated by the experiments. \n  - The authors effectively showcase the performance of their approach across various experiments.  \n\n- Weaknesses/Questions:  \n\n  - In my opinion the methodology is not appropriately justified. While it is clearly stated that the authors will try to approximate the iterates of the IPFP, they do not clearly define what it means to do so in the presence of observations. To be more precise, the reference process in the IPFP should be the posterior of the original reference process conditioned on the sparse observations. Then, by writing the consecutive iterates the authors should be able to obtain theoretical updates that involve the observations and their likelihood. I believe that this has not been done in this paper, which makes evaluating the soundness of the methodology rather difficult. \n  - In Algorithm 1, resampling is performed only on the states $x^i _{t_k}$ where $t_k$ is an observation timestep. This is quite strange as I would expect the resampling of the complete trajectory $x^i _{0:t_k}$. It seems to me that resampling only the states $x^i _{t_k}$ and not the complete trajectory up to step $t_k$ leads to weird and inconsistent trajectories. Again, I think that the authors should write down the theoretical updates (which are not the same as in [1] due to the presence of observations) and then use it to justify their practical implementation. \n  - Since the observation likelihood is linear in the states, the optimal kernel (in the sense of the auxiliary particle filter) is available in closed form and leads to trajectories that better match the observations. Have the authors considered such an approach for sampling and resampling of the states at the observed timesteps?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes a novel method for learning a diffusion process with constrained trajectories where the constraint is defined by a set of observations at sparse time steps. More precisely, the authors approximate the iterates of the iterative proportional fitting procedure (IPFP), following [1], while at the same time incorporating the observations. In order to achieve the latter, they introduce  a resampling step in order to select the samples that match best the observations. As the resampling step is discrete and non-differentiable in nature, the authors use the differentiable resampling procedure of [2]. The resulting loss which incorporates resampled samples then becomes fully differentiable. \n\n[1] Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion Schrödinger bridge with applications to score-based generative modeling. In Advances in Neural Information Processing Systems 34\n(NeurIPS), pp. 17695–17709. Curran Associates, Inc., 2021.\n\n\n[2] Adrien Corenflos, James Thornton, George Deligiannidis, and Arnaud Doucet. Differentiable particle filtering via entropy-regularized optimal transport. In Proceedings of the 38th International Conference on Machine Learning, ICML 2021, volume 139 of Proceedings of Machine Learning Research, pp. 2100–2111. PMLR, 2021.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I recommend that the authors write down a rigorous justification of their methodology by clearly writing the IPFP iterates when the reference process is the joint posterior of a hidden Markov model.", "broader_impact_concerns": "There are no ethical concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "OCT0ioIGz4", "reviewer": "Reviewer_xkTm", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths:**\n\n- extremely well written paper\n- provides connection to several related methods published recently\n- targets a rather challenging problem: solving a Schroedinger bridge problem is sufficiently challenging it self. Here the authors consider additional intermediate constraints that are often relevant for applied problems with intermediate observations.\n- they provide well defined theoretical formulation and analysis for the proposed problem both in Supplement and main text.\n\n\n\n**Weaknesses:**\n\nI cannot really find weaknesses for the proposed framework. I would consider some comment on the computational demands of the method would be helpful, and some discussion regarding the convergence of the method to the optimal solution.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes the **Iterative Smoothing Bridge** method for solving dynamic Schroedinger bridge problems with intermediate constraints, i.e. not only distribution constraints at the beginning and end of the considered time interval but also during _that_ interval.The proposed approach employs the iterative pass idea from\nthe Iterative Proportional Fitting procedure [1] together with differentiable particle filtering of [2], and show that their framework can be understood as an iterative version of the recent optimal control framework of [3] with distributional initial and terminal constraints.\nThey demonstrate their method on a series of toy and applied experiments (bird migration, image generation, single cell RNA sequence ).\n\n**I found the paper extremely well written, providing sufficient background and connections to the existing literature, targeting a challenging problem, and therefore I strongly recommend the paper for publication.**\n\n\n\n\n**References:**\n\n[1] Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion Schrödinger bridge with\napplications to score-based generative modeling. In Advances in Neural Information Processing Systems 34\n(NeurIPS), pp. 17695–17709. Curran Associates, Inc., 2021.\n\n[2] Adrien Corenflos, James Thornton, George Deligiannidis, and Arnaud Doucet. Differentiable particle filtering\nvia entropy-regularized optimal transport. In Proceedings of the 38th International Conference on Machine\nLearning, ICML 2021, volume 139 of Proceedings of Machine Learning Research, pp. 2100–2111. PMLR,\n2021\n\n[3] Dimitra Maoutsa and Manfred Opper. Deterministic particle flows for constraining stochastic nonlinear\nsystems. Physical Review Research, 4(4), 2022.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See weaknesses.", "broader_impact_concerns": "None.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "AN3ljIDJcC", "reviewer": "Reviewer_9PRD", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper provides an interesting methodology and considers an interesting problem. Constraining generative modelling with observed data is a crucial problem for constrained generation, inverse problems. Also I found the \"dynamic\" approach that allowed the use of filtering interesting.\n\nThe weaker aspects of the paper is that, I think,\n\n1) The paper does not position itself clearly in the literature. Reading the abstract, it is not very clear that, for example, the paper is about generative modelling. It reads more like a classical sampling-optimal transport paper.\n\n2) The utility of the algorithm is not very clear for many problems. Considering something like score-based models, it is rarely need to condition on an intermediate stage (which is, for images, a noisy image). To me, it looked like the paper is more related to the classical filtering approaches with sparse observations.\n\n3) Some of the approaches in this paper may not scale well. For example, it is surprising that there is no curse of dimensionality problem due to the use of particle filters.\n\nSome questions/requests below.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper is about defining stochastic (generative) processes conditioned on observations on the path to constrain/restrict the generation process. The paper achieves this by introducing Iterative Smoothing Bridge (ISB). The method combines ideas from Bayesian filtering to define this constrained generation processes.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Requests/questions:\n\n1) The paper can use a clarification of its purpose. Is the main aim here \"generative modelling\" and constraining it throughout the generation process? While the optimal transport tools is useful for this purpose, I think it is important to clearly mention the applications in the abstract.\n\n2) Connected to point 1), as I said, this is sort of not very intuitive for many generation examples. The authors provide one example in their figure 9 for MNIST to condition the images to generate the digit 8 which looks interesting. I would also like to see a multimodal example : Can you condition ambiguously, for example the upper half of 8 which could be also 9. Is the sampler generate diverse samples in this sense?\n\n3) Going into the details of the method: The authors seems to (1) evolve forward particle filtering trajectories (I assume this is particle filtering?), (2) update the drift function $\\phi$, (3) then run a filter again backwards in time to assimilate observations this time using the reverse process, (4) update $\\theta$. Here, it is interesting to discuss regular problems of particle filters. Typically, this is not expected to work really well for high dimensional states (but looking at MNIST, it looks like it is working). Is the use of differentiable resampling key here to success of this method? Did you look at these filter's efficiency in terms of ESS?\n\n4) Overall, I'd ask authors to clarify the method further. For example Steps 2 and 4 are not that clear from the explanations. The choice of losses seem arbitrary. Since the notation is also convoluted, the intuition gets lost here. Please provide more intuitive explanations for Steps 2 and 4 for the readers.\n\n5) I would also like a discussion of the differences between this approach and general data assimilation methods which deals with the same problem. There is lots of work in the DA literature for these kinds of sparsely observed dynamical systems and incorporation of data. What are the connections/relationships?", "broader_impact_concerns": "The paper's methodology is about generative modelling -- therefore all risks in this area applies to this paper. There was no Broader Impact statement.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper has received mixed feedback from the reviewers:

One of the reviewers is very enthusiastic about the paper.

A second reviewer agrees that the paper has some merits but believes that the methodology lacks novelty.

The last reviewer found the proposed method interesting but pointed out that the method does not have a strong theoretical foundation.

After reviewing the paper, I concur with the last reviewer. While the methodology presented in the paper seems intuitively reasonable, it's unclear whether it effectively addresses the problem as stated by the authors, i.e., "solving control problems
under constraints on both the initial and terminal distribution and sparse observational data constraints". From my perspective, the proposed method appears to be a generative model for time series that aims to incorporate constraints, but I'm not entirely convinced that it directly solves the original optimal transport problem with constraints.

As a result, I have decided to accept the paper with minor revisions. The authors should acknowledge that their method is primarily a generative model inspired by stochastic optimal control theory, but the overall connection between the two is unclear.
Decision: Accept (poster)