Title: Diffusion Schrödinger Bridge Matching
Venue: NeurIPS 2023
Accepted: yes
Keywords: diffusion Schrödinger bridge

bridge matching

optimal transport
Abstract: Solving transport problems, i.e. finding a map transporting one given distribution to another, has numerous applications in machine learning. Novel mass transport methods motivated by generative modeling have recently been proposed, e.g. Denoising Diffusion Models (DDMs) and Flow Matching Models (FMMs) implement such a transport through a Stochastic Differential Equation (SDE) or an Ordinary Differential Equation (ODE). However, while it is desirable in many applications to approximate the deterministic dynamic Optimal Transport (OT) map which admits attractive properties, DDMs and FMMs are not guaranteed to provide transports close to the OT map. In contrast, Schrödinger bridges (SBs) compute stochastic dynamic mappings which recover entropy-regularized versions of OT. Unfortunately, existing numerical methods approximating SBs either scale poorly with dimension or accumulate errors across iterations. In this work, we introduce Iterative Markovian Fitting (IMF), a new methodology for solving SB problems, and Diffusion Schrödinger Bridge Matching (DSBM), a novel numerical algorithm for computing IMF iterates. DSBM significantly improves over previous SB numerics and recovers as special/limiting cases various recent transport methods. We demonstrate the performance of DSBM on a variety of problems.

Diffusion Schrödinger Bridge Matching

Abstract

Solving transport problems, i.e. finding a map transporting one given distribution to another, has numerous applications in machine learning. Novel mass transport methods motivated by generative modeling have recently been proposed, e.g. Denoising Diffusion Models (DDMs) and Flow Matching Models (FMMs) implement such a transport through a Stochastic Differential Equation (SDE) or an Ordinary Differential Equation (ODE). However, while it is desirable in many applications to approximate the deterministic dynamic Optimal Transport (OT) map which admits attractive properties, DDMs and FMMs are not guaranteed to provide transports close to the OT map. In contrast, Schrödinger bridges (SBs) compute stochastic dynamic mappings which recover entropy-regularized versions of OT. Unfortunately, existing numerical methods approximating SBs either scale poorly with dimension or accumulate errors across iterations. In this work, we introduce Iterative Markovian Fitting (IMF), a new methodology for solving SB problems, and Diffusion Schrödinger Bridge Matching (DSBM), a novel numerical algorithm for computing IMF iterates. DSBM significantly improves over previous SB numerics and recovers as special/limiting cases various recent transport methods. We demonstrate the performance of DSBM on a variety of problems.

1 Introduction

Mass transport problems are ubiquitous in machine learning (Peyré and Cuturi, 2019). For discrete measures, the Optimal Transport (OT) map can be computed exactly but is computationally intensive. In a landmark paper, Cuturi (2013) showed that an entropy-regularized version of OT can be computed more efficiently using the Sinkhorn algorithm (Sinkhorn, 1967). This has enabled the use of OT techniques in a variety of applications ranging from biology (Bunne et al., 2022) to shape correspondence (Feydy et al., 2017). However, applications involving high-dimensional continuous distributions and/or large datasets remain challenging for these techniques.

One of such data-rich applications is generative modeling, a central transport problem in machine learning which requires designing a deterministic or stochastic mapping transporting a reference 'noise' distribution to the data distribution. For example, Generative Adversarial Networks (Goodfellow et al., 2014) define a static, deterministic transport map, while Denoising Diffusion Models (DDMs) (Song et al., 2021b; Ho et al., 2020) build a dynamic, stochastic transport map by simulating a Stochastic Differential Equation (SDE), whose drift is learned using score matching (Hyvärinen, 2005; Vincent, 2011). The excellent performances of DDMs have motivated recent developments of Bridge Matching and Flow Matching models, which are dynamic transport maps using SDEs (Song et al., 2021a; Peluchetti, 2021; Liu, 2022; Albergo et al., 2023) or ODEs (Albergo and VandenEijnden, 2023; Heitz et al., 2023; Lipman et al., 2023; Liu et al., 2023b). Compared to DDMs, Bridge and Flow Matching methods do not rely on a forward 'noising' diffusion converging to the reference distribution in infinite time, and are also more generally applicable as they can approximate transport maps between two general distributions based on their samples. Nonetheless, these transport maps

DSBM

Denoising

Diffusion

Flow

Matching

are not necessarily close to the OT map minimizing the Wasserstein-2 metric, which is appealing for its many attractive properties (Peyré and Cuturi, 2019; Villani, 2009).

In contrast, the Schrödinger Bridge (SB) problem is a dynamic version of entropy-regularized OT (EOT) (Föllmer, 1988; Léonard, 2014b). The SB is the finite-time diffusion which admits as initial and terminal distributions the two distributions of interest and is the closest in Kullback-Leibler divergence to a reference diffusion. Numerous methods to approximate SBs numerically have been proposed, see e.g. (Bernton et al., 2019; Chen et al., 2016; Finlay et al., 2020; Caluya and Halder, 2021; Pavon et al., 2021), but these techniques tend to be restricted to low-dimensional settings. Recently, novel techniques using diffusion-based ideas have been proposed in (De Bortoli et al., 2021; Vargas et al., 2021; Chen et al., 2022) based on Iterative Proportional Fitting (IPF) (Fortet, 1940; Kullback, 1968; Rüschendorf and Thomsen, 1993), a continuous state-space extension of the Sinkhorn algorithm (Essid and Pavon, 2019). These approaches have been shown to scale better empirically, but numerical errors tend to accumulate over iterations (Fernandes et al., 2021).

In this paper, our contributions are three-fold. First, we introduce Iterative Markovian Fitting (IMF), a new procedure to compute SBs which alternates between projecting on the space of Markov processes and on the reciprocal class , i.e. the measures which have the same bridge as the reference measure of SB (Léonard et al., 2014). We establish various theoretical results for IMF. Contrary to IPF, the IMF iterates always preserve the initial and terminal distributions. The differences between IPF and IMF are presented in Table 1. Second, we propose Diffusion Schrödinger Bridge Matching (DSBM), a novel algorithm approximating numerically the SB solution derived from IMF. DSBM requires at each iteration solving a simple regression problem in the spirit of Bridge and Flow Matching, and does not suffer from the time-discretization and 'forgetting' issues of previous DSB techniques (De Bortoli et al., 2021; Vargas et al., 2021; Chen et al., 2022). Finally, we demonstrate the performance of DSBM on a variety of transport tasks. 2

Notations. We denote by P ( C ) , the space of path measures , i.e. P ( C ) = P (C([0 , T ] , R d )) where T > 0 . The subset of Markov path measures associated with an SDE of the form d X t = v t ( X t )d t + σ t d B t , with σ, v locally Lipschitz, is denoted M . For any Q ∈ M , the reciprocal class of Q is denoted R ( Q ) , see Definition 3. We also denote Q t its marginal distribution at time t , Q s,t the joint distribution at times s and t , Q s | t the conditional distribution at time s given state at time t , and Q | 0 ,T ∈ P ( C ) its diffusion bridge . Unless specified otherwise, all gradient operators ∇ are w.r.t. the variable x t with time index t . Let ( X , X ) and ( Y , Y ) be probability spaces. Given a Markov kernel K : X × Y → [0 , 1] and a probability measure µ defined on X , we write µ K the probability measure on Y such that for any A ∈ Y we have µ K( A ) = ∫ X K( x, A )d µ ( x ) . In particular, for any joint distribution Π 0 ,T over R d × R d , we denote the mixture of bridges measure as Π = Π 0 ,T Q | 0 ,T ∈ P ( C ) , which is short for Π( · ) = ∫ R d × R d Q | 0 ,T ( ·| x 0 , x T )Π 0 ,T (d x 0 , d x T ) .

2 Dynamic Mass Transport Techniques

2.1 Denoising Diffusion and Bridge Matching Models

Denoising Diffusion Models (Song et al., 2021b; Ho et al., 2020) are a popular class of generative models. They define a forward noising process Q ∈ M using the SDE d X t = -1 2 X t d t +d B t on the time-interval [0 , T ] , where X 0 ∈ R d is drawn from the data distribution π 0 and ( B t ) t ∈ [0 ,T ] is a d -

dimensional Brownian motion. This diffusion 3 converges towards the standard Gaussian distribution N(0 , Id) as T →∞ . A generative model is given by its time-reversal ( Y t ) t ∈ [0 ,T ] = ( X T -t ) t ∈ [0 ,T ] , where Y 0 ∼ Q T and d Y t = { 1 2 Y t + ∇ log Q T -t ( Y t ) } d t + d B t (Anderson, 1982; Haussmann and Pardoux, 1986). In practice, ( Y t ) t ∈ [0 ,T ] is initialized with Y 0 ∼ π T = N(0 , Id) , and the Stein score ∇ log Q t ( x t ) = E Q 0 | t [ ∇ log Q t | 0 ( X t | X 0 ) | X t = x t ] is approximated using a neural network s θ ( t, x t ) minimizing the denoising score matching loss E Q 0 ,t [ ‖∇ log Q t | 0 ( X t | X 0 ) -s θ ( t, X t ) ‖ 2 ] .

An alternative to considering the time-reversal of a forward noising process is to 'build bridges' between the two distributions and learn a mimicking diffusion process. This approach generalizes DDMs and allows for more flexible choices of sampling processes. We call this framework Bridge Matching and adopt a presentation similar to Peluchetti (2021); Liu et al. (2022b), where π T is the data distribution. 4 We denote Q ∈ M the path measure associated with the following process

Consider now the distribution of this process pinned down at an initial and terminal point x 0 , x T , denoted Q | 0 ,T ( ·| x 0 , x T ) . Under mild assumptions, the pinned process Q | 0 ,T ( ·| x 0 , x T ) is a diffusion bridge and is given by

which satisfies X 0 ,T T = x T using Doob h -transform theory (Rogers and Williams, 2000). Next, we define an independent coupling Π 0 ,T = π 0 ⊗ π T , and let Π = Π 0 ,T Q | 0 ,T . This path measure Π is a mixture of bridges . We aim to find a Markov diffusion d Y t = { f t ( Y t ) + v t ( Y t ) } d t + σ t d B t on [0 , T ] which admits the same marginals as Π ; i.e. for any t ∈ [0 , T ] , Y t ∼ Π t , so Y T ∼ π T . For such v t , a generative model for sampling data distribution π T is obtained by simulating ( Y t ) t ∈ [0 ,T ] . It can be verified that indeed Y t ∼ Π t for v glyph[star] t ( x t ) = σ 2 t E Π T | t [ ∇ log Q T | t ( X T | X t ) | X t = x t ] . We present the theory behind this idea more formally using Markovian projections in Section 3.1. In practice, we do not have access to v glyph[star] t and it is learned using neural networks with regression loss

For f t = 0 and σ t = σ , Q | 0 ,T is a Brownian Bridge and we have

with ( B t -t T B T ) ∼ N(0 , t (1 -t T ) Id) . The regression loss (3) associated with (4) is given by

Letting σ → 0 , we recover Flow Matching models (see Appendix A.1 for further details).

2.2 Schrödinger Bridges and Optimal Transport

The Schrödinger Bridge (SB) problem (Schrödinger, 1932) consists in finding a path measure P SB ∈ P ( C ) such that

where Q ∈ P ( C ) is a reference path measure. In what follows, we consider Q defined by the diffusion process (1) which is Markov, and without loss of generality, we assume Q 0 = π 0 . Hence P SB is the path measure closest to Q in terms of Kullback-Leibler divergence which satisfies the initial and terminal constraints P SB 0 = π 0 and P SB T = π T .

Another crucial property of P SB is that it can also be defined as a mixture of bridges P SB = Π SB 0 ,T Q | 0 ,T , where Π SB 0 ,T = argmin Π 0 ,T { KL(Π 0 ,T | Q 0 ,T ) : Π 0 = π 0 , Π T = π T } is the solution of the static SB problem (Léonard, 2014b). In particular, for Q associated with ( σ B t ) t ∈ [0 ,T ] we have

where H ( µ ) denotes the entropy, i.e. Π SB 0 ,T is the solution of the entropy-regularized OT problem. In this case, the SB can also be obtained theoretically by solving the following problem (Dai Pra, 1991)

Then P SB is given by the SDE with drift v SB initialized with X 0 ∼ π 0 . For σ = 0 , we recover the classical OT problem and the Benamou-Brenier formula (Benamou and Brenier, 2000).

A common approach to solve (6) is the Iterative Proportional Fitting (IPF) method (Fortet, 1940; Kullback, 1968; Rüschendorf, 1995) defining a sequence of path measures ( ˜ P n ) n ∈ N where

with initialization ˜ P 0 = Q . This procedure alternates between projections on the set of path measures with given initial distribution π 0 and terminal distribution π T . It can be shown (De Bortoli et al., 2021) that ( ˜ P n ) n ∈ N are associated with diffusions and that for any n ∈ N , ˜ P 2 n +1 is the time-reversal of ˜ P 2 n with initialization π T , and ˜ P 2 n +2 is the time-reversal of ˜ P 2 n +1 with initialization π 0 . Leveraging this property, De Bortoli et al. (2021) proposed Diffusion Schrödinger Bridge (DSB), an algorithm which learns the time-reversals iteratively. In particular, DDMs can be seen as the first iteration of DSB.

3 Iterative Markovian Fitting

3.1 Markovian Projection and Reciprocal Projection

Markovian Projection. Projecting on Markov measures is a key ingredient in our methodology and in the Bridge Matching framework. This concept was introduced multiple times in the literature (Gyöngy, 1986; Peluchetti, 2021; Liu et al., 2022b). In particular, we focus on Markovian projection of path measures given by a mixture of bridges Π = Π 0 ,T Q | 0 ,T ∈ P ( C ) .

Definition 1. Assume that Q is given by (1) and that for any ( x 0 , x T ) ∈ R d , Q | 0 ,T ( ·| x 0 , x T ) is associated with ( X 0 ,T t ) t ∈ [0 ,T ] given by d X 0 ,T t = { f t ( X 0 ,T t )+ σ 2 t ∇ log Q T | t ( x T | X 0 ,T t ) } d t + σ t d B t , with σ : [0 , T ] → (0 , + ∞ ) . Then, when it is well-defined, we introduce the Markovian projection of Π , M glyph[star] = proj M (Π) ∈ M , which is associated with the SDE

Note that in our definition σ t > 0 so ∇ log Q T | t ( x T | x t ) is well-defined, but Flow Matching can be recovered as the deterministic case in the limit σ t = σ → 0 . In the following proposition, we show that the Markovian projection is indeed a projection for the reverse Kullback-Leibler divergence, and that it preserves marginals of Π t .

Proposition 2. Assume that σ t > 0 . Let M glyph[star] = proj M (Π) . Then, under mild assumptions, we have M glyph[star] = argmin M { KL(Π | M ) : M ∈ M} ,

In addition, we have that for any t ∈ [0 , T ] , M glyph[star] t = Π t . In particular, M glyph[star] T = Π T .

glyph[negationslash]

Definition 3. Π ∈ P ( C ) is in the reciprocal class R ( Q ) of Q ∈ M if Π = Π 0 ,T Q | 0 ,T . We define the reciprocal projection of P ∈ P ( C ) as Π glyph[star] = proj R ( Q ) ( P ) = P 0 ,T Q | 0 ,T .

Reciprocal Projection. While the Markovian projection ensures that the obtained measure is Markov, the associated bridge measure is not preserved in general, i.e. proj M (Π) | 0 ,T = Π | 0 ,T = Q | 0 ,T . Measures with same bridge as Q are said to be in its reciprocal class (Léonard et al., 2014).

Similarly to Proposition 2, we have the following result, which justifies the term reciprocal projection. Proposition 4. Let P ∈ P ( C ) , Π glyph[star] = proj R ( Q ) ( P ) . Then, Π glyph[star] = argmin Π { KL( P | Π) : Π ∈ R ( Q ) } .

The reciprocal projection Π glyph[star] of a Markov path measure M does not preserve the Markov property in general. In fact, the Schrödinger Bridge is the unique path measure which satisfies the initial and terminal conditions, is Markov and is in the reciprocal class of Q , see (Léonard, 2014b).

Proposition 5. Let P be a Markov measure in the reciprocal class of Q such that P 0 = π 0 , P T = π T . Then, under assumptions on Q , π 0 and π T , P is unique and is equal to the Schrödinger Bridge P SB .

3.2 Iterative Markovian Fitting

Based on Proposition 5, we propose a novel methodology called Iterative Markovian Fitting (IMF) to solve Schrödinger Bridges. We consider a sequence ( P n ) n ∈ N such that

with P 0 such that P 0 0 = π 0 , P 0 T = π T and P 0 ∈ R ( Q ) . These updates correspond to alternatively performing Markovian projections and reciprocal projections.

Combining Proposition 2 and Definition 3, we get that for any n ∈ N , P n 0 = π 0 and P n T = π T . This property is in contrast to the IPF algorithm (7) for which the marginals at the initial and final times are not preserved. We highlight this duality between IPF (7) and IMF (8) in Table 1.

We conclude this section with a theoretical analysis of IMF. First, we start by showing a Pythagorean theorem for both the Markovian projection and the reciprocal projection.

Lemma 6. Under mild assumptions, if M ∈ M , Π ∈ R ( Q ) and KL(Π | M ) < + ∞ , we have

If KL( M | Π) < + ∞ , we have

Using Lemma 6, we have the following proposition.

Proposition 7. Under mild assumptions, we have KL( P n +1 | P SB ) ≤ KL( P n | P SB ) < ∞ , and lim n → + ∞ KL( P n | P n +1 ) = 0 .

Hence, for the IMF sequence ( P n ) n ∈ N , the Markov path measures ( P 2 n +1 ) n ∈ N are getting closer to the reciprocal class, while the reciprocal path measures ( P 2 n +2 ) n ∈ N are getting closer to the set of Markov measures. Proposition 7 should be compared with (Rüschendorf, 1995, Proposition 2.1, Equation (2.16)) which shows that, for the IPF sequence ( ˜ P n ) n ∈ N , we have lim n → + ∞ KL( ˜ P n +1 | ˜ P n ) = 0 . This result is similar to Proposition 7 but for the forward Kullback-Leibler divergence.

Using Proposition 7, we finally prove the convergence of the IMF sequence ( P n ) n ∈ N to the Schrödinger Bridge. This result was first shown in the concurrent work (Peluchetti, 2023, Theorem 2). We present a simpler proof in Appendix C.6.

Theorem 8. Under mild assumptions, the IMF sequence ( P n ) n ∈ N admits a unique fixed point P glyph[star] = P SB , and lim n → + ∞ KL( P n | P glyph[star] ) = 0 .

4 Diffusion Schrödinger Bridge Matching

In this section, we present Diffusion Schrödinger Bridge Matching (DSBM), a practical algorithm for solving the SB problem obtained by combining the IMF procedure with Bridge Matching.

Iterative Markovian Fitting in practice. IMF alternatively projects on the Markov class M and the reciprocal class R ( Q ) . We denote M n +1 = P 2 n +1 ∈ M and Π n = P 2 n ∈ R ( Q ) . Assuming we know how to sample from the bridge Q | 0 ,T given the initial and terminal conditions, sampling from the reciprocal projection proj R ( Q ) ( M ) is simple: First, sample ( X 0 , X T ) from the joint distribution M 0 ,T . 5 Then, sample from the bridge Q | 0 ,T ( ·| X 0 , X T ) . The bottleneck of IMF is in the computation of Markovian projections. By Definition 1, M glyph[star] = proj M (Π) is associated with the process

where { v θ : θ ∈ Θ } is a parametric family of functions, usually given by a neural network. The optimal v θ glyph[star] ( t, x t ) = σ 2 t E Π T | t [ ∇ log Q T | t ( X T | X t ) | X t = x t ] for any t ∈ [0 , T ] and x t ∈ R d .

With the above two procedures for computing proj R ( Q ) ( M ) and proj M (Π) , we can now describe a numerical method implementing IMF (8). Let Π 0 = Π 0 0 ,T Q | 0 ,T where Π 0 0 = π 0 , Π 0 T = π T . Learn M 1 ≈ proj M (Π 0 ) given by (9) with v θ glyph[star] given by (10). Next, sample from Π 1 = proj R ( Q ) ( M 1 ) = M 1 0 ,T Q | 0 ,T by sampling from M 1 0 ,T and reconstructing the bridge Q | 0 ,T . We iterate the process to obtain a sequence (Π n , M n +1 ) n ∈ N . In practice, this algorithm performs poorly (see Figure 3), since the approximate minimization (10) for computing M n +1 may not admit M n +1 T = π T exactly as in Proposition 2. Instead, we incur a bias between M n +1 T and π T which accumulates for each n ∈ N .

To mitigate this problem, we alternate between a forward Markovian projection and a backward Markovian projection. This procedure is justified by the following proposition.

Proposition 9. Assume that Π = Π 0 ,T Q | 0 ,T with Q associated with d X t = f t ( X t )d t + σ t d B t . Under mild conditions, the Markovian projection M glyph[star] = proj M (Π) is associated with both

In Proposition 9, (11) is the definition of the Markovian projection, see Definition 1. However, (12) is an equivalent representation as a time-reversal . In practice, ( Y t ) t ∈ [0 ,T ] is approximated with

The optimal v φ glyph[star] ( t, x t ) = σ 2 t E Π 0 | t [ ∇ log Q t | 0 ( X t | X 0 ) | X t = x t ] for any t ∈ [0 , T ] and x t ∈ R d .

Algorithm 1 Diffusion Schrödinger Bridge Matching

1: Input: Joint distribution Π 0 0 ,T , tractable bridge Q | 0 ,T , number of outer iterations N ∈ N .

3: for n ∈ { 0 , . . . , N -1 } do

2: Let Π 0 = Π 0 0 ,T Q | 0 ,T .

4: Learn v φ glyph[star] using (14) with Π = Π 2 n .

5: Let M 2 n +1 be given by (13).

6:

Let

Π

2

n

+1

=

M

2

n

+1

Q

0

,T

.

8: Let M 2 n +2 be given by (9).

0 ,T | 7: Learn v θ glyph[star] using (10) with Π = Π 2 n +1 .

9: Let Π 2 n +2 = M 2 n +2 0 ,T Q | 0 ,T .

11: Output: v θ glyph[star] , v φ glyph[star]

10: end for

Note that X 0 ∼ π 0 in the forward projection, while Y 0 ∼ π T in the backward projection. Therefore, using the backward projection removes the bias on π T accumulated from the forward projection. Leveraging the time-symmetry of the Markovian projection and alternating between (13) and (9) yields the DSBM methodology summarized in Algorithm 1.

It is also possible to learn both the forward and backward processes at each step, and enforce that the backward and forward processes match. We explore this in Appendix G.

Initialization coupling. We now relate Algorithm 1 to the classical IPF and practical algorithms such as DSB (De Bortoli et al., 2021). Instead of initializing DSBM with Π 0 0 ,T given by a coupling between π 0 , π T , if we initialize it by Π 0 0 ,T = Q 0 ,T where Q 0 = π 0 and Q T | 0 is given by the reference process defined in (1), then DSBM also recovers the IPF iterates used in DSB.

Proposition 10. Suppose the families of functions { v θ : θ ∈ Θ } and { v φ : φ ∈ Φ } are rich enough so that they can model the optimal vector fields. Let (Π n , M n +1 ) n ∈ N be the optimal DSBM sequence in Algorithm 1 initialized with Π 0 0 ,T = Q 0 ,T , and let ( ˜ P n ) n ∈ N be the optimal DSB sequence given by the IPF iterates in (7) . Then for any n ∈ N , n ≥ 1 , we have M n = ˜ P n .

We will thus call DSBM-IPF, the DSBM algorithm initialized with the joint distribution given by the forward reference process Π 0 0 ,T = Q 0 ,T ; and DSBM-IMF, the DSBM algorithm initialized with an independent coupling Π 0 0 ,T = π 0 ⊗ π T . However, the training procedure of DSBM-IPF is very different from the one of (De Bortoli et al., 2021; Chen et al., 2022). In existing works, ˜ P n +1 is

obtained as the time-reversal of ˜ P n which requires full trajectories from ˜ P n , see e.g. (De Bortoli et al., 2021, Proposition 6). In contrast, in Algorithm 1 we only use the coupling M n 0 ,T to create the bridge measure Π n = M n 0 ,T Q | 0 ,T . By doing so, (i) the losses (10) and (14) can be easily evaluated at any time t ∈ [0 , T ] ; (ii) the trajectory caching procedure in DSBM is more computationally and memory efficient; (iii) while every IPF iteration ˜ P n is also supposed to be in R ( Q ) , in practice one can observe a forgetting of the bridge Q | 0 ,T (Fernandes et al., 2021). In DSBM, this effect is countered by explicit projections on the reciprocal class. See Appendix F for more details.

Probability flow ODE. At equilibrium of DSBM, we have that ( Y t ) t ∈ [0 ,T ] given by (13) is the time reversal of ( X t ) t ∈ [0 ,T ] given by (9) and are both associated with the optimal Schrödinger Bridge path measure P glyph[star] . As a result, we have that v φ glyph[star] ( t, x ) = -v θ glyph[star] ( t, x ) + σ 2 t ∇ log P glyph[star] t ( x ) . Hence, a probability flow ( Z glyph[star] t ) t ∈ [0 ,T ] such that Law( Z glyph[star] t ) = P glyph[star] t for any t ∈ [0 , T ] is given by

See also De Bortoli et al. (2021); Chen et al. (2022) for derivation of this result. Note however that the path measure induced by ( Z glyph[star] t ) t ∈ [0 ,T ] does not correspond to P glyph[star] ; in particular, ( Z glyph[star] 0 , Z glyph[star] T ) is not an entropic OT plan. However, since for any t ∈ [0 , T ] , Z glyph[star] t has marginal distribution P glyph[star] t , we can compute the log-likelihood of the model (Song et al., 2021b; Huang et al., 2021).

5 Related Work

Markovian projection and Bridge Matching. The concept of Markovian projection has been rediscovered multiple times (Krylov, 1984; Gyöngy, 1986; Dupire, 1994). In the machine learning context, this was first proposed by Peluchetti (2021) to define Bridge Matching models. More recently, Liu et al. (2022b) derived theoretical properties of the Markovian projection in Proposition 2, first part of Lemma 6, and applied Bridge Matching for learning data on discrete and constrained domains.

Bridge and Flow Matching. Flow Matching corresponds to deterministic bridges with deterministic samplers (ODEs) and has been under active study (Liu et al., 2023b; Liu, 2022; Lipman et al., 2023; Albergo and Vanden-Eijnden, 2023; Heitz et al., 2023; Pooladian et al., 2023; Tong et al., 2023). Denoising Diffusion Implicit Models (DDIM) (Song et al., 2021a) can also be formulated as a discrete-time version of Flow Matching, see Liu et al. (2023b). These models have been extended to the Riemannian setting by Chen and Lipman (2023). Recently, Albergo et al. (2023) studied the influence of stochasticity in the bridge, through the concept of stochastic interpolants. Liu et al. (2023a); Delbracio and Milanfar (2023) used Bridge Matching to perform image restoration tasks and noted benefits of stochasticity empirically. Closely related to our work is the Rectified Flow algorithm of Liu et al. (2023b), which corresponds to an iterative Flow Matching procedure in order to improve the straightness of the flow and thus eases its simulation. An iterative rectifying procedure using stochastic interpolants is also proposed in (Albergo et al., 2023, Section 3.5). Our proposed DSBM-IMF algorithm is closest to Rectified Flow, which can be seen as the deterministic limiting case of DSBM-IMF as σ → 0 . However, there are a few important theoretical and practical differences. Most notably, we adopt the SDE approach which is crucial for the validity of Proposition 5 as well as for the empirical performance of DSBM. We discuss further distinctions between DSBM and Rectified Flow in Appendix A.3.

Diffusion Schrödinger Bridge. Schrödinger Bridges (Schrödinger, 1932) are ubiquitous in probability theory (Léonard, 2014b) and stochastic control (Dai Pra, 1991; Chen et al., 2021). More recently, they have been used for generative modeling: De Bortoli et al. (2021) introduced the DSB algorithm and Vargas et al. (2021); Chen et al. (2022) introduced similar algorithms. The case of Dirac delta terminal distribution was investigated by Wang et al. (2021). These methods were later extended to solve conditional simulation and more general control problems (Shi et al., 2022; Thornton et al., 2022; Liu et al., 2022a; Chen et al., 2023; Tamir et al., 2023). In Somnath et al. (2023), SBs are learned using one Bridge Matching iteration, assuming access to the true Schrödinger static coupling. Our proposed method DSBM-IPF is closest to DSB, but with improved continous-time training and projections on the reciprocal class which mitigate two limitations of DSB. Concurrently with our work, Peluchetti (2023) independently introduced the DSBM-IMF approach (named IDBM therein).

6 Experiments

2D Experiments. Wefirst show our proposed methods can generate correct samples and learn lower kinetic energy transport maps in some 2D examples. We compare our method DSBM with flow-based methods including Flow Matching (FM) (Lipman et al., 2023), Conditional Flow Matching (CFM), OT-CFM (Tong et al., 2023), and Rectified Flow (RF) (Liu et al., 2023b); and other SB methods including DSB (De Bortoli et al., 2021) and SB-CFM (Tong et al., 2023). OT-CFM and SB-CFM utilizes sample-based mini-batch OT or EOT solvers (Fatras et al., 2021; Flamary et al., 2021) to define an approximate OT or SB static coupling ˜ Π OT 0 ,T or ˜ Π SB 0 ,T , see also Pooladian et al. (2023); Stromme (2023). We can also utilize this idea in the DSBM-IMF framework, which corresponds to using the initialization coupling Π 0 0 ,T = ˜ Π SB 0 ,T in Algorithm 1. This approximate SB coupling ˜ Π SB 0 ,T also satisfies ˜ Π SB 0 = π 0 and ˜ Π SB T = π T but can provide a better initialization than the independent coupling Π 0 0 ,T = π 0 ⊗ π T . We name this approach DSBM-IMF+. The rest of the methods do not use OT solvers. DSB and DSBM directly learn the EOT map as the solution of the diffusion process.

In Table 2, we show the 2-Wasserstein distance between the true and generated samples, as well as the integrated path energy defined as E [ ∫ T 0 || v ( t, Z t ) || 2 d t ] where v is the learned drift along the ODE trajectory Z t . For direct comparability, we report for DSBM using its probability flow ODE. Lower path energies represent shorter (and potentially easier to integrate) trajectories. We find that in this low dimensional setting, OT-CFM performs the best by utilizing OT solvers, but DSBM outperforms FM and CFM when OT solvers are not used. Further, DSBM outperforms DSB on all datasets, suggesting DSBM solves the SB problem with higher accuracy. The results also show that among SB methods, DSBM-IMF+ can achieve lower sampling error than DSBM-IPF and DSBM-IMF. It also performs better than SB-CFM on 3 of the datasets and achieve lower path energy on all datasets. Finally, we find Rectified Flow achieves lower sampling error than DSBM except for the moons-8gaussians task, for which DSBM is significantly more accurate. Since RF can be informally seen as DSBM in the case σ → 0 , this suggests the optimal σ varies for each task and between generative and general transfer tasks. Figure 2 visualizes how σ affects the straightness and sample quality of learned transport maps between two mixture distributions.

High-Dimensional Gaussian Experiment. We next perform the Gaussian transport experiment in De Bortoli et al. (2021) with dimension d = 50 to verify the scalability of our proposed approach. The true SB can be computed analytically in this case (Bunne et al., 2023). In Figure 3, we plot the convergence of the learned mean E [ X 0 ] , variance Var( X 0 ) , and covariance Cov( X 0 , X T ) between times 0 , T . We also consider RF and a related baseline IMF-b, which performs IMF numerically but only in the backward direction. All methods converge approximately to the correct mean. However,

0=0.0

0=0.1

0=0.3

0=1.0

0=3.0

0

5

10

15

20

Iteration

-

0

.

10

-

0

.

05

0

.

00

Mean

DSB

IMF-b

DSBM-IPF

DSBM-IMF

RF

0

5

10

15

20

Iteration

0

.

9

1

.

0

1

.

1

0

5

10

15

20

Iteration

0

.

50

0

.

55

0

.

60

0

.

65

Covariance

(a) OT-CFM

(b) DSB

(c) DSBM-IPF

0

KL

×

10

-

3

d

= 5

d

= 20

d

= 50

DSB

3.26±1.60

13.0±3.49

32.8±1.28

SB-CFM

1.45±0.73

12.3±1.47

49.4±3.91

DSBM-IPF

1.23±0.23

4.42±0.76

8.75±0.87

DSBM-IMF

1.34±0.51 5.05±0.95 9.76±1.67

DSB

200

100

50

20

10

DSBM-IPF

DSBM-IMF

BM

CFM

OT-CFM

RF

50

100

150

200

250

Training step (×1000)

the variance estimates become inaccurate for RF and IMF-b. Among SB methods, DSB and IMF-b also gave inaccurate SB covariance estimates as the number of iteration increases. On the other hand, DSBM does not suffer from this issue. In Table 3, we further quantify the accuracy and compare with SB-CFM (Tong et al., 2023) by computing the KL divergence between the marginal distributions of the learned process P t and the true SB P SB t . Our proposed methods achieve similar KL divergence as SB-CFM in dimension d = 5 , but are much more accurate in higher dimensions.

MNIST, EMNIST transfer. We test our method for domain transfer between MNIST digits and EMNIST letters as in De Bortoli et al. (2021). We compare DSBM as a direct substitute of DSB, and also with Bridge Matching (BM) (Peluchetti, 2021; Liu et al., 2022b), CFM, OT-CFM and RF. We plot some output samples from different algorithms in Figure 4 and the convergence of FID score in Figure 5. We find that OT-CFM becomes less applicable in higher dimensions and produces samples of worse quality (Figure 4a). On the other hand, image quality deteriorates during training of DSB and RF. DSBM achieves higher quality samples visually, and does not suffer from deterioration. It is also about 30% more efficient than DSB in terms of runtime.

CelebA transfer. Next, we evaluate and perform some ablations of our method on a transfer task on the CelebA 64 × 64 dataset. We consider the images given by the tokens male / old and female / young . In Figures 6 and 7, we show that as σ increases, the quality of the images (as measured by the FID score) increases until σ is too high, but the alignment (as measured by LPIPS) between the generated image and the original sample decreases. Additionally, we investigate the dependency between σ and image dimension in Figure 8. In particular, for the same σ = 1 , the outputs of DSBM for CelebA 128 × 128 are better aligned with the original data than for CelebA 64 × 64 . This is in agreement with the observations of Chen (2023); Hoogeboom et al. (2023) that the noise schedule in diffusion models should scale with the resolution.

AFHQtransfer. Wedemonstrate the scalability of our method on an additional transfer experiment on the AFHQ 512 × 512 dataset between the classes cat and wild . The results are shown in Figure 9. On this higher-dimensional problem, we observe that DSBM can also generate realistic samples which are similar to the input.

Unpaired Fluid Flows Downscaling. Finally, we apply DSBM to perform downscaling of geophysical fluid dynamics, i.e. super-resolution of low resolution spatial data. We use the dataset in (Bischoff and Deck, 2023), which consists of unpaired low ( 64 × 64 ) and high ( 512 × 512 ) resolution fields. As shown in Figure 10, DSBM is able to learn high resolution reconstructions by only slightly noising the low resolution input. In contrast, Bischoff and Deck (2023) use two diffusion models in forward and backward directions (Diffusion-fb) based on Meng et al. (2022), which improves over the Random baseline. Figure 11 shows that DSBM-IPF and DSBM-IMF achieve much lower glyph[lscript] 2

FID

10

-

2

10

-

1

10

0

10

1

30

40

50

60

70

10

-

2

10

-

1

10

0

10

1

0

.

05

0

.

10

0

.

15

0

.

20

(a)

(b)

2

4

8

16

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

1

.

2

L2 of high and low res

Supersaturation

2

4

8

16

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

Vorticity

Random

Diffusion-fb

DSBM-IPF

DSBM-IMF

Data wavenumber

distances for all frequency classes in the dataset than Diffusion-fb (and thus Random), indicating DSBM is able to reconstruct high resolution fields consistent with the low resolution source.

7 Discussion

In this work, we introduce IMF, a new methodology for learning Schrödinger Bridges. IMF is an alternative to the classical IPF and can be interpreted as its dual. Building on this new framework, we present two practical algorithms, DSBM-IPF and DSBM-IMF, for learning SBs. These algorithms mitigate the time-discretization and bias accumulation issues of existing methods. However, DSBM still has some limitations. First, our results suggest DSBM is most effective for solving general transport problems. For generative modeling, we only find minor improvements compared to Bridge and Flow Matching on CIFAR-10 (see Appendix I.6). Second, while DSBM is more efficient than DSB, it still requires sampling from the learned process during the caching step. Finally, the EOT problem becomes more difficult to solve numerically for small values of σ .

In future work, we would like to further investigate the differences between DSBM-IMF and DSBMIPF. IMF also appears useful for developing a better understanding of the Rectified Flow algorithm (Liu et al., 2023b), as IMF minimizes a clear objective (6) and Rectified Flow can be seen as a limiting case of it. Finally, Rectified Flow has also been extended to solve OT problems with general convex costs by Liu (2022), and it would be interesting to derive a SB version of this extension.

References

Albergo, M. S., Boffi, N. M., and Vanden-Eijnden, E. (2023). Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.08797 .

Albergo, M. S. and Vanden-Eijnden, E. (2023). Building normalizing flows with stochastic interpolants. In International Conference on Learning Representations .

Anderson, B. D. (1982). Reverse-time diffusion equation models. Stochastic Processes and their Applications , 12(3):313-326.

Banerjee, A., Guo, X., and Wang, H. (2005). On the optimality of conditional expectation as a Bregman predictor. IEEE Transactions on Information Theory , 51(7):2664-2669.

Barczy, M. and Kern, P. (2013). Representations of multidimensional linear process bridges. Random Operators and Stochastic Equations , 21(2):159-189.

Benamou, J.-D. and Brenier, Y. (2000). A computational fluid mechanics solution to the MongeKantorovich mass transfer problem. Numerische Mathematik , 84(3):375-393.

Bernton, E., Heng, J., Doucet, A., and Jacob, P. E. (2019). Schrödinger bridge samplers. arXiv preprint arXiv:1912.13170 .

Bischoff, T. and Deck, K. (2023). Unpaired downscaling of fluid flows with diffusion bridges. arXiv preprint arXiv:2305.01822 .

Bogachev, V. I., Krasovitskii, T. I., and Shaposhnikov, S. V. (2021). On uniqueness of probability solutions of the Fokker-Planck-Kolmogorov equation. Sbornik: Mathematics , 212(6):745.

Bunne, C., Hsieh, Y.-P., Cuturi, M., and Krause, A. (2023). The Schrödinger bridge between Gaussian measures has a closed form. In International Conference on Artificial Intelligence and Statistics .

Bunne, C., Papaxanthos, L., Krause, A., and Cuturi, M. (2022). Proximal optimal transport modeling of population dynamics. In International Conference on Artificial Intelligence and Statistics , pages 6511-6528. PMLR.

Caluya, K. F. and Halder, A. (2021). Wasserstein proximal algorithms for the Schrödinger bridge problem: Density control with nonlinear drift. IEEE Transactions on Automatic Control , 67(3):11631178.

Chen, R. T. and Lipman, Y. (2023). Riemannian flow matching on general geometries. arXiv preprint arXiv:2302.03660 .

Chen, T. (2023). On the importance of noise scheduling for diffusion models. arXiv preprint arXiv:2301.10972 .

Chen, T., Liu, G.-H., Tao, M., and Theodorou, E. A. (2023). Deep momentum multi-marginal Schrödinger bridge. arXiv preprint arXiv:2303.01751 .

Chen, T., Liu, G.-H., and Theodorou, E. A. (2022). Likelihood training of Schrödinger bridge using forward-backward SDEs theory. In International Conference on Learning Representations .

Chen, Y., Georgiou, T., and Pavon, M. (2016). Entropic and displacement interpolation: a computational approach using the Hilbert metric. SIAM Journal on Applied Mathematics , 76(6):2375-2396.

Chen, Y., Georgiou, T. T., and Pavon, M. (2021). Optimal transport in systems and control. Annual Review of Control, Robotics, and Autonomous Systems , 4.

Choi, Y., Uh, Y., Yoo, J., and Ha, J.-W. (2020). Stargan v2: Diverse image synthesis for multiple domains. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition .

Chung, K. L. and Walsh, J. B. (2006). Markov processes, Brownian motion, and Time Symmetry , volume 249. Springer Science & Business Media.

Csiszár, I. (1975). I-divergence geometry of probability distributions and minimization problems. The Annals of Probability , 3(1):146-158.

Cuturi, M. (2013). Sinkhorn distances: Lightspeed computation of optimal transport. In Advances in Neural Information Processing Systems .

Dai Pra, P. (1991). A stochastic control approach to reciprocal diffusion processes. Applied Mathematics and Optimization , 23(1):313-329.

De Bortoli, V., Thornton, J., Heng, J., and Doucet, A. (2021). Diffusion Schrödinger bridge with applications to score-based generative modeling. In Advances in Neural Information Processing Systems .

Delbracio, M. and Milanfar, P. (2023). Inversion by direct iteration: An alternative to denoising diffusion for image restoration. arXiv preprint arXiv:2303.11435 .

Dupire, B. (1994). Pricing with a smile. Risk , 7(1):18-20.

Essid, M. and Pavon, M. (2019). Traversing the Schrödinger bridge strait: Robert Fortet's marvelous proof redux. Journal of Optimization Theory and Applications , 181(1):23-60.

Fatras, K., Zine, Y., Majewski, S., Flamary, R., Gribonval, R., and Courty, N. (2021). Minibatch optimal transport distances; analysis and applications. arXiv preprint arXiv:2101.01792 .

Fernandes, D. L., Vargas, F., Ek, C. H., and Campbell, N. D. (2021). Shooting Schrödinger's cat. In Fourth Symposium on Advances in Approximate Bayesian Inference .

Feydy, J., Charlier, B., Vialard, F.-X., and Peyré, G. (2017). Optimal transport for diffeomorphic registration. In Medical Image Computing and Computer Assisted Intervention- MICCAI , pages 291-299.

Finlay, C., Gerolin, A., Oberman, A. M., and Pooladian, A.-A. (2020). Learning normalizing flows from entropy-Kantorovich potentials. arXiv preprint arXiv:2006.06033 .

Flamary, R., Courty, N., Gramfort, A., Alaya, M. Z., Boisbunon, A., Chambon, S., Chapel, L., Corenflos, A., Fatras, K., Fournier, N., Gautheron, L., Gayraud, N. T., Janati, H., Rakotomamonjy, A., Redko, I., Rolet, A., Schutz, A., Seguy, V ., Sutherland, D. J., Tavenard, R., Tong, A., and Vayer, T. (2021). Pot: Python optimal transport. Journal of Machine Learning Research , 22(78):1-8.

Föllmer, H. (1988). Random fields and diffusion processes. In École d'Été de Probabilités de Saint-Flour XV-XVII, 1985-87 , pages 101-203. Springer.

Fortet, R. (1940). Résolution d'un système d'équations de M. Schrödinger. Journal de Mathématiques Pures et Appliqués , 1:83-105.

Goodfellow, I. J., Pouget-Abadie, J., Mirza, M., Xu, B., Warde-Farley, D., Ozair, S., Courville, A., and Bengio, Y. (2014). Generative adversarial networks. arXiv preprint arXiv:1406.2661 .

Gyöngy, I. (1986). Mimicking the one-dimensional marginal distributions of processes having an Itô differential. Probability Theory and Related Fields , 71:501-516.

Haussmann, U. G. and Pardoux, E. (1986). Time reversal of diffusions. The Annals of Probability , 14(4):1188-1205.

Heitz, E., Belcour, L., and Chambon, T. (2023). Iterative α -(de)blending: a minimalist deterministic diffusion model. In SIGGGRAPH .

Heng, J., Bortoli, V . D., Doucet, A., and Thornton, J. (2021). Simulating diffusion bridges with score matching. arXiv preprint arXiv:2111.07243 .

Ho, J., Jain, A., and Abbeel, P. (2020). Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems .

Hoogeboom, E., Heek, J., and Salimans, T. (2023). simple diffusion: End-to-end diffusion for high resolution images. arXiv preprint arXiv:2301.11093 .

Huang, C.-W., Lim, J. H., and Courville, A. C. (2021). A variational perspective on diffusion-based generative models and score matching. In Advances in Neural Information Processing Systems .

Hyvärinen, A. (2005). Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research , 6(4).

Krylov, N. (1984). Once more about the connection between elliptic operators and Itô's stochastic equations. In Statistics and Control of Stochastic Processes, Steklov Seminar , pages 214-229.

Kullback, S. (1968). Probability densities with given marginals. The Annals of Mathematical Statistics , 39(4):1236-1243.

Kullback, S. (1997). Information Theory and Statistics . Dover Publications, Inc., Mineola, NY. Reprint of the second (1968) edition.

Léonard, C. (2012). Girsanov theory under a finite entropy condition. In Séminaire de Probabilités XLIV , pages 429-465. Springer.

Léonard, C. (2014a). Some properties of path measures. In Séminaire de Probabilités XLVI , pages 207-230. Springer.

Léonard, C. (2014b). A survey of the Schrödinger problem and some of its connections with optimal transport. Discrete & Continuous Dynamical Systems-A , 34(4):1533-1574.

Léonard, C., Rœlly, S., Zambrini, J.-C., et al. (2014). Reciprocal processes. a measure-theoretical point of view. Probability Surveys , 11:237-269.

Lipman, Y., Chen, R. T., Ben-Hamu, H., Nickel, M., and Le, M. (2023). Flow matching for generative modeling. In International Conference on Learning Representations .

Liu, G.-H., Chen, T., So, O., and Theodorou, E. A. (2022a). Deep generalized Schrödinger bridge. In Advances in Neural Information Processing Systems .

Liu, G.-H., Vahdat, A., Huang, D.-A., Theodorou, E. A., Nie, W., and Anandkumar, A. (2023a). I 2 SB: Image-to-image Schrödinger bridge. arXiv preprint arXiv:2302.05872 .

Liu, Q. (2022). Rectified flow: A marginal preserving approach to optimal transport. arXiv preprint arXiv:2209.14577 .

Liu, X., Gong, C., and Liu, Q. (2023b). Flow straight and fast: Learning to generate and transfer data with rectified flow. In International Conference on Learning Representations .

Liu, X., Wu, L., Ye, M., and Liu, Q. (2022b). Let us build bridges: Understanding and extending diffusion generative models. arXiv preprint arXiv:2208.14699 .

Liu, Z., Luo, P., Wang, X., and Tang, X. (2015). Deep learning face attributes in the wild. In International Conference on Computer Vision .

Meng, C., He, Y., Song, Y., Song, J., Wu, J., Zhu, J.-Y., and Ermon, S. (2022). SDEdit: Guided image synthesis and editing with stochastic differential equations. In International Conference on Learning Representations .

Palmowski, Z. and Rolski, T. (2002). A technique for exponential change of measure for Markov processes. Bernoulli , pages 767-785.

Pavon, M., Trigila, G., and Tabak, E. G. (2021). The data-driven Schrödinger bridge. Communications on Pure and Applied Mathematics , 74:1545-1573.

Peluchetti, S. (2021). Non-denoising forward-time diffusions. https://openreview.net/forum? id=oVfIKuhqfC .

Peluchetti, S. (2023). Diffusion bridge mixture transports, Schrödinger bridge problems and generative modeling. arXiv preprint arXiv:2304.00917 .

Peyré, G. and Cuturi, M. (2019). Computational optimal transport. Foundations and Trends® in Machine Learning , 11(5-6):355-607.

Pooladian, A.-A., Ben-Hamu, H., Domingo-Enrich, C., Amos, B., Lipman, Y., and Chen, R. (2023). Multisample flow matching: Straightening flows with minibatch couplings. In International Conference on Machine Learning .

Rogers, L. C. G. and Williams, D. (2000). Diffusions, Markov processes, and Martingales. Vol. 2 . Cambridge Mathematical Library. Cambridge University Press, Cambridge. Itô calculus, Reprint of the second (1994) edition.

Rüschendorf, L. (1995). Convergence of the iterative proportional fitting procedure. The Annals of Statistics , 23(4):1160-1174.

Rüschendorf, L. and Thomsen, W. (1993). Note on the Schrödinger equation and i-projections. Statistics & Probability letters , 17(5):369-375.

Schrödinger, E. (1932). Sur la théorie relativiste de l'électron et l'interprétation de la mécanique quantique. Annales de l'Institut Henri Poincaré , 2(4):269-310.

Shi, Y., De Bortoli, V., Deligiannidis, G., and Doucet, A. (2022). Conditional simulation using diffusion Schrödinger bridges. In Uncertainty in Artificial Intelligence .

Sinkhorn, R. (1967). Diagonal equivalence to matrices with prescribed row and column sums. The American Mathematical Monthly , 74(4):402-405.

Somnath, V. R., Pariset, M., Hsieh, Y.-P., Martinez, M. R., Krause, A., and Bunne, C. (2023). Aligned diffusion Schrödinger bridges. arXiv preprint arXiv:2302.11419 .

Song, J., Meng, C., and Ermon, S. (2021a). Denoising diffusion implicit models. In International Conference on Learning Representations .

Song, K.-U. (2022). Applying regularized Schrödinger-bridge-based stochastic process in generative modeling. arXiv preprint arXiv:2208.07131 .

Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. (2021b). Scorebased generative modeling through stochastic differential equations. In International Conference on Learning Representations .

Stromme, A. (2023). Sampling from a Schrödinger bridge. In International Conference on Artificial Intelligence and Statistics .

Tamir, E., Trapp, M., and Solin, A. (2023). Transport with support: Data-conditional diffusion bridges. arXiv preprint arXiv:2301.13636 .

Thornton, J., Hutchinson, M., Mathieu, E., De Bortoli, V., Teh, Y. W., and Doucet, A. (2022). Riemannian diffusion Schrödinger bridge. arXiv preprint arXiv:2207.03024 .

Tong, A., Malkin, N., Huguet, G., Zhang, Y., Rector-Brooks, J., Fatras, K., Wolf, G., and Bengio, Y. (2023). Conditional flow matching: Simulation-free dynamic optimal transport. arXiv preprint arXiv:2302.00482 .

van Erven, T. and Harremoes, P. (2014). Rényi divergence and kullback-leibler divergence. IEEE Transactions on Information Theory , 60(7):3797-3820.

Vargas, F., Thodoroff, P., Lamacraft, A., and Lawrence, N. (2021). Solving Schrödinger bridges via maximum likelihood. Entropy , 23(9):1134.

Villani, C. (2009). Optimal transport , volume 338 of Grundlehren der Mathematischen Wissenschaften [Fundamental Principles of Mathematical Sciences] . Springer-Verlag, Berlin. Old and new.

Vincent, P. (2011). A connection between score matching and denoising autoencoders. Neural Computation , 23(7):1661-1674.

Wang, G., Jiao, Y., Xu, Q., Wang, Y., and Yang, C. (2021). Deep generative learning via Schrödinger bridge. In International Conference on Machine Learning .

Outline of the Appendix

In Appendix A, we first clarify the relationship between different methods in the existing literature and our proposed DSBM framework. In Appendix B, we focus on the family of linear SDEs, and draw a link between the parameterization of bridges in this paper and the stochastic interpolant in Albergo et al. (2023). In Appendix C, we give proofs for results in the main text. In Appendix D, we present additional theoretical results for IMF in the Gaussian case. In Appendix E, we derive the discrete-time version of Markovian projection. In Appendix F, we explain the benefits of DSBM compared to DSB in more detail. In Appendix G, we describe a method for learning the forward and backward processes jointly and propose a consistency loss between the forward and backward processes. In Appendix H, we present additional methodological details for a practical scaling of the loss function to reduce variance, similar to standard Denoising Diffusion Models. In Appendix I, we give further details for all experiments and additional experimental results. Finally, we discuss broader impacts of our work in Appendix J.

A Discussion of Existing Works

A.1 Bridge Matching and Flow Matching Models

In this section, we clarify the relationship between variants of Flow Matching and show that they are equivalent under some conditions. We follow the nomenclature of Tong et al. (2023). We refer to the algorithm originally proposed in Lipman et al. (2023) using linear probability paths and described in (Tong et al., 2023, Section 4.1) as Flow Matching (FM), and the algorithm proposed in (Tong et al., 2023, Section 4.2) as Conditional Flow Matching (CFM). There is a small constant parameter σ min in both algorithms, which controls the smoothing of the modeled distribution. We consider the case σ min = 0 . Then CFM recovers exactly the 1st iteration of Rectified Flow (Liu et al., 2023b). Furthermore, FM, CFM and the 1st iteration of Rectified Flow are all equivalent when performing generative modeling with a standard Gaussian π 0 . We refer to them collectively as Flow Matching models (FMMs) as they only differ in the smoothing method. We also present them all under the Bridge Matching framework. These models can also be interpreted in the context of stochastic interpolants (Albergo and Vanden-Eijnden, 2023; Albergo et al., 2023). Finally, we present recent applications of Bridge Matching and show that some of the objectives in Somnath et al. (2023); Liu et al. (2023a); Delbracio and Milanfar (2023) are identical.

Flow Matching and Conditional Flow Matching. In Flow Matching (FM), the objective (Lipman et al., 2023, Equation (21)) is

where Π t,T is given by π T ( X T )N( X t ; t T X T , (1 -t T ) 2 ) .

In Conditional Flow Matching (CFM), X 0 ,T t = t T X T +(1 -t T ) X 0 , with X 0 ∼ N(0 , Id) and the objective (Tong et al., 2023, Equation (16)) is given by

This is the same as (Liu et al., 2023b, Equation (1)). Furthermore, ( X T -X 0 ) /T = (1 -t T )( X T -X 0 ) / ( T -t ) = ( X T -X 0 ,T t ) / ( T -t ) , so the CFM objective is equivalent to

The optimal v θ ( t, x t ) = ( E Π T | t [ X T | X t = x t ] -x t ) / ( T -t ) . In the case of generative modeling, π 0 is a standard Gaussian distribution and Π 0 ,T is given by N( X 0 ; 0 , Id) π T ( X T ) . Thus, Π t,T is also given by π T ( X T )N( X 0 ,T t ; t T X T , (1 -t T ) 2 ) . Therefore, the FM (Lipman et al., 2023) and CFM (Tong et al., 2023) objectives are exactly the same. However, CFM is also applicable when π 0 is not Gaussian distributed, so CFM is a generalized version of FM 6 .

Stochastic Interpolant. In (Albergo and Vanden-Eijnden, 2023; Albergo et al., 2023), the concept of stochastic interpolant is introduced. In Albergo and Vanden-Eijnden (2023), the interpolation is deterministic (not necessarily linear), of the form I t ( x 0 , x T ) = α ( t ) x 0 + β ( t ) x T , while in Albergo et al. (2023), the interpolation is stochastic given by I t ( x 0 , x T ) = α ( t ) x 0 + β ( t ) x T + γ ( t ) Z for Z ∼ N(0 , Id) . In Albergo and Vanden-Eijnden (2023), an ODE is learned and the associated velocity field v θ is obtained by minimizing the following objective (Albergo and Vanden-Eijnden, 2023, Equation (9))

Link with Bridge Matching. When Q is associated with the Brownian motion ( σ B t ) t ∈ [0 ,T ] and σ → 0 in Bridge Matching, we recover the same objective (5) as the Flow Matching objective (16), since ∇ log Q T | t ( X T | X t ) = ( X T -X t ) / ( σ 2 ( T -t )) . Bridge Matching can also be applied to general distributions π 0 , π T ; i.e. π 0 does not have to be restricted to a Gaussian. Therefore, Bridge Matching is a generalized version of Flow Matching, see also (Liu et al., 2022b, Equation (10)).

Inverse problems and interpolation. Somnath et al. (2023) and Liu et al. (2023a) present Bridge Matching algorithms between aligned data ( X 0 , X T ) ∼ Π 0 ,T . The objectives (Somnath et al., 2023, Equation (8)) and (Liu et al., 2023a, Equation (12)) are equivalent to the Bridge Matching objective (3). The main difference between Liu et al. (2023a) and Somnath et al. (2023) resides in the choice of Π 0 ,T . In the case of Somnath et al. (2023), this choice is motivated by the access to aligned data with applications in biology assuming they are distributed as the true Schrödinger static coupling, i.e. Π 0 ,T = Π SB 0 ,T . In the case of Liu et al. (2023a), Π 0 ,T corresponds to a pairing between clean and corrupted images, e.g. with Π 0 = π 0 the distribution of clean images and Π T = π T the distribution of corrupted images obtained from the clean images using the degradation kernel Π T | 0 .

Finally, in (Delbracio and Milanfar, 2023, Equation (5)) the authors consider a reconstruction process of the form where here we have replaced F ( x t , t ) by E Π 0 | t [ X 0 | X t = x t ] . This is justified if the ‖ · ‖ p norm in (Delbracio and Milanfar, 2023, Equation (4)) is replaced by ‖·‖ 2 2 (or any Bregman Loss Function, see Banerjee et al. (2005)). In Delbracio and Milanfar (2023), Π 0 ,T corresponds to the joint distribution of clean and corrupted images as in Liu et al. (2023a). Exchanging the role of Π 0 and Π T , (17) can be rewritten equivalently as

Wethus obtain the optimal Flow Matching vector field v θ ( t, x t ) = ( E Π T | t [ X T | X t = x t ] -x t ) / ( T -t ) in (16). Note that Delbracio and Milanfar (2023) also incorporates a stochastic version of their objective (Delbracio and Milanfar, 2023, Equation (7)). It remains an open question whether this objective can be understood as a special instance of the Bridge Matching framework.

A.2 On DSBM and Existing Works

In this section, we show that the DSBM framework recovers the above existing algorithms for different choices of bridges Q | 0 ,T and couplings Π 0 0 ,T in Algorithm 1. For the independent coupling Π 0 0 ,T = π 0 ⊗ π T and Brownian bridge Q | 0 ,T (4) with diffusion parameter σ t = σ , the loss function (10) recovers the Brownian Bridge Matching loss (5). Letting σ → 0 , we recover Flow Matching (Lipman et al., 2023). In this case, further iterations repeating lines 7-9 in Algorithm 1 (with only forward projections) recover Rectified Flow (Liu et al., 2023b). If the coupling Π 0 0 ,T is given by an estimation of the OT map between π 0 and π T , then the first iteration recovers OT-CFM (Tong et al., 2023; Pooladian et al., 2023). Finally, for general bridges Q | 0 ,T , if we are given the optimal Schrödinger Bridge static coupling Π 0 0 ,T = Π SB 0 ,T , then the DSBM procedure converges in one iteration and we recover Somnath et al. (2023).

A.3 DSBM and Rectified Flow

We discuss the differences in more detail between our proposed DSBM method and Rectified Flow. Both Rectified Flow and DSBM are general frameworks for building transport maps between

two general distributions π 0 , π T . However, there are a few important theoretical and practical differences. Firstly, we adopt the SDE approach as opposed to the ODE approach in Rectified Flow. This distinction is crucial in theory, as Proposition 5, which guarantees the uniqueness of the characterization of SB, is valid only when σ t > 0 . Consequently, Rectified Flow is not guaranteed to converge to the dynamic optimal transport solution (see e.g. counterexample in Liu (2022)). In a following work, Liu (2022) established formal connections between Rectified Flow and OT when restricting the class of vector fields to gradient fields. In DSBM, the connection to OT is obtained by considering its entropy-regularized version. Furthermore, by adopting the SDE approach, we observe significant improvements of sample quality in our experiments when performing transport between two general distributions. This is in line with the theoretical analysis in Albergo et al. (2023). On the other hand, while Bridge Matching also achieves high sample quality using the SDE approach, the transported samples are much more dissimilar to the input data (see e.g. Figures 14, 15, 19). Lastly, Rectified Flow also performs Markovian projections iteratively, but only in the forward direction. Consequently, the bias in the learned marginals P n T is accumulated and cannot be corrected in later iterations, i.e. the first iteration of RF will achieve the most accurate marginal P 1 T . Subsequent iterations can improve the straightness of the flow, but at the cost of sampling accuracy of P n T . We observe in practice that this becomes particularly problematic if the first iteration of Rectified Flow (which is equivalent to CFM) fails to provide a good transport and learn an accurate P 1 T , e.g. in the case of moons-8gaussians (Table 2), Gaussian transport (Figure 3), and MNIST, EMNIST transfer (Figure 5 and Figure 13). As Rectified Flow cannot recover from this issue, we observe the accuracy of P n T only deteriorates in further iterations as n increases. In our methodology, we leverage Proposition 9 to perform forward and backward Bridge Matching, and we observe that the marginal accuracy is able to improve with iteration.

B The Design Space of Brownian Bridges

B.1 Relationship to Stochastic Interpolants

From stochastic interpolants to Brownian bridges. In this section, we draw a link between our parameterization of bridges and the one used in Albergo et al. (2023). In Albergo et al. (2023), a stochastic interpolant is defined as

where Z ∼ N(0 , Id) . Since their methodology and analysis mainly relies on the probability flow, they work with (18), which is easier to analyse. In our setting, as we deal mostly with diffusions, it is natural to parameterize Brownian bridges as follows

The goal of this section is to derive explicit formulas between the parameters ¯ α t , ¯ β t and ¯ γ t of (18) and the parameters α t , β t and γ t of (19). Consider ( X t ) t ∈ [0 ,T ] given by (19). We have that for any t ∈ [0 , T ]

where A t = ∫ t 0 α s d s . Therefore, we have that

Using this relationship, we get that the Markovian projection, see Definition 1, is given by

We have that

Using (21), we get that

In Albergo et al. (2023), it is shown that ∇ log M glyph[star] t ( x t ) = -E Π 0 ,T | t [ Z | X t = x t ] / ¯ γ t , where M glyph[star] is the Markovian projection. The probability flow associated with ( X glyph[star] t ) t ∈ [0 ,T ] is given by

Hence, we recover (Albergo et al., 2023, Theorem 2.6).

Non-Markov path measures. A natural question is whether (19) arises as the bridge measure of some Markov measure. For instance, if Q is associated with ( x 0 + B t ) t ∈ [0 ,T ] , then pinning the process at x T at time T , we get that the associated bridge measure Q | 0 ,T is given by

Therefore, we recover (19) with α t = β t = 1 T -t and γ t = 1 . Using (20), we get that ¯ α t = 1 -t T , ¯ β t = t T and ¯ γ 2 t = ( T -t ) t/T . We recover (4), upon noting that B t -t T B T is Gaussian with zero mean and variance ( T -t ) t/T .

More generally, we consider a Markov measure Q associated with ( X t ) t ∈ [0 ,T ] such that

We now derive the associated bridge measure Q | 0 ,T :

with Λ t = ∫ t 0 a s d s . We have that

Therefore, combining this result and (2), we get that Q | 0 ,T is associated with

In that case ( a t , c t ) t ∈ [0 ,T ] entirely parameterize ( α t , β t , γ t ) t ∈ [0 ,T ] . Hence, in the Ornstein-Uhlenbeck setting, if Q | 0 ,T is the bridge of a Markov measure, it is fully parameterized by two functions while in the non-Markov setting it is parameterized by three functions.

In this paper, we present our framework in the Markovian setting as the Schrödinger Bridge problem is usually defined with respect to Markov reference measures. However, our methodology could be extended in a straightforward fashion to the non-Markovian setting. This would allow for a further exploration of the design space of DSBM.

B.2 Linear SDE and Bridge Matching

In this section, we study further the diffusion bridge of linear SDEs. Arbitrary Markov measures can be chosen to build bridges; however, we want to be able to compute some representations of the bridge in an explicit way. More precisely, denoting ( X 0 ,T t ) t ∈ [0 ,T ] the diffusion bridge with x 0 , x T the initial and final condition, we want to have access to the following:

integral sampler : we want to have a formula to sample X 0 ,T t for any t ∈ [0 , T ] without having to run a stochastic process forward or backward.

forward sampler : we want to have a forward SDE for X 0 ,T t with explicit coefficients, which might depend on x T , running in a forward fashion terminating at x T .

backward sampler : we want to have a backward SDE for Y 0 ,T t = X 0 ,T T -t with explicit coefficients, which might depend on x 0 , running in a backward fashion terminating at x 0 .

We focus on linear SDEs of the form d X t = -αβ t X t d t + σβ 1 / 2 t d B t , which are particularly amenable, where ( β t ) t ∈ [0 ,T ] is a schedule with β ∈ C([0 , T ] , (0 , + ∞ )) .

B.2.1 Brownian motion

First, we consider the Brownian motion setting and let ( X t ) t ∈ [0 ,T ] be associated with Q with d X t = β 1 / 2 t d B t . We consider ( X 0 ,T t ) t ∈ [0 ,T ] conditioned at both ends X 0 ,T 0 = x 0 and X 0 ,T T = x T . First, using (Barczy and Kern, 2013, Theorem 3.3), we have that for any t ∈ [0 , T ]

with R ( s, t ) = ∫ t s β u d u = σ 2 ( B t -B s ) , where for any t ∈ [0 , T ] , B t = ∫ t 0 β s d s . Therefore, we get that

(22) defines the integral sampler . In addition, using (Barczy and Kern, 2013, Theorem 3.2), we have

where γ ( s, t ) = σ 2 ( B ( t ) -B ( s )) . Therefore, we get that

(23) defines the forward sampler . Finally, we derive the backward sampler by considering the time-reversal of the forward unconditional process (initialized at x 0 ). Following Haussmann and Pardoux (1986),

In addition, we have that

Hence, we get that for any t ∈ [0 , T ] and x ∈ R d

Combining this result and (24), we get

Combining (22), (23) and (25), we get

with

B.2.2 Ornstein-Uhlenbeck

glyph[negationslash]

Second, we consider the Ornstein-Uhlenbeck setting and let ( X t ) t ∈ [0 ,T ] with d X t = -αβ t X t d t + σβ 1 / 2 t d B t , with α = 0 . We consider ( X 0 ,T t ) t ∈ [0 ,T ] , the stochastic process ( X t ) t ∈ [0 ,T ] conditioned at both ends X 0 ,T 0 = x 0 and X 0 ,T T = x T . First, using (Barczy and Kern, 2013, Theorem 3.3), we have that for any t ∈ [0 , T ]

with R ( s, t ) = exp[ α ( B ( t ) -B ( s ))] γ ( s, t ) , with γ ( s, t ) = ∫ t s σ 2 β ( u ) exp[ -2 α ( B ( t ) -B ( u ))]d u . In particular, we have

Therefore, we get that

(27) defines the integral sampler . In addition, using (Barczy and Kern, 2013, Theorem 3.2) and (26), we have X 0 ,T 0 = x 0 and

In the formula, coth is the hyperbolic cotangent function defined as coth( x ) = 1 tanh( x ) = cosh( x ) sinh( x ) and csch is the hyperbolic cosecant function defined as csch( x ) = 1 sinh( x ) . Combining this result and (26), we get that

(28) defines the forward sampler . Finally, we derive the backward sampler by considering the time-reversal of the forward unconditional process (initialized at x 0 ). Following Haussmann and Pardoux (1986),

In addition, we have that

Hence, we get that for any t ∈ [0 , T ] and x ∈ R d

Combining this result and (29), we get Y 0 ,T 0 = x T and

Therefore

Combining (27), (28) and (30), we get with

Using that tanh( x ) ∼ x and sinh( x ) ∼ x for x → 0 , we recover the Brownian motion setting by letting α → 0 . Note that Albergo et al. (2023) show that given an integral sampler , which does not necessarily comes from a Markovian process, a forward sampler with the same marginals can be defined, although it does not necessarily satisfies the fact that the paths have the same distribution.

C Proofs

C.1 Proof of Proposition 2

We refer the reader to Chung and Walsh (2006); Rogers and Williams (2000) for an introduction to Doob h -transform. Our theoretical treatment of the Doob h -transform closely follows Palmowski and Rolski (2002).

First, we introduce the infinitesimal generator A given for any f ∈ C ∞ c ([0 , T ] × R d , R ) , t ∈ [0 , T ] and x ∈ R d by

The following assumption ensures that the diffusion associated with Q as well as its Markovian projections are well-defined.

A1. f , σ and ( t, x t ) ↦→ E Π T | t [ ∇ log Q T | t ( X T | X t ) | X t = x t ] are locally Lipschitz and there exist C > 0 , ψ ∈ C([0 , T ] , R + ) such that for any t ∈ [0 , T ] and x 0 , x t ∈ R d , we have

We consider the following assumption, which will ensure that we can apply Doob h -transform techniques.

A2. For any x 0 ∈ R d , Π T | 0 is absolutely continuous w.r.t. Q T | 0 . For any x 0 ∈ R d , let ϕ T | 0 be given for any x T ∈ R d by ϕ T | 0 ( x T | x 0 ) = dΠ T | 0 ( x T | x 0 ) / d Q T | 0 ( x T | x 0 ) and assume that for any x 0 ∈ R d , x T ↦→ ϕ T | 0 ( x T | x 0 ) is bounded. For any x 0 ∈ R d , let ϕ t | 0 given for any x t ∈ R d and t ∈ [0 , T ] by

Finally, we assume that for any x 0 ∈ R d , ( t, x t ) ↦→ 1 /ϕ t | 0 ( x t | x 0 ) and ( t, x t ) ↦→A ϕ t | 0 ( x t | x 0 ) are bounded.

This means that for any x 0 ∈ R d , ( t, x t ) ↦→ ϕ t ( x t | x 0 ) is a good function in the sense of (Palmowski and Rolski, 2002, Proposition 3.2). Note here that these assumptions could be relaxed on a case-bycase basis. We leave this study for future work.

The following lemma is a direct consequence of A 2 and (32). It ensures that the h -function ϕ t | 0 satisfies the backward Kolmogorov equation.

Lemma 11. Assume A 2 . Then, ϕ ∈ C 1 , 2 ([0 , T ) × R d , R ) and A ϕ | 0 = 0 .

Using (31), we have that for any x 0 ∈ R d and f ∈ C ∞ c ([0 , T ] × R d , R ) , t ∈ [0 , T ] and x t ∈ R d

Finally, we consider the following assumption, which will ensure that the Doob h -transform is well-defined.

A3. For any x 0 ∈ R d , there exists C ≥ 0 such that for any t ∈ [0 , T ] and x t ∈ R d , ‖∇ log ϕ t | 0 ( x t | x 0 ) ‖ ≤ C (1 + ‖ x 0 ‖ + ‖ x t ‖ ) .

We are now ready to state and prove Proposition 2. Note that the Markovian projection is defined in Definition 1. Finally, we define M the space of path measures such that P ∈ M if P is associated with d X t = { f t ( X t ) + v t ( X t ) } d t + σ t d B t , with σ, v locally Lipschitz. This restriction of Markov measures allows us to apply the entropic version of the Girsanov theorem (Léonard, 2012). It has no impact on our methodology.

Proposition. Assume A 1 , A 2 , A 3 . Let M glyph[star] = proj M (Π) . Then,

In addition, we have that for any t ∈ [0 , T ] , M = Π . In particular, M = Π

2 0 | | glyph[star] t t glyph[star] T T .

Proof. First, we recall that Π is given by Π = Q ϕ 0 ,T with ϕ 0 ,T = dΠ 0 ,T d Q 0 ,T . In particular, we have Π | 0 = Q | 0 ϕ T | 0 , where ϕ T | 0 = dΠ T | 0 d Q T | 0 . Therefore, using Lemma 11, (Palmowski and Rolski, 2002, Lemma 3.1, Lemma 4.1), the remark following (Palmowski and Rolski, 2002, Lemma 4.1), A 1, A 2 and A 3, we get that Π | 0 is Markov and associated with the distribution of ( X t ) t ∈ [0 ,T ] given for any t ∈ [0 , T ] by where for any t ∈ [0 , T ] , x 0 , x t ∈ R d we recall that

First, we have that for any t ∈ [0 , T ] , x t , x 0 ∈ R d

Therefore, we get that for any t ∈ [0 , T ] and x t , x 0 ∈ R d

In addition, we have the following identity for any t ∈ [0 , T ] , x 0 , x t , x T ∈ R d

Using (34), this result and (35), we get that for any t ∈ [0 , T ] and x 0 , x t ∈ R d

Hence, combining this result and (33), we get

Let M be Markov defined by d X t = { f t ( X t ) + v t ( X t ) } d t + σ t d B t , such that KL(Π | M ) < + ∞ with σ, v locally Lipschitz. Using (Léonard, 2012, Theorem 2.3), we get that

| | t. In addition, we have that for any t ∈ [0 , T ] ,

where v glyph[star] t ( x t ) = σ 2 t E Π T | t [ ∇ log Q T | t ( X T | X t ) | X t = x t ] which concludes the first part of the proof. For the second part of the proof, we show that for any t ∈ [0 , T ] , we have M glyph[star] t = Π t . First, we have that M glyph[star] t and Π t satisfy the same Fokker-Planck equation, see (Peluchetti, 2021, Theorem 2) for instance. We conclude by uniqueness of the solutions of the Fokker-Planck equation under A 1 and A 3, see Bogachev et al. (2021) for instance.

C.2 Proof of Proposition 4

Proof. By the additive property of KL divergence (Léonard, 2014a), KL( P | Π) = KL( P 0 ,T | Π 0 ,T ) + E P 0 ,T [KL( P | 0 ,T | Π | 0 ,T )] . Restricting Π | 0 ,T = Q | 0 ,T directly gives that the KL minimizer is Π glyph[star] with Π glyph[star] 0 ,T = P 0 ,T , and thus Π glyph[star] = P 0 ,T Q | 0 ,T which we recall is short for Π glyph[star] ( · ) = ∫ R d × R d Q | 0 ,T ( ·| x 0 , x T ) P 0 ,T (d x 0 , d x T ) .

C.3 Proof of Proposition 5

This result is a direct consequence of (Léonard, 2014b, Theorem 2.12), which we recall here for completeness.

Proposition. Assume that Q ∈ M , that Q 0 = Q T = ¯ Q , that for any x 0 , x T ∈ R d , d Q 0 ,T / d( ¯ Q ⊗ ¯ Q )( x 0 , x T ) ≥ exp[ -A ( x 0 ) -A ( x T )] with A ≥ 0 measurable, ∫ R d × R d exp[ -B ( x 0 ) -B ( x T )]d Q ( x 0 , x T ) < + ∞ with B ≥ 0 measurable. Assume that there exists t 0 ∈ (0 , T ) and X measurable such that Q t 0 ( X ) > 0 and for all x ∈ X , Q 0 ,T glyph[lessmuch] Q 0 ,T | t 0 ( ·| X t 0 = x ) . In addition, assume that KL( π 0 | ¯ Q ) < + ∞ , KL( π T | ¯ Q ) < + ∞ , ∫ R d ( A + B )( x 0 )d π 0 ( x 0 ) < + ∞ , ∫ R d ( A + B )( x T )d π T ( x T ) < + ∞ .

Then there exists a unique Schrödinger Bridge P SB . In addition let P be a Markov measure in the reciprocal class of Q such that P 0 = π 0 and P T = π T . Assume that KL( P | Q ) < + ∞ . Then P is the unique Schrödinger Bridge P SB .

Proof. The first part of the proof is a consequence of (Léonard, 2014b, Theorem 2.12(a)). The second part is a consequence of (Léonard, 2014b, Theorem 2.12(b)) and (Léonard et al., 2014, Theorem 2.14).

C.4 Proof of Lemma 6

Lemma. Let M ∈ M and Π ∈ R ( Q ) and assume A 1 , A 2 , A 3 . If KL(Π | M ) < + ∞ and KL(proj M (Π) | M ) < + ∞ we have

For any P ∈ P ( C ) , if KL( P | Π) < + ∞ , we have

Proof. We start with the proof of (36). Similarly to Proposition 2, where we have M ∈ M to ensure that we can apply (Léonard, 2012, Theorem 2.3), we get

In addition, we have

Finally, using Proposition 2, we have that

Using this result, we have

which concludes the first part of the proof.

For (37), define Π glyph[star] = proj R ( Q ) ( P ) = P 0 ,T Q | 0 ,T . Using (Csiszár, 1975, Equation 2.6), we have

which concludes the proof.

C.5 Proof of Proposition 7

Proposition. Assume that the conditions of Proposition 5 and Lemma 6 apply for P n for every n ∈ N and for the Schrödinger Bridge P SB , we have KL( P n +1 | P SB ) ≤ KL( P n | P SB ) < ∞ , and lim n → + ∞ KL( P n | P n +1 ) = 0 .

Proof. We follow the technique of Rüschendorf (1995) but for the reverse Kullback-Leibler divergence. Applying Lemma 6, we get for any N ∈ N

C.6 Proof of Theorem 8

Theorem. Assume that the conditions of Proposition 5 and Lemma 6 apply for P n for every n ∈ N and for the Schrödinger Bridge P SB , the IMF sequence ( P n ) n ∈ N admits a unique fixed point P glyph[star] = P SB , and lim n → + ∞ KL( P n | P glyph[star] ) = 0 .

Proof. By Proposition 7, KL( P n | P SB ) ≤ KL( P 0 | P SB ) < ∞ for all n ∈ N . Now, using the coercivity of KL( ·| P SB ) (this is where our analysis differs from the one of the IPF), we have that the IMF sequence ( P n ) n ∈ N and its subsequences ( M n +1 ) n ∈ N and (Π n ) n ∈ N are subsets of { P ∈ P ( C ) : KL( P | P SB ) ≤ KL( P 0 | P SB ) } which is (relatively) compact. Thus, ( M n +1 ) n ∈ N contains a convergent subsequence M n j → M glyph[star] as j → ∞ , and (Π n j ) j ∈ N contains a further convergent subsequence Π n j k → Π glyph[star] weakly as k → ∞ . As the Markov and the reciprocal classes are closed under weak convergence, M glyph[star] ∈ M and Π glyph[star] ∈ R ( Q ) . Now, by the lower semicontinuity of KL divergence in the weak topology (van Erven and Harremoes, 2014, Theorem 19), 0 ≤ KL( M glyph[star] | Π glyph[star] ) ≤ lim inf k →∞ KL( M n j k | Π n j k ) = 0 . Hence, M glyph[star] = Π glyph[star] which we denote as P glyph[star] ∈ M ∩ R ( Q ) . Also, P glyph[star] satisfies P glyph[star] 0 = π 0 and P glyph[star] T = π T as is satisfied by all P n . By Proposition 5, P glyph[star] is the unique Schrödinger Bridge P SB . Finally, lim n → + ∞ KL( P n | P glyph[star] ) = 0 follows using lim k →∞ KL( M n j k | P glyph[star] ) = 0 and the monotonicity of KL( P n | P SB ) by Proposition 7.

C.7 Proof of Proposition 9

Proof. The proof is similar to the one of Proposition 2.

In particular, the time-reversal of Q | 0 ,T ( ·| x 0 , x T ) is associated with

One can view both (11) (12) as SDEs with drift defined as the conditional expectation of the drift of (2) (38) under Π 0 ,T | t in the forward and backward directions respectively.

C.8 Proof of Proposition 10

Proof. We proceed by induction. Firstly, for Π 0 0 ,T = Q 0 ,T , Π 0 = ˜ P 0 = Q at initialization. We can also define M 0 = Q , such that M 0 = ˜ P 0 and Π n = proj R ( Q ) ( M n ) for all n ∈ N . By (De Bortoli et al., 2021, Section 3.5), the optimal DSB sequence ˜ P n is Markov and ˜ P n = ˜ P n 0 ,T Q | 0 ,T , where ˜ P n 0 ,T is the IPF sequence of the static SB problem. In other words, ˜ P n ∈ M∩R ( Q ) .

Suppose M 2 n +1 = ˜ P 2 n +1 . By definition, M 2 n +2 0 = ˜ P 2 n +2 0 = π 0 , i.e. both forward processes are initialized at π 0 . In DSB, by De Bortoli et al. (2021), ˜ P 2 n +2 is defined as the time-reversal of ˜ P 2 n +1 , such that ˜ P 2 n +2 | 0 = ˜ P 2 n +1 | 0 . Hence, ˜ P 2 n +2 = π 0 ˜ P 2 n +1 | 0 .

In DSBM, we first perform reciprocal projection Π 2 n +1 = proj R ( Q ) ( M 2 n +1 ) = M 2 n +1 0 ,T Q | 0 ,T . Since M 2 n +1 = ˜ P 2 n +1 ∈ R ( Q ) , however, we have that Π 2 n +1 = M 2 n +1 . Furthermore, since M 2 n +1 = ˜ P 2 n +1 ∈ M , proj M (Π 2 n +1 ) = proj M ( M 2 n +1 ) = M 2 n +1 . Thus, M 2 n +2 given by (9) is such that M 2 n +2 0 = π 0 and M 2 n +2 | 0 = proj M (Π 2 n +1 ) | 0 = M 2 n +1 | 0 . We conclude that M 2 n +2 = π 0 M 2 n +1 | 0 = ˜ P 2 n +2 . Similar arguments holds for the the reverse projection (13). Therefore, M n = ˜ P n for all n ∈ N .

C.9 The set of Markov measures is not convex

The result of Lemma 6 should be compared with the information geometry result of (Csiszár, 1975, Theorem 2.2), which states that if C is a convex set and P ∈ C , then under mild conditions, KL( P | Q ) = KL( P | proj C ( Q )) + KL(proj C ( Q ) | Q ) , where proj C ( Q ) = argmin P { KL( P | Q ) : P ∈ C} is the projection of Q on C . Note that, contrary to Lemma 6, (Csiszár, 1975, Theorem 2.2) is given for the forward Kullback-Leibler divergence whereas Lemma 6 is given for the reverse KL divergence. In addition, (Csiszár, 1975, Theorem 2.2) requires the projection set C to be convex which is not satisfied for the space of Markov measures M . We give a simple counter-example proving that the set of Markov measures is not convex.

Let p 1 ( x 0 , x 1 , x 2 ) = p 1 ( x 0 ) p 1 ( x 1 | x 0 ) p 1 ( x 2 | x 1 ) and p 2 ( x 0 , x 1 , x 2 ) = p 2 ( x 0 ) p 2 ( x 1 | x 0 ) p 2 ( x 2 | x 1 ) on { 0 , 1 } 3 such that

Additionally, we set

Finally, we set q = (1 / 2) p 1 +(1 / 2) p 2 . Consider q ( x 2 = 1 | x 1 = 1 , x 0 = 1) = q ( x 2 = 1 , x 1 = 1 , x 0 = 1) /q ( x 1 = 1 , x 0 = 1) and q ( x 2 = 1 | x 1 = 1) = q ( x 2 = 1 , x 1 = 1) /q ( x 1 = 1) . Let

glyph[negationslash]

glyph[negationslash]

q is Markov if and only if ∆ = 0 . Therefore q is not Markov as soon as α 0 = β 0 and α 2 = β 2 .

C.10 Impact of the resolution on the entropic regularization

In this section, we show how the resolution of the input affects the entropic regularization. First, we consider the Schrödinger bridge problem between π 0 and π 1 ( T = 1 ) with Q associated with ( σ B t ) t ∈ [0 , 1] . In that case the static Schrödinger Bridge is given by P glyph[star] 0 , 1 such that where ε = 2 σ 2 . We now describe the solution of a higher dimensional problem given by the upsampling of the marginals. Let f ∈ N with f ≥ 1 and down : R fd → R d such that for any k ∈ { 1 , . . . , d } , and x ∈ R fd , down( x ) = ¯ x with ¯ x ∈ R d and ¯ x k = x kf . We also denote up : R d → R fd such that for any k ∈ { 1 , . . . , d } , and ¯ x ∈ R d , up(¯ x ) = x with x ∈ R fd and ¯ x k = x ( k -1) f + j , for j ∈ { 1 , . . . , d } . Note that down ◦ up = Id .

We denote π up 0 = up # π 0 and π up 1 = up # π 1 . We extend up and down to R d × R d and R fd × R fd by letting for any ¯ x, ¯ y ∈ R d and x, y ∈ R fd

First, note that since up ◦ down = Id on up( R d ) we get that for any P supported on up( R d ) × up( R d ) , using (Kullback, 1997, Theorem 4.1)

Second let P such that P 0 = π up 0 and P 1 = π up 1 . Then, P is supported on up( R d ) × up( R d ) . Therefore, for any P such that P 0 = π up 0 and P 1 = π up 1

Finally, for any P such that P 0 = π up 0 and P 1 = π up 1 we have

Therefore, for any P such that P 0 = π up 0 and P 1 = π up 1 we have

Therefore, we have the following result.

Proposition 12. Let ε > 0 . P glyph[star] is the solution of the static Schrödinger bridge with marginals π up 0 , π up 1 and regularization ε if and only if down # P glyph[star] is the solution of the static Schrödinger bridge with marginals π 0 , π 1 and regularization ε/f 2 .

This means in particular that the Schrödinger Bridge is not invariant via upsampling. In Appendix I.4, we confirm these results visually upon noting that for the same σ , running DSBM at resolution 64 × 64 and 128 × 128 gives different results, even after downsampling of the 128 × 128 results.

D Convergence of IMF in the Gaussian setting

In this section, we study the IMF in an one-dimensional Gaussian case. We consider T = 1 , Π 0 = Π 1 = N(0 , (1 / 2 β 2 )) and Q associated with ( σ B t ) t ∈ [0 , 1] where σ > 0 . In what follows, we let

where c ∈ [0 , 1] . We also denote ¯ σ 2 = 2 σ 2 β 2 . We start with the following result which gives an explicit expression of some marginals of the reciprocal projection.

Lemma 13. Let Π 0 0 , 1 = N(0 , Σ 0 ) with Σ 0 given by (39) . Let Π 0 = Π 0 0 , 1 Q | 0 , 1 . For any t ∈ [0 , 1] , we have that Π 0 0 , 1 ,t = N(0 , Σ) with

where we have

Proof. Let t ∈ [0 , 1] and X 0 , 1 t ∼ Π 0 t | 0 , 1 = Q t | 0 , 1 . Using (Barczy and Kern, 2013, Theorem 3.3), we get that

with Z ∼ N(0 , Id) independent from ( X 0 , X 1 ) . Hence, we get that E [ X 0 , 1 t ] = 0 , and

Similarly, we get that Cov( X 1 , X 0 , 1 t ) = a 1 -t k / (2 β 2 ) . Finally, we get that

which concludes the proof.

Leveraging Lemma 13, we can give an explicit expression of the drift term in the Markovian projection.

Lemma 14. Let Π 0 0 , 1 = N(0 , Σ 0 ) with Σ 0 given by (39) . Let Π 0 = Π 0 0 , 1 Q | 0 , 1 . For any t ∈ [0 , 1] and x t ∈ R d , we have that

Hence, the Markovian projection of Π 0 , denoted M 1 is associated with ( X t ) t ∈ [0 , 1] with

Proof. Using (Barczy and Kern, 2013, Theorem 3.2), we get that Q | 0 , 1 is associated with

where for any t ∈ [0 , 1) , we have and x t ∈ R d , σ 2 ∇ log Q 1 | t ( x 1 | x t ) = ( x 1 -x t ) / (1 -t ) . Therefore, we get that M 1 is associated with ( X t ) t ∈ [0 , 1] such that

Therefore, we get that

Using Lemma 13, we have that for any t ∈ [0 , 1] and x t ∈ R d , E Π 0 1 | t [ X 1 | X t = x t ] = a 1 -t k /b t k x t . In addition, we have for any t ∈ [0 , 1]

which concludes the proof.

Note that since σ > 0 and c 2 ∈ [0 , 1] , we get that for any t ∈ [0 , 1] , 1+ t (1 -t )(2( c 2 -1) + ¯ σ 2 ) > 0 and therefore the drift is well-defined, smooth and sublinear. In particular, (40) admits a unique strong solution. In what follows, we denote G : [0 , 1] × [0 , 1] → R given for any t ∈ [0 , 1] by

We have the following useful lemma.

Lemma 15. Let c ∈ [0 , 1] , ¯ σ > 0 and p = 2 c 2 + ¯ σ 2 . We distinguish three cases:

(a) If p < 2 , then

(b) If p = 2 , then

(c) If p > 2 , then

This lemma is useful combined with the following proposition, which gives the analytical update formula of Gaussian IMF covariances as a function of G (1 , c 2 ) .

Proposition 16. Let Π 0 0 , 1 = N(0 , Σ 0 ) with Σ 0 given by (39) . Then Π 1 0 , 1 = N(0 , Σ 1 ) with

with f : [0 , 1] → [0 , 1] given for any c ∈ [0 , 1] by

Proof. We have that X 1 = exp[ G (1 , c 2 )] X 0 + M 1 , where M 1 is a Gaussian random variable with zero mean independent from X 0 . Therefore, we get that Cov( X 0 , X 1 ) = exp[ G (1 , c 2 )] / (2 β 2 ) . In addition, we have that E [ X 2 1 ] = E [ X 2 0 ] = 1 / 2 β 2 , which concludes the proof.

Iterating the procedure in Proposition 16, we obtain a sequence of IMF covariances ( c 2 n ) n ∈ N satisfying c 2 n +1 = f ( c 2 n ) . Finally, we show that this iterative procedure recovers the true SB coupling Π SB 0 , 1 = N(0 , Σ SB ) as a fixed point. The formula of Σ SB is given e.g. in (Bunne et al., 2023, Equation (2)) which we use below.

Proposition 17. Let Π SB 0 , 1 = N(0 , Σ SB ) be the true static SB solution, with

Then Π SB 0 , 1 is a fixed point of the iterative procedure in Proposition 16, i.e. f ( c 2 SB ) = c 2 SB .

Proof. By straightforward calculations, p SB = 2 c 2 SB + ¯ σ 2 = √ 4 + ¯ σ 4 . If ¯ σ = 0 , p SB = 2 and thus G (1 , c 2 SB ) = -¯ σ 2 / 2 = 0 . Hence f ( c 2 SB ) = exp( G (1 , c 2 SB )) = 1 = c 2 SB . If ¯ σ > 0 , we have p SB > 2 and G (1 , c 2 SB ) = -tanh -1 (¯ σ 2 / √ 4 + ¯ σ 4 ) . Hence, f ( c 2 SB ) = exp( G (1 , c 2 SB )) = 1 2 ( √ 4 + ¯ σ 4 -¯ σ 2 ) = c 2 SB . We thus correctly recover the true static SB solution Π SB 0 , 1 as fixed point of IMF.

We visualize the convergence of this fixed point procedure for a variety of parameter settings in Figure 12. The convergence appears to be very fast in only two or three iterations.

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

c

2

n

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

c

2

n

+1

Convergence of IMF

Update function

Identity

True SB solution

Fixed point iterations

(a)

¯

σ

= 0

.

1

, c

2

0

= 0

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

c

2

n

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

c

2

n

+1

Convergence of IMF

(b)

¯

σ

= 1

, c

2

0

= 0

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

c

2

n

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

c

2

n

+1

Convergence of IMF

(c)

¯

σ

= 1

, c

2

0

= 1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

c

2

n

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

c

2

n

+1

Convergence of IMF

(d)

¯

σ

= 5

, c

2

0

= 1

E Discrete-Time Markovian Projection

We derive in this section a discrete-time version of the Markovian projection and show that, in some limiting case, we recover the continuous-time projection. In the discrete case, we let

We consider a Markovian measure p given by p ( x 0: N ) = p ( x 0 ) ∏ N -1 k =0 p k +1 | k ( x k +1 | x k ) . Now let us compute KL( π | p ) . We have

In what follows, we denote

We have the following proposition.

Proposition 18. The minimizer p k +1 | k of L k +1 is given by

If p 0 = q 0 , then for any k ∈ { 0 , . . . , N -1 } , p k = π k . In addition, assume that p k +1 | k ( x k +1 | x k ) = exp[ -‖ x k +1 -x k -γf ( x k ) ‖ 2 / (2 γ )] / (2 π γ ) -d/ 2 and q k +1 | k,N ( x k +1 | x k , x N ) = exp[ -‖ x k +1 -x k -γf ( x k , x N ) ‖ 2 / (2 γ )] / (2 π γ ) d/ 2 . Finally, assume that ‖ x k +1 -x k ‖ ≤ γ 1 / 2 . Then, we have that

Proof. The proofs of (41) and (42) are straightforward and left to the reader. We now prove that if p 0 = π 0 , then for any k ∈ { 1 , . . . , N } , p k = π k . First, we have that for any k ∈ { 0 , . . . , N -1 } ,

Assume now that p k = π k , then we have

which concludes the proof.

In particular, in the previous proposition, if f ( x k , x N ) = ∇ log q ( x N | x k ) , i.e. we have a discretization of the bridge then f ( x k ) = ∫ R d ∇ log q ( x N | x K ) π ( x N | x k )d x N , which recovers the Markovian projection in continuous-time.

F Comparing DSBM-IPF and DSB

Weanalyze further the differences between DSBM-IPF proposed here and DSB proposed in De Bortoli et al. (2021) and related algorithms in Vargas et al. (2021); Chen et al. (2022). All algorithms solve the SB problem using the IPF iterates. However, DSB-type algorithms solve for the IPF iterates using time-reversals , whereas DSBM solves for the iterates using Markovian and reciprocal projections . A comparison between these two methodologies is made in Section 4.

We investigate here further benefit (iii) of DSBM in Section 4, i.e. the benefit of explicitly projecting onto the reciprocal class of Q . Intuitively speaking, we directly incorporate the reference measure Q in the training procedure as our inductive bias. More formally, suppose we have at the current

glyph[negationslash]

IPF iteration M 2 n (Markov diffusion in the forward direction) and want to learn M 2 n +1 (Markov diffusion in the backward direction). Due to training error and the forgetting issue (Fernandes et al., 2021), however, M 2 n no longer has the correct bridge Q | 0 ,T . Now suppose we first perform IMF for M 2 n and learn M 2 n,glyph[star] in the forward direction. That is to say, we repeat alternative reciprocal and Markovian projections and obtain a sequence ( M 2 n,m ) m ∈ N in the forward direction converging to M 2 n,glyph[star] . Then M 2 n,glyph[star] now has the correct bridge Q | 0 ,T by Proposition 7, since M 2 n,glyph[star] is the SB between M 2 n 0 and M 2 n T . Theoretically, M 2 n t = M 2 n,glyph[star] t for t = 0 , T , but due to training error accumulating it may be that M 2 n T = M 2 n,glyph[star] T . However, M 2 n 0 = M 2 n,glyph[star] 0 , since M 2 n,glyph[star] is in the forward direction and starting from samples from π 0 . As a result, we can obtain a Markov forward diffusion M 2 n,glyph[star] , which has M 2 n,glyph[star] 0 = π 0 and the correct bridge Q | 0 ,T . These are the same set of properties that the reference measure Q has. As a result, replacing Q with M 2 n,glyph[star] in (6) results in the same SB solution. Consequently, now continuing the IPF iterations from M 2 n,glyph[star] 0 , it is as if we restart IPF afresh using a modified SB problem

If M 2 n,glyph[star] is closer than Q to M glyph[star] in the sense of KL divergence, then we obtain a better initialization of the IPF procedure. As proposed in Algorithm 1, DSBM performs the Markovian and reciprocal projection only once before switching between the forward and backward directions. However, it is still beneficial compared to DSB with less bias accumulation in the bridge.

Algorithmically, one main difference between DSB-type algorithms and DSBM due to the above distinction occurs in the trajectory caching step. In DSB, a fixed discretization of SDE needs to be chosen, and all intermediate samples from the discretized Euler-Maruyama simulation of the SDE need to be saved. Furthermore, a second set of drift evaluation needs to be performed for all datapoints in the trajectory (De Bortoli et al., 2021, Equations (12), (13)). The IPML algorithm in Vargas et al. (2021) is also similar to DSB, but Gaussian processes are used to fit the drifts of forward and backward SDEs instead of neural networks. In Chen et al. (2022), the implicit score matching loss is used instead, but all intermediate points in the SDE trajectory also need to be saved. On the contrary, DSBM does not require intermediate samples during trajectory caching and only retains the joint samples at times 0 , T . Then, the intermediate trajectories are reconstructed using the reference bridge Q | 0 ,T .

G Joint Training of Forward and Backward Processes

We recall below the DSBM algorithm given in Algorithm 1.

Algorithm 2 Diffusion Schrödinger Bridge Matching

Our main observation comes from Proposition 9. In particular, under mild assumptions, we have that the Markovian projection M = proj M (Π) is associated with

Considering the following losses,

If the families of functions { v θ : θ ∈ Θ } and { v φ : θ ∈ Φ } are rich enough, we have for any t ∈ [0 , T ] and x t ∈ R d , v θ glyph[star] ( t, x t ) = σ 2 t E Π T | t [ ∇ log Q T | t ( X T | X t ) | X t = x t ] and v φ glyph[star] ( t, x t ) = σ 2 t E Π 0 | t [ ∇ log Q t | 0 ( X t | X 0 ) | X t = x t ] . In practice, this means that the Markovian projection can be computed in a forward or backward fashion equivalently.

Therefore, given a coupling Π = Π 2 n , we can update both v θ and v φ . This means that we train the forward and backward model jointly . We then consider M 2 n +1 b associated with (13) and M 2 n +1 f associated with (9). Note that if the families of functions { v θ : θ ∈ Θ } and { v φ : θ ∈ Φ } are rich enough then M 2 n +1 f = M 2 n +1 b .

Mixture from forward and backward. Once we have obtained both the forward update and the backward update, our next task is to define the new mixture of bridge Π 2 n +1 . In Algorithm 1, since we train only the backward model M 2 n +1 = M 2 n +1 b , we define Π 2 n +1 = M 2 n +1 0 ,T Q | 0 ,T . In the case of joint training , we have access to M 2 n +1 b and M 2 n +1 f . One way to define a new mixture of bridge is to compute Π 2 n +1 = 1 2 ( M 2 n +1 b, 0 ,T Q | 0 ,T + M 2 n +1 f, 0 ,T ) Q | 0 ,T . This choice ensures that in the case where M 2 n +1 f = M 2 n +1 b we have

It also ensures that all the steps in the joint DSBM training algorithms are symmetric. We leave the study of an optimal combination of M 2 n +1 f and M 2 n +1 b for future work.

Consistency loss. In addition to the losses (43) and (44), we also consider an additional consistency loss. A similar idea was explored in Song (2022). In DSB (De Bortoli et al., 2021; Chen et al., 2022) and DSBM, see Algorithm 1, the processes parameterized by v θ (forward) and v φ backward are identical only at equilibrium . Thus imposing the forward and the backward processes match at each step of DSB or DSBM would lead to some bias. However, this is not the case in the joint training setting. Indeed, in that case, we have M 2 n +1 f = M 2 n +1 b if the families are rich enough. Therefore, we get that

is the time reversal of

Computing the time-reversal of (45), we have

Identifying (47) and (46), we get that for any t ∈ [0 , T ] and x t ∈ R d

We highlight that letting σ t → 0 for any t ∈ [0 , T ] , we get that v θ = -v φ , which confirms that the time-reversal of an ODE is simply given by flipping the sign of the velocity. Therefore, we propose the following loss which links the parameters θ and φ

Leveraging tools from implicit score matching (Hyvärinen, 2005) and the divergence theorem, we get that

where C ≥ 0 is a constant which does not depend on θ and φ . Alternatively, (48) shows that

We thus also have the following denoising score matching consistency loss

The advantage of this DSM loss is that it does not rely on any divergence computation.

Below, we recall the two losses used to estimate the Markovian projection (43) and (44)

The complete loss we consider in the joint training of the algorithm is of the form

where λ > 0 is an additional regularization parameter. We now state a version of DSBM which performs joint training in Algorithm 3.

Algorithm 3 Diffusion Schrödinger Bridge Matching (Joint Training)

1: Input: Coupling Π 0 0 ,T , tractable bridge Q | 0 ,T , N ∈ N

3: for n ∈ { 0 , . . . , N -1 } do

2: Let Π 0 = Π 0 0 ,T Q | 0 ,T .

4: Learn v φ glyph[star] , v θ glyph[star] using (49) with Π = Π n .

6: Let M n +1 b be given by (13).

5: Let M n +1 f be given by (9).

7: Let M n +1 = 1 2 ( M n +1 f + M n +1 b ) .

9: end for

8: Let Π n +1 = M n +1 0 ,T Q | 0 ,T .

10: Output: v θ glyph[star] , v φ glyph[star]

H Loss Scaling

Similar to the loss weighting in standard diffusion models (Song et al., 2021b; Ho et al., 2020), we derive a similar weighting to reduce the variance of our objective. We focus on the forward direction of Markovian projection in this case, and the backward case can be derived similarly. Our forward loss in the DSBM framework is given by (10), where the inner expectation is given by

Letting Q | 0 ,T be a Brownian bridge with diffusion parameter σ and assuming T = 1 , this becomes

with X 0 ,T t = t X 1 + (1 -t ) X 0 + σ √ t (1 -t ) Z . When t ≈ 1 , we see that the regression target is dominated by the noise term σ √ t/ (1 -t ) Z which needs to be predicted based on information contained within X 0 ,T t . The loss will have an approximate scale of σ 2 t/ (1 -t ) when t ≈ 1 which will be very large. To avoid these large values affecting gradient descent, we can downweight the loss by 1 + σ 2 t/ (1 -t ) (we can add 1 to effectively cause no loss scaling when t is close to 0 )

Similar arguments can be applied to the backward loss (14)

which we then downweight by 1 + σ 2 (1 -t ) /t

I Experiments

In this section, we present further details of the experiment setups as well as further experiment results. In all experiments, we use Brownian motion for the reference measure Q with corresponding Brownian bridge (4) and T = 1 . We use the Adam optimizer with learning rate 10 -4 and SiLU activations unless specified otherwise. The experiments are run on computing clusters with a mixture of both CPU and GPU resources.

I.1 2D Experiments

For the 2D experiments, we closely follow Tong et al. (2023) and the released code 7 , and use the same synthetic datasets and the 2-Wasserstein distance between the test set and samples simulated using probability flow ODE as the evaluation metric. However, we use 10000 samples in the test set since we find the 2-Wasserstein distance can vary greatly with only 1000 samples (which can be as high as 0.3 even between two set of samples both drawn from the ground truth distribution). We use a simple MLP with 3 hidden layers and 256 hidden units to parameterize the forward and backward drift networks. We use batch size 128 and 20 diffusion steps with uniform schedule at sampling time. Each outer iteration is trained for 10000 steps and we train for 20 outer iterations. As the initial coupling is more optimal in DSBM-IMF+, we reduce the number of outer iterations to 4 and train for 50000 steps in each outer iteration. For flow methods, we train for 200000 steps in total. For Table 2 we use σ t = 1 in all cases, except the moons-8gaussians dataset where we use σ t = 5 . Note that FM cannot be used for the moons-8gaussians task since it requires a Gaussian source, but CFM is applicable. The experiments are run using 1 CPU and take approximately 200 minutes (for both training and testing).

I.1.1 Variance of the reference measure Q

We comment further on the effect of σ t in the reference path measure Q . We assume a timehomogeneous σ t = σ for simplicity. In Figure 2, we vary σ and visualize the learned transport for a 3gaussians problem of transporting between two Gaussian mixtures. In Table 4 we show the 2-Wasserstein distance between the test set and generated samples for this 3gaussians problem as well as the moons-8gaussians problem. We find that large values of σ result in increasingly curved transport paths, and correspondingly reduced performance when σ is excessively large. Conversely, we also find reduced performance when σ is excessively small. We conjecture this is due to increased optimization difficulty and bias accumulation. Firstly, the EOT problem becomes more difficult to solve as σ is taken to 0 , which would require a higher number of outer iterations. Further, the introduction of noise also decreases optimization difficulty by smoothing the intermediate marginals between the two terminal distributions. The benefit of setting σ > 0 and using a stochastic sampler was also observed in Albergo et al. (2023); Delbracio and Milanfar (2023). Finally, we conjecture setting σ > 0 could also increase the diversity of sampled couplings and may alleviate some bias accumulation issues in the outer iterations. When σ = 0 , these issues result in the artifacts observed in the transferred samples (marked using yellow points) in Figure 2. The appropriate value for σ depends on the spatial scaling of the problem as shown in Table 4, where the optimum σ is larger for the larger scale moons-8gaussians problem.

I.2 Gaussian Experiment

Similar to the 2D experiments, we use a simple MLP with 2 hidden layers and 256 hidden units to parameterize the forward and backward drift networks. This is a smaller network compared to the 'large' network in De Bortoli et al. (2021). We use batch size 128 and 20 diffusion sampling steps with uniform time schedule at inference time. Each outer iteration is trained for 10000 steps and we train for 20 outer iterations. The experiments are also run using 1 CPU, and for the case d = 50 finish in approximately 200 minutes. We note that DSB, IMF-b, and DSBM methods all solve for the Schrödinger Bridge transport, whereas Rectified Flow does not and so is not plotted on the covariance convergence plot in Figure 3 since it is not comparable.

For Table 3, we assume the marginals of the learned process P t are also independently Gaussian distributed in each dimension. We thus estimate the KL divergence from the sample mean and variance of each dimension of P t using the analytic KL formula between Gaussian distributions.

I.3 MNIST Transfer Experiment

We follow De Bortoli et al. (2021) closely for the setup of this experiment. We use the set of first 5 letters of EMNIST (A-E, a-e) such that both domains have 10 classes in total. We use the same U-Net architecture, batch size 128 and 30 diffusion sampling steps. The network size is approximately 6.6 million parameters. Each outer iteration is trained for 5000 steps. We refresh the cache dataloader every 1000 steps with 10000 new samples. Contrary to De Bortoli et al. (2021), in our experiments we find that we obtain better sampling quality for both DSB and DSBM using a uniform noising schedule. We simply choose σ t = 1 for all t and T = 1 in our experiments.

For DSBM-IPF, we train for at most 50 outer iterations (i.e. 250000 total number of steps). For DSBM-IMF and Rectified Flow, since their first iteration corresponds to Bridge Matching and Flow Matching respectively, we first pretrain the forward and backward networks for 100000 steps using the Bridge Matching or Flow Matching losses. DSBM-IMF then switches to iterative training with 5000 steps per outer iteration, the same as DSBM-IPF. For Rectified Flow, we train for 50000 steps per outer iteration. The experiments are performed using 2 GPUs and take approximately one day.

We provide further experiment results in Figures 13, 14 and 15. In Figure 13, we show samples generated using different algorithms and at different points of convergence. Samples generated using CFM and Rectified Flow with 2 rectified steps in (a) and (b) appear to be less clear and identifiable. OT-CFM improves upon CFM slightly in (c), but many samples still appear to be unclear. For DSB, the algorithm has not converged after 10 iterations, and many samples in (d) still appear to be letters. After 20 iterations, there are still letter-like samples in Figure 4b, and the digit classes also appear to be unbalanced with many instances of digit '0'. After 30 iterations, however, the sample quality of DSB becomes very poor in (e). On the other hand, as shown in (f)-(j), we observe that DSBM-IPF converges faster than DSB, with more accurate samples even in iterations 10 and 20, and the sample quality continues to improve until the end of training after 50 iterations.

We present some additional trajectory samples at the end of training in Figures 14 and 15 in both the forward and backward directions. We observe DSBM is able to transfer samples across the two domains faithfully, and the output samples preserve interesting similarities compared to the input, whereas Bridge Matching preserves much less similarity. The Mean Squared Distance (MSD) between the initial and final samples also confirm that DSBM methods transfer more closely to the original inputs, and DSBM-IMF further achieves the best FID score out of the three methods.

I.4 CelebA Transfer Experiment

For CelebA (Liu et al., 2015), we test DSBM on resized 64 × 64 and 128 × 128 resolution images. We evaluate the dependency of the results with respect to σ > 0 on CelebA 64 × 64 . We also showcase the scalibility of our method by training DSBM on CelebA 128 × 128 . In both cases, the dataset is split between male/old and female/young . We gather 20000 samples of each class on which we perform classical data augmentation such as horizontal flipping.

For our ablation study w.r.t. the value of σ > 0 , we run DSBM for 20 iterations (note that the loss and the quality was still improving after 20 DSBM iterations but by stopping the runs early we were allowed to draw comparisons with more values of σ ). We use a U-Net architecture with 4 resolution levels and 2 residual blocks per resolution level. The batch size is fixed to 64 and the EMA rate is

(a) CFM

(b) Rectified Flow @ 3

(c) OT-CFM

(d) DSB @ 10

(e) DSB @ 30

内

e

a

6

5

b

2

3

2

2

S

9955391

C

3

8

8399

3

C

8

33

0

0428642

0

5

560

a

Y

5

b

58

67

4

e

2

3

9

2

78530277

Q

2

6

b

6

5

3

9

277%680

6

220

3

5

7

3

2

3

2

3

f

2

9

2

0

S

了

2

4

6

3

d

8

3

7

3

4

cko

3/7

2

Y

8

G

3

2

3

3

4

7

2

O

O

S

2

3

1

7

了

2

e

6

6

5

2

七

8

2

b

5

5

9

5

8

2

3

d

3

8

2

9

2

5

2

9

3

6

3

8

6

3

G

2

4

2

3

3

X

2

4

4

6

2

3

5

5

4

5

5

？

3

5

3

6

6

5

D

h

7

238

aq

8

8

2

2

3

8

S

8S

0

2

2

4

3

3

3

8

4

7

6

6

6

2

2

2

0

8

&

9

C

2

3248516

G

8

2

3

Q

7

932462

H

Q

6

$

9

800008586

6

39600346

6

22335758

g

053C3327.39

5

92788379

(f) DSBM-IPF @ 10

(g) DSBM-IPF @ 20

(h) DSBM-IPF @ 30

(i) DSBM-IPF @ 40

(j) DSBM-IPF @ 50

fixed to 0.999. We refresh the cache dataloader every 500 steps with 10000 new samples. The SDE sampler is chosen to be the modified Euler-Maruyama sampler, see Heng et al. (2021) for instance, with a constant schedule for the stepsizes. We use 100 sampling steps at inference time and to refresh the cache. For each outer DSBM iteration we train the model for 20000 iterations.

We provide additional transfer results in resolution 128 × 128 in Figure 16. We do not change the training setting for this experiment.

I.5 AFHQ Transfer Experiment

For AFHQ (Choi et al., 2020), we test DSBM between classes cat and wild with 512 × 512 resolution images. Each class contains approximately 5000 samples. We first pretrain the networks using Bridge Matching for 100000 steps, then run DSBM for 20 iterations with 25000 steps per outer iteration. We follow Liu et al. (2023b) and use the same U-Net architecture 8 . The batch size is 4 and the EMA rate is 0.999. We choose σ 2 = 5 and again we use 100 sampling steps with constant stepsizes.

I.6 CIFAR-10 Generative Modeling Experiment

We also test our method in the standard generative modeling framework on the CIFAR-10 dataset. We again use a U-Net architecture with 4 resolution levels and 2 residual blocks per resolution level. The network size is approximately 39.6 million parameters. The batch size is fixed to 128 and the EMA rate is fixed to 0.9999. The AdamW optimizer is used for this task. For DSBM-IMF and Rectified Flow, again we first pretrain the networks using the Bridge Matching or Flow Matching losses, then switch to DSBM-IMF or RF training with 100000 steps per outer iteration. We find that 1 or 2 additional outer iterations appear sufficiently effective on this task, and additional outer iterations can cause sample quality to drop. For our main experiment, the pretraining stage ran for approximately 6 days, and DSBM-IMF ran for approximately 4 additional days using 4 V100 GPUs.

The best results during training for each method are reported in Table 5 and Figure 17, where we compute the FID score between 50000 samples in the CIFAR-10 training set and 50000 generated samples following standard practice for this task. We observe that DSBM-IMF can clearly improve upon Bridge Matching at the same value of σ , which suggests that further outer iterations in DSBM is beneficial for improving sample quality. This is contrary to Rectified Flow, which causes the FID score to worsen compared to Flow Matching after only 1 rectified iteration. However, as σ increases,

(c) DSBM-IMF

FID

17.14

MSD

0.579

FID

15.27

MSD

0.354

FID MSD

10.59

0.375

b

AcDDbbc?b

b

CbCAeB

e

S

7

00

3

5

6

S

482

3

3

2

9

9

3

4

2477673z

B

A

B99J

d

d

b

8

782302

C

BdAD

bcdc

C

e

e

e

C

e

59355

4

5

B:

CCC

C

C

C

3

5

7

3

C

abBD

b

e

D

2

E

CdbeEA

8

3

d

6

3

x

4

3

65

9

2

ccCBaccbe

99540

b

ACDDbbC?b

B

b

CbCAeB

A

B99J

d

d

e

b

七

8

5

4

b

7

2

9

3

6

89

L

O

3

3

7

8

7n

603~054

F

b

bcdc

BdAD

C

e

e

e

e

C

80752

4

2

8

3

385542

B

JUJ

C

C

6

2

7

2

C

abB7

b

e

7

E

CdbeEA

5

S

力

2

5

G

h1OS

3

ccCBaccbe

B

278260

89

S

cACDdbbc?b

be

b

CbCAeB

e

7

5

A

3

6

7

S

6

8

n

O

3

7

5

9

B

A

B99J

d

d

b

9

7

>

6

5

6

3

G

qb

BdAD

C

e

C

+

8

6

3

B

c

b

bcdc

CcC

C

C

e

e

C

e

A

3

6

2

S

S

3

8

7

53

5

0sy

C

abBD

b

e

2

5

7

V

2

E

CdbeEAc

a

3

7

9

3

a

2

S

813

8

ccCBaccbe

B

3

6

7138

9

&

FID

9.402

MSD

0.561

FID

13.59

MSD

0.359

FID

8.990

MSD

0.375

FID

9.402

MSD

0.561

FID

13.59

MSD

0.359

(b) DSBM-IPF

FID

8.990

MSD

0.375

70731

2

4

b

9

C

4

BAEbdee

G

6

B

ABAceace

d

1

3

A

5

BaedL

e

a

E

9

4

9

C

A

B

AEBdao

BabbbE

e

B6

e

3

F

6

h

3

d

a

beAeae

9

）

3

2

9

B

a

5952p

B

13

G

2

门

h

7

e

B

Bebceeee

4

6

g

3

4

6

9

aeedabbeb

C

70731

2

h

CdbB

A

Cd

9

C

b

C

dBd

C

b

d

6

Q

d

b

b

h

C

C

A

C

P

1

3

d

b

d

J

P

CAAe

bA

7

3

1

5

4

9

6

a

e

E

d

b

LDA

B

P

6

d

3

F

6

h

3

AC

b

d

db

2

9

）

3

2

9

C

C

aba

e

J

d

d

13

G

门

6

h

7

e

D

b

aceA

b

b

6

g

3

4

6

9

d

bdeaedbA

7073

2

eredaeacd

q

6

dc&cca

B

b

中

b

DCCCA

3

q

e

Q（）

C

b

B

a

d

C

D

edd

C

bcc

a

713

3

8

2

5

6

5

子

A

3

2

d

D

e

CaaAee

C

ERLO

b

cbA

C

C

C

d

a

9

2

7

h

7

3y

B

cacenaBl

D

B

6

9

3

4

G

9

B

bdeaeanca

FID

17.14

MSD

0.579

FID MSD

10.59

0.375

20

a

we observe the FID score worsens for both Bridge Matching and DSBM as more stochasticity is introduced in the sampler. The best result of DSBM-IMF is obtained using σ 2 = 0 . 2 and is slightly better than FM (i.e. with σ 2 = 0 ) using 100 Euler steps. On the other hand, using the dopri5 ODE solver, FM achieves a FID of 4.055 with on average 148 integration steps. In Figure 17, we observe that both RF and DSBM-IMF are very effective in improving sampling quality at low number of diffusion steps, i.e. low number of function evaluations (NFEs), compared to Bridge and Flow Matching as well as OT-CFM which improves upon CFM slightly. DSBM-IMF also achieves lower FID score than RF as the NFEs are taken higher. Additional strategies such as distillation and fast SDE solvers can also be useful for improving few-step sampling quality further.

I.7 Fluid Flows Experiment

We use the fluid flows dataset 9 from Bischoff and Deck (2023). The dataset consists of unpaired low ( 64 × 64 ) and high ( 512 × 512 ) resolution fields, as well as a context field with local information for the high resolution field. The data fields consist of two channels representing supersaturation and vorticity. The context field is dependent on the wavenumber k x = k y ∈ { 1 , 2 , 4 , 8 , 16 } , which

10

0

10

1

10

2

10

3

NFE

10

1

10

2

FID

BM

CFM

OT-CFM

RF

DSBM-IMF

10

2

10

3

NFE

4

.

5

5

.

0

5

.

5

6

.

0

FID

BM

CFM

OT-CFM

RF

DSBM-IMF

(a) Low-res

(b) Diffusion-fb

(c) Bridge Matching

(d) DSBM-IPF

(e) DSBM-IMF

specifies the frequency of the saturation specific humidity modulation. We follow Bischoff and Deck (2023) for data processing and network architecture, which is given by a U-Net with an additional spatial mean bypass network given by an MLP. The network size is approximately 11.3 million parameters in total. We train using batch size 4, learning rate 2 × 10 -4 , for 5000 steps per outer iteration and N = 12 outer iterations. We refresh the cache dataloader every 2500 steps with 1250 new samples. The training takes approximately 20 hours on a single RTX GPU. We used σ 2 = 0 . 3 and sample using 30 diffusion steps without finetuning these parameters. For the Diffusion-fb method in Bischoff and Deck (2023), we use the released code 10 without modifying any parameters.

We visualize intermediate and final reconstruction samples for different algorithms in Figure 19. We see that DSBM-IPF and DSBM-IMF provide consistent samples with the low resolution source, whereas Diffusion-fb and Bridge Matching produce dissimilar samples. We also follow Bischoff and Deck (2023) for a more refined statistical analysis in Figures 20, 21, 22. DSBM-IMF achieves comparable performance as Diffusion-fb in terms of these statistical profiles, and can be comparatively more accurate e.g. in the tails of the distributions in Figure 20, and for the case k x = k y = 4 in Figure 21 for which the power spectrum of supersaturation is correctly captured by DSBM-IMF but not by other methods. Comparing this analysis with Figure 11, DSBM-IMF is also significantly more accurate in terms of conditional consistency than Diffusion-fb. On the other hand, DSBM-IPF appears less accurate in terms of these unconditional statistics than Diffusion-fb and DSBM-IMF, but achieves lower glyph[lscript] 2 distances from the input sources in Figure 11. This suggests that DSBM-IPF and DSBM-IMF exhibit different empirical biases before convergence, and DSBM-IMF is more preferable when the accuracy of the samples are important. This is in line with IMF theory as the marginals π 0 , π T are preserved in IMF but not in IPF.

J Broader Impact

Our work focuses on theoretical and methodological research and is intended to bring closer the fields of generative modeling and optimal transport. It can be useful for learning transport maps between general distributions with high accuracy and high scalability, which can have useful applications

-

25

-

20

-

15

-

10

-

5

0

5

10

-

4

10

-

3

10

-

2

10

-

1

10

0

Density

k

x

=

k

y

= 2

True high-res

True low-res

Diffusion-fb

DSBM-IPF

DSBM-IMF

-

25

-

20

-

15

-

10

-

5

0

5

10

-

4

10

-

3

10

-

2

10

-

1

10

0

k

x

=

k

y

= 4

-

25

-

20

-

15

-

10

-

5

0

5

10

-

4

10

-

3

10

-

2

10

-

1

10

0

k

x

=

k

y

= 8

-

25

-

20

-

15

-

10

-

5

0

5

10

-

4

10

-

3

10

-

2

10

-

1

10

0

k

x

=

k

y

= 16

Supersaturation

-

20

-

10

0

10

20

10

-

4

10

-

3

10

-

2

10

-

1

10

0

Density

k

x

=

k

y

= 2

True high-res

True low-res

Diffusion-fb

DSBM-IPF

DSBM-IMF

-

20

-

10

0

10

20

10

-

4

10

-

3

10

-

2

10

-

1

10

0

k

x

=

k

y

= 4

-

20

-

10

0

10

20

10

-

4

10

-

3

10

-

2

10

-

1

10

0

k

x

=

k

y

= 8

-

20

-

10

0

10

20

10

-

4

10

-

3

10

-

2

10

-

1

10

0

k

x

=

k

y

= 16

Vorticity

10

0

10

1

10

2

10

-

6

10

-

5

10

-

4

10

-

3

10

-

2

10

-

1

10

0

Density

k

x

=

k

y

= 2

True high-res

True low-res

Diffusion-fb

DSBM-IPF

DSBM-IMF

10

0

10

1

10

2

10

-

6

10

-

5

10

-

4

10

-

3

10

-

2

10

-

1

10

0

k

x

=

k

y

= 4

10

0

10

1

10

2

10

-

6

10

-

5

10

-

4

10

-

3

10

-

2

10

-

1

10

0

k

x

=

k

y

= 8

10

0

10

1

10

2

10

-

6

10

-

5

10

-

4

10

-

3

10

-

2

10

-

1

10

0

k

x

=

k

y

= 16

Supersaturation

10

0

10

1

10

2

10

-

6

10

-

5

10

-

4

10

-

3

10

-

2

10

-

1

10

0

Density

k

x

=

k

y

= 2

True high-res

True low-res

Diffusion-fb

DSBM-IPF

DSBM-IMF

10

0

10

1

10

2

10

-

6

10

-

5

10

-

4

10

-

3

10

-

2

10

-

1

10

0

k

x

=

k

y

= 4

10

0

10

1

10

2

10

-

6

10

-

5

10

-

4

10

-

3

10

-

2

10

-

1

10

0

k

x

=

k

y

= 8

10

0

10

1

10

2

10

-

6

10

-

5

10

-

4

10

-

3

10

-

2

10

-

1

10

0

k

x

=

k

y

= 16

Vorticity

-

10

-

8

-

6

-

4

-

2

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

Density

k

x

=

k

y

= 2

True high-res

True low-res

Diffusion-fb

DSBM-IPF

DSBM-IMF

-

10

-

8

-

6

-

4

-

2

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

k

x

=

k

y

= 4

-

10

-

8

-

6

-

4

-

2

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

k

x

=

k

y

= 8

-

10

-

8

-

6

-

4

-

2

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

k

x

=

k

y

= 16

Supersaturation

in machine learning, but also natural science areas such as physics, biology and geosciences in which optimal transport maps with theoretical guarantees are appealing. Our fluid flows experiment demonstrates such potentials. However, as is the case for generative models as a whole, intentional malicious use could cause detrimental societal impacts.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper formulates a new framework of a probabilistic model called Diffusion Schrödinger Bridge Matching (DSBM) and an iterative algorithm which is called iterative Markovian fitting (IMF). Firstly, the paper has contribution of pointing out a relationship between famous generative models (score matching and flow matching) in the form of DSBM. IMF is also novel and might be more suitable to be trained in neural networks than well-established algorithms such as IPF. To achieve these, the authors utilize a tractable bridge structure (Brownian bridge), which enables preserving key properties, i.e., densities of marginal distribution at the start and end of the process. The proposed algorithm often shows better generation results in certain tasks, such as image transfer.\n\nSTRENGTHS:\n1. The paper is well-written despite dealing with a complex subject. As far as I understand, most of the arguments seem to be theoretically sound.\n2. SB algorithms have suffered from numerical errors and computational complexity. Since the IMF approach is akin to standard matching algorithms for neural networks, the provided empirical results show that it scales well. \n3. The paper provides a novel perspective on the diffusion bridge problem with a novel algorithm. As a result, DSBM-IMF shows promising results in certain tasks.\n\nWEAKNESSES:\n1. Significance. In Tables 2, 3 and Figure 10, It is difficult for me to observe that DSBM-IMF excels DSMB-IPF. Therefore, I would say significance of the proposed algorithm is not clearly shown in the experiments. The paper needs more empirical evidence of the actual benefits of choosing this algorithm, such as performance boost and scalability, in more challenging, or large-scale settings.\n2. Limited applications. I believe the actual implementation of the algorithm heavily relies on the well-known Brownian bridge. Although I think this is actually a valid contraint for certain applications such as image generation, at the same time it means that the algorithm is not always applicable to general optimal transport problems. Therefore, I recommend putting more effort on analyzing more classes of $\\mathbb{Q}$ in the experiment (if there are any), or putting clear justification on certain classes of $\\mathbb{Q}$ used in the experiment.\n\nQUESTIONS:\n1. Does $f_t$ is predefined in all tasks? \n2. Is $\\sigma_t$ constant and **not** scheduled? The authors of [1] mentioned time-symmetric scheduling the diffusion term, as suggested by prior SB models.\n3. Could you explain how the IMF scheme is fundamentally better than IPF algorithms? \n\n\n\n[1] Guan-Horng Liu et al. (2023), I2SB: Image-to-Image Schrödinger Bridge.\n\nLIMITATIONS:\nAlthough I think overall paper is good, but I also think the proposed experimental results are not quite sufficient; the paper need more empirical evidence to show the excellence of the model. I hope to see response from the authors for this point.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nFlow matching and alpha blending have achieved tremendous attention in the matching problems. Although the straight trajectories based on $X_t = (1-\\alpha)X_0+\\alpha X_1$ yield **fast inference**, it doesn't necessarily mean they are efficient in score estimations in training. To make the training theoretically more efficient, rectified flow leverages the optimal transport ideas to minimize the transport cost. However, the stochastic alternative is still lacking (stochastic interpolant may not be optimal in terms of optimal transport). To tackle this issue, the authors proposed the stochastic version of flow matching/ alpha blending with entropic optimized transport. A novel scheme IMF that resembles the IPF iteration is proposed.\n\nSTRENGTHS:\n1. Despite the advances of stochastic interpolant, the transport from marginals may not be optimal in general. The **optimal-transport**-based stochastic version of flow matching and alpha blending is **still missing** in the community.\n\n2. In my understanding, the proposal of the IMF algorithm is the key/ most important contribution of this paper.  \n\n3. solid experiments and comparisons are studied, which verify the promising potential of DSBM.\n\n4. Connections to flow matching, rectified flow, and stochastic interpolant are studied, which makes the authors easier to understand the benefit of this paper.\n\nWEAKNESSES:\n1. Clearer clarification between IPF and IMF would be required to facilitate the understanding. IPF is based on forward K; IMF is based on backward KL, this is good in math, but more detailed information is lacking. for example, if we solely check proposition 2, we can derive almost the same loss based on forward KL (which leads to the IPF algorithm) except that the underlying measure is not based on $\\Pi_{0, t}$. So **what is the real benefit/ advantages of simulating $\\Pi_{0, t}$ to train the loss in proposition 2 would clarify our concerns**\n\n2. Algorithm 1 is not clear enough. The authors may consider providing us more details in the learning of $\\upsilon_{\\phi^{\\star}}$ in Alg 1. For example, in the training of alpha-blending (line 60-64 in https://github.com/tchambon/IADB/blob/main/iadb.py), given some samples of $x_1$ and $x_0$, very simple code would be used to train the models. However, when it comes to DSBM, how is it implemented? pseudocode would work if it clarifies our concerns.\n\n3. My concern is that this algorithm may be quite expensive.\n\n\n\n\nMinor \n\n1. missing one related work:\n\n    Discussions on the accumulations of IPF errors based on general optimal-transport cost functions studied in [1] would be appreciated.\n\n    [1] Provably Convergent Schrödinger Bridge with Applications to Probabilistic Time Series Imputation. ICML'23.\n\n\n2. writing of Equation between line 637 and line 638 is weird, because given X_t=x_t and x_T, what is the randomness there and why do we need expectation w.r.t. $\\Pi_{T|t}$?  \n\n3. derivation of eq after line 113 is not clear to me. why $E_{\\Pi_{0, T}}[||X_0-X_T||]^2$? is there any reference for that?\n\n4. first part of section 4 in the main paper, the need the alternating schemes is still not well motivated\n\n5. what is IMF-b?\n\nQUESTIONS:\n1. DDMs can be seen as the first iteration of DSM, which implies that we can use DDMs model weights to initialize the model weights of DSM. Does it mean that empirically we can also use weights from stochastic interpolant / the stochastic version of flow matching to train the first step of DSBM?\n\n2. In line 993 in the appendix, the authors mentioned that \"The advantage of this DSM loss is that it does **not rely on any divergence computation**\". I believe this is quite an important part of scalability, but why it is only emphasized in the joint training instead of stating it in the main paper?\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe submission suggests a numerically effective approach for solving the Schroedinger Bridge (SB) problem and illustrates its potential for generative modeling. The approach generalizes in some sense some recent flow matching methods. Furthermore, it provides a more efficient (no full trajectory caching) algorithm with less error accumulations that better remembers the prior reference measure for solving SB problems compared to previous works. Empirical results indicate for instance that it can yield better generative performance for MNIST-EMNIST transfer compared to previous work, or better reconstructions for downscaling geophysical fluid dynamics.\n\nSTRENGTHS:\nOriginality: The proposed methodology and algorithm is novel as far as I am aware. The work provides additional contributions compared to cited concurrent work [Peluchetti (2023)] that shares some similar ideas. Related literature and approaches are discussed in detail.\n\nQuality: Claims are adequately supported by proofs, and numerical experiments seem to support this. \n\nClarity: The paper is well written. The presentation is quite rigorous. \n\nSignificance: In my opinion, this work is of interest to the community as the method (i) generalizes recent flow matching approaches with non-degenerate noise and (ii) improves experimentally previous Diffusion Schrödinger Bridge approaches and (iii) provides some theoretical guarantees for solving the SP problem (under idealized assumptions).\n\nWEAKNESSES:\nIt was not clear to me how the computational cost of the suggested method compares to related works per iteration or diffusion step (e.g. in Figures 5 and 16).\n\nQUESTIONS:\nIn Figure 16, it appears that FID can be improved by using more than 100 steps. Does DBSM obtain better FID values compared to Bridge matching for more than 100 steps?\n\nIt was not clear to me what to take away from some parts of the appendix, such as as E Discrete-Time Markovian Projection or the Brownian/OU bridge representations in B.2. Are the results in B.2 helpful for saying anything about how to specify the SDE in the forward process $X$ for the suggested approach?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces Iterative Markovian Fitting (IMF) as a new method to compute Schrödinger Bridges (SBs), which are dynamic versions of entropy-regularized optimal transport. IMF alternates between projecting on the space of Markov processes and the reciprocal class, and it preserves the initial and terminal distributions. The paper also proposes Diffusion Schrödinger Bridge Matching (DSBM) as an algorithm to approximate SB solutions derived from IMF. DSBM overcomes issues of previous techniques and solves a simple regression problem at each iteration. The performance of DSBM is demonstrated in various transport tasks. The paper provides theoretical results, notations, and definitions related to the topics discussed.\n\nSTRENGTHS:\n1, The paper introduces Iterative Markovian Fitting (IMF) as a new procedure for computing Schrödinger Bridges (SBs). This approach provides a fresh perspective and contributes to the field of optimal transport and generative modeling.\n2, The paper establishes various theoretical results for IMF, demonstrating the validity and effectiveness of the proposed method. These results enhance the understanding of SBs and their applications in machine learning.\n\nWEAKNESSES:\n1， While the paper mentions that IMF and DSBM address scalability issues, it lacks a detailed analysis of their scalability in terms of computational resources and dataset sizes. Understanding the scalability limitations of the proposed methods is crucial for their applicability to real-world, large-scale problems.\n2, The paper assumes locally Lipschitz continuity of drift and limited settings for the proposed methods. These assumptions may limit the applicability of IMF and DSBM to certain scenarios. A discussion of the limitations imposed by these assumptions would provide a more comprehensive view of the proposed methods.\n\nQUESTIONS:\n1, How does the Iterative Markovian Fitting (IMF) method differ from the traditional Iterative Proportional Fitting (IPF) approach in computing Schrödinger Bridges (SBs)?\n2, Can you provide more insights into the theoretical results established for IMF in the paper?\n3, How does Diffusion Schrödinger Bridge Matching (DSBM) overcome the time-discretization and \"forgetting\" issues of previous diffusion-based techniques?\n\nLIMITATIONS:\nAs stated in \"Weakness\"", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes IMF, an iterative method for solving Schrodinger Bridge problems where the solution at each iteration preserves the correct marginal distributions at times 0 and T. This is in contrast to the existing method IPF which only satisfies this in the limit (and is very difficult to reach in practice).\n\nSTRENGTHS:\n- An interesting marginal-preserving algorithm for solving Schrodinger bridge problems. The popular method IPF can have convergence problems in practice, while this seems to fix some of that.\n\nWEAKNESSES:\n- Presentation. While I eventually understood the resulting IMF algorithm, I felt the presentation can be much more simplified. The algorithm is actually quite simple and closely related to priors works on SB and OT. While I understand the need to prove convergence and showcase novelty, the heavily front-loaded mathematical notation made reading on the first pass very difficult. I imagine the current paper may be off-putting for many readers who are not extremely familiar with related works. \n\n\n  - Novelty. The proposed algorithm seems very similar to Rectified Flow but just with an additional noise parameter.\n\n\n  - Limitation. The proposed algorithm only works for the standard formulation of Schrodinger bridge with quadratic costs. However, in this setting, it isn't clear to me why the extra stochasticity is useful, as opposed to solving deterministic maps (optimal transport problems).\n\nQUESTIONS:\n- \"Most notably, we adopt the SDE approach which is crucial for the validity of Proposition 5\" - Can the authors expand on what breaks when there is no stochasticity? From what I understand, RF should also preserve marginal distributions and converge to the optimal transport solution, which is discussed more in [1]. The solutions are also unique similar to any Flow Matching-based approach.\n\n\n  - All of the experiments in the main paper are currently only for transferring between two data distributions, but being able to find optimal transport maps is useful for generative modeling as well. I think it could really make the paper stronger by having experiments in generative modeling (Gaussian -> data) and showing some empirical metrics from e.g. [2] such as sampling cost and consistency, since these are important problems in generative modeling at the moment. [Edit: I found the experiments in I.5. I see that sampling efficiency is on par with RF but Table 6 seems pretty good! Although it does look like if we extrapolate to NFE > 100, the other methods are descending quite a bit faster in FID...]\n\n\n  - Can the authors expand on the differences between the two proposed methods, DSBM-IPF and DSBM-IMF? If I understood correctly, the main difference between DSBM-IPF and DSBM-IMF is only in the first iteration, where the pairs (x0, xT) are sampled from (either an OU process Q or the independent joint distribution). And the real difference is that Q has a slight bias at time T because it may not have reached its stationary distribution. Practically, this difference seems very negligible, yet there seems to be differences in empirical results (e.g. Table 3 and Figure 5). Theoretically, since Q does not match the correct marginal at T, this is also problematic for the proof of convergence to SB right?\n\n\n[1] \"Rectified Flow: A Marginal Preserving Approach to Optimal Transport\" https://arxiv.org/abs/2209.14577\n\n\n[2] \"Multisample Flow Matching: Straightening Flows with Minibatch Couplings\" https://arxiv.org/abs/2304.14772\n\nLIMITATIONS:\nThe authors do mention limitations such as marginal improvement, computational issues, and difficulties when sigma --> 0. I think these are sufficient.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-08"}