Title: Schrodinger Bridge Flow for Unpaired Data Translation
Venue: NeurIPS 2024
Accepted: yes
Keywords: unpaired data translation

diffusion models

Schrodinger bridge

bridge matching

stochastic interpolant

flow matching
Abstract: Mass transport problems arise in many areas of machine learning whereby one wants to compute a map transporting one distribution to another. Generative modeling techniques like Generative Adversarial Networks (GANs) and Denoising Diffusion Models (DMMs) have been successfully adapted to solve such transport problems, resulting in CycleGAN and Bridge Matching respectively. However, these methods do not approximate Optimal Transport (OT) maps, which are known to have desirable properties. Existing techniques approximating OT maps for high-dimensional data-rich problems, including DDMs-based Rectified Flow and Schrodinger bridge procedures, require fully training a DDM-type model at each iteration, or use mini-batch techniques which can introduce significant errors. We propose a novel algorithm to compute the Schrodinger bridge, a dynamic entropy-regularized version of OT, that eliminates the need to train multiple DDMs-like models. This algorithm corresponds to a discretization of a flow of path measures, referred to as the Schrodinger Bridge Flow, whose only stationary point is the Schrodinger bridge. We demonstrate the performance of our algorithm on a variety of unpaired data translation tasks.

Schrödinger Bridge Flow for Unpaired Data Translation

Abstract

Mass transport problems arise in many areas of machine learning whereby one wants to compute a map transporting one distribution to another. Generative modeling techniques like Generative Adversarial Networks (GANs) and Denoising Diffusion Models (DDMs) have been successfully adapted to solve such transport problems, resulting in CycleGAN and Bridge Matching respectively. However, these methods do not approximate Optimal Transport (OT) maps, which are known to have desirable properties. Existing techniques approximating OT maps for high-dimensional data-rich problems, such as DDM-based Rectified Flow and Schrödinger Bridge procedures, require fully training a DDM-type model at each iteration, or use mini-batch techniques which can introduce significant errors. We propose a novel algorithm to compute the Schrödinger Bridge, a dynamic entropyregularised version of OT, that eliminates the need to train multiple DDM-like models. This algorithm corresponds to a discretisation of a flow of path measures, which we call the Schrödinger Bridge Flow, whose only stationary point is the Schrödinger Bridge. We demonstrate the performance of our algorithm on a variety of unpaired data translation tasks.

1 Introduction

The problem of finding a map to transport one probability distribution to another one has numerous applications in machine learning. In particular, it is at the core of generative modeling where the idea is to transform a noise distribution into the data distribution, and is also central to transfer learning tasks such as image-to-image translation. For discrete probability distributions, it is possible to compute the Optimal Transport (OT) map but this is computationally expensive (Peyré et al., 2019). By showing that an entropy-regularised version of OT, the Entropic OT (EOT), could be computed much more efficiently using the Sinkhorn algorithm, Cuturi (2013) has enabled transport ideas to be used in numerous applications (Ge et al., 2021; Zhou et al., 2022). However, the computational complexity of Sinkhorn algorithm is quadratic in the sample size, which makes its application to very large datasets impractical. Mini-batch versions have been proposed, see e.g. (Genevay et al., 2018), but tend to introduce significant errors in high dimensions (Sommerfeld et al., 2019).

In the context of generative modeling, Denoising Diffusion Models (DDMs) (Song et al., 2021a; Ho et al., 2020) have shown impressive performance in a variety of domains. DDMs define a forward process progressively noising the data, and sample generation is achieved by approximating the timereversal of this diffusion. In order to leverage the iterative refinement properties of DDMs in the OT setting, methods exploiting the equivalence between the static versions of (E)OT and their dynamic counterparts (Benamou and Brenier, 2000; Léonard, 2014) have been developed. A procedure to approximate the dynamic OT is considered by Liu et al. (2023b), while techniques to approximate the dynamic equivalent to EOT, the Schrödinger Bridge (SB), have been proposed in (De Bortoli et al., 2021; Vargas et al., 2021; Chen et al., 2022; Peluchetti, 2023; Shi et al., 2023). These techniques are

Arnaud Doucet Google DeepMind

expensive however, as they require training multiple DDM-type models. Mini-batch versions of OT and Sinkhorn (Pooladian et al., 2023; Tong et al., 2024b) combined with bridge or flow matching have also been proposed to approximate the OT path and SB, but they optimise a minibatch OT objective that can introduce significant errors in high dimensions: the error in Wasserstein1 distance is of order O ( B -1 / (2 d ) ) , where d is the dimension of the problem and B the minibatch size, see (Sommerfeld et al., 2019, Corollary 1).

In this paper, we propose a novel approach to computing the SB. Similarly to Iterative Markovian Fitting (IMF) and its practical implementation, Diffusion Schrödinger Bridge Matching (DSBM) (Shi et al., 2023; Peluchetti, 2023), it leverages the fact that the SB is the only Markov process with prescribed marginals at the endpoints which is in the reciprocal class of the Brownian motion, i.e. it has the same bridge as the Brownian motion (Léonard, 2014); see Section 2 for more details on Markov processes and the reciprocal class. Compared to DSBM, our approach is easier to implement as it does not require caching samples, alternating between optimising two different losses, and, optionally, uses one neural network instead of two. In Section 3, we start by introducing a flow of path measures whose time-discretisation yields a family of algorithms called α -IMF and presented in Section 4. Notably, we show that α -IMF converges to the Schrödinger Bridge for any α ∈ (0 , 1] . Additionally, for a special value of the discretisation stepsize α = 1 , we recover the IMF procedure (Peluchetti, 2023; Shi et al., 2023), while α < 1 corresponds to online versions of IMF. We implement a parametric version of the α -IMF as an online DSBM procedure, called α -DSBM. We illustrate the efficiency of our approach in unpaired image-to-image translation settings in Section 6.

Notation. We denote the space of path measures by P ( C ) , i.e. P ( C ) = P (C([0 , 1] , R d )) , where C([0 , 1] , R d ) is the space of continuous functions from [0 , 1] to R d . The subset of Markov path measures associated with a diffusion of the form d X t = v t ( X t )d t + σ t d B t , with σ, v locally Lipschitz, is denoted M . For Q induced by ( √ ε B t ) t ∈ [0 , 1] , with ε > 0 and ( B t ) t ≥ 0 a d -dimensional Brownian motion, the reciprocal class of Q is denoted R ( Q ) , see Definition 2.1. For any P ∈ P ( C ) , we denote by P t its marginal distribution at time t , P s,t the joint distribution at times s, t , P s | t the conditional distribution at time s given the state at time t , and P | 0 , 1 ∈ P ( C ) the distribution of the path on time interval (0 , 1) given its endpoints; e.g. Q | 0 , 1 is a scaled Brownian bridge. Unless specified otherwise, all gradient operators ∇ are w.r.t. the variable x t with time index t . Given probability spaces ( X , X ) and ( Y , Y ) , a Markov kernel K : X ×Y → [0 , 1] , and a probability measure µ defined on X , we write µ K for 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 , 1 over R d × R d , we denote the mixture of bridges measure as Π = Π 0 , 1 P | 0 , 1 ∈ P ( C ) , which is short for Π( · ) = ∫ R d × R d P | 0 , 1 ( ·| x 0 , x 1 )dΠ 0 , 1 ( x 0 , x 1 ) . Finally, we define the Kullback-Leibler (KL) divergence between two probability measures π 0 , π 1 ∈ P ( X ) as KL( π 0 | π 1 ) = ∫ X log((d π 0 / d π 1 )( x ))d π 0 ( x ) if π 0 is absolutely continuous w.r.t. π 1 and KL( π 0 | π 1 ) = + ∞ otherwise.

2 Optimal Transport and Schrödinger Bridge

Unpaired Transfer and Optimal Transport. Given unpaired data samples from π 0 and π 1 , where π 0 , π 1 are two distributions on R d , we are interested in designing a transport map from π 0 to π 1 . This corresponds to an unpaired data transfer task . We can formulate this problem as finding a distribution Π on R d × R d with marginals Π 0 = π 0 and Π 1 = π 1 so that if X 0 ∼ π 0 then X 1 | X 0 ∼ Π 1 | 0 ( ·| X 0 ) satisfies X 1 ∼ π 1 . Among an infinite number of such so-called coupling distributions Π , we are here interested in finding the Entropic Optimal Transport (EOT) coupling Π ⋆ defined as

where H(Π) is the differential entropy of Π and ε > 0 is a regularisation hyperparameter (Peyré et al., 2019). For ε = 0 , we recover the standard OT.

In order to leverage the recent advances in generative modeling, and in particular the concept of iterative refinement central to DDMs, we turn to a dynamic formulation of EOT known as the Schrödinger Bridge problem (Léonard, 2014). It is defined as follows: find P ⋆ ∈ P ( C ) such that

with Q ∈ P ( C ) induced by a scaled d -dimensional Brownian motion ( √ ε B t ) t ∈ [0 , 1] . The term dynamic here refers to the fact that (2) is defined on path measures, i.e. on (stochastic) processes, in contrast to the static problem (1) which is defined on measures on the space R d × R d . In Section 3, we show that solving (2) is equivalent to optimising the vector field of a stochastic process using objectives similar to the ones of bridge matching (Peluchetti, 2021; Albergo and Vanden-Eijnden, 2023; Lipman et al., 2023; Liu et al., 2023a). Under mild assumptions, it can be shown that P ⋆ 0 , 1 = Π ⋆ , see e.g. (Léonard, 2014; Pavon et al., 2021). Hence solving (1) reduces to solving (2). Once we have found P ⋆ associated with ( X ⋆ t ) t ∈ [0 , 1] , we can sample from P ⋆ by first sampling X ⋆ 0 ∼ π 0 and then sampling the trajectory ( X ⋆ t ) t ∈ (0 , 1] which yields ( X ⋆ 0 , X ⋆ 1 ) ∼ Π ⋆ .

Reciprocal and Markov projections. To introduce our methodology, it is necessary to recall the notions of reciprocal and Markov projections. We refer to Shi et al. (2023) for more details. For practitioners, a more intuitive explanation of these projections is given in Appendix E.

Definition 2.1 (Reciprocal projection): P ∈ P ( C ) is in the reciprocal class R ( Q ) of Q if P = P 0 , 1 Q | 0 , 1 . We define the reciprocal projection of P ∈ P ( C ) as P ⋆ = proj R ( Q ) ( P ) = P 0 , 1 Q | 0 , 1 . We will write proj R instead of proj R ( Q ) to simplify notation.

In other words, P is in the reciprocal class of Q if the conditional distribution of a path given its endpoints is identical under P and Q , see (Rœlly, 2013). Sampling from the reciprocal projection of P can be achieved by sampling a path ( X t ) t ∈ [0 , 1] from P , keeping only the values of the endpoints, say X 0 , X 1 , and then sampling a new value for the bridge ( X t ) t ∈ (0 , 1) from Q | 0 , 1 .

Definition 2.2 (Markov projection): Assume that Q is induced by ( √ ε B t ) t ∈ [0 , 1] for ε > 0 . Then, when it is well-defined, for any P ∈ R ( Q ) , the Markovian projection M = proj M ( P ) ∈ M is the path measure induced by the diffusion ( X ⋆ t ) t ∈ [0 , 1] with for any t ∈ [0 , 1] d X ⋆ t = v ⋆ t ( X ⋆ t )d t + √ ε d B t , v ⋆ t ( x t ) = ( E P 1 | t [ X 1 | X t = x t ] -x t ) / (1 -t ) , X ⋆ 0 ∼ P 0 .

In practice, implementing a Markovian projection requires solving a regression problem to approximate E P 1 | t [ X 1 | X t = x t ] , similar to the one appearing in bridge matching and flow matching. One key property of the Markovian projection is that M t = P t for all t ∈ [0 , 1] , i.e. the Markovian projection preserves the marginals; see (Peluchetti, 2021) for instance.

Iterative Markovian Fitting. Leveraging the reciprocal and Markovian projections, Peluchetti (2023) and Shi et al. (2023) concurrently introduced IMF. Starting from ˆ P 0 = ( π 0 ⊗ π 1 ) Q | 0 , 1 , a measure where endpoints are sampled independently from π 0 and π 1 and then interpolated using a (scaled) Brownian bridge, they define a sequence of path measures ( P n , ˆ P n ) n ∈ N where P n = proj M ( ˆ P n ) and ˆ P n +1 = proj R ( P n ) . This ensures that P n 0 = π 0 , P n 1 = π 1 for all n , and it can be shown that the sequence ( P n ) n ∈ N converges to the SB, see (Peluchetti, 2023, Theorem 2). The practical implementation of this algorithm proposed by Shi et al. (2023) is called DSBM. Implementing DSBM poses challenges, as each Markovian projection requires training a neural network to approximate the relevant conditional expectations by minimising a bridge matching loss. Furthermore, in practice, generated model samples are stored in a cache in order to train the next iterations of DSBM. This introduces additional hyperparameters that require tuning. In Section 3 we propose α -IMF, an algorithm which can be interpreted as the discretisation of a flow of path measures . This leads to α -DSBM, an algorithm that is computationally much more efficient than DSBM as it does not rely on a Markovian projection at each step.

3 Schrödinger Bridge flow

We will now introduce a flow of path measures ( P s ) s ≥ 0 , and show that the time-discretisation of this flow with an appropriate stepsize α ∈ (0 , 1] yields a family of procedures called α -IMF, which all converge to the Schrödinger Bridge. While α = 1 yields the classical IMF, α ∈ (0 , 1) yields an incremental version of IMF. In Section 4 we show that α -IMF can be implemented as an online version of DSBM.

3.1 A flow of path measures

Let ( P s , ˆ P s ) s ≥ 0 be a flow of path measures defined for any s ≥ 0 by

which we assume is well-defined. Note that for any s ≥ 0 , P s is Markov while ˆ P s is in the reciprocal class of Q . Crucially, the only fixed point of (3) is the SB. Indeed, let ¯ P be a fixed point of ( P s ) s ≥ 0 in (3). Then, we have that ¯ P = proj R (proj M ( ¯ P )) . Hence, we get ¯ P = proj R (proj M ( . . . (proj R (proj M ( ¯ P ))) . . . )) . Hence, under mild assumptions, ¯ P is a limit point of IMF and therefore ¯ P is the SB P ⋆ given by (2), see (Peluchetti, 2023, Theorem 2).

Next, for any α ∈ (0 , 1] , we define the following discretisation of (3) called α -IMF:

and P n = proj M ( ˆ P n ) . Note that for any n ∈ N , ˆ P n ∈ R ( Q ) . This recovers the IMF procedure (Shi et al., 2023; Peluchetti, 2023) when α = 1 . Using the definition of the sequence ( ˆ P n ) n ∈ N , it is possible to analyse the sequence

( P n ) n ∈ N using the properties of the KL divergence as well as the Pythagorean identities derived in (Shi et al., 2023; Peluchetti, 2023). We first introduce some assumptions on the Schrödinger Bridge problem. We recall that the differential entropy of a probability measure π is given by

if π admits a density with respect to the Lebesgue measure and + ∞ otherwise. Recall that Q is associated with ( √ ε B t ) t ∈ [0 , 1] and assume that Q 0 = Leb . Let π 0 , π 1 ∈ P ( R d ) such that

for i ∈ { 0 , 1 } . Under these assumptions, we can use the characterisation of the SB as the only path measure that preserves π 0 , π 1 , and is both Markov and in the reciprocal class of Q (see e.g. (Léonard, 2014, Theorem 2.12)). We get the following result.

Theorem 3.1 (Convergence of α -IMF): Let α ∈ (0 , 1] and ( P n , ˆ P n ) n ∈ N defined by (4) . Under mild assumptions, we have that lim n → + ∞ P n = P ⋆ , where P ⋆ is the solution of the Schrödinger Bridge problem (2) .

3.2 Discretisation and non-parametric loss

We show here that α -IMF is associated with an incremental version of DSBM for α ∈ (0 , 1) .

Iterative Markovian Fitting. For any v : [0 , 1] × R d → R d , we introduce the loss function

where we recall that Q is induced by ( ε B t ) t ∈ [0 , 1] for some ε > 0 . This loss was already considered in (Peluchetti, 2021; Lipman et al., 2023; Liu et al., 2023a; Liu, 2022; Shi et al., 2023). We also define the path measure P v ∈ P ( C ) associated with

Consider first the sequence ( v n ) n ∈ N defined by

Using Definition 2.2, we have that P v n +1 = proj M (proj R ( P v n )) , which corresponds to P n +1 in the IMF sequence. Therefore we have that lim n → + ∞ P v n = P ⋆ under mild assumptions (Peluchetti, 2023, Theorem 2).

Functional gradient descent. We now introduce a relaxation of (7), where, instead of considering the argmin , we update the vector field with one gradient step. To define this relaxation, we recall that for a functional F : F → R , where F is an appropriate function space, its functional derivative (Courant and Hilbert, 2008) with reference measure µ is denoted ∇ µ F and is given for any ϕ ∈ F , when it exists, by

Initialised with v 0 t ( x ) = ( E ˆ P 0 1 | t [ X 1 | X t = x ] -x ) / (1 -t ) , where ˆ P 0 = ( π 0 ⊗ π 1 ) Q | 0 , 1 , we now introduce a sequence of vector fields ( v n ) n ∈ N . This corresponds to training a bridge matching model (see e.g. Liu et al. (2023a); Albergo et al. (2023)), giving P v 0 = proj M ( ˆ P 0 ) . Then for n ∈ N , let

with δ n > 0 and µ n ∈ P ( C ) . The parameters ( δ n , µ n ) n ∈ N will be made explicit in Proposition 3.2. We emphasize that, in contrast to the IMF procedure, in the online update (8) we do not need to solve a Markovian projection problem at every step; instead we simply take a gradient step on the loss (5).

Connection with α -IMF. The following proposition shows that ( P v n ) n ∈ N defined by (8) is associated with α -IMF defined in (4).

Proposition 3.2 (Non-parametric updates are α -IMF): Let α ∈ (0 , 1] , ( P n , ˆ P n ) n ∈ N as in (4) , δ n = α and µ n = (1 -α ) ˆ P n + α proj R ( P n ) . Then, under mild assumptions, we have P v n = P n for all n ∈ N .

Combining Theorem 3.1 to Proposition 3.2, we get that lim n → + ∞ P v n = P ⋆ , i.e. the non-parametric procedure converges to the SB.

4 α -Diffusion Schrödinger Bridge Matching

From DSBM to α -DSBM. In Section 3, we introduced α -IMF, a scheme which defines a sequence of path measures converging to the SB for all α ∈ (0 , 1] . For α = 1 , this corresponds to the IMF, whose practical DSBM implementation (Shi et al., 2023) requires repeatedly solving an expensive minimisation problem (7). In contrast, for α < 1 we are only required to take one (non-parametric) gradient step to update the vector field, see (8). This suggests the following practical implementation of α -IMF, called α -DSBM: First, pretrain a bridge matching model so that for t ∈ [0 , 1] and x ∈ R d , v θ t ( x ) = ( E ˆ P 0 1 | t [ X 1 | X t = x ] -x ) / (1 -t ) , where ˆ P 0 = ( π 0 ⊗ π 1 ) Q | 0 , 1 . Then, perform the parametric version of the update (8):

where P ¯ θ is a stop-gradient version of P v θ . In Appendix D.2, we give a theoretical justification for this parametric equivalent of (5) and (8) by showing that, as α → 0 , the update on the velocity fields v θ given by (9) corresponds to a direction of descent for the non-parametric loss (8) on average. Once again, we emphasize that if we replace the gradient step in (9) with the minimisation θ ← argmin θ L( θ, P ¯ θ ) , we recover DSBM.

Bidirectional online procedure. As with DSBM, directly implementing (9) leads to error quickly accumulating, see Appendix I for details. One way to circumvent this error accumulation issue is to consider a bidirectional procedure, in which we train both a forward and a backward

model. This is possible because the Markovian projection coincides for forward and backward path measures, see (Shi et al., 2023, Proposition 9). This suggests considering the loss L ( v , v , P , P ) = ∫ 1 0 L t ( v t , v t , P , P )d t , which is an extension of (5), where

Similarly to (6), we define P v , P v , associated with ( X t ) t ∈ [0 , 1] and ( Y 1 -t ) t ∈ [0 , 1] respectively, which are defined by forward and backward SDEs

Similarly to (8), we define non-parametric updates for any n ∈ N , t ∈ [0 , 1] and x ∈ R d

We have the following proposition which ensures our bidirectional procedure is still valid and that the results of Proposition 3.2 still hold.

Proposition 4.1 (Bidirectional updates): Let α ∈ (0 , 1] . For any n ∈ N , define ( P n , ˆ P n ) n ∈ N by (4) . Then, under mild assumption and assuming that δ n = α and µ n = (1 -α ) ˆ P n + α proj R ( P n ) , we have that for any n ∈ N , P v n, = P v n, = P n .

In Appendix I, we show that in the Gaussian setting the bidirectional procedure (4.1) does not accumulate error when the vector field is approximated, while the unidirectional one (8) does.

Vector field parameterisation. Contrary to existing procedures (Shi et al., 2023; Peluchetti, 2023; Liu, 2022), we do not parameterise v and v using two separate networks. Instead, we consider an additional input s ∈ { 0 , 1 } such that v θ (1 , · ) ≈ v and v θ (0 , · ) ≈ v . This allows us to substantially reduce the number of parameters in the model. The conditioning on s in the network is detailed in Appendix K. Before stating our full algorithm in Algorithm 1, we introduce a batched parametric version of (10). For ease of notation, we write Interp t for the operation corresponding to sampling from Q t | 0 , 1 , i.e.

Weare now ready to introduce the batched parametric version of (10). For a given batch of inputs X 1: B 0 and X 1: B 1 , timesteps t ∼ Unif([0 , 1]) ⊗ B , and X t = Interp t ( X 0 , X 1 , Z ) with Z ∼ N (0 , Id) ⊗ B , we compute the empirical forward and backward losses as

We present the resulting α -DSBM in Algorithm 1. Note that in this algorithm, we maintain an Exponential Moving Average (EMA) of model parameters, as is common in diffusion models (Nichol and Dhariwal, 2021). During the finetuning stage, when we generate samples to use as model's inputs, we then have a choice of sampling using the EMA or non-EMA parameters. At test time, we always sample using the EMA parameters, as it is known to improve the visual quality (Song and Ermon, 2020). In Algorithm 1, we specify α ∈ (0 , 1] as a stepsize parameter. In practice, we use Adam (Kingma and Ba, 2015) for optimization, thus the choice of α is implicit and adaptive throughout the training. To emphasize the importance of the parameter α , we sweep over its value with an explicit solver SGD in a toy setting, see Appendix K.2. We refer to Appendix K for more details on our experimental setup.

Algorithm 1 α -Diffusion Schrödinger Bridge Matching

1: Input: datasets π 0 and π 1 , entropic regularisation ε , number of pretraining and finetuning steps N pretraining and N finetuning , batch size B and half batch size b = B/ 2 , EMA decay γ , initial parameters θ and initial EMA parameters θ EMA = θ , α ∈ (0 , 1]

2: for n ∈ { 1 , . . . , N pretraining } do

3: Sample ( X 0 , X 1 ) ∼ ( π 0 ⊗ π 1 ) ⊗ B

5: Update θ with a gradient step on 1 2 [ ℓ ( t 1: b , X 1: b 1 , X 1: b t ) + ℓ ( t b +1: B , X b +1: B 0 , X b +1: B t )]

4: Sample t ∼ Unif([0 , 1]) ⊗ B and Z ∼ N (0 , Id) ⊗ B and compute X t = Interp t ( X 0 , X 1 , Z )

6: Update EMA parameters: θ EMA = γθ EMA +(1 -γ ) θ

7: end for

8: for n ∈ { 1 , . . . , N finetuning } do

9: Sample ( X 0 , X 1 ) ∼ ( π 0 ⊗ π 1 ) ⊗ b

10: Sample ˆ X 1 solving forward SDE (11)-(fwd) with v θ EMA (1 , · ) or v θ (1 , · ) starting from X 0

11: Sample ˆ X 0 solving backward SDE (11)-(bwd) with v θ EMA (0 , · ) or v θ (0 , · ) starting from X 1

12: Sample t ∼ Unif([0 , 1]) ⊗ b and Z ∼ N (0 , Id) ⊗ b and compute X t = Interp t ( ˆ X 0 , X 1 , Z )

14: Update θ with a gradient step on 1 2 [ ℓ ( t , X 1 , X t ) + ℓ ( t , X 0 , X t )] and stepsize α

13: Sample t ∼ Unif([0 , 1]) ⊗ b and Z ∼ N (0 , Id) ⊗ b and compute X t = Interp t ( X 0 , ˆ X 1 , Z )

15: Update EMA parameters: θ EMA = γθ EMA +(1 -γ ) θ

16: end for

17: Output: ( θ, θ EMA ) parameters of the finetuned model

5 Related work

Solving Schrödinger Bridge problems. Schrödinger Bridges (Schrödinger, 1932) have been thoroughly studied through the lens of probability theory (Léonard, 2014) and stochastic control (Dai Pra, 1991; Chen et al., 2021). They recently found applications in generative modeling and related fields leveraging recent advances in diffusion models (De Bortoli et al., 2021; Vargas et al., 2021; Chen et al., 2022). Extensions of these methods to other machine learning problems and modalities were studied in (Shi et al., 2022; Thornton et al., 2022; Liu et al., 2022; Chen et al., 2023; Tamir et al., 2023). Shi et al. (2023); Peluchetti (2023) concurrently introduced the DSBM algorithm which relies on a new procedure called IMF, while the DSB algorithm introduced in (De Bortoli et al., 2021) is based on the standard Iterative Proportional Fitting (IPF) scheme. Neklyudov et al. (2023a,b); Liu et al. (2022) generalise DSBM to arbitrary cost functions, albeit at the expense of having to learn the reciprocal projection which is no longer given by a Brownian bridge. These new methodologies translate to improved numerics when compared to their IPF counterparts, but they remain reliant on alternating between the optimisation of two losses. Finally, we note that the Schrödinger Bridge flow and the α -IMF procedure can be linked to the Sinkhorn flow recently introduced by Karimi et al. (2024), see Appendix H.1 for a detailed discussion.

Sampling-free methodologies. Sampling-free methodologies have been proposed to solve OT related objectives. In (Liu et al., 2023a; Somnath et al., 2023; Diefenbacher et al., 2024; Cao et al., 2024), the authors perform one step of DSBM, i.e. only consider the pretraining stage of our algorithm. While the obtained bridge might enjoy transport properties, it does not solve an OT problem. In another line of work, Pooladian et al. (2023); Tong et al. (2024a,b); Eyring et al. (2024) have proposed simulation-free methods to minimise OT objectives. However, they target not the OT problem, but a minibatch version of it which coincides with OT only in the limit of infinite batch size, see (Pooladian et al., 2023, Theorem 4.2). Other sampling-free methods to solve the Schrödinger Bridge problem include Kim et al. (2024); Gushchin et al. (2024b) both of which rely on adversarial losses to solve the OT problem. In (De Bortoli et al., 2021; Vargas et al., 2021; Liu et al., 2022; Shi et al., 2023; Peluchetti, 2023) the adversarial objective is dropped and instead the procedure requires alternating objectives during training and is not sampling-free. We also highlight the line of work of Korotin et al. (2024); Gushchin et al. (2024a) in which the Schrödinger Bridge potentials are parameterised with mixtures of Gaussians, allowing for fast training in small dimensions. Finally, recently Deng et al. (2024) introduced a variation on Schrödinger Bridge for generative modeling, which while still not sampling-free, does not require learning a forward process.

0.8

distance

Iterative @1K

Base

Iterative@2.5K

Online

0.6

Iterative @5K

0.4

Normalized

0.2

10k20k30k40k50k60k70k80k90k100k

Trainingiterations

1.0

0.9

alue

val

0.8

ariance

0.7

Covar

0.6

0.5

Iterative@1K

一

Base

Iterative@2.5K—

Online

0.4

Iterative@5K

Optimum

10k

20k

30k

40k

50k

Trainingiterations

6 Experiments

In this section, we illustrate the efficiency of α -DSBM on different tasks. In Section 6.1, we compare α -DSBM to DSBM in a Gaussian setting where the EOT coupling is tractable and show that α -DSBM recovers the solution faster than DSBM. In Section 6.2, we illustrate the scalability of our method through a range of unpaired image translation experiments.

6.1 Gaussian case

We compare α -DSBM to DSBM in the Gaussian setting where π 0 = N (0 , σ 2 0 Id) , π 1 = N (0 , σ 2 1 Id) and Q is associated with ( √ ε B t ) t ∈ [0 , 1] with √ ε = 0 . 5 . In this case, the EOT coupling is N (0 , Σ ⋆ ) , with Σ ⋆ given by

with Id being a d × d identity matrix. We consider d = 50 , σ 0 = σ 1 = 1 , resulting in σ 2 ⋆ ≈ 0 . 88 . To showcase the robustness of α -DSBM, we consider the initial coupling P 0 , 1 , where ( X 0 , X 1 ) ∼ P 0 , 1 , X 0 ∼ N (0 , Id) , X 1 = -X 0 , and let ˆ P 0 = P 0 , 1 Q | 0 , 1 . In this setting, the base model, i.e. bridge matching, significantly underestimates the true covariance σ 2 ⋆ , as shown in Section 6.1. Additionally, the figure illustrates that online finetuning approaches the true solution faster than the original iterative DSBM finetuning. For the latter, we can set how often we alternate between updating the forward and backward networks, and as this frequency increases, the behaviour approaches that of the online finetuning.

Full covariance Gaussian case. Let π 0 = N ( µ 0 , Σ 0 ) , π 1 = N ( µ 1 , Σ 1 ) with Σ i = Id + 1 2 Z i Z ⊤ i for i ∈ { 0 , 1 } and Z 0 , Z 1 independent d × d matrices with unit Gaussian entries. We also set µ 0 = µ 1 = 0 . We consider the Entropic Optimal Transport (EOT) with regularization σ = 0 . 5 and d = 3 , given by

with C ⋆ = 1 2 [Σ 1 / 2 0 D ⋆ Σ -1 / 2 0 -σ 2 Id] , with D ⋆ = (4Σ 1 / 2 0 Σ 1 Σ 1 / 2 0 + σ 4 Id) 1 / 2 . Let normFrob = ∥ A -B ∥ Fro / ∥ A ∥ Fro be the normalized Frobenius distance between matrices A and B . The results are presented in Section 6.1 and confirm those presented in the original manuscript considering a diagonal covariance.

base

finetuned

101

0.2

MSD

0.4

0.6

0.8

0.00 0.14 0.22 0.37 0.61 1.00 1.65 2.72 4.48 7.39

VE

initialsamples

√E=0.0

V=0.25

√=0.75

VE=2.0

VE=5.0

6.2 Image datasets

Similarly to Shi et al. (2023), we apply our method to image translation problems, such as MNIST digits to EMNIST letters (LeCun and Cortes, 2010; Cohen et al., 2017), Wild to Cat domains from the Animal Faces-HQ (AFHQ) dataset (Choi et al., 2020), downsampled to 64 × 64 and 256 × 256 resolutions and CelebA 64 × 64 .

The whole training procedure can be framed as a two-stage process: first, we train a base model on the true data samples, performing bridge matching (Peluchetti, 2021; Albergo and Vanden-Eijnden, 2023; Lipman et al., 2023; Liu et al., 2023a), and then we finetune this model. We compare models that combine different vector field parameterisations (two networks vs. one bidirectional net), finetuning methods (iterative vs. online), and sample generation strategies during the finetuning stage.

Following the established practice (Choi et al., 2020), we evaluate our models using FID (Heusel et al., 2017) for visual quality, and mean squared distance (MSD) or LPIPS (Zhang et al., 2018) for alignment. It is important to note that for image translation tasks at hand, FID scores are not ideal, as FID was designed for natural RGB images, which is not the case for MNIST. It is also not well suited for small sample sizes as it is the case with AFHQ, where the test set in each domain has fewer than 500 examples. Thus quantitative results in Table 1 should be interpreted cautiously, and we recommend a visual inspection of samples to complement these quantitative measures, especially for the AFHQ models. Samples from the models along with the training and evaluation protocols are given in Appendix K.

Compared to the iterative DSBM, our online finetuning α -DSBM reduces the number of tunable hyperparameters, i.e. inner and outer iterations, refresh rate and the size of the cache for storing generated samples. This simplifies implementation and makes the algorithm more practical. The primary remaining hyperparameter, the variance of a Brownian motion ε , requires careful tuning as it influences the trade-off between the visual quality and alignment, as was also observed in Shi et al. (2023). An appropriate ε needs to balance the two: setting ε too low results in poor visual quality, while high values of ε cause poorly aligned and oversmoothed samples. Figure 3 illustrates how FID and MSD metrics vary with ε for the case of MNIST. Additionally, it demonstrates the impact of ε on the generated samples for the AFHQ-64 model.

We run α -DSBM on CelebA with image size 64 × 64 with σ = 2 . 0 . We do not change the training hyper-parameters compared to AFHQ. Visual results are reported in Figure 5 and Figure 6. In Figure 5, we show the influence of σ during the pretraining. The visual quality of the transfer is much lower for σ = 0 than for σ = 2 . 0 . The case σ = 0 corresponds to the first step of Rectified Flow (i.e. Flow Matching). Given the poor quality of the samples, we do not perform finetuning with σ = 0 . In Figure 6, we compare the visual quality and alignment of DSBM and α -DSBM after 4000 training steps, corresponding to two outer DSBM iterations. In this case DSBM is trained with a bidirectional network and both procedures consist of finetuning the pretrained model obtained with σ = 2 . 0 . We note that the alignment is better in the case of α -DSBM.

(a) Cat → Wild

(b) Wild → Cat

7 Discussion

In this paper we have introduced α -Diffusion Schrödinger Bridge Matching ( α -DSBM), a new methodology to solve Entropic Optimal Transport problems. α -DSBM is an improved version of DSBM, which does not require training multiple DDM-type models. We have shown that a non-parametric version of this method recovers the Schrödinger Bridge (SB). In addition, α -DSBM is easier to implement than existing SB methodologies while exhibiting similar performance. We illustrated the efficiency of our algorithm on a variety of unpaired transfer tasks.

While α -DSBM solves one of the most critical limitations of DBSM, namely the alternative optimisation, several issues remain to be addressed in order for the method to scale comparably to generative DDMs. In particular, the method is not sampling-free, as during training it requires sampling from the model from the previous iteration to obtain the training data for the current iteration. While it seems difficult to derive a completely sampling-free method to solve SB problems without resorting to the Minibatch OT approximation, there is still room for improvement.

厦

2

22

四

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. International Conference on Learning Representations .

Ambrosio, L., Gigli, N., and Savaré, G. (2008). Gradient Flows in Metric Spaces and in the Space of Probability Measures . Lectures in Mathematics ETH Zürich. Birkhäuser Verlag, Basel, second edition.

Bauschke, H. H. and Kruk, S. G. (2004). Reflection-projection method for convex feasibility problems with an obtuse cone. Journal of Optimization Theory and Applications , 120(3):503-531.

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

Black, K., Janner, M., Du, Y., Kostrikov, I., and Levine, S. (2023). Training diffusion models with reinforcement learning. arXiv preprint arXiv:2305.13301 .

Brekelmans, R. and Neklyudov, K. (2023). On Schrödinger bridge matching and expectation maximization. In NeurIPS 2023 Workshop Optimal Transport and Machine Learning .

Brock, A., Donahue, J., and Simonyan, K. (2019). Large scale GAN training for high fidelity natural image synthesis. In International Conference on Learning Representations .

Cao, Z., Wu, X., and Deng, L.-J. (2024). Neural shrödinger bridge matching for pansharpening. arXiv preprint arXiv:2404.11416 .

Chen, T., Liu, G.-H., Tao, M., and Theodorou, E. (2023). Deep momentum multi-marginal Schrödinger bridge. Advances in Neural Information Processing Systems .

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. T., and Pavon, M. (2021). Optimal transport in systems and control. Annual Review of Control, Robotics, and Autonomous Systems , 4.

Chen, Y., Goldstein, M., Hua, M., Albergo, M. S., Boffi, N. M., and Vanden-Eijnden, E. (2024). Probabilistic forecasting with stochastic interpolants and Föllmer processes. arXiv preprint arXiv:2403.13724 .

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 .

Chong, M. J. and Forsyth, D. (2020). Effectively unbiased FID and inception score and where to find them. In IEEE/CVF Conference on Computer Vision and Pattern Recognition .

Cohen, G., Afshar, S., Tapson, J., and van Schaik, A. (2017). EMNIST: an extension of MNIST to handwritten letters. arXiv preprint arXiv:1702.05373 .

Courant, R. and Hilbert, D. (2008). Methods of Mathematical Physics: Partial Differential Equations . John Wiley & Sons.

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.

Daras, G., Dagan, Y., Dimakis, A., and Daskalakis, C. (2023a). Consistent diffusion models: Mitigating sampling drift by learning to be consistent. In Advances in Neural Information Processing Systems .

Daras, G., Shah, K., Dagan, Y., Gollakota, A., Dimakis, A., and Klivans, A. (2023b). Ambient diffusion: Learning clean distributions from corrupted data. In Advances in Neural Information Processing Systems .

De Bortoli, V., Hutchinson, M., Wirnsberger, P., and Doucet, A. (2024). Target score matching. arXiv preprint arXiv:2402.08667 .

De Bortoli, V., Liu, G.-H., Chen, T., Theodorou, E. A., and Nie, W. (2023). Augmented bridge matching. arXiv preprint arXiv:2311.06978 .

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 .

De Lange, M., Aljundi, R., Masana, M., Parisot, S., Jia, X., Leonardis, A., Slabaugh, G., and Tuytelaars, T. (2021). A continual learning survey: Defying forgetting in classification tasks. IEEE Transactions on Pattern Analysis and Machine Intelligence , 44(7):3366-3385.

Deng, W., Luo, W., Tan, Y., Biloš, M., Chen, Y., Nevmyvaka, Y., and Chen, R. T. (2024). Variational Schrödinger diffusion models. arXiv preprint arXiv:2405.04795 .

Dhariwal, P. and Nichol, A. Q. (2021). Diffusion models beat GANs on image synthesis. In Advances in Neural Information Processing Systems .

Diefenbacher, S., Liu, G.-H., Mikuni, V., Nachman, B., and Nie, W. (2024). Improving generative model-based unfolding with Schrödinger bridges. Physical Review D , 109(7):076011.

Dupuis, P. and Ellis, R. S. (2011). A Weak Convergence Approach to the Theory of Large Deviations . John Wiley & Sons.

Eyring, L., Klein, D., Uscidda, T., Palla, G., Kilbertus, N., Akata, Z., and Theis, F. (2024). Unbalancedness in neural Monge maps improves unpaired domain translation. In International Conference on Learning Representations .

Fan, Y., Watkins, O., Du, Y., Liu, H., Ryu, M., Boutilier, C., Abbeel, P., Ghavamzadeh, M., Lee, K., and Lee, K. (2024). Reinforcement learning for fine-tuning text-to-image diffusion models. In Advances in Neural Information Processing Systems .

Ge, Z., Liu, S., Li, Z., Yoshie, O., and Sun, J. (2021). Ota: Optimal transport assignment for object detection. In IEEE/CVF Conference on Computer Vision and Pattern Recognition .

Genevay, A., Peyre, G., and Cuturi, M. (2018). Learning generative models with Sinkhorn divergences. In Artificial Intelligence and Statistics .

Gushchin, N., Kholkin, S., Burnaev, E., and Korotin, A. (2024a). Light and optimal schrödinger bridge matching. arXiv preprint arXiv:2402.03207 .

Gushchin, N., Kolesov, A., Korotin, A., Vetrov, D. P., and Burnaev, E. (2024b). Entropic neural optimal transport via diffusion processes. In Advances in Neural Information Processing Systems .

Heusel, M., Ramsauer, H., Unterthiner, T., Nessler, B., and Hochreiter, S. (2017). GANs trained by a two time-scale update rule converge to a local Nash equilibrium. In Advances in Neural Information Processing Systems .

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. In International Conference on Machine Learning .

Hu, E. J., Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., Wang, L., Chen, W., et al. (2021). LoRA: Low-rank adaptation of large language models. In International Conference on Learning Representations .

Hudson, D. A., Zoran, D., Malinowski, M., Lampinen, A. K., Jaegle, A., McClelland, J. L., Matthey, L., Hill, F., and Lerchner, A. (2023). Soda: Bottleneck diffusion models for representation learning. arXiv preprint arXiv:2311.17901 .

Karimi, M. R., Hsieh, Y.-P., and Krause, A. (2024). Sinkhorn flow as mirror flow: A continuous-time framework for generalizing the Sinkhorn algorithm. In International Conference on Artificial Intelligence and Statistics .

Karras, T., Aittala, M., Aila, T., and Laine, S. (2022). Elucidating the design space of diffusion-based generative models. In Advances in Neural Information Processing Systems .

Kim, B., Kwon, G., Kim, K., and Ye, J. C. (2024). Unpaired image-to-image translation via neural schrödinger bridge. In International Conference on Learning Representations .

Kingma, D. P. and Ba, J. (2015). Adam: A method for stochastic optimization. In International Conference on Learning Representations .

Korotin, A., Gushchin, N., and Burnaev, E. (2024). Light Schrödinger bridge. In International Conference on Learning Representations .

LeCun, Y. and Cortes, C. (2010). MNIST handwritten digit database.

Lee, K., Liu, H., Ryu, M., Watkins, O., Du, Y., Boutilier, C., Abbeel, P., Ghavamzadeh, M., and Gu, S. S. (2023). Aligning text-to-image models using human feedback. arXiv preprint arXiv:2302.12192 .

Léonard, C. (2014). 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. Q., 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. (2022). 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). I2sb: image-to-image Schrödinger bridge. In International Conference on Machine Learning .

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 .

Masip, S., Rodriguez, P., Tuytelaars, T., and van de Ven, G. M. (2023). Continual learning of diffusion models with generative distillation. arXiv preprint arXiv:2311.14028 .

Mnih, V., Kavukcuoglu, K., Silver, D., Rusu, A. A., Veness, J., Bellemare, M. G., Graves, A., Riedmiller, M., Fidjeland, A. K., Ostrovski, G., et al. (2015). Human-level control through deep reinforcement learning. Nature , 518(7540):529-533.

Neal, R. M. and Hinton, G. E. (1998). A view of the EM algorithm that justifies incremental, sparse, and other variants. In Learning in Graphical Models , pages 355-368. Springer.

Neklyudov, K., Brekelmans, R., Severo, D., and Makhzani, A. (2023a). Action matching: Learning stochastic dynamics from samples. In International Conference on Machine Learning .

Neklyudov, K., Brekelmans, R., Tong, A., Atanackovic, L., Liu, Q., and Makhzani, A. (2023b). A computational framework for solving Wasserstein Lagrangian flows. arXiv preprint arXiv:2310.10649 .

Nichol, A. Q. and Dhariwal, P. (2021). Improved denoising diffusion probabilistic models. In International Conference on Machine Learning .

Parisi, G. I., Kemker, R., Part, J. L., Kanan, C., and Wermter, S. (2019). Continual lifelong learning with neural networks: A review. Neural networks , 113:54-71.

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. Journal of Machine Learning Research , 24:1-51.

Perez, E., Strub, F., de Vries, H., Dumoulin, V., and Courville, A. C. (2018). Film: Visual reasoning with a general conditioning layer. In AAAI .

Peyré, G., Cuturi, M., et al. (2019). Computational optimal transport: With applications to data science. 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. T. (2023). Multisample flow matching: Straightening flows with minibatch couplings. In International Conference on Learning Representations .

Rafailov, R., Sharma, A., Mitchell, E., Manning, C. D., Ermon, S., and Finn, C. (2024). Direct preference optimization: Your language model is secretly a reward model. In Advances in Neural Information Processing Systems .

Rœlly, S. (2013). Reciprocal processes: a stochastic analysis approach. In Modern Stochastics and Applications , pages 53-67. Springer.

Ronneberger, O., Fischer, P., and Brox, T. (2015). U-net: Convolutional networks for biomedical image segmentation. In Navab, N., Hornegger, J., Wells, W. M., and Frangi, A. F., editors, Medical Image Computing and Computer-Assisted Intervention - MICCAI 2015 , pages 234-241, Cham. Springer International Publishing.

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 ., Campbell, A., and Doucet, A. (2023). Diffusion Schrödinger bridge matching. In Advances in Neural Information Processing Systems .

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

Smith, J. S., Hsu, Y .-C., Zhang, L., Hua, T., Kira, Z., Shen, Y ., and Jin, H. (2023). Continual diffusion: Continual customization of text-to-image diffusion with c-lora. arXiv preprint arXiv:2304.06027 .

Sommerfeld, M., Schrieber, J., Zemel, Y., and Munk, A. (2019). Optimal transport: Fast probabilistic approximation with exact solvers. Journal of Machine Learning Research , 20(105):1-23.

Somnath, V. R., Pariset, M., Hsieh, Y.-P., Martinez, M. R., Krause, A., and Bunne, C. (2023). Aligned diffusion Schrödinger bridges. In Uncertainty in Artificial Intelligence .

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

Song, Y. and Ermon, S. (2020). Improved techniques for training score-based generative models. In Advances in Neural Information Processing Systems .

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 .

Su, X., Song, J., Meng, C., and Ermon, S. (2023). Dual diffusion implicit bridges for image-to-image translation. In International Conference on Learning Representations .

Tamir, E., Trapp, M., and Solin, A. (2023). Transport with support: Data-conditional diffusion bridges. Transactions on Machine Learning Research .

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., Fatras, K., Malkin, N., Huguet, G., Zhang, Y., Rector-Brooks, J., Wolf, G., and Bengio, Y. (2024a). Improving and generalizing flow-based generative models with minibatch optimal transport. Transactions on Machine Learning Research .

Tong, A., Malkin, N., Fatras, K., Atanackovic, L., Zhang, Y., Huguet, G., Wolf, G., and Bengio, Y. (2024b). Simulation-free Schrödinger bridges via score and flow matching. In International Conference on Artificial Intelligence and Statistics .

Vargas, F., Padhy, S., Blessing, D., and Nüsken, N. (2024). Transport meets variational inference: Controlled Monte Carlo diffusions. In International Conference on Learning Representations .

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

Yang, K., Tao, J., Lyu, J., Ge, C., Chen, J., Li, Q., Shen, W., Zhu, X., and Li, X. (2023). Using human feedback to fine-tune diffusion models without any reward model. arXiv preprint arXiv:2311.13231 .

Zaj ˛ ac, M., Deja, K., Kuzina, A., Tomczak, J. M., Trzci´ nski, T., Shkurti, F., and Miło´ s, P. (2023). Exploring continual learning of diffusion models. arXiv preprint arXiv:2303.15342 .

Zhang, R., Isola, P., Efros, A. A., Shechtman, E., and Wang, O. (2018). The unreasonable effectiveness of deep features as a perceptual metric. In Proceedings of the IEEE conference on Computer Vision and Pattern Recognition , pages 586-595.

Zhou, T., Wang, W., Konukoglu, E., and Van Gool, L. (2022). Rethinking semantic segmentation: A prototype view. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition .

A Appendix organisation

The supplementary material is organised as follows. First in Appendix B, we analyze an Euclidean counterpart to the α -IMF sequence identified in Section 3 and the associated flow. In Appendix C, we show that the Markovian projection can be recovered as the parameterisation of the vector field that minimises the accumulation of errors, extending the results of Chen et al. (2024). Theoretical results are gathered in Appendix D. In particular in Appendix D.1 we show that the proposed non-parametric method coincides with the α -IMF and prove the convergence of the α -IMF. In Appendix D.2, we show the connection between the non-parametric and the parametric updates. In Appendix E, we provide more background on DSBM and propose an extension of the DSBM methodology. Consistency losses similar to (Daras et al., 2023b; De Bortoli et al., 2024) are proposed in Appendix F. Model stitching procedures are described in Appendix G. We comment on extended related work in Appendix H. In particular we draw connections with Sinkhorn flows (Karimi et al., 2024), Reinforcement Learning policies, Expectation-Maximisation schemes following (Brekelmans and Neklyudov, 2023) and comment on finetuning of diffusion models. In Appendix I, we investigate the accumulation of bias in a Gaussian setting and compare forward-forward and forward-backward methods. In Appendix J, we derive the preconditioning of loss following the principles of (Karras et al., 2022) in the case of bridge matching. Additional results and experimental details are presented in Appendix K.

B Euclidean flow and iterative procedure

In this section, we study a simplified counterpart of the Schrödinger flow and of DSBM in a Euclidean setting. The goal of this section is to draw some conclusions in the Euclidean case which also remain true empirically when analyzing the Schrödinger Bridge problem.

We consider the set A 1 = { ( x, y ) ∈ R 2 : y ≥ x } and the set A 2 = { ( x, y ) ∈ R 2 : y ≤ 0 } . Loosely speaking, one can identify A 1 with the reciprocal class R ( Q ) and A 2 with the set of Markov path measures proj M . In that case, we have that for any ( x, y ) ∈ R 2 , proj A 1 (( x, y )) = (( x + y ) / 2 , ( x + y ) / 2) if ( x, y ) / ∈ A 1 and otherwise, proj A 1 (( x, y )) = ( x, y ) . In addition, we have that for any ( x, y ) ∈ R 2 , proj A 2 (( x, y )) = ( x, 0) if ( x, y ) / ∈ A 2 and proj A 2 (( x, y )) = ( x, y ) otherwise. We consider the following flow ( x t , y t ) t ≥ 0 given by

Let ( x 0 , y 0 ) / ∈ A 1 and ( x 0 , y 0 ) / ∈ A 2 . Denote T the explosion time of ( x t , y t ) , i.e. for any t ≥ T we have that ( x t , y t ) = ∞ , where R 2 ∪ {∞} is the one-point compactification of R 2 . Finally, denote τ ≤ T such that for any t ∈ [0 , τ ] , ( x t , y t ) ̸∈ A 1 and ( x t , y t ) ̸∈ A 2 . Then, we have

Hence, we have that x t = x 0 exp[ -t/ 2] for any t ∈ [0 , τ ] and y t = x 0 exp[ -t/ 2] + ( x 0 exp[ -t/ 2]) 2 ( y 0 -x 0 ) /x 2 0 . Therefore, we get that τ = T = + ∞ and we have that for any t ≥ 0

Hence, (( x t , y t )) t ≥ 0 converges exponentially fast to (0 , 0) with rate 1 / 2 .

We now investigate the rate of convergence of the alternate projection scheme, i.e. the Euclidean equivalent of DSBM. We define (( x n , y n )) n ∈ N such that for any n ∈ N ,

Hence, we get that x n = x 0 2 -n and therefore (( x n , y n )) n ∈ N converges exponentially fast to (0 , 0) . Note that this procedure corresponds to a discretisation of the flow (( x t , y t )) n ∈ N with stepsize α = 1 .

More generally, we define for any α ∈ (0 , 1] , (( x α n , y α n )) n ∈ N such that for any n ∈ N ,

Hence, we get that x n = x 0 2 -n and therefore (( x n , y n )) n ∈ N converges exponentially fast to (0 , 0) . It can be shown that for any n ∈ N , x α n = x α 0 (1 -α/ 2) n and in addition, we have that

Therefore, we get that

We now analyse the complexity of the different discretisations assuming that the cost of discretising the flow with stepsize α ∈ (0 , 1] is C α . In that case in order to reach the threshold value ε , i.e. | x α n | ≤ ε , we get a total cost C α n = O (log(1 /ε ) C α / log(1 / (1 -α/ 2))) , where we have neglected the terms that do not depend on log(1 /ε ) . Hence, if C α is constant then the choice α = 1 is the best possible one in the range α ∈ (0 , 1] . Otherwise, one has to consider the ratio C α / log(1 / (1 -α/ 2)) , where the lower is the better. The flow procedure and the iterative based one are presented in Figure 1.

Based on this simplified Euclidean experiment, we draw some conclusions which also remain true in our setting, see Appendix K for more experimental details. First, we have that different discretisations of the flow yield different convergence rates. Large stepsizes incur faster convergence. This suggests to choose α = 1 . However, if the cost of choosing α = 1 is too high then one might turn to alternative schemes with α ∈ (0 , 1) assuming that C α < C 1 in that case. To draw a parallel with our setting, in the case of DSBM (case α = 1 ), we need to solve the projection subproblem at each step which incurs a great cost. On the other hand, one step of the online algorithm only requires sampling once from the model and performing one gradient step.

C Minimisation of errors and Markovian projection

For a given non-Markovian (stochastic) interpolant process (see definition below), there exist an infinite number of Markov processes admitting the same marginals (Albergo and Vanden-Eijnden, 2023). In this section, when it is well-defined, we show that the Markovian projection corresponds to the process which minimises an error measure (defined further) in case one has access to the oracle of x t ↦→ E [ X 1 | X t = x t ] .

Stochastic Interpolant. We first start by recalling the framework of Albergo and Vanden-Eijnden (2023). Consider a coupling Π between π 0 and π 1 , one builds a (stochastic) flow between π 0 and π 1 using the following interpolation procedure

where α 1 = β 0 = γ 0 = γ 1 = 0 and α 0 = β 1 = 1 . This defines a non-Markovian process. We denote by π t the induced unconditional distribution of X t . Let us now consider the Markov process ( X ε t ) given by

where ( B t ) t ∈ [0 , 1] is a d -dimensional Brownian motion and ε t is an additional hyperparameter. It can then be shown that ( X ε t ) t ∈ [0 , 1] satisfies that X ε t ∼ π t for all t ∈ [0 , 1] ; see e.g. (Albergo et al., 2023, Theorem 2.8, Corollary 2.10). Hence ( X ε t ) t ∈ [0 , 1] is a (stochastic) flow mapping π 0 onto π 1 . Note that ( X ε t ) t ∈ [0 , 1] in (14) can be rewritten as

In the specific case where α t = 1 -t , β t = t and γ t = σ 0 √ t (1 -t ) then (15) becomes

This corresponds to the marginal distribution of the bridge associated with ( σ 0 B t ) t ∈ [0 , 1] . In this case, (15) becomes √

for ε 2 t = (2 γ t )( ˙ γ t -γ t ˙ α t /α t ) = 2 σ 2 0 . In Proposition C.1, we will show that this choice of ( ε t ) t ∈ [0 , 1] is optimal in some sense.

Consider ( ˆ X ε t ) t ∈ [0 , 1] given by

with ˆ X ε 0 ∼ π 0 and where ˆ Z ( t, x ) is an approximation of E [ Z | X t = x ] . We have the following result.

Proposition C.1 (Optimality and stochastic interpolant): Denote P ε , respectively ˆ P ε the path measures associated with ( X ε t ) t ∈ [0 , 1] and ( ˆ X t ) t ∈ [0 , 1] respectively. Consider ℓ ( ε ) = KL( P ε | ˆ P ε ) . Let ε ⋆ = argmin ε ℓ ( ε ) . Then we have

In particular, if α t = t , β t = 1 -t and γ t = σ 0 √ t (1 -t ) , then ε t = √ 2 σ 0 . The value ( ε ⋆ t ) t ∈ [0 , 1] corresponds to Markovian projection when it is well defined.

Proof. We have that for any

where ∆ t = ∥ ˆ Z ( t, X ε t ) -E [ Z | X t = X ε t ] ∥ 2 and the expectation is w.r.t. P ε . We have that KL( P ε | ˆ P ε ) = 0 if ε = ε ⋆ , which concludes the proof.

Proposition C.1 is related to (Chen et al., 2024, Section 3.4). Therein it is noticed that, in the case of Augmented Bridge matching (De Bortoli et al., 2023), the choice of ε t does not affect the joint distribution of ( X ε 0 , X ε 1 ) . The authors then select ( ε t ) so as to minimise an approximation error. They show that, in that case, they recover the Föllmer process.

We now show that Proposition C.1 can be further strengthened to establish that ε ⋆ is also the optimal value if we interpolate between π s and π 1 , or π 0 and π s , for any s ∈ [0 , 1] and π s the distribution of X s . Consider in this context for any s, t ∈ [0 , 1] with t ≥ s , γ t /γ s ≥ α t ≥ α s the following interpolation model.

where X s ∼ π s , X 1 ∼ π 1 and Z ∼ N (0 , Id) . Assume that α t = 1 -t , β t = t and γ t = σ 0 √ t (1 -t ) for any t ∈ [0 , 1] then (16) corresponds to the Brownian bridge associated with ( σ 0 B t ) t ∈ [0 , 1] with endpoints X s at time s and X 1 at time 1 . We have the following proposition.

Proposition C.2 (Stochastic interpolant with intermediate time points): Define ( X ε t,s ) t ∈ [ s, 1] given by

with γ t,s = √ γ 2 t -α 2 t γ 2 s /α 2 s . Then for any t ∈ [ s, 1] , X ε t,s and X t defined by (16) have the same distribution.

Proof. We let s ∈ [0 , 1] and X 1 , X s ∈ R d . From (16), we have directly that for any t ∈ [ s, 1]

where Z ∼ N (0 , Id) . In addition, rearranging (16), we also have that

Hence, by combining these two expressions, we get that

It follows that ( X t,s ) t ∈ [ s, 1] given by

is such that for any t ∈ [ s, 1] the same distribution as X t defined by (16). Then, we conclude similarly to (Albergo et al., 2023, Theorem 2.8, Corollary 2.10).

We now consider the following approximate version of (17)

Similarly to Proposition C.1 we consider the best choice of ε to minimise the interpolation cost.

Proposition C.3 (Optimality and stochastic interpolant): Let s ∈ [0 , 1] . Denote P ε , respectively ˆ P ε , the path measure associated with ( X ε t,s ) t ∈ [0 , 1] , respectively ( ˆ X t,s ) t ∈ [0 , 1] . Consider ℓ ( ε ) = KL( P ε | ˆ P ε ) . Let ε ⋆ = argmin ε ℓ ( ε ) . Then we have

In particular, ε ⋆ does not depend on s ∈ [0 , 1] and for every s 1 , s 2 ∈ [0 , 1] with s 1 ≤ s 2 , we have that ( X ε t,s 1 ) t ∈ [ s 2 , 1] and ( X ε t,s 2 ) t ∈ [ s 2 , 1] coincide.

Proof. Similarly to Proposition C.1, we get first that for any s, t ∈ [0 , 1] with s ≤ t

Second, we have that for any s, t ∈ [0 , 1] with s ≤ t

Third, we have that

Combining (18), (19) and (20), we can conclude.

D Theoretical results

In this section, we prove the main theoretical results of the paper. In Appendix D.1, we first prove the convergence of the α -IMF sequence, i.e. we prove Theorem 3.1. Second, we show that the non-parametric updates (8) correspond to the α -IMF sequence, i.e. we prove Proposition 3.2. In Appendix D.2, we link the non-parametric updates to the parametric updates.

D.1 Non-parametric sequence and convergence

Let Q ∈ P ( C ) be associated with ( √ ε B t ) t ∈ [0 , 1] , where ( B t ) t ∈ [0 , 1] is a d -dimensional Brownian motion and ε > 0 . In this section, we abuse notation and denote P ( C ) the set of path measures which are not necessarily probability path measures. In particular, we will consider Q ∈ P ( C ) associated with ( √ ε B t ) t ∈ [0 , 1] with Q 0 = Leb . In that case, the Kullback-Leibler divergence is still well-defined and we refer to (Léonard, 2014) for more details. We recall that we have defined ( P n , ˆ P n ) n ∈ N for any n ∈ N and α ∈ (0 , 1] by

In addition, for any n ∈ N , t ∈ [0 , 1) and x ∈ R d we have defined

where

We define ( P v n ) n ∈ N associated with (27), where for any suitable vector field v , P v is associated with

where ( B t ) t ∈ [0 , 1] is a d -dimensional Brownian motion.

In order to rigorously prove Proposition D.1 detailed further, we introduce P 2 ( C ) , such that P ∈ P 2 ( C ) if P ∈ P ( C ) and for

Note that if P ∈ P 2 ( C ) then we have that for any t ∈ [0 , 1]

In addition, we recall that ϕ ∈ L 2 ( µ ) for µ ∈ P ( R d ) if ϕ : R d → R d and

Then for any t ∈ [0 , 1) , we define L t : A 2 → R given for any ( v, P ) ∈ A 2 by (21).

Proposition D.1 (Non-parametric updates are α -IMF): Let α ∈ (0 , 1] , ( P n , ˆ P n ) n ∈ N as in (4) , δ n = α and µ n = (1 -α ) ˆ P n + α proj R ( P n ) . Assume that for any n ∈ N , P v n is well-defined. Then, for any n ∈ N , P v n = P n .

Proof. First, we have that for any t ∈ [0 , 1) , v, P ∈ A 2 and ϕ ∈ L 2 ( P t ) we have

Hence, we have that

Assume that for some n ∈ N we have that for any t ∈ [0 , 1) and x t ∈ R d , we have v k t ( x t ) = ( E ˆ P k [ X 1 | X t = x t ] -x t ) / (1 -t ) . We are going to show that for any t ∈ [0 , 1) and x t ∈ R d , we have v n +1 t ( x t ) = ( E ˆ P n +1 [ X 1 | X t = x t ] -x t ) / (1 -t ) . For any t ∈ [0 , 1) and x t ∈ R d , we denote

Since we have that δ n = α and µ n = (1 -α ) ˆ P n + α proj R ( P n ) , we obtain for any t ∈ [0 , 1] and x t ∈ R d

so that

Finally, we define

Therefore, combining (8) with (23), (24), (22), we get that for any t ∈ [0 , 1) and x t ∈ R d

Hence, we have that for any t ∈ [0 , 1) and x t ∈ R d , v n +1 t ( x t ) = ( E ˆ P n +1 [ X 1 | X t = x t ] -x t ) / (1 -t ) . Since, for any t ∈ [0 , 1) and x t ∈ R d , v 0 t ( x t ) = ( E ˆ P 0 [ X 1 | X t = x t ] -x t / (1 -t ) by definition, we get that for any n ∈ N , t ∈ [0 , 1) and x t ∈ R d , v n t ( x t ) = ( E ˆ P n [ X 1 | X t = x t ] -x t ) / (1 -t ) . Using, Definition 2.2, we get that P v n = proj M ( ˆ P n ) , which concludes the proof.

Before stating our convergence theorem, we show the following result which is a direct consequence of (Léonard, 2014, Theorem 2.12) and (Léonard et al., 2014, Theorem 2.14). We recall that the differential entropy of a probability measure π is given by

if π admits a density with respect to the Lebesgue measure and + ∞ otherwise.

Lemma D.2 (Characterisation of Schrödinger Bridge): Recall that Q is associated with ( √ ε B t ) t ∈ [0 , 1] and assume that Q 0 = Leb . Let π 0 , π 1 ∈ P ( R d ) such that

for i ∈ { 0 , 1 } . Let P ⋆ such that P ⋆ is Markov, P ⋆ ∈ R ( Q ) , P ⋆ 0 = π 0 and P ⋆ 1 = π 1 . Then P ⋆ is the Schrödinger Bridge, i.e. the unique solution to (2) .

Proof. First, we have that Q 0 , 1 is equivalent to Leb ⊗ Leb . Indeed, we have that for any x 0 , x 1 ∈ R d

Similarly, we have that for any t ∈ (0 , 1) and x t ∈ R d , Q 0 , 1 | t ( ·| x t ) is equivalent to Leb ⊗ Leb . Indeed, we have that for any t ∈ (0 , 1) and x 0 , x t , x 1 ∈ R d

Hence, for any t ∈ (0 , 1) and x t ∈ R d , Q 0 , 1 | t ( ·| x t ) is equivalent to Q 0 , 1 . Since P ⋆ is Markov and P ⋆ ∈ R ( Q ) we get that there exist φ ◦ 0 and φ ⋆ 1 which are Lebesgue measurable such that for any ω ∈ C we have that

Second we verify that the conditions (i)-(vii) of (Léonard, 2014, Theorem 2.12) are satisfied. First, Q is Markov and hence (i) is satisfied. Then, (ii) is satisfied since for any t ∈ (0 , 1) and x t ∈ R d ,

Q 0 , 1 | t ( ·| x t ) is equivalent to Q 0 , 1 . We have that Q 0 = Q 1 = Leb and (iii) is satisfied. We have that for any x 0 , x 1 ∈ R d

Hence, (iv) is satisfied and we let A : R d → R + be given for any x ∈ R d by A ( x ) = ∥ x ∥ 2 /ε . In addition, we have that for any x 0 , x 1 ∈ R d

Hence, (v) is satisfied and we let B : R d → R + given for any x ∈ R d by B ( x ) = ∥ x ∥ 2 /ε . By assumption (vi) and (vii) are satisfied. We conclude the proof upon using (Léonard, 2014, Theorem 2.12-(b)) and (25).

We are now ready to state our main convergence result.

Proposition D.3 (Convergence of α -IMF): Let α ∈ (0 , 1] and ( P n , ˆ P n ) n ∈ N defined by (4) . Under mild assumptions, we have that lim n → + ∞ P n = P ⋆ , where P ⋆ is the solution of the Schrödinger Bridge problem (2) .

Proof. Using the convexity of the Kullback-Leibler divergence with respect to its first argument (see e.g. (Dupuis and Ellis, 2011)), the data processing inequality (see e.g. (Ambrosio et al., 2008, Lemma 9.4.5)), the fact that the Schrödinger Bridge is Markov and in the reciprocal class of Q (see e.g. (Léonard, 2014, Theorem 2.12) and (Léonard et al., 2014, Theorem 3.2)), and the Pythagorean theorem for the Markovian projection (Shi et al., 2023, Lemma 6), we have that for any n ∈ N

Therefore, we get that

Hence, it follows that

So we obtain lim n → + ∞ KL( ˆ P n | proj M ( ˆ P n )) = 0 . In addition, using (26) we have that KL( ˆ P n | P ⋆ ) ≤ KL( ˆ P 0 | P ⋆ ) < + ∞ for all n ∈ N . Using (Shi et al., 2023, Lemma 6), we also get that KL(proj M ( ˆ P n ) | P ⋆ ) ≤ KL( ˆ P 0 | P ⋆ ) < + ∞ for any n ∈ N . Hence both the sequences ( ˆ P n ) n ∈ N and ( P n ) n ∈ N = (proj M ( ˆ P n )) n ∈ N are relatively compact in P ( C ) . Let ¯ P ∈ P ( C ) be an adherent point to the sequence ( ˆ P n ) n ∈ N and φ : N → N increasing such that lim n → + ∞ P φ ( n ) = ¯ P . Similarly, let ϕ : N → N increasing such that ( ϕ ( n )) n ∈ N is a subsequence of ( φ ( n )) n ∈ N such that lim n → + ∞ proj M ( ˆ P n ) = ¯ P ′ , with ¯ P ′ and adherent point to the sequence (proj M ( ˆ P n )) n ∈ N . Using the lower semi-continuity of the Kullback-Leibler divergence in both arguments (Dupuis and Ellis, 2011), we get that

Since the set of Markov measures and the set of reciprocal measures w.r.t. Q are both closed, we have that ¯ P is Markov and in the reciprocal class of Q . Since we also have that ¯ P 0 = π 0 and ¯ P 1 = π 1 , we get that ¯ P = P ⋆ using Appendix D.1. Since every adherent point of ( ˆ P n ) n ∈ N is P ⋆ , we have that lim n → + ∞ ˆ P n = P ⋆ . Similarly, using that lim n → + ∞ KL( ˆ P n | proj M ( ˆ P n )) = 0 and again the lower semi-continuity of the Kullback-Leibler divergence in both arguments, we get that every adherent point of (proj M ( ˆ P n )) n ∈ N is P ⋆ . Hence, we have that lim n → + ∞ P n = P ⋆ , which concludes the proof.

D.2 From parametric to non-parametric.

In this section, we show that the parametric updates considered in (9) are a preconditioned version of the non-parametric updates considered in (8). We first recall the non-parametric loss

and the parametric loss

The non-parametric sequence ( v n ) n ∈ N is given by (27), i.e. we have for any n ∈ N , t ∈ [0 , 1] and x ∈ R d

Similarly the sequence of parametric updates is given for any n ∈ N , t ∈ [0 , 1] and x ∈ R d by

We recall that P ¯ θ n is a stop gradient version of P v ¯ θn . We are going to show that on average the parametric algorithm yields a direction of descent for the non-parametric loss. We assume that the set of parameters Θ is an open subset of R p for some p ∈ N . For any t ∈ [0 , 1] and x ∈ R d we assume that θ ↦→ v θ t ( x ) is twice continuously differentiable and denote D θ v θ t ( x ) ∈ R d × p its Jacobian and D 2 θ v θ t ( x ) its Hessian. For any θ ∈ Θ , we denote

We show the following result.

Proposition D.4 (Velocity field parametric update): Assume that there exists C > 0 such that for any θ ∈ Θ and x ∈ R d

where h θ = ∇ θ L( θ, P ¯ θ ) . We have that for any n ∈ N , t ∈ [0 , 1) and x ∈ R d

Proof. First, we have that for any µ ∈ P ( C )

Therefore, using (22), we get that

Let θ ∈ Θ and denote θ ′ = θ -α ∇ θ L( θ, P ¯ θ ) . Using a Taylor expansion, we get that for any θ ∈ Θ , we have that

Since the functional gradient is not applied on the second coordinate, we can drop the stop gradient operator and therefore we have for any θ ∈ Θ

Combining this result with (28), we conclude the proof.

The corresponding update on the velocity field is given for any n ∈ N , t ∈ [0 , 1] and x ∈ R d by

We immediately have the following corollary.

Proposition D.5 (Parametric direction of descent): For any n ∈ N , if

̸

then we have

E Background material on DSBM and extensions

In this section we recall some basics on Markovian and reciprocal projections in Appendix E.1. We explain the link between the concept of iterative refinement and Schrödinger Bridges in Appendix E.2. Then, we briefly present Diffusion Schrödinger Bridge Matching (DSBM) (Shi et al., 2023) in Appendix E.3 and propose some new extensions in Appendix E.4.

E.1 Markov and reciprocal projections in practice

In this section, we recall the definition of the reciprocal and Markov projection. We provide more details on how these different projections can be performed and illustrate them on simple examples.

Markov projection. First, we recall the definition of the Markovian projection.

Definition E.1 (Markov projection): Assume that Q is induced by ( √ ε B t ) t ∈ [0 , 1] for ε > 0 . Then, when it is well-defined, for any P ∈ R ( Q ) , the Markovian projection M = proj M ( P ) ∈ M is the path measure induced by the diffusion

In Figure 7 and Figure 8, we illustrate the effect of the Markovian projection, following the example of (Liu, 2022). We consider two distributions π 0 and π 1 such that

In Figure 7, we display samples from the distributions π 0 and π 1 as well as trajectories from the path measure P = ( π 0 ⊗ π 1 ) Q | 0 , 1 . Practically, this means that we sample X 0 ∼ π 0 and X 1 ∼ π 1 independently and then consider a Brownian bridge between X 0 and X 1 . The SDE associated with the Brownian bridge with scale ε > 0 is given for any t ∈ [0 , 1] by

Note that the measure P = ( π 0 ⊗ π 1 ) Q | 0 , 1 is in the reciprocal class, i.e. P ∈ R ( Q ) .

2

1

0

Stochastic process

Data

-1

-2

2

-1

0

1

2

Next in Figure 8, we display samples from the distributions π 0 and π 1 as well as trajectories from the path measure P ⋆ = proj M ( P ) . Note that in Figure 8, contrary to Figure 7, we observe less crossings between the trajectories. Indeed in the limit case where ε → 0 the Markov measures P ⋆ is an ODE with regular coefficients and therefore admits a unique solution for every starting point in the space so no crossing is possible. In particular, note that most of the trajectories starting from the upper-left Gaussian end at the upper-right Gaussian. Similarly, most of the trajectories starting from the lower-left Gaussian end at the lower-right Gaussian.

2

1

0

Stochastic process

Data

1

-2

-2

1

0

1

2

o

In practice, computing the Markov projection involves finding the optimal drift v ⋆ t . This optimal drift is the minimizer of a regression problem, see (Shi et al., 2023) for more details. Hence, computing the Markovian projection requires training a neural network to define a vector field.

Reciprocal projection. First, we recall the definition the reciprocal projection.

Definition E.2 (Reciprocal projection): P ∈ P ( C ) is in the reciprocal class R ( Q ) of Q if P = P 0 , 1 Q | 0 , 1 . We define the reciprocal projection of P ∈ P ( C ) as P ⋆ = proj R ( Q ) ( P ) = P 0 , 1 Q | 0 , 1 . We will write proj R instead of proj R ( Q ) to simplify notation.

To sample from P ⋆ = proj R ( P ) , we only need to sample ( X 0 , X 1 ) ∼ P 0 , 1 and then to sample from the Brownian bridge conditioned on ( X 0 , X 1 ) . This means that in order to sample X ⋆ t ∼ P ⋆ t , we only need to sample ( X 0 , X 1 ) ∼ P 0 , 1 and then compute

with Z ∼ N (0 , Id) . In particular, sampling from P ⋆ = proj R ( P ) does not require training any neural network. However, in practice, in order to obtain samples ( X 0 , X 1 ) ∼ P , we have that P is associated with an SDE and therefore obtaining ( X 0 , X 1 ) requires unrolling the SDE associated with P . In Algorithm 1, the measure P is associated with an SDE with parametric drift v θ .

In Figure 9, we continue our study of the example of (Liu, 2022) that we used to explain the concept of Markovian projection. We consider the path measure M obtained as the Markov projection of P = ( π 0 ⊗ π 1 ) Q | 0 , 1 . In Figure 9, we display samples from the distributions π 0 and π 1 as well as trajectories from the path measure P ⋆ = M 0 , 1 Q | 0 , 1 . In order to sample from P ⋆ we first sample ( X 0 , X 1 ) ∼ M 0 , 1 . This involves unrolling the SDE associated with M . Once we have access to samples ( X 0 , X 1 ) , we draw trajectories from the Brownian bridge following the SDE (29). We can also sample from any time t without having to unroll the SDE (29) by simply sampling from (30). This is what is done in Algorithm 1.

2

1

0

Stochastic process

Data

-1

-2

-2

-1

0

1

2

E.2 Iterative refinement and Schrödinger Bridge

The Schrödinger Bridge problem (2) can be solved leveraging techniques from diffusion models and bridge matching. De Bortoli et al. (2021); Vargas et al. (2021) consider an alternating projection algorithm, corresponding to a dynamic version of the celebrated Sinkhorn algorithm. Peluchetti (2023); Shi et al. (2023) introduce the Iterative Markovian Fitting procedure which corresponds to perform an alternating projection algorithm on the class of Markov processes and the reciprocal class of the Brownian motion. It can be shown that the solution of this iterative algorithm converges to the Schrödinger Bridge under mild assumptions, see (Peluchetti, 2023; Shi et al., 2023). We highlight that in the case where ε → 0 then DSBM is equivalent to the Rectified Flow algorithm (Liu et al., 2023b). One of the main limitation of those previously introduced procedures which provably converge to the solution of the Schrödinger Bridge problem is that they rely on these expensive iterative solvers and requires to consider two networks, one parameterising the forward process π 0 → π 1 and one parameterising the backward π 1 → π 0 .

E.3 Diffusion Schrödinger Bridge Matching

Diffusion Schrödinger Bridge Matching corresponds to the practical implementation of the Iterative Markovian Fitting procedure proposed in Shi et al. (2023); Peluchetti (2023). The IMF procedure alternates between projecting on the Markov class M and the reciprocal class R Q . In what follows, we denote M n +1 = P 2 n +1 ∈ M and Π n = P 2 n ∈ R ( Q ) . We also recall that Q is a (rescaled) Brownian motion associated with ( σ 0 B t ) t ∈ [0 , 1] and that therefore sampling from Q | 0 , 1 ( ·| x 0 , x 1 ) corresponds to sampling from

We recall that the main computational bottleneck of the DSBM lies in the approximation of the Markovian projection. Indeed, using (Shi et al., 2023, Definition 1, Proposition 2), we have that

M ⋆ = proj M (Π) is associated with the process

We also have using (Shi et al., 2023, Proposition 2) that M ⋆ can be approximated using M θ ⋆ given by

where { v θ : θ ∈ Θ } is a parametric family of functions, usually given by a neural network.

Hence, since we can approximate proj R ( Q ) ( M ) and proj M (Π) we can approximate the IMF procedure. This is the DSBM algorithm introduced in Shi et al. (2023); Peluchetti (2023). We describe the first few iterations. Let Π 0 = Π 0 0 , 1 Q | 0 , 1 where Π 0 0 = π 0 , Π 0 1 = π 1 . Learn M 1 ≈ proj M (Π 0 ) given by (31) with v θ ⋆ given by (32). Next, sample from Π 1 = proj R ( Q ) ( M 1 ) = M 1 0 , 1 Q | 0 , 1 by sampling from M 1 0 , 1 and reconstructing the bridge Q | 0 , 1 . Upon iterating the previous procedure, we obtain a sequence (Π n , M n +1 ) n ∈ N . To mitigate the bias accumulation problem caused by approximating only the forward process, we alternate between a forward Markovian projection and a backward Markovian projection. We give more details on the advantage of using a forward-backward parameterisation instead of a forward-forward in Appendix I. This procedure is valid using (Shi et al., 2023, Proposition 9). The optimal backward process is approximated with

We recall the full DSBM algorithm in Algorithm 2.

Algorithm 2 Diffusion Schrödinger Bridge Matching

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

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

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

4: Learn v ϕ ⋆ using (34) with Π = Π 2 n .

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

6: Let Π 2 n +1 = M 2 n +1 0 , 1 Q | 0 , 1 .

7: Learn v θ ⋆ using (5) with Π = Π 2 n +1 .

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

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

10: end for

11: Output: v θ ⋆ , v ϕ ⋆

E.4 A Reflection-projection extension

First, we consider a reflection-projection method similar to the one investigated in Bauschke and Kruk (2004). We recall that the DSBM algorithm is associated with a sequence ( P n ) n ∈ N such that for any n ∈ N

In a reflection-projection scheme, one of the projection is replaced by a reflection. As noted in Bauschke and Kruk (2004), this can yield faster convergence rates in practice. We consider the sequence ( P n ) n ∈ N such that for any n ∈ N

In what follows, we make the assumption that 2 P n +1 / 2 -P n is a probability measure, even though it is not clear if this path measure is non-negative. However, even by making this strong assumption, we show that we can recover DSBM in Algorithm 4. By considering a relaxation of the reflectionprojection scheme.

First, note that for any n ∈ N , P n | 0 is Markov, see De Bortoli et al. (2023) for instance. Hence, we assume that P n is associated with

Estimating P n +1 . First, we compute v n +1 t, 0 assuming that we can sample from P n and P n +1 / 2 . Since P n +1 is in the reciprocal class, we have that P n +1 is associated with

We refer to De Bortoli et al. (2023) for a proof of this fact. Hence, using (35), we have that

Next, we turn to the estimation of P n +3 / 2 .

Estimating P n +3 / 2 . Next, we assume that for any n ∈ N , P n +1 / 2 is associated with

Using (36), we have that v n +1 t is given by

We note also that using (Shi et al., 2023, Proposition 2) and the fact P n is in the reciprocal class of Q , we also have that

Hence, assuming that we can sample from P n and P n +1 / 2 then we can estimate v n +1 t, 0 and v n +1 t , i.e. sample from P n +1 and P n +3 / 2 . Note that the losses (37) and (38) only differ by the conditioning with respect to the initial condition x 0 and therefore the optimisation can be conducted in parallel. We are now able to propose the following projection-reflection algorithm, see Algorithm 3.

Algorithm 3 Reflection Diffusion Schrödinger Bridge Matching

Note that in Algorithm 3 we only consider the optimisation of a forward process but similarly to Algorithm 2, one can construct a forward backward extension to alleviate some of the bias accumulation problems. Finally, we can interpolate between DSBM and this new reflection algorithm and DSBM by introducing an hyperparameter α ≥ 0 and consider the following extension given in Algorithm 4

Algorithm 4 Reflection Diffusion Schrödinger Bridge Matching

Using different values of α ≥ 0 in Algorithm 4, we recover different existing algorithms. If α = 1 , we recover DSBM Shi et al. (2023). Finally, if α = 1 / 2 , we recover the reflection algorithm Algorithm 3.

F Consistency in Schrödinger Bridge

The idea of training both the forward and the backward jointly was mentioned in (Shi et al., 2023, Section G). However, it was still assumed that, while being trained jointly, the forward and backward vector fields were obtained using an argmin operation, see (Shi et al., 2023, Equation (43), (44)). In addition, in (Shi et al., 2023, Section G) a consistency loss was proposed in order to enforce that the forward and backward processes match, see (Shi et al., 2023, Equation (49)). In this section, we leverage new results from Daras et al. (2023a); De Bortoli et al. (2024) in order to enforce the internal consistency of the model.

First, note that for any ( X t ) t ∈ [0 , 1] associated with P ∈ R ( Q ) we have for any 0 ≤ t 0 ≤ t ≤ t 1 ≤ 1 that where

Let p t be the density of X t with respect to the Lebesgue measure, we have that for any 0 ≤ t 0 ≤ t ≤ t 1 ≤ 1 and x t ∈ R d

Using the change of variable x t 0 → x t 0 + x t and x t 1 → x t 1 + x t we get that for any 0 ≤ t 0 ≤ t ≤ t 1 ≤ 1 and x t ∈ R d

This identity for the score has already been presented in a bridge matching context in (De Bortoli et al., 2024, Section 3.3). Let P ∈ R ( Q ) then we have that proj M ( P ) is such that for any t ∈ [0 , 1] , proj M ( P ) t = P t , see (Shi et al., 2023, Proposition 2). We have that for any t ∈ [0 , 1] , v t ( x ) + v 1 -t ( x ) = σ 2 0 ∇ log p t ( x ) . Combining this result and (39), this suggests considering the following consistency loss

Similarly to (13), we can consider an empirical version of (40).

G Model stitching

In Algorithm 1, the finetuning stage requires a pretrained bridge matching model interpolating between π 0 and π 1 (lines 2-7). However, for large datasets with complex distributions π 0 and π 1 , e.g. ImageNet, training this bridge model from scratch can be computationally expensive. To improve efficiency, we can leverage existing diffusion models targeting π 0 and π 1 . Specifically, we assume access to generative models transferring between N (0 , Id) and π 0 , and between N (0 , Id) and π 1 . In the rest of this section, we show how one can adapt Algorithm 1 to this setting. We then comment on the link between the proposed algorithm and Dual Diffusion Implicit Bridges (Su et al., 2023).

Setting. For simplicity, assume that we have two pretrained diffusion models for π 0 and π 1 . We describe our procedure for π 0 . Consider a forward process of the form X t = X 0 + σ t Z , with Z ∼ N (0 , Id) , where σ t is a hyperparameter. Note that we could have considered an interpolant of the form X t = α t X 0 + σ t Z instead, see Song et al. (2021a) for instance.

We assume that the model X t = X 0 + σ t Z is associated with the forward diffusion model

where we assume that g t ≥ 0 for all t ∈ [0 , 1] . Note that we have that for any t ∈ [0 , 1] , σ 2 t = ∫ t 0 g 2 s d s . In particular, we have that for any s, t ∈ [0 , 1] with s ≤ t

Our goal is to solve the following Entropic Optimal Transport problem

where ε > 0 is some entropic regularisation. We assume that ε > 0 is fixed and assume that there exists t ′ ∈ [0 , 1] such that σ 2 t ′ = ε/ 2 . We now consider a dynamic version of (42) with

where Q is associated with ( X t ) t ∈ [0 ,t ′ ] (41). Note that contrary to the setting presented in the main paper, here we do not consider the integration between time 0 and 1 but between time 0 and t ′ . It can be shown that for any t ∈ [0 , t ′ ] , ( X t ) t ∈ [0 ,t ′ ] associated with Q t | 0 ,t ′ is given by

Solving (43) is equivalent to solving (42). We now propose an algorithm to solve (43). It corresponds to the finetuning stage of Algorithm 1 with a specific initialisation, similar to DSBM-IPF in Shi et al. (2023).

By v ϕ , we denote a DDM model associated with π 1 :

Similarly, v θ denotes a diffusion model associated with π 0 :

In analogy to Equation (11), the two equations above correspond to the forward and backward SDEs.

For a given batch of inputs X 1: B 0 and X 1: B 1 , timesteps t ∼ Unif([0 , t ′ ]) ⊗ B , and interpolations X t and X t , we compute the empirical forward and backward losses as the following modification of Equation (13):

Algorithm 5 corresponds to an online version of DSBM-IPF (Shi et al., 2023) with the initialisation given by two generative models. In Algorithm 5, we finetune the trained vector fields to solve the interpolation task. At inference time, the SDE associated with vector field v θ interpolates between π 1 → π 0 , while the SDE associated with the vector field v ϕ interpolates between π 0 → π 1 .

Algorithm 5 α -Diffusion Schrödinger Bridge Matching for DDM finetuning

1: Input: datasets π 0 and π 1 , number finetuning steps N finetuning , batch size B , DDMparameters ϕ and θ .

2: for n ∈ { 1 , . . . , N finetuning } do

3: Sample ( X 0 , X 1 ) ∼ ( π 0 ⊗ π 1 ) ⊗ B , t ∼ Unif([0 , 1]) , Z 1: B ∼ N (0 , Id) ⊗ B

4: Sample ˆ X t ′ by solving (44) starting from X 0

6: Sample t ∼ Unif([0 , t ′ ]) ⊗ B , Z ∼ N (0 , Id) ⊗ B , and compute X t = Interp t ( X 0 , ˆ X t ′ , Z )

5: Sample ˆ X t ′ by solving (45) starting from X 1

7: Sample t ∼ Unif([0 , t ′ ]) ⊗ B , Z ∼ N (0 , Id) ⊗ B , and compute X t = Interp t ( X 1 , ˆ X t ′ , Z )

8: Update θ with gradient step on ℓ ( θ ; t , X 0 , X t )

9: Update ϕ with gradient step on ℓ ( ϕ ; t , X 1 , X t )

10: end for

11: Output: θ, ϕ parameters of the finetuned models

Our model stitching approach is related to Dual Diffusion Implicit Bridges (DDIB) (Su et al., 2023), which uses pretrained diffusion models, but without further finetuning. As highlighted in Shi et al. (2023), DDIB is inferior to DSBM in terms of quality and alignement of the samples.

H Extended related work

We highlight links between our proposed flow and Sinkhorn flows in Appendix H.1. We draw connection between our practical approach and Reinforcement Learning in Appendix H.2. We discuss how α -IMF is related to (incremental) Expectation-Maximisation in Appendix H.3. Finally, we discuss how our algorithm can be seen as an instance of continual learning in Appendix H.5.

H.1 Links with Sinkhorn flow

In this section, we discuss the links between our approach and the Sinkhorn flow introduced by Karimi et al. (2024). We start by recalling how Sinkhorn flows are defined and then discuss how they are related to our approach.

γ -Sinkhorn and Sinkhorn flows. We first consider the static EOT problem and recall the Sinkhorn procedure, also called Iterative Proportional Fitting. We define a sequence of coupling ( ¯ Π n , Π n ) n ∈ N , i.e. for any n ∈ N , Π n ∈ P ( R d × R d ) . We let Π 0 = Q 0 , 1 and we consider for any n ∈ N ,

In Karimi et al. (2024), the authors generalise (46) by introducing an extra hyperparameter γ ∈ (0 , 1] and defining

Using (Karimi et al., 2024, Lemma 2), we have that for any γ ∈ (0 , 1] , any n ∈ N and any x 0 , x 1 ∈ R d

with f 0 γ = g 0 γ = 0 and for any n ∈ N , γ ∈ (0 , 1] and x 1 ∈ R d

In addition, using (Karimi et al., 2024, Equation (9)) we have that for any n ∈ N , γ ∈ (0 , 1] and x 0 ∈ R d

When letting γ → 0 , (48) and (49) suggest to consider for any s ≥ 0 , x 0 , x 1 ∈ R d

where for any s ≥ 0 , x 1 ∈ R d

Comparison with Schrödinger Bridge flows. In order to compare our approach with the one of Karimi et al. (2024), we start by rewriting the γ -Sinkhorn algorithm defined by (47). To do so, we introduce the projection on the measures with fixed marginal.

Definition H.1 (Projection on marginals): Let Π ∈ P ( R d × R d ) and π 0 ∈ P ( R d ) , we define proj 0 ,π 0 (Π) as follows

Similarly, for any Π ∈ P ( R d × R d ) and π 1 ∈ P ( R d ) , we define proj 1 ,π 1 (Π) as follows

With these definitions, we have that for any n ∈ N , ¯ Π n +1 = proj 0 ,π 0 (proj 1 ,π 1 ( ¯ Π n )) , with ( ¯ Π n ) n ∈ N the original Sinkhorn sequence defined by (46). Similarly, we have that the original Iterative Markovian Fitting (IMF) sequence ( ˆ P n ) n ∈ N as defined in (4) with α = 1 satisfies for any n ∈ N , ˆ P n +1 = proj R ( Q ) (proj M ( P n )) . The analogy between the Sinkhorn iterates and the IMF sequence was already highlighted in Shi et al. (2023); Peluchetti (2023) and further studied in Brekelmans and Neklyudov (2023). We know show that similarly, we can draw an analogy between the sequences defined in (4) with α ∈ (0 , 1) and the sequences obtained in γ -Sinkhorn. To do so, we start by introducing for any Π , ˜ Π ∈ P ( R d × R d )

With this notation, we can now rewrite (47) for any n ∈ N as

Now, we are going to see that (50) is linked with the discretisation of the path measure flow described in (4). Recall that for any suitable v , we define the path measure P v associated with

We define

Similarly, for any µ ∈ P ( C ) , we define

Next, we define the sequence of path measures ( ¯ P n ) n ∈ N such that for any n ∈ N

Now, if we denote ( v n ) n ∈ N the sequence such that for any n ∈ N , ¯ P n = P v n then we have that for any n ∈ N , t ∈ [0 , 1] and x ∈ R d

Recall that ( P n ) n ∈ N given by (4) is associated with ( v n ) n ∈ N such that for any n ∈ N , t ∈ [0 , 1] and x ∈ R d

see Proposition 3.2. Therefore, the only difference between (52) and (51) is that (52) is an explicit update whereas (51) is an implicit update. We summarise the differences between γ -Sinkhorn and the discretisation we introduce in Table 2.

H.2 Links with Reinforcement Learning

In this section, we draw some connection between Algorithm 1 and self-play in Reinforcement Learning. In particular, we introduce a generalisation of Algorithm 1 which uses the concept of replay buffer commonly used in Reinforcement learning, see Mnih et al. (2015) for instance.

We first present a generalisation of Algorithm 1 called Replay Buffer Diffusion Schrödinger Bridge Matching Algorithm 6. We define a buffer B as a collection of samples { ( X k 0 , X k 1 ) } N k =1 , where

N ∈ N is the size of the buffer equipped with two functions Add and Sample . We have that Add : Ω × ⊔ k ∈ N ( R 2 d ) k × ( R 2 d ) N → ( R 2 d ) N , where Ω is a probability space. In practice Add takes a random number (the function can be stochastic), any number of proposed samples as well as the current buffer. As an output Add returns the updated buffer. We also define Sample : Ω × N × ( R 2 d ) N → ⊔ k ∈ N ( R 2 d ) k . This function takes a random number (the function can be stochastic), a natural number k representing the number of samples to return as well as the current buffer. As an output Sample returns a batch of k samples from the buffer.

Algorithm 6 Replay Buffer Diffusion Schrödinger Bridge Matching

1: Input: π 0 , π 1 , ε (entropic regularisation), N pretraining (number of pretraining steps), N finetuning (number of finetuning steps), B (batch size), γ (EMA parameter), θ (initial parameters), B fwd (forward buffer), B bwd (backward buffer)

In Algorithm 6, we allow for more flexibility than the online procedure by leveraging the concept of replay buffer originally introduced in Reinforcement Learning Mnih et al. (2015). The concept of replay buffer has been used previously in Schrödinger Bridge works, with the notion of cache where every n refresh steps a cache is emptied and filled with new samples. If n refresh = 1 , N = B for both B fwd and B fwd we have that for any ω ∈ Ω and ( X 1: B 0 , X 1: B 1 ) ∈ ( R 2 d ) B

This means that the Add simply fills the buffer with the new samples while Sample just return the whole current buffer. In that case we recover Algorithm 1. For more general update rules, the replay buffers B fwd and B bwd allow us to collect previous samples and therefore to keep a memory of the past experiences. In future work, we plan to investigate popular choice in experience replay and their impact on the performance of Algorithm 6.

H.3 Links with Expectation Maximisation

In this section, we make a connection between DSBM and the Expectation Maximisation (EM) algorithm, and show that the discretisation of the Schrödinger Flow proposed in Algorithm 1 corresponds to some incremental version of an idealised algorithm, as discussed in Neal and Hinton (1998). We would like to emphasize that the link between the EM algorithm and Diffusion Schrödinger Bridge based methodologies was already highlighted by Vargas et al. (2024); Brekelmans and Neklyudov (2023). Below, we follow the framework of Brekelmans and Neklyudov (2023) and recall the following definitions.

Definition H.2 (Projections and maximisations): Let A be a subset of P ( C ) . Then, for any P ∈ P ( C ) , when it is well-defined, we define its E -projection on A as P ⋆ = argmin Q ∈ A KL( Q | P ) . Similarly, for any P ∈ P ( C ) , when it is well-defined, we define its M -projection on A as P ⋆ = argmin Q ∈ A KL( P | Q ) .

In Brekelmans and Neklyudov (2023), the authors choose M -projection because this corresponds to the Maximisation step in an EM algorithm while the E -projection corresponds to the expectation step in the EM algorithm. In Brekelmans and Neklyudov (2023), the authors highlight that the Iterative Proportional Fitting procedure is a Expectation-Expectation procedure, i.e. the alternating projections are both E -projections. In contrast, the Iterative Markovian Fitting procedure is a MaximisationMaximisation procedure, i.e. the alternating projections are both M -projections. In particular, we can define the following sequence of path measures ( P n ) n ∈ N , where for any n ∈ N we have

In addition, we have that

since we have that P n = P v n ⋆ . Hence, our online procedure Algorithm 1, which corresponds to the discretisation of the flow of path measures (3) can be rewritten as

Therefore, our proposed algorithm can be seen as an incremental version of the MaximisationMaximisation algorithm associated with DSBM instead of an incremental version of the ExpectationMaximisation algorithm discussed in (Neal and Hinton, 1998).

H.4 Links with finetuning of diffusion models

Algorithm 1 can be seen as a method to finetune bridge matching. Finetuning of diffusion models and flow matching procedures is an active research area. Most of the existing methodologies optimise for an external cost after a pretraining phase. These procedures rely on Reinforcement Learning strategies (Lee et al., 2023; Black et al., 2023; Fan et al., 2024). Recently Direct Preference optimisation (DPO) (Rafailov et al., 2024) has been applied to the finetuning of diffusion models in (Yang et al., 2023; Rafailov et al., 2024). Our approach departs from these works as the objective we minimise is given by the EOT cost. However all of these approaches involve some level of self-play, i.e. are not simulation free.

H.5 Links with continual learning

Continual learning develops techniques to train models when the dataset changes during the training, usually to solve different tasks De Lange et al. (2021); Parisi et al. (2019); Zaj ˛ ac et al. (2023). In the context of diffusion models, continual learning has been investigated in Masip et al. (2023); Zaj ˛ ac et al. (2023); Smith et al. (2023). In (Masip et al., 2023), the authors consider a weighted loss between a diffusion model loss and a distillation loss which ensures some consistency between the model being trained and the previous task model. Similarly to our approach this distillation loss is not simulation-free but, contrary to our loss, the clean samples are not obtained by unrolling the diffusion model but by applying a one-step prediction operator. In (Zaj ˛ ac et al., 2023), consider different replay buffer techniques to train continual diffusion models and observe that experience replay with a small coefficient can bring improvements. Finally, in (Smith et al., 2023), the authors consider the continual training of a text-to-image diffusion model with LoRA (Hu et al., 2021).

I Forward-Forward, Forward-Backward and accumulation of error

In this section, we investigate how error accumulates in the context of DSBM. In practice, we observe similar conclusions in the case of the online version of DSBM. We compare two methods: one which only trains a forward model and one which trains a forward and a backward model.

In what follows, we assume that π 0 = π 1 = N (0 , Id) , we also assume that Q is associated with ( √ 2 B t ) t ∈ [0 , 1] . We recall that for any t ∈ [0 , 1] , we have that

We are going to consider to approximate schemes to implement IMF.

Forward-forward. First, we consider the following sequence of path measures ( P n ) n ∈ N . We set P 0 = ( π 0 ⊗ π 1 ) Q | 0 , 1 . For any n ∈ N , we define P 2 n +2 = P 2 n +1 0 , 1 Q | 0 , 1 , i.e. P 2 n +2 = proj R ( P 2 n +1 ) . In addition, we define P 2 n +1 = proj ε, M ( P 2 n ) such that P 2 n +1 is associated with ( X t ) t ∈ [0 , 1] where for any t ∈ [0 , 1]

with ε ∈ R . Recall that if we define ¯ P 2 n +1 = proj M ( P 2 n ) we have that for any t ∈ [0 , 1] , ¯ P 2 n +1 is associated with ( X t ) t ∈ [0 , 1] where for any t ∈ [0 , 1]

Hence, proj ε, M corresponds to making an error of order x ↦→ εx on the estimated velocity field. Doing so, we now longer have that for any n ∈ N , P n 1 = π 1 . In what follows, we are going to show how the error accumulates for the sequence P n 0 , 1 .

Before stating Proposition I.1, we introduce f : R 4 → R such that for any c 0 , 0 , c 1 , 1 , c 0 , 1 > 0 and t ∈ [0 , 1]

We define F ( c 0 , 0 , c 1 , 1 , c 0 , 1 , ε, t ) = 2 ∫ t 0 f ( c 0 , 0 , c 1 , 1 , c 0 , 1 , s )d s +2 εt . Finally, we define

as well as

Proposition I.1 (Forward-Forward updates): For any n ∈ N , we have that P 2 n +1 0 , 1 = N (0 , Σ n +1 Id) where

Proof. Let P = ( P 0 , 1 ) Q | 0 , 1 where P 0 , 1 is a Gaussian random variable with zero mean and covariance matrix Σ ∈ R 2 d × 2 d such that

where Id is the d -dimensional identity matrix and we assume that c 0 , 1 , c 1 , 1 > 0 . We denote P ⋆ = proj ε, M ( P ) . We have that P 1 | t is a Gaussian random variable with zero mean. We now compute its covariance matrix. First, we have that

and for any n ∈ N

We also have that

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

Hence, we have that for any t ∈ [0 , 1] and x t ∈ R d

So it follows that

Note that if we set c 0 , 1 = c 2 and c 1 , 1 = 1 , we recover (Shi et al., 2023, Lemma 13) with σ = 2 . Denote P ⋆ = proj ε, M ( P ) . Combining (54) and (53) we get that P ⋆ is associated with ( X t ) t ∈ [0 , 1] such that for any t ∈ [0 , 1] we have

Hence, we get that

where Z ∼ N (0 , Id) is independent from X 0 and for any t ∈ [0 , 1] , c 0 , 1 , c 1 , 1 , ε > 0 we have

In addition, we define

Hence, we have that

Therefore, since for any n ∈ N , we have that P 2 n +1 = proj ε, M ( P 2 n ) and P 2 n +2 = P 2 n +1 0 , 1 Q | 0 , 1 , we define ( c n 0 , 1 , c n 1 , 1 ) n ∈ N such that for any n ∈ N

Note that for any n ∈ N , we have that

Since P 0 = ( π 0 ⊗ π 1 ) Q | 0 , 1 we get that c 0 0 , 1 = 0 and c 1 , 1 = 1 . We have that for any n ∈ N

which concludes the proof.

Forward-backward. Next, we consider the following sequences of path measures ( P n, ) n ∈ N and ( P n, ) n ∈ N . We set P 0 , = P 0 , = ( π 0 ⊗ π 1 ) Q | 0 , 1 . For any n ∈ N , we define P 2 n +2 , = P 2 n +1 , 0 , 1 Q | 0 , 1 and P 2 n +2 , = P 2 n +1 , 0 , 1 Q | 0 , 1 , i.e. P 2 n +2 , = proj R ( P 2 n +1 , ) and P 2 n +2 , = proj R ( P 2 n +1 , ) . In addition, we define P 2 n +1 , = proj ε, M ( P 2 n, ) such that for any t ∈ [0 , 1] , P 2 n +1 , is associated with ( X t ) t ∈ [0 , 1] where

with ε ∈ R . Similarly, we define P 2 n +1 , = proj ε, M ( P 2 n, ) such that for any t ∈ [0 , 1] , P 2 n +1 , is associated with ( Y 1 -t ) t ∈ [0 , 1] where

Proposition I.2 (Forward-Backward updates): For any n ∈ N , we have that P 2 n +1 , 0 , 1 = N (0 , Σ n +1 , Id) and P 2 n +1 , 0 , 1 = N (0 , Σ n +1 , Id) where

and for any n ∈ N

The proof is similar to the one of Proposition I.1.

Proof. Let P = ( P 0 , 1 ) Q | 0 , 1 where P 0 , 1 is a Gaussian random variable with zero mean and covariance matrix Σ ∈ R 2 d × 2 d such that

where Id is the d -dimensional identity matrix and c 0 , 1 , c 0 , 0 > 0 . We denote P ⋆ = proj ε, M ( P ) . We have that P 1 | t is a Gaussian random variable with zero mean. We now compute its covariance matrix. First, we have that

We also have that

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

Hence, we have that for any t ∈ [0 , 1] and x t ∈ R d

Finally, we have that for any t ∈ [0 , 1] and x t ∈ R d

Note that if we set c 0 , 1 = c 2 and c 0 , 0 = 1 , we recover (Shi et al., 2023, Lemma 13) with σ = 2 . Denote P ⋆ = proj ε, M ( P ) . Combining (55) and (53) we get that P ⋆ is associated with ( X t ) t ∈ [0 , 1] such that for any t ∈ [0 , 1] we have

Hence, we get that

where Z ∼ N (0 , Id) is independent from X 0 and for any t ∈ [0 , 1] , c 0 , 1 , c 1 , 1 , ε > 0 we have

In addition, we define

Hence, we have that

Remember that P 0 , = P 0 , = ( π 0 ⊗ π 1 ) Q | 0 , 1 . In addition, for any n ∈ N , we have P 2 n +2 , = P 2 n +1 , 0 , 1 Q | 0 , 1 and P 2 n +2 , = P 2 n +1 , 0 , 1 Q | 0 , 1 , i.e. P 2 n +2 , = proj R ( P 2 n +1 , ) and P 2 n +2 , = proj R ( P 2 n +1 , ) . In addition, we also have P 2 n +1 , = proj ε, M ( P 2 n, ) and P 2 n +1 , = proj ε, M ( P 2 n, ) . We also define ( c n, 0 , 1 , c n, 1 , 1 ) n ∈ N such that for any n ∈ N

Finally, we define ( c n, 0 , 1 , c n, 1 , 1 ) n ∈ N such that for any n ∈ N

Using this definition and (56) we get that for any n ∈ N

In addition, we have that c n, 0 , 1 = c n, 0 , 1 = 0 and c 0 , 1 , 1 = c 0 , 0 , 0 = 1 . This concludes the proof.

Error accumulation. In Proposition I.1 and Proposition I.2, we derive the sequences corresponding to the evolution of the variance and the covariance throughout the DSBM iterations in forwardforward mode or forward-backward mode. In what follows, we showcase the behavior of these sequences for different values of ε > 0 . We recall that ε corresponds to the error made in the Markov projection, i.e. proj M is replaced by proj ε, M in the forward-forward mode and proj M is replaced by proj ε, M and proj ε, M in the forward-backward mode. First, if we consider the perfect scenario, i.e. ε = 0 , then we observe that both the forward-forward mode and the forward-backward mode satisfy that E P 2 n [ X 1 X ⊤ 1 ] = Id , see Figure 10 and Figure 11. Additionally, we can show that in the perfect scenario, i.e. ε = 0 , then both the forward-forward mode and the forward-backward mode satisfy that lim n → + ∞ E P 2 n [ X 1 X ⊤ 0 ] = ( √ 2 -1)Id , see Figure 10 and Figure 11. However, as ε increases the behavior between the forward-forward sequence and the forward-backward sequence significantly differs. More precisely, the error explodes as ε increases along the DSBM iteration for the forward-forward mode. On the contrary, in the forward-backward mode, the error remains bounded along the DSBM iterations, see Figure 10 and Figure 11.

1014

0.00e+00

1.00e-09

1.00e-08

1011

1.00e-07

1.00e-06

1.00e-05

10#

1.00e-04

1.00e-03

1.00e-02

105

1.00e+00

1.00e-01

10²

101

104

10

20

OE

40

50

09

70

08

90

0

0.00e+00

1.00e-09

1.00e-08

1.00e-07

1.00e-06

1.00e-05

1.00e-04

1.00e-03

1.00e-02

1.00e-01

1.00e+00

10

20

30

40

50

60

70

80

90

1039

0.00e+00

1.00e-09

1033

1.00e-08

1.00e-07

1.00e-06

1027

1.00e-05

1.00e-04

1021

1.00e-03

1.00e-02

1.00e-01

105

1.00e+00

109

10²

103

0

10

20

30

40

50

60

70

80

90

0

10

20

30

40

50

60

70

80

90

J Preconditioning of the loss function

In this section, we provide details on the scaling of the loss function we implement when training our online version of DSBM. We adapt the method of (Karras et al., 2022, Appendix B.2) to the case of bridge matching. We only present our derivations in the case of the forward training of the online version of DSBM, i.e. (9). The preconditioning of the loss described in this setting can be readily extended to the forward-backward loss we consider in practice, i.e. the parametric version of (10).

We consider the following objective function for any t ∈ [0 , 1]

We also define for any t ∈ [0 , 1] and x t ∈ R d , v θ t ( x t ) = c o t nn θ t ( c i t x t ) + c s t x t . Hence, c i t is an input scaling function, c o t is an output scaling function and c s t is a skip-connection function. During the training of the online version of DSBM, P will be given by P n , where P n = P v θn , where the sequence ( θ n ) n ∈ N is given by (9). Here, we apply the principles of Karras et al. (2022) to the case where P = ( π 0 ⊗ π 1 ) Q | 0 , 1 , i.e. at initialisation of the sequence. In what follows, we assume that E π 0 [ ∥ X 0 ∥ 2 ] = E π 1 [ ∥ X 1 ∥ 2 ] = d . Note that our considerations can be generalised to E π 0 [ ∥ X 0 ∥ 2 ] = σ 2 0 d and E π 1 [ ∥ X 1 ∥ 2 ] = σ 2 1 d . We also have that

Using (58), we have that for any t ∈ [0 , 1]

We set c i t so that E [ ∥ c i t X t ∥ 2 ] = d for every t ∈ [0 , 1] . Hence, we have that for any t ∈ [0 , 1]

0.00e+00

1.00e-09

1.00e-08

1.00e-07

1.00e-06

1.00e-05

1.00e-04

1.00e-03

1.00e-02

1.00e-01

1.00e+00

100

101

10-2

103

104

105

101

102

103

104

105

106

2.00

1.75

1.50

1.25

4.0

5.0

6.0

7.0

8.0

9.0

0.0

1.0

2.0

3.0

10.0

6

2

0.0

2.0

3.0

4.0

5.0

6.0

7.0

8.0

9.0

10.0

1.0

20

25

15

0.0

1.0

2.0

3.0

4.0

5.0

6.0

7.0

8.0

9.0

10.0

1.00

D

10

0.75

0.50

5

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

0.9

Next, we rewrite (57). For any t ∈ [0 , 1] we have that

Hence, we get that for any t ∈ [0 , 1] , T t = [ X 1 -(1+ c s t (1 -t )) X t 1 -t ] /c o t is the target of the network in the regression loss. We are going to fix c o t and c s t such that i) E [ ∥ T t ∥ 2 ] = d , ii) c o t is as small as possible in order not to minimise the error propagation made by the neural network. Using (58), we have that for any t ∈ [0 , 1]

Hence, we get that for any t ∈ [0 , 1]

We now minimise ( c o t ) 2 with respect to (1 + c s t (1 -t )) . We get that

Hence, we get that c s t = t/ [(1 -t )((1 -t ) 2 + t 2 + εt (1 -t ))] -1 / (1 -t ) . With that choice, we get that for any t ∈ [0 , 1]

In Karras et al. (2022), the weighting function λ t is set so that the weight in front of the regression loss is equal to one for all times t ∈ [0 , 1] . Hence, Karras et al. (2022) suggests to set λ t = 1 / ( c o t ) 2 . However, in practice we observe better results by letting λ t = 1 . This means that the effective weight is given by 1 / ( c o t ) 2 . Therefore, for any t ∈ [0 , 1] we have

K Experimental details

In this section, we delve deeper into the specifics of each experiment, implementation details, and share additional results.

We consider two ways of parameterising the vector fields: as in DSBM, we can use two separate neural networks to approximate the forward and backward vector fields, or we can use a single neural network that is conditioned on the direction. In the latter case, we do the conditioning in a similar fashion to how DDM's neural networks, U-Nets or MLPs, are conditioned on time embeddings. After all, if we work with continuous time variables t ∈ [0 , 1] , then the direction signal s ∈ { 0 , 1 }

can be thought of as a target time. Thus, we perform the same initial transformations on t and s , i.e. computing sinusoidal embeddings followed by a 2-layer MLP, and use the concatenated outputs in adaptive group normalisation layers (Dhariwal and Nichol, 2021; Hudson et al., 2023; Perez et al., 2018).

To optimise our networks, we use Adam (Kingma and Ba, 2015) with β = (0 . 9 , 0 . 999) , and we modify the gradients to keep their global norm below 1.0. We re-initialise the optimiser's state when the finetuning phase starts.

All image samples in the paper are generated using EMA parameters as it has been known to increase the visual quality of resulting images (Song and Ermon, 2020). Sampling is also the integral part of DSBM's finetuning stage, both iterative and online. Here, we have two options: sample with EMA or non-EMA parameters. The non-EMA sampling might be easier to implement, while EMA sampling results in a more stable training and slightly better quality, e.g. see AFHQ samples in Figure 23 and Figure 24 for comparison.

For every model used in the paper, we provide hyperparameters in Table 3.

K.1 2D Experiments

In addition to the experiments presented in the main text, we test our models in the simplest 2D data settings used in Tong et al. (2024a) and Shi et al. (2023). Note, that low-dimensional datasets might not be the ideal showcase for α -DSBM given that one can successfully employ less computationally demanding techniques based on minibatch-OT methods (Tong et al., 2024b).

The results of our bidirectional model finetuned with online updates are given in Table 4. During finetuning, we generate samples using 100 Euler-Maruyama steps to solve the forward and backward SDEs. At test time, we solve the forward probability flow ODE (PF-ODE) given by:

To evaluate model fit, we compute 2-Wasserstein distance between the true and generated samples (generated with 20 Euler steps). Additionally, we estimate path energy as a measure of trajectory simplicity: E X 0 ∼ π 0 [ ∫ 1 0 ∥ v θ ( t, X t ) ∥ 2 d t ] where v θ ( t, X t ) is the drift of PF-ODE in (59), and the integral is approximated using 100 steps. We have made a deliberate effort to closely replicate the experimental setup of Shi et al. (2023) to ensure the comparability of our results. However, as

illustrated in Figure 13, 2-Wasserstein distance can be very noisy even with 10K samples in the test set. To mitigate this variance, we averaged the 2-Wasserstein distance across five random sets of 10K samples per run, and then averaged these results across multiple runs. Despite these measures, we recommend a future redesign of these 2D experiments to facilitate more robust comparisons between methods.

K.2 Gaussian data

To parameterise the forward and backward drifts, we use a 2-layer MLP network with 256 hidden units. To process time variables, we compute sinusoidal time embeddings, followed by a 2-layer MLP with 256 hidden units and 50 output units. The resulting time embeddings are then concatenated with X t , so the drift networks receive 100-dimensional input vectors.

For iterative DSBM finetuning, we perform 40K steps with varying number of outer iterations, i.e. when we switch between training the forward and the backward networks. Alternating every 5K steps, corresponds to 8 outer DSBM iteration. Similarly, changing the direction every 1K steps, leads to 40 outer iterations.

We do not have a cache dataloader like in the original DSBM implementation 2 , thus we generate training samples on the fly by sampling either from the forward or the backward network. For this simple task, we also do not use EMA.

During training and evaluation, we use Euler-Maruyama method with 100 equidistant time steps between 0 and 1. The covariance is evaluated using 10K samples.

Additional comparison α -DSBM vs OT bridge-matching. We consider the scalar Gaussian setting. We highlight the dependence of OT bridge matching (256/16/8) on the batch size, as the mini-batch OT coupling can be far from the true OT coupling as the batch size decreases. All experiments are run with similar compute and architecture/training hyperparameters, see Table 5.

Ablation of the hyperparameter α . Instead of letting α be determined by Adam and adaptive for α -DSBM, we explicitly set it by using Stochastic Gradient Descent (SGD) with learning parameter α ; see Figure 14. We also ran online α -DSBM with α = 10 -1 but the training diverges in this case.

distance

1.5

Online@Adam

Online @1e-3

Online @1e-2

Base

1.0

Normalized

0.5

10k20k30k40k50k60k70k80k90k100k

Trainingiterations

K.3 MNIST ↔ EMNIST transfer

We closely follow the setup of Shi et al. (2023) and De Bortoli et al. (2021), and train the models to transfer between 10 EMNIST letters, A-E and a-e, and 10 MNIST digits (CC BY-ND 4.0 license). We use the same U-Net architecture with hyperparameters given in Table 3.

For DSBM finetuning, we perform 30 outer iterations, i.e. alternating between training the forward and the backward networks, while at each outer iteration a network is trained for 5000 steps. We do not have a cache dataloader and generate training samples on the fly by sampling either from the forward or the backward network with EMA parameters.

During training and evaluation, we use Euler-Maruyama method with 30 equidistant time steps between 0 and 1. For evaluation, we compute FID based on the whole MNIST training set of 60000 examples and a set of 4000 samples that were initialised from each test image in the EMNIST dataset. MSD is computed between 4000 initial EMNIST test examples and their corresponding MNIST samples.

In Figures 15-18, we provide forward and backward samples, i.e. EMNIST → MNIST and MNIST → EMNIST, from models that differ in parameterisation, finetuning methods, and sampling strategy. For all the models above, we used ε = 1 . Figure 19 illustrated the behaviour of the samples when we sweep over the ε hyperparameter.

Pretraining a bidirectional model on 4 v3 TPUs takes 1 hour, while the online finetuning stage requires 4 hours on 16 v3 TPUs. The number of pretraining and finetuning steps is chosen to match the experimental setup of Shi et al. (2023).

K.4 AFHQ: Cat ↔ Wild

We consider the problem of image translation between Cat and Wild domains of AFHQ (Choi et al. (2020); CC BY-NC 4.0 DEED licence) as introduced by Shi et al. (2023). Each domain has approximately 5000 samples in the training set, and around 500 samples in the test set. We resize the original 512 × 512 images to 64 × 64 or 256 × 256 resolutions.

aaBEeeebBA

7

97B1

b

7

2

93583213

3

2

3585215

5

G

baEEdEdac

B

6

一

8

6

8

7

万

α

5

2

LlLO0h

a

A

2

8

9

b

8

3

3

3

5

S

S

3

5

2

S

6

εaBe

c

eCeAbc.

Addceec

d

b

e

C

a

8

D

9

4

h

D

3

7

S

a

2

3

9

6

S

S

O

3

9

S

0

b

5

S

3

5

D

eeCBE

a

e

9

h

2

h

3

3

6393

9

3

?

7

3

2

6

C

<sece

D

9

8

4

D

2

S

7

8

3

3

2

3

8

6

8

D

derDeBd

a

9

5

7

2

8

D

5

5

7

4

7

D

7

5

59

9

d

DbERdCace

832

3

4

9

2

8893

2

7

b

8843

7

e

bdeBDdB<

g

7

b

55

h

S

5775

9

1

(a) Initial EMNIST letters (b) Bridge matching: FID=6.02, MSD=0.564

(c) Bridge matching + DSBM finetuning: FID=5.25, MSD=0.345

(d) Bridge matching + online finetuning: FID=4.28, MSD=0.368

aaBEeee6BA507976291

7

2

93589015

2

9358961

中

h

baEEdεdac

B

2

984b9

2

8

o

7

6842

2

6

6

６4

Ceabebeca

A

2

b

9

N

S

5

3

2

1

36

S

6

4

2

Q

529

6

6

G

LAddceec

b

C

8

S

S

6

3

4055

3

3

O

4

3

G

7

3

EaBe

eCeAbc de

P

大

4

h

S

q

3

b

了

7

23

8

5

4

3

b

S

7

3

8

S

e

9

y

h

3

3

7

S

6

3

9

D

3

9

9

C

<ieceDB

7

6

4

4

7

3

98

6

2

3

３9

y

8

6

2

8

DdecdeBd

a

67

8

9

D

4

55

9

8

9

5

98

7

dDbEBdcac

3734

9

3

Zh8890

一

3

3

6

0883321

了

e

bdeBndB

967

a

64S57y591

（

S5751

(a) Initial EMNIST letters (b) Bridge matching: FID=6.33, MSD=0.572

(c) Bridge matching + online non-EMA finetuning: FID=4.57, MSD=0.369

(d) Bridge matching + online finetuning: FID=4.39, MSD=0.387

41078127

6

P

b

ApdlcDa6ldsAaEced6B

ddcaccaabA

6

4

1773379

DBecAbaC2bbcdaBeaa

AncdCeCaeA

6

69948

9

EcAcccaec

ab6AarddAc

ab6baAddbc

3

3

Q

G

9

6

eEeOACh

BecaAACEbA

B

BecaAACAbA

7

2

6

6

acAbeda

6

P

acccodcaAr

B

Ccccdbaag

6

5

9

ecae6decA

Aadbbcaea

P

S

b

C

6bcaBb

P

7

5

8

1

aacecdBbe

A

Deeaddde

E

b

S

8411

cceBdca

C

bopqt9pp

EedbcAcda

1子

6

tgaeece

4

ocd

eAeαedbcc

a

Z

4497.

E

beBaaDbA

6

ICdeacBecd

CdAaCBeCd

a

(a) Initial MNIST digits (b) Bridge matching: FID=7.50, MSD=0.553

(c) Bridge matching + DSBM finetuning: FID=3.56, MSD=0.330

(d) Bridge matching + online finetuning: FID=3.67, MSD=0.357

4107812/

5

dApdEeDbcldbcaecaab6ldbcaeeadbc

6

1773379

eBecebBce

6A

cceBeae

A

6

JteBeacA

6

569948

9

EeAcceaec

a

b

6qede

te

d

C

baeteeA

C

3

3

Q

9

6

eEeOAeh

E

a

B

E

eadcc

EE

A

ecaAcceE

AdA

7

2

6

6

acAbeda

51

R

P

EcCCA&aC

cccAsaa

y

6

5

9

3P90

CA

d

b

bbaAda

aeapear

deBabBb=b

7

5

8

1

e2eebCb

de

8

deBabdbe

S

8

C

ceedcac

D

Cdb/ddca

CddrAddaa

71子

6

ElaeacacdlbcBaebqcd

eABaebacdc

Z

4497

E

brgaabcabddsacBbcd

(a) Initial MNIST digits (b) Bridge matching: FID=7.97, MSD=0.572

(c) Bridge matching + online non-EMA finetuning: FID=4.16, MSD=0.370

(d) Bridge matching + online finetuning: FID=3.97, MSD=0.392

Our U-Net (Ronneberger et al., 2015) implementation is based on Ho et al. (2020) with a few improvements suggested in Dhariwal and Nichol (2021); Song et al. (2021b) such as rescaling of skip connections by 1 / √ 2 , using residual blocks from BigGAN (Brock et al., 2019), and convolution-based up- and downsampling. Hyperparameters are given in Table 3. Compared to the straightforward parameterisation of the vector fields, we obtained slightly better results using EDM preconditioning Karras et al. (2022), which we derive in Appendix J for the case of bridge matching. During training, we use horizontal flips as a way to augment the data.

During training and evaluation, we use Euler-Maruyama method with 100 equidistant time steps between 0 and 1. When evaluating the quality of Cat → Wild transfer, we compute FID based on the whole training set of 4576 examples in the Wild domain and a set of 480 samples that were initialised from test images in the Cat domain. LPIPS and MSD are computed between 480 initial Cat images and Wild samples from the model. The same procedure is followed when evaluating in the reverse direction from Wild to Cat. Given that train, and especially the test sets are small, the quantitative results for AFHQ are likely unreliable (Chong and Forsyth, 2020). In Figure 22 we provide samples from the models finetuned either with an iterative or an online method. While their FID scores are different, the samples look similar between the two models.

As we discussed in the main text, hyperparameter ε trades off the visual quality and alignment of the samples in the resulting transfer models. In Figure 20, we provide AFHQ 64 × 64 samples for pretrained and finetuned models with different values of ε . In addition to its relation to EOT, from a DDM perspective, ε can be seen as the controlling factor of the noise schedule. As observed by Hoogeboom et al. (2023), noise schedules should be adjusted for different image sizes by shifting the noise schedule of some reference resolution where it is proven to be successful. In our case, if we find a good value of ε for 64 × 64 images, then a shifted ε for the 256 × 256 resolution can be computed as ε 256 = ε 64 ( 256 64 ) 2 . Thus, if we choose √ ε = 0 . 75 for AFHQ-64, then for AFHQ-256, we can expect √ ε = 3 . 0 to also work well. Samples from an AFHQ-256 model trained with √ ε = 3 . 0 are given in Figure 27.

On 16 v3 TPUs, the bidirectional base and finetuned AFHQ-64 models take 4 and 14 hours to train, respectively. For AFHQ-256, the base model trains for 15 hours, and finetuning takes an additional 37 hours. While we did not experiment with varying pretraining and fine-tuning iterations, these training times suggest that a longer pretraining stage followed by fewer fine-tuning steps may be desirable.

VE=0.00initial samples

6BAba

6BAba

b

BAb

V∈=0.14

3

6

333

60

VE=0.22

S

4

ZZ

VE=0.37

6

VE=0.61

5

VE=1.00

4

VE=1.65

4

7

VE=2.72

5

5

h

2

VE=7.39

5

6

5

(a) Base model

(b) Finetuned model

initial samples

VE=0.0

VE=0.25

VE=0.5

VE=0.75

VE=1.0

VE=2.0

VE=3.0

VE=5.0

(a) Base model

(b) Finetuned model

initial samples

initial samples

NFE=1

NFE=2

NFE=2

NFE=10

NFE=50

NFE=100

NFE=100

NFE=150

NFE=150

NFE=200

NFE=200

(a) Base model

(b) Finetuned model

(a) Forward: Cat → Wild

(b) Backward: Wild → Cat

(a) Forward: Cat → Wild

(b) Backward: Wild → Cat

(a) Forward: Cat → Wild with inputs from Wild.

(b) Backward: Wild

→

Cat with inputs from Cat.

initial samples

VE=0.0

VE=0.25

VE=0.5

VE=0.75

VE=1.0

VE=2.0

VE=3.0

VE=5.0

(a) Forward: Cat

→

Wild

(b) Backward: Wild → Cat

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: our main theoretical and experimental contributions are claimed in the abstract and demonstrated in the paper. We summarize our main contributions hereafter. Theoretically, we identify a new family of sequence of path measures related to the IMF algorithm, called α -IMF. We show that these sequences correspond to non-parametric updates. We then introduce a parametric update that corresponds to an online version of the DSBM algorithm. We show that our procedure retains the favorable properties of DSBM while not requiring the expensive repeated inner minimisation procedure of DSBM.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: The limitations are addressed in the discussion section. The main limitation of our algorithm is that it is not a sampling free methodology. In future work, we would like to see how to mitigate the fact that our algorithm depends on some self-play.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: All theoretical results are proven in the supplementary material, see Appendix D.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: Full experimental details are provided in Appendix K.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [No]

Justification: Due to IP restrictions, we cannot share the codebase used for this paper. However, we plan to release some notebooks in order to reproduce experiments in a small scale setting.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimiser, etc.) necessary to understand the results?

Answer: [Yes]

Justification: Full experimental details are provided in Appendix K.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: All metrics are computed using multiple random seeds and error bars are provided.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: Full details on the compute requirements are given in Appendix K.

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: After careful review of the NeurIPS Code of Ethics, we can ensure that the research presented in this paper conforms with the Code of Ethics in every respect. Indeed, we see no immediate safety, security, discrimination, surveillance, deception, harassment, environment, human rights or bias and fairness concerns to our work. In addition, we release details and documentation regarding the datasets and models used. We disclose essential details for reproducibility and have ensured that our work is legally compliant.

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: This paper addresses the problem of unpaired dataset translation and proposes an improvement to the DSBM methodology. As the current paper is mostly theoretical and methodological we do not see immediate societal impact of this work and therefore do not discuss these issues. However, we acknowledge that large scale implementation of our algorithm might suffer from the same societal biases as generative models. We hope to address the limitations of such models when turning to more experimental work.

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: The paper poses no such risks.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We have referenced the license of the datasets we use and cite the original papers that produced the code packages and datasets that we use in that paper.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: The paper does not release new assets.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes $\\alpha$-IMF, which is an incremental version of IMF method in DSBM paper. Specifically, this work demonstrate the convergence properties of $\\alpha$-IMF and implement it through an online learning method. Here, $\\alpha$ is implicitly reflected. Moreover, the functional flow interpretation provides a basis for online learning, allowing for online learning as opposed to the iterative approach of DSBM. From an engineering perspective, the authors suggest parameterizing forward and backward control simultaneously by adding a forward/backward indicator (0 or 1) as an input variable to the neural network. The method shows feasible results on toy (Gaussian) data and image-to-image translation tasks.\n\nSTRENGTHS:\n-\tThis work is theoretically fruitful. The introduction of $\\alpha$-IMF leaded to a more expansive theoretical framework. This flow-based perspective enables online learning. \n-\tReduced the sensitivity to hyperparameters such as the iteration number, the number of phase in DSBM.\n-\tExpands the interpretation of the Schrödinger bridge by demonstrating connections with reinforcement learning (RL) and Sinkhorn flow.\n\nWEAKNESSES:\n- The experiment on practical data is insufficient and the performance improvement is incremental. The paper only presents real-world I2I on Cat <-> Wild. Moreover, the performance of bidirectional online learning improves only incrementally. Totally, it needs more evaluation on other kinds of dataset, e.g. Male <-> Female / Handbags <-> Shoes. \n- It would be beneficial to quantify the advantages of bidirectional learning or online learning over DSBM. For instance, comparing the actual wall-clock time or GPU memory usage would be useful. \n- The current algorithm reflects $\\alpha$ implicitly, and so there is no analysis provided on $\\alpha$. If there were any intuition or insights into the role of $\\alpha$ (step size) or if an algorithm could be developed to explicitly control $\\alpha$, it would offer more valuable insights.\n\nQUESTIONS:\n- If we perform this algorithm with full gradient descent algorithm, is it possible to explicitly control $\\alpha$?\n- Can we develop exact algorithm that reflect $\\alpha$ explicitly? If so, can we obtain some experimental results that compares $\\alpha$?\n\nLIMITATIONS:\nThe authors adequately addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors consider developing a new algorithm for a Schrödinger Bridge problem of translating between two probability distributions. Motivated by the fact that current Schrödinger Bridge (SB) approaches either use mini-batch optimal transport techniques or require training diffusion at every iteration, the authors propose discretizing the Schrödinger Bridge Flow (SB Flow) of path measures that converge to SB. They call this discretization $\\alpha$-Iterative Markovian Fitting ($\\alpha$-IMF) procedure. Using the fact that Markovian path measures characterizing the considered flow can be parametrized by diffusion vector field $v$, authors propose a non-parametric scheme for updating $v_n$ corresponding to $n$ discretization step of SB Flow and prove its correctness. To implement the procedure in practice, authors prove a parametric counterpart to the non-parametric update scheme that allows for optimization over a vector field $v_{\\theta}$ parametrized as a neural network called $\\alpha$ Diffusion Schrödinger Bridge Matching ($\\alpha$-DSBM). Furthermore, the authors computationally validate the proposed algorithm on both vector and image state spaces. The authors compare their $\\alpha$-DSBM algorithm to other DSBM versions [3] in the Gaussian setting. Image state space experiments were held on image-to-image translation tasks, i.e., MNIST $\\leftrightarrow$ EMNIST, Wild $\\leftrightarrow$ Cat on AFHQ 64/256, with quantitative evaluation and comparison with the original DSBM algorithm [3].\n\nSTRENGTHS:\n- The overall theory is novel and interesting. The authors prove that the proposed $\\alpha$-IMF procedure with non-parametric updates converges to the solution of SB.\n- The parametric counterpart to $\\alpha$-IMF ($\\alpha$-DSBM), which can be used with neural networks, is proposed and empirically verified.\n- The possibility of online $\\alpha$-DSBM finetuning seems beneficial in terms of computational resources compared to regular DSBM iterative training.\n- A sufficient comparison of the proposed $\\alpha$-DSBM with the previously known DSBM algorithm [3] is presented for a wide set of hyperparameter $\\epsilon$.\n- The Appendix presents an extensive additional theoretical review of the proposed concept and its connections to existing research.\n- The study on the parameterization of diffusion in both directions by one neural network with varied conditional vectors is proposed and tested.\n\nWEAKNESSES:\n- The proposed $\\alpha$-DSBM algorithm [3] is, in some sense, an advancement of the previously known DSBM algorithm [3]. It is expected to be more efficient because it removes the necessity to learn a distinct model at each iteration. However, the paper lacks a study of image quality boost given the same computational budget. Thus, it is not clear whether the authors achieved the goal of developing a more efficient algorithm.\n- The convergence of parametric $\\alpha$-DSBM is not proved. While it may be a hard theoretical problem to prove, at least a more extensive empirical evaluation of the setup with the known ground truth solution would be beneficial. The presented setup with a scalar covariance matrix is too simple. It would be a more solid argument to evaluate the Gaussian setup with a full covariance matrix [6] or use a mixture benchmark [7] and compare it with other neural network methods. \n- The FID metric is similar for DSBM and $\\alpha$-DSBM for the AFHQ experiment in Table 1 considering measured standard deviation. Thus again highlighting the lack of study with the same computational budget to justify that $\\alpha$-DSBM is more efficient.\nThe chosen AFHQ dataset may not be the best choice considering the small test set, which may introduce bias in measuring the FID metric. In turn, the visual comparison is hard since most images produced by DSBM and $\\alpha$-DSBM and given in Figure 19 are similar, and it is hard to deduce which algorithm is better. A bigger dataset such as Celeba [5] may help to solve this issue. At least a bigger dataset would give a good estimate of FID metrics. \n- It seems that some hyperparameters for DSBM AFHQ experiments are missing. The number of pre-training iterations, grad updates per IMF iteration, and IMF iterations are not specified. What is the training time and computational budget? \n- There is no study on *cornerstone* hyperparameter $\\alpha$. What is the value of $\\alpha$ used in the experiments? Is it equal to the learning rate of the diffusion model parameters?\n- The absence of code may cause trouble with reproducing the presented results.\n\nQUESTIONS:\n- Since the AFHQ dataset is quite small, experiments on unpaired male-to-female translation on the Celeba [5] dataset could provide a more meaningful comparison (as it was in DSBM[3]).\n- How do DSBM and $\\alpha$-DSBM compare in terms of computational budget? Could authors provide a quantitative comparison?\n- How does variation of $\\alpha$ affect speed of convergence and deviation from the SB solution?\n- Can authors provide experimental results of $\\alpha$-DSBM on more complex Gaussian distributions or multimodal low dimensional data, i.e., gaussian mixtures? \n- Is it possible to extend the proposed approach to GSBM [4], where, as far as my understanding goes, authors change the reciprocal projection step compared to DSBM [3]?\n- Can this method with small $\\epsilon$ can be applied to the generation problem? This could drastically reduce the computational budget of such and be applied in modern text-to-image generative models [2] similar to the work [1]. It would be great to see such experimental results. \n- Is it possible to perform *parametric* updates for Sinkhorn flow in a similar to $\\alpha$-DSBM way?\n\nOverall, the paper looks very promising. I will consider increasing my score if the authors address some of the listed weaknesses and questions, especially related to the same computational budget comparison, and evaluate their method on one of the listed more complex setups with the known ground truth distributions.\n\n[1] Liu, Xingchao, et al. \"Instaflow: One step is enough for high-quality diffusion-based text-to-image generation.\" The Twelfth International Conference on Learning Representations. 2023.\n\n[2] Esser, Patrick, et al. \"Scaling rectified flow transformers for high-resolution image synthesis.\" Forty-first International Conference on Machine Learning. 2024.\n\n[3] Shi, Yuyang, et al. \"Diffusion Schrödinger bridge matching.\" Advances in Neural Information Processing Systems 36 (2024).\n\n[4] Liu, Guan-Horng, et al. \"Generalized Schrödinger Bridge Matching.\" The Twelfth International Conference on Learning Representations.\n\n[5] Liu Z. et al. Deep learning face attributes in the wild //Proceedings of the IEEE international conference on computer vision. – 2015. – С. 3730-3738.\n\n[6] Hicham Janati, Boris Muzellec, Gabriel Peyré, and Marco Cuturi. Entropic optimal transport between unbalanced gaussian measures has a closed form. Advances in neural information processing systems, 33:10468–10479, 2020.\n\n[7] Gushchin N. et al. Building the bridge of schrödinger: A continuous entropic optimal transport benchmark //Advances in Neural Information Processing Systems. – 2023. – Т. 36. – С. 18932-18963.\n\nLIMITATIONS:\nThe authors adequately addressed the limitations and potential negative societal impact of their work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work introduces $\\alpha$-DSBM, a new way of training DSBM-like models, which does not require a Markovian projection at each step and eliminates the need to train multiple models. The main advantage over previous DSBM-based approaches is that $\\alpha$-DSBM only needs to train a single model with a single model, thus exhibiting a more stable training procedure.\n\nThe authors thoroughly contextualize their work within related work and provide detailed theoretical derivations to motivate their approach. Empirically $\\alpha$-DSBM is validated by comparing to existing DSBM-based approaches across different unpaired image translation tasks as well as some toy data examples.\n\nSTRENGTHS:\n- The new $\\alpha$-DSBM is introduced based on thorough theoretical motivation and derivations.\n- The authors manage to present and contextualize their method very clearly within the scope of related work across the different bridge matching approaches (Appendix E) and highlight connections accordingly (e.g. connection to the reflow procedure).\n- In general, the paper is well-written and, given the complexity of the topic, quite clear to follow.\n- Empirical validation includes important ablations giving further insights into the hyperparameter choice of e.g. $\\varepsilon$.\n\nWEAKNESSES:\n- While the experimental section of the paper thoroughly analyzes and compares DSBM and its different flavors, including the proposed method, a comparison with other competing methods could further support $\\alpha$-DSBM through empirical evidence. Specifically, a comparison to bridge/flow matching with mini-batch OT similar to [1, 2], a highly optimized CycleGAN as in [3], and adversarial-based OT methods like [4] could be good candidates for further comparisons. The authors also mention these as competing methods.\n\n\n[1] Alexander Tong and Kilian Fatras and Nikolay Malkin and Guillaume Huguet and Yanlei Zhang and Jarrid Rector-Brooks and Guy Wolf and Yoshua Bengio. \"# Improving and generalizing flow-based generative models with minibatch optimal transport\". IN TMLR 2024.\n\n[2] Luca Eyring and Dominik Klein and Théo Uscidda and Giovanni Palla and Niki Kilbertus and Zeynep Akata and Fabian Theis. \"Unbalancedness in Neural Monge Maps Improves Unpaired Domain Translation\". In ICLR 2024.\n\n[3] Dmitrii Torbunov and Yi Huang and Huan-Hsin Tseng and Haiwang Yu and Jin Huang and Shinjae Yoo and Meifeng Lin and Brett Viren and Yihui Ren. \"UVCGAN v2: An Improved Cycle-Consistent GAN for Unpaired Image-to-Image Translation\". In Arxiv 2023.\n\n[4] Beomsu Kim and Gihyun Kwon and Kwanyoung Kim and Jong Chul Ye. \"Unpaired Image-to-Image Translation via Neural Schrödinger Bridge\". In ICLR 2024.\n\nQUESTIONS:\n- How does $\\alpha$-DSBM perform compared to bridge matching with mini-batch OT sampling empirically? While it is clear that these methods introduce significant errors because of the mini-batch approximation, it is unclear how this affects results empirically compared to $\\alpha$-DSBM. I think this is the most important additional competing work (apart from DSBM), as these methods share their overall goal.\n- Does it make sense to combine $\\alpha$-DSBM with mini-batch OT sampling for the initial pretraining? How would this impact $\\alpha$-DSBM?\n\nLIMITATIONS:\n- The main limitation of $\\alpha$-DSBM is that new data needs to be generated with the current model during its fine-tuning, making the procedure not simulation-free and, thus, more expensive. This is sufficiently addressed in the paper.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a novel algorithm for mass transport problems, aiming to compute maps that transport one distribution to another. This paper introduces the Schrödinger Bridge Flow algorithm, a dynamic entropy-regularized version of OT, eliminating the need to train multiple DDM-type models. The algorithm discretizes a flow of path measures, ensuring the Schrödinger Bridge as the only stationary point. The paper demonstrates the algorithm's effectiveness on various unpaired data translation tasks, showcasing its potential to solve Entropic Optimal Transport problems with improved computational efficiency and practical implementation compared to existing methodologies.\n\nSTRENGTHS:\n1.This paper is eloquently written, and its ideas are easy to follow.\n2.While I did not strictly follow all the proofs provided by the authors, I found them to be generally detailed, clear, and easy to understand.\n3.The authors provide both quantitative and qualitative results on several datasets, along with a relatively detailed discussion and visualization. However, it is worth noting that the figures in the appendix appear to have low resolution, and the positioning of the legends is somewhat odd.\n4.I understand that the convergence of the parametric methodology may require several strong assumptions. Therefore, it is reasonable that the authors have not provided further proof of convergence.\n\nWEAKNESSES:\n1.Assumptions in Theorem 3.1: Theorem 3.1 only mentions \"under mild assumptions\" without specifying them in the main text. It is necessary to state these assumptions clearly in the main body of the paper and provide references to justify their rationality. While a more detailed explanation can be provided in the appendix, the main text should at least include a summary of the assumptions to ensure the reader understands the conditions under which the theorem holds.\n2.Experimental Results and Comparisons: The experimental results on some datasets are not particularly impressive, even though the authors claim that their method approximates OT maps more accurately than existing methods. The results in Table 1 do not show significant progress or state-of-the-art (SOTA) performance on certain datasets. Additionally, it is confusing that the authors explicitly mentioned Rectified Flow, Flow Matching, and OT-Flow Matching in the literature review, but did not include comparisons with these methods in the experiments. For instance, Rectified Flow has conducted experiments on mass transport, and other related methods, such as Denoising Diffusion Bridge Models from ICLR 2024, have tackled similar tasks. A direct comparison of the proposed method's performance with these existing methods would provide a clearer picture of its advantages and limitations.\n3.Discussion of Related Work: The paper could benefit from a more comprehensive discussion of related work, including diffusion models, flow matching, and the application of Wasserstein gradient flow-based models in mass transport tasks. This would help in better contextualizing the authors' innovations and contributions. By situating their work more clearly within the existing literature, the authors could highlight the unique aspects and potential advantages of their proposed method, enhancing the readers' understanding of its significance.\n\nQUESTIONS:\n1.See Weaknesses\n2.I am curious about the complexity and cost of the training method proposed by the authors. Could the authors provide more detailed quantitative data, such as training time comparisons with related methods?\n\nLIMITATIONS:\nAlthough the authors proposed a novel model to address mass transport problems and provided some theoretical proofs, the method is simulation-based. I believe that the training cost is higher compared to simulation-free methods, and the model performance does not achieve state-of-the-art (SOTA) results. Additionally, there is a lack of direct comparison with existing SOTA methods.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}