Title: Entropic Neural Optimal Transport via Diffusion Processes
Venue: NeurIPS 2023
Accepted: yes
Keywords: Optimal transport

Schrödinger Bridge

Entropy regularized OT

Neural Networks

Unpaired Learning
Abstract: We propose a novel neural algorithm for the fundamental problem of computing the entropic optimal transport (EOT) plan between probability distributions which are accessible by samples. Our algorithm is based on the saddle point reformulation of the dynamic version of EOT which is known as the Schrödinger Bridge problem. In contrast to the prior methods for large-scale EOT, our algorithm is end-to-end and consists of a single learning step, has fast inference procedure, and allows handling small values of the entropy regularization coefficient which is of particular importance in some applied problems. Empirically, we show the performance of the method on several large-scale EOT tasks. The code for the ENOT solver can be found at https://github.com/ngushchin/EntropicNeuralOptimalTransport

Entropic Neural Optimal Transport via Diffusion Processes

Abstract

We propose a novel neural algorithm for the fundamental problem of computing the entropic optimal transport (EOT) plan between continuous probability distributions which are accessible by samples. Our algorithm is based on the saddle point reformulation of the dynamic version of EOT which is known as the Schrödinger Bridge problem. In contrast to the prior methods for large-scale EOT, our algorithm is end-to-end and consists of a single learning step, has fast inference procedure, and allows handling small values of the entropy regularization coefficient which is of particular importance in some applied problems. Empirically, we show the performance of the method on several large-scale EOT tasks. The code for the ENOT solver can be found at https://github.com/ngushchin/EntropicNeuralOptimalTransport .

t=0.2

t=0.3

t=0.4

t=0.5

t=0.6

t=0.7

t=0.8

t=0.9

t=1.0

0=

3

1

=3

1 Introduction

Optimal transport (OT) plans are a fundamental family of alignments between probability distributions. The majority of scalable neural algorithms to compute OT are based on the dual formulations of OT, see [28] for a survey. Despite the success of such formulations in generative modeling [42, 31], these dual form approaches can hardly be generalized to the popular entropic OT [12]. This is due to the numerical instability of the dual entropic OT problem [14] which appears for small entropy

regularization values which are suitable for downstream generative modeling tasks. At the same time, entropic OT is useful as it allows to learn one-to-many stochastic mappings with tunable level of sample diversity. This is particularly important for ill-posed problems such as super-resolution [37].

Contributions . We propose a saddle-point reformulation of the Entropic OT problem via using its dynamic counterpart known as the Schrödinger Bridge problem ( M 4.1). Based on our new reformulation, we propose a novel end-to-end neural algorithm to solve the related entropic OT problem for a pair of continuous distributions accessible by samples ( M 4.2). Unlike many predecessors, our method allows handling small entropy coefficients. This enables practical applications to data → data mapping tasks that require slight variability in the learned maps. Furthermore, we provide an error analysis for solving the suggested saddle point optimization problem through duality gaps which are the errors of solving the inner and outer optimization problems ( M 4.3). Empirically, we illustrate the performance of the method on several toy and large-scale EOT tasks ( M 5).

2 Background

Optimal Transport (OT) and Schrödinger Bridge (SB) problems imply finding an efficient way to transform some initial distribution P 0 to target distribution P 1 . While the solution of OT only gives the information about which part of P 0 is transformed to which part of P 1 , SB implies finding a stochastic process that describes the entire evolution from P 0 to P 1 . Below we give an introduction to OT and SB problems and show how they are related. For a detailed overview of OT, we refer to [49, 41] and of SB - to [32, 11].

2.1 Optimal Transport (OT)

Kantorovich's OT formulation (with the quadratic cost). We consider D -dimensional Euclidean spaces X , Y and use P 2 ( X ) = P 2 ( Y ) to denote the respective sets of Borel probability distributions on them which have finite second moment. For two distributions P 0 ∈ P 2 ( X ) , P 1 ∈ P 2 ( Y ) , consider the following minimization problem:

where Π( P 0 , P 1 ) ⊂ P 2 ( X × Y ) is the set of probability distributions on X × Y with marginals P 0 and P 1 . Such distributions π ∈ Π( P 0 , P 1 ) are called the transport plans between P 0 and P 1 . The set Π( P 0 , P 1 ) is non-empty as it always contains the trivial plan P 0 × P 1 . A minimizer π ∗ of (1) always exists and is called an OT plan. If P 0 is absolutely continuous, then π ∗ is deterministic: its conditional distributions are degenerate, i.e., π ∗ ( ·| x ) = δ T ∗ ( x ) for some T ∗ : X → Y (OT map).

Entropic OT formulation. We use H ( π ) to denote the differential entropy of distribution π and KL ( π || π ′ ) to denote the Kullback-Leibler divergence between distributions π and π ′ . Two most popular entropic OT formulations regularize (1) with the entropy H ( π ) or KL-divergence between plan π and the trivial plan P 0 × P 1 , respectively ( ϵ > 0 ):

Since π ∈ Π( P 0 , P 1 ) , it holds that KL ( π || P 0 × P 1 )= -H ( π )+ H ( P 0 )+ H ( P 1 ) , i.e., both formulations are equal up to an additive constant when P 0 and P 1 are absolutely continuous. The minimizer of (2) is unique since the functional is strictly convex in π thanks to the strict convexity of H ( π ) . This unique minimizer π ∗ is called the entropic OT plan.

2.2 Schrödinger Bridge (SB)

SB with the Wiener prior. Let Ω be the space of R D valued functions of time t ∈ [0 , 1] describing some trajectories in R D , which start at time t = 0 and end at time t = 1 . We use P (Ω) to denote the set of probability distributions on Ω . We use dW t to denote the differential of the standard Wiener process. Let W ϵ ∈ P (Ω) be the Wiener process with the variance ϵ which starts at P 0 . This diffusion process can be represented via the following stochastic differential equation (SDE):

We use KL ( T || Q ) to denote the Kullback-Leibler divergence between stochastic processes T and Q . The Schrödinger Bridge problem was initially proposed in 1931/1932 by Erwin Schrödinger [44]. It can be formulated as follows [11, Problem 4.1]:

where F ( P 0 , P 1 ) ⊂ P (Ω) is the set of probability distributions on Ω having marginal distributions P 0 and P 1 at t = 0 and t = 1 , respectively. Thus, the Schrödinger Bridge problem implies finding a stochastic process T with marginal distributions P 0 and P 1 at times t = 0 and t = 1 , respectively, which has minimal KL divergence with the prior process W ϵ .

Link to OT problem. Here we recall how Scrödinger Bridge problem (5) relates to entropic OT problem (2). This relation is well known (see, e.g., [32] or [11, Problem 4.2]), and we discuss it in detail because our proposed approach ( M 4) is based on it. Let π T denote the joint distribution of a stochastic process T at time moments t = 0 , 1 and π T 0 , π T 1 denote its marginal distributions at time moments t = 0 , 1 , respectively. Let T | x,y denote the stochastic processes T conditioned on values x, y at times t = 0 , 1 , respectively. One may decompose KL ( T || W ϵ ) as [48, Appendix C]:

i.e., KL divergence between T and W ϵ is a sum of two terms: the first represents the similarity of the processes at start and finish times t = 0 and t = 1 , while the second term represents the similarity of the processes for intermediate times t ∈ (0 , 1) conditioned on the values at t = 0 , 1 . For the first term, it holds (see Appendix A or [11, Eqs 4.7-4.9]):

where C is a constant which depends only on P 0 and ϵ . In [32, Proposition 2.3], the authors show that if T ∗ is the solution to (5), then T ∗ | x,y = W ∗ | x,y . Hence, one may optimize (5) over processes T for which T | x,y = W ϵ | x,y for every x, y and set the last term in (6) to zero. In this case:

i.e., it suffices to optimize only over joint distributions π T at time moments t = 0 , 1 . Hence, minimizing (8) is equivalent (up to an additive constant C ) to solving EOT (2) divided by ϵ , and their respective solutions π T ∗ and π ∗ coincide. Thus, SB problem (5) can be simplified to the entropic OT problem (2) with entropy coefficient ϵ .

Dual form of EOT . Entropic OT problem (8) has several dual formulations. Here we recall the one which is particularly useful to derive our algorithm. The dual formulation follows from the weak OT theory [7, Theorem 1.3] and we explain it in detail in the proof of Lemma B.3:

where β C ( x ) def = inf ν ∈P 2 ( Y ) { KL ( ν || π W ϵ ( ·| x ) ) -∫ Y β ( y ) dν ( y ) } . The sup here is taken over β belonging to the set of functions

i.e., β should be continuous with mild boundness assumptions.

Dynamic SB problem (DSB) . It is known that the solution to the SB problem (5) belongs to the class D ( P 0 ) of finite-energy diffusions T f [32, Proposition 4.1] which are given by:

where f : R D × [0 , 1] → R D is the drift function. The last inequality in (9) means that T f is a finite-energy diffusion. Hence, optimizing only over finite-energy diffusions rather than all possible

processes is enough to solve SB problem (5). For finite-energy diffusions, it is possible to rewrite the optimization objective. One can show that KL ( T f || W ϵ ) between processes T f and W ϵ is [40]:

By substituting (10) in (5), SB problem reduces to the following problem [11, Problem 4.3]:

where D ( P 0 , P 1 ) ⊂ P (Ω) is the set of finite-energy diffusion on Ω having marginal distributions P 0 and P 1 at t = 0 and t = 1 , respectively. Note that D ( P 0 , P 1 ) ⊂ F ( P 0 , P 1 ) . Since the problem (11) is the equivalent reformulation of (5), its optimal value also equals (8). However, solving this problem, as well as (8) is challenging as it is still hard to satisfy the boundary constraints.

3 Related Work

In this section, we overview the existing methods to compute the OT plan or map. To avoid any confusion, we emphasize that popular Wasserstein GANs [5] compute only the OT cost but not the OT plan and, consequently, are out of scope of the discussion.

Discrete OT. The majority of algorithms in computational OT are designed for the discrete setting where the inputs P 0 , P 1 have finite supports. In particular, the usage of entropic regularization (2) allows to establish efficient methods [13] to compute the entropic OT plan between discrete distributions with the support size up to 10 5 -10 6 points, see [41] for a survey. For larger support sizes, such methods are typically computationally intractable.

Continuous OT. Continuous methods imply computing the OT plan between distributions P 0 , P 1 which are accessible by empirical samples. Discrete methods "as-is" are not applicable to the continuous setup in high dimensions because they only do a stochastic matching between the train samples and do not provide out-of-sample estimation. In contrast, continuous methods employ neural networks to explicitly or implicitly learn the OT plan or map. As a result, these methods can map unseen samples from P 0 to P 1 according to the learned OT plan or map.

There exists many methods to compute OT plans [51, 36, 38, 26, 28, 27, 42, 16, 19, 18] but they consider only unregularized OT (1) rather than the entropic one (2). In particular, they mostly focus on computing the deterministic OT plan (map) which may not exist. Recent works [31, 30] design algorithms to compute OT plans for weak OT [20, 7]. Although weak OT technically subsumes entropic OT , these works do not cover the entropic OT (2) because there is no simple way to estimate the entropy from samples. Below we discuss methods specifically for EOT (2) and DSB (11).

3.1 Continuous Entropic OT

In LSOT [45], the authors solve the dual problem to entropic OT (2). The dual potentials are then used to compute the barycentric projection x ↦→ ∫ Y y dπ ∗ ( y | x ) , i.e., the first conditional moment of the entropic OT plan. This strategy may yield a deterministic approximation of π ∗ for small ϵ but does not recover the entire plan itself.

In [14, Figure 3], the authors show that the barycentric projection leads to the averaging artifacts which make it impractical in downstream tasks such as the unpaired image super-resolution. To solve these issues, the authors propose a method called SCONES. It recovers the entire conditional distribution π ∗ ( y | x ) of the OT plan π ∗ from the dual potentials. Unfortunately, this is costly. During the training phase, the method requires learning a score-based model for the distribution P 1 . More importantly, during the inference phase, one has to run the Langevin dynamic to sample from π ∗ ( y | x ) .

The optimization of the above-mentioned entropic approaches requires evaluating the exponent of large values which are proportional to ϵ -1 [45, Eq. 7]. Due to this fact, those methods are unstable for small values ϵ in (2). In contrast, our proposed method ( M 4) resolves this issue: technically, it works even for ϵ = 0 .

3.2 Approaches to Compute Schrödinger Bridges

Existing approaches to solve DSB mainly focus on generative modeling applications ( noise → data ). For example, FB-SDE [10] utilizes data likelihood maximization to optimize the parameters of

forward and backward SDEs for learning the bridge. On the other hand, MLE-SB [48] and DiffSB [15] employ the iterative proportional fitting technique to learn the DSB. Another method proposed by [50] involves solving the Schrodinger Bridge only between the Dirac delta distribution and real data to solve the data generation problem.

Recent studies [34, 33] have indicated that the Schrödinger Bridge can be considered as a specific instance of a more general problem known as the Mean-Field Game [1]. These studies have also suggested novel algorithms for resolving the Mean-Field Game problem. However, the approach presented by [33] cannot be directly applied to address the hard distribution constraints imposed on the start and final probability distribution as in the Schrödinger Bridge problem (see Appendix H). The approach suggested by [34] coincides with that proposed in FB-SDE [10] for the SB problem.

4 The Algorithm

This section presents our novel neural network-based algorithm to recover the solution T f ∗ of the DSB problem (11) and the solution π ∗ of the related EOT problem (8). In M 4.1, we theoretically derive the proposed saddle point optimization objective. In M 4.2 we provide and describe the practical learning procedure to optimize it. In M 4.3, we perform the error analysis via duality gaps. In Appendix B, we provide proofs of all theorems and lemmas.

4.1 Saddle Point Reformulation of EOT via DSB

For two distributions P 0 and P 1 accessible by finite empirical samples, solving entropic OT (2) "as-is" is non-trivial. Indeed, one has to (a) enforce the marginal constraints π ∈ Π( P 0 , P 1 ) and (b) estimate the entropy H ( π ) from empirical samples which is challenging. Our idea below is to employ the relation of EOT with DSB to derive an optimization objective which in practice can recover the entropic plan avoiding the above-mentioned issues. First, we introduce the functional

This functional can be viewed as the Lagrangian for DSB (11) with the relaxed constraint dπ T f 1 ( y ) = d P 1 ( y ) , and function β : Y → R ( potential ) playing the role of the Langrange multiplier.

Theorem 4.1 (Relaxed DSB formulation) . Consider the following saddle point optimization problem:

where sup is taken over potentials β ∈ C b, 2 ( Y ) and inf is taken over diffusion processes T f ∈ D ( P 0 ) . Then for every optimal pair ( β ∗ , T f ∗ ) for (13), i.e.,

it holds that the process T ∗ f is the solution to SB (11).

Corollary 4.2 (Entropic OT as relaxed DSB) . If ( β ∗ , T f ∗ ) solves (13) , then π T f ∗ is the EOT plan (2).

Our results above show that by solving (13), one immediately recovers the optimal process T f ∗ in DSB (11) and the optimal EOT plan π ∗ = π T ∗ f . The notable benefits of considering (13) instead of (2), (3) and (11) is that (a) it is as an optimization problem over ( β, T f ) without the constraint dπ T f 1 ( y ) = d P 1 ( y ) , and (b) objective (12) admits Monte-Carlo estimates by using random samples from P 0 , T f , P 1 . In M 4.2 below, we describe the straightforward practical procedure to optimize (13) with stochastic gradient methods and neural nets.

Relation to prior works. In the field of neural OT, there exist so many maximin reformulations of OT (classic [28, 42, 16, 19, 22], weak [31, 30], general [6]) resembling our (13) that a reader may naturally wonder (1) why not to apply them to solve entropic OT? (2) how does our reformulation differ from all of them? It is indeed true that, e.g., algorithm from [31, 6] mathematically covers the entropic case (2). Yet in practice it requires estimation of entropy from samples for which there is no easy way and the authors do not consider this case. These methods can be hardly applied to EOT.

In contrast to the prior works, we deal with entropic OT through its connection with the Schrödinger Bridge. Our max-min reformulation is for DSB and it avoids computing the entropy term H ( π )

in EOT. This term is replaced by the energy of the process E T f [ ∫ 1 0 || f ( X t , t ) || 2 dt ] which can be straightforwardly estimated from the samples of T f , allowing to establish a computational algorithm.

4.2 Practical Optimization Procedure

To solve (12), we parametrize drift function f ( x, t ) of the process T f and potential β ( y ) 4 by neural nets f θ : R D × [0 , 1] → R D and β ϕ : R D → R . We consider the following maximin problem:

We use standard Euler-Maruyama (Eul-Mar) simulation (Algorithm 2 in Appendix C) for sampling from the stochastic process T f by solving its SDE (9). To estimate the value of E T f [ ∫ 1 0 || f ( X t , t ) || 2 dt ] in (14), we utilize the mean value of || f ( x, t ) || 2 over time t of trajectory X t that is obtained during the simulation by Euler-Maruyama algorithm (Appendix C). We train f θ and β ϕ by optimizing (12) with the stochastic gradient ascent-descent by sampling random batches from P 0 , P 1 . The optimization procedure is detailed in Algorithm 1. We use f n,m to denote the drift at time step n for the m -th object of the input sample batch. We use the averaged of the drifts as an estimate of ∫ 1 0 || f ( X t , t ) || 2 dt in the training objective.

Remark. For the image tasks ( M 5.3, M 5.4), we find out that using a slightly different parametrization of T f considerably improves the quality of our Algorithm, see Appendix F.

It should be noted that the term E T f [ ∫ 1 0 || f ( X t , t ) || 2 dt ] is not multiplied by 1 2 ϵ in the algorithm since this does not affect the optimal T f ∗ solving the inner optimization problem, see Appendix D.

Algorithm 1: Entropic Neural OT (ENOT)

Input:

samples from distributions P 0 , P 1 ; Wiener prior noise variance ϵ ≥ 0 ; drift network f θ : R D × [0 , 1] → R D ; beta network β ϕ : R D → R ; number of steps N for Eul-Mar (App C); number of inner iterations K f .

Output: drift f ∗ θ of T f ∗ θ solving DSB (11). repeat

until converged ;

Relation to GANs . At the first glance, our method might look like a typical GAN as it solves a maximin problem with the "discriminator" β ϕ and SDE "generator" with the drift f θ . Unlike GANs, in our saddle point objective (12), optimization of "generator" T f and "discriminator" β are swapped, i.e., "generator" is adversarial to "discriminator", not vise versa, as in GANs. For further discussion of differences between saddle point objectives of neural OT/GANs, see [31, M 4.3], [42, M 4.3], [16].

4.3 Error Bounds via Duality Gaps

Our algorithm solves a maximin optimization problem and recovers some approximate solution ( ˆ β, T ˆ f ) . Given such a pair, it is natural to wonder how close is the recovered T ˆ f to the optimal T f ∗ . Our next result sheds light on this question via bounding the error with the duality gaps .

Theorem 4.3 (Error analysis via duality gaps) . Consider a pair ( ˆ β , T ˆ f ). Define the duality gaps, i.e., errors of solving inner and outer optimization problems by:

Then it holds that

where we use ρ TV ( · , · ) to denote the total variation norm (between the processes or plans).

(a)

x

∼

P

0

,

y

∼

P

1

(b) ENOT

(ours)

,

ϵ

= 0

(c) ENOT

(ours)

,

ϵ

= 0

.

01

(d) ENOT

(ours)

,

ϵ

= 0

.

1

Fitted distribution

2

Fitteddistribution

2

Fitted distribution

0

0

岛

Inputdistribution

2

-1

0

1

2

0

2

Targetdistribution

2

Trajectories

2

Target data

1

D

Q

Trajectories

Trajectories

2

Targetdata

2

Targetdata

-1

0

1

2

-2

-1

0

1

2

Relation to prior works . There exist deceptively similar results, see [38, Theorem 3.6], [42, Theorem 4.3], [16, Theorem 4], [6, Theorem 3]. None of them are relevant to our EOT/DSB case.

In [38], [42], [16], the authors consider maximin reformulations of unregularized OT (1), i.e., nonentropic. Their result requires the potential β ( f, ψ in their notation) to be a convex function which in practice means that one has to employ ICNNs [4] which have poor expressiveness [29, 17, 26]. Our result is free from such assumptions on β . In [6], the authors consider general OT problem [39] and require the general cost functional to be strongly convex (in some norm). Their results also do not apply to our case as the (negative) entropy which we consider is not strongly convex.

5 Experimental Illustrations

In this section, we qualitatively and quantitatively illustrate the performance of our algorithm in several entropic OT tasks. Our proofs apply only to EOT ( ϵ > 0 ), but for completeness, we also present results ϵ = 0 , i.e., unregularized case (1). Furthermore, we test our algorithm with ϵ = 0 on the Wasserstein-2 Benchmark [28], see Appendix J. We also demonstrate the extension of our algorithm to costs other than the squared Euclidean distance in Appendix I. The implementation details are given in Appendices E, F and G. The code is written in PyTorch and is publicly available at

5.1 Toy 2D experiments

Here we give qualitative examples of our algorithm's performance on toy 2D pairs of distributions. We consider two pairs P 0 , P 1 : Gaussian → Swiss Roll , Gaussian → Mixture of 8 Gaussians . We provide qualitative results in Figure 2 and Figure 5 (Appendix E), respectively. In both cases, we provide solutions of the problem for ϵ = 0 , 0 . 01 , 0 . 1 and sample trajectories. For ϵ = 0 , all the trajectories are straight lines as they represent solutions for non-regularized OT (1), see [43, M 5.4]. For bigger ϵ , trajectories, as expected, become more noisy and less straight.

5.2 High-dimensional Gaussians

For general continuous distributions P 0 , P 1 , the ground truth solution of entropic OT (2) and DSB (11) is unknown. This makes it challenging to assess how well does our algorithm solve these problems. Fortunately, when P 0 and P 1 are Gaussians, there exist closed form solutions of these related problems, see [25] and [9]. Thus, to quantify the performance of our algorithm, we consider entropic OT problems in dimensions D ∈ { 2 , 16 , 64 , 128 } with ϵ = 1 for Gaussian P 0 = N (0 , Σ 0 ) , P 1 = N (0 , Σ 1 ) . We pick Σ 0 , Σ 1 at random: their eigenvectors are uniformly distributed on the unit sphere and eigenvalues are sampled from the loguniform distribution on [ -log 2 , log 2] .

Metrics . We evaluate (a) how precisely our algorithm fits the target distribution P 1 on R D ; (b) how well it recovers the entropic OT plan π ∗ which is a Gaussian distribution on R D × R D ; (c) how accurate are the learned marginal distributions at intermediate times t = 0 , 1 10 , . . . , 1 .

In each of the above-mentioned cases, we compute the BW 2 2 -UVP [29, M 5] between the learned and the ground truth distributions. For two distributions ̂ χ and χ , it is the Wasserstein-2 distance between their Gaussian approximations which is further normalized by the variance of the distribution χ :

We estimate the metric by using 10 5 samples.

Baselines . We compare our method ENOT with LSOT [45], SCONES [14], MLE-SB [48], DiffSB[15], FB-SDE [10] (two algorithms, A and J). Results are given in Tables 1, 2 and 3. LSOT and SCONES solve EOT without solving SB, hence there are no results in Table 3 for these methods. Our method achieves low BW 2 2 -UVP values indicating that it recovers the ground truth process and entropic plan fairly well. The competitive methods have good results in low dimensions, however they perform worse in high dimensions. Importantly, our methods scores the best results in recovering the OT plan (Table 2) and the marginal distributions of the Schrodinger bridge (Table 3). To illustrate the stable convergence of ENOT, we provide the plot of BW 2 2 -UVP between the learned plan and the ground truth plan for ENOT during training in Figure 6 (Appendix E).

5.3 Colored MNIST

In this section, we test how the entropy parameter ϵ affects the stochasticity of the learned plan in higher dimensions. For this, we consider the entropic OT problem between colorized MNIST digits of classes " 2 " ( P 0 ) and " 3 " ( P 1 ).

Effect of parameter ϵ . For ϵ = 0 , 1 , 10 , we learn our Algorithm 1 on the train sets of digits "2" and "3". We show the translated test images in Figures 3a, 3b and 3c, respectively. When ϵ = 0 , there is no diversity in generated "3" samples (Figure 3a), the color remains since the map tried to minimally change the image in the RGB pixel space. When ϵ = 1 , some slight diversity in the shape of "3" appears but the color of the input "2" is still roughly preserved (Figure 3c). For higher ϵ , the diversity of generated samples becomes clear (Figure 3c). In particular, the color of "3" starts to slightly deviate from the input "2". That is, increasing the value ϵ of the entropy term in (2) expectedly leads to bigger stochasticity in the plan. We add the conditional LPIPS variance [24, Table 1] of generated samples for test datasets by ENOT to show how diversity changes for different ϵ (Table 4). We provide examples of trajectories learned by ENOT in Figure 7 (Appendix F).

Metrics and baselines. We compare our method ENOT with SCONES [14], and DiffSB [15] as these are the only methods which the respective authors applied for data → data tasks. To evaluate the results, we use the FID metric [23] which is the Bures-Wasserstein (Freschet) distance between the

(a) ENOT

(ours)

samples,

ϵ

= 0

, FID:

6

.

0

(b) ENOT

(ours)

samples,

ϵ

= 1

, FID:

6

.

28

(c) ENOT

(ours)

samples,

ϵ

=10

, FID:

6

.

9

(d) DiffSB [15] samples,

ϵ

= 1

, FID:93

(e) DiffSB [15] samples,

ϵ

= 10

, FID:105

(f) DOT samples,

ϵ

= 0

, FID:N/A

(g) DOT samples,

ϵ

= 1

, FID:N/A

(h) DOT samples,

ϵ

= 10

, FID:N/A

(i) SCONES [14] samples

ϵ

= 25

, FID:

14

.

73

(j) SCONES [14] samples

ϵ

= 100

, FID:

14

.

22

X~Po

X~Po

X~Po

X~Po

X~Po

2

3

3

3

2

3

3

3

2

3

3

3

2

3

3

2

2

3

3

3

2

3

3

3

2

3

3

3

2

3

3

3

2

3

3

3

2

3

3

3

2

3

3

3

2

3

3

3

2

3

3

3

2

3

2

2

2

3

3

3

2

3

3

3

2

3

3

3

2

3

3

3

2

2

W

2

3

3

2

X~Po

DOT(X)

X~P。

DOT(X) DOT(X) DOT(X)

X~P

DOT(X)DOT(X)

DOT(X)

S(X)

S(X)

S(X)

S(X)

S(X)

S(X)

S(X)

S(X)

2

3

3

3

2

3

3

3

2

3

3

3

3

3

3

2

3

3

3

2

3

3

3

2

3

了

3

2

3

3

3

2

3

3

3

2

3

3

3

2

3

3

3

2

3

3

3

2

3

E

3

2

3

3

3

2

3

3

3

2

3

3

3

2

3

3

3

2

3

3

3

2

3

3

3

2

3

E

3

distributions after extracting features using the InceptionV3 model [47]. We measure test FID for every method and present the results and qualitative examples in Figure 3. There are no results for SCONESwith ϵ = 0 , 1 , 10 , since it is not applicable for such reasonably small ϵ due to computational instabilities, see [14, M 5.1]. DiffSB [15] can be applied for small regularization ϵ , so we test ϵ = 1 , 10 . By the construction, this algorithm is not suitable for ϵ = 0 .

Our ENOT method outperforms the baselines in FID. DiffSB [15] yield very high FID. This is presumably due to instabilities of DiffSB which the authors report in their sequel paper [46]. SCONES yields reasonable quality but due to high ϵ = 25 , 100 the shape and color of the generated images "3" starts to deviate from those of their respective inputs "2".

For completeness, we provide the results of the stochastic matching of the test parts of the datasets by the discrete OT for ϵ = 0 and EOT [12] for ϵ = 1 , 10 (Figures 3f, 3g, 3h). This is not the out-of-sample estimation, obtained samples "3" are just test samples of " 3 " (this setup is unfair ). Discrete OT is not a competitor here as it does not generate new samples and uses target test samples. Still it gives a rough estimate what to expect from the learned plans for increasing ϵ .

5.4 Unpaired Super-resolution of Celeba Faces

For the large-scale evaluation, we adopt the experimental setup of SCONES [14]. We consider the problem of unpaired image super-resolution for the 64 × 64 aligned faces of CelebA dataset [35].

We do the unpaired train-test split as follows: we split the dataset into 3 parts: 90k (train A1), 90k (train B1), 20k (test C1) samples. For each part we do 2 × bilinear downsample and then 2 × bilinear upsample to degrade images but keep the original size. As a result, we obtain degraded parts A0, B0, C0. For training in the unpaired setup, we use parts A0 (degraded faces, P 0 ) and B1 (clean faces, P 1 ). For testing, we use the hold-out part C0 (unseen samples) with C1 considered as the reference.

We train our model with ϵ =0 , 1 , 10 to and test how it restores C1 (Figure 4a) from C0 images (Figure 4b) and present the qualitative results in Figures 4f, 4g, 4h. We provide examples of trajectories learned by ENOT in Figure 1.

Metrics and baselines. To quantify the results, as in [14], we compute the FID score [23] between the sets of mapped C0 images and C1 images (Table 5). The FID of ENOT is better than FID values of the other methods, but increases with ϵ probably due to the increasing variance of gradients during training . As in M 5.1 and M 5.3, the diversity of samples increases with ϵ . Our method works with small values of ϵ and provides reasonable amount of diversity in the mapped samples which grows with ϵ (Table 4). As the baseline among other methods for EOT we consider only SCONES, as it is the only EOT/DSB algorithm which has been applied to data → data task at 64 × 64 resolution. At the same time, we emphasize that SCONES is not applicable for small ϵ due to instabilities, see [14, M 5.1]. This makes it impractical , as due to high ϵ , its produces up-scaled images (Figures 4c) are

(a)

(b)

(c) SCONES,

ϵ

= 100

(d) ICNN

(e) AugCycleGAN

X~Po

Y~Po

S(X)

S(X)

S(X)

S(X)

Tf(X)

Tf(X)

Tf(X)

Tf(X)

Tf(X)

Tf(X)

Tf(X)

Tf(X)

Tf(X)

Tf(X)

Tf(X)

(f) ENOT

(ours)

,

ϵ

= 0

(g) ENOT

(ours)

,

ϵ

= 1

(h) ENOT

(ours)

,

ϵ

= 10

nearly random and do not reflect the attributes of the input images (Figure 4a). We do not provide results for DiffSB [15] since it already performs bad on Colored MNIST ( M 5.3) and the authors also did not consider any image-to-image apart of grayscale 28x28 images.

For completeness, we present results on this setup for other methods, which do not solve EOT: ICNNbased OT [38] and AugCycleGAN [2]. ICNN (4d) learns a deterministic map. AugCycleGAN (4e) learns a stochastic map, but the generated samples differ only by brightness.

6 Discussion

Potential impact. There is a lack of scalable algorithms for learning continuous entropic OT plans which may be used in data → data practical tasks requiring control of the diversity of generated samples. We hope that our results provide a new direction for research towards establishing scalable and efficient methods for entropic OT by using its connection with SB.

Potential social impact. Like other popular methods for generating images, our method can be used to simplify the work of designers with digital images and create new products based on it. At the same time, our method may be used for creating fake images just like the other generative models.

Limitations. To simulate the trajectories following SDE (9), we use the Euler-Maruyama scheme. It is straightforward but may be imprecise when the number of steps is small or the noise variance ϵ is high. As a result, for large ϵ , our Algorithm 1 may be computationally heavy due to the necessity to backpropagate through a large computational graph obtained via the simulation. Employing time and memory efficient SDE integration schemes is a promising avenue for the future work.

ACKNOWLEDGEMENTS. This work was partially supported by Skoltech NGP program (SkoltechMIT joint project).

References

Yves Achdou, Pierre Cardaliaguet, François Delarue, Alessio Porretta, and Filippo Santambrogio. Mean Field Games: Cetraro, Italy 2019 , volume 2281. Springer Nature, 2021.

Amjad Almahairi, Sai Rajeshwar, Alessandro Sordoni, Philip Bachman, and Aaron Courville. Augmented cyclegan: Learning many-to-many mappings from unpaired data. In International Conference on Machine Learning , pages 195-204. PMLR, 2018.

Brandon Amos. On amortizing convex conjugates for optimal transport. In The Eleventh International Conference on Learning Representations , 2022.

Brandon Amos, Lei Xu, and J Zico Kolter. Input convex neural networks. In Proceedings of the 34th International Conference on Machine Learning-Volume 70 , pages 146-155. JMLR. org, 2017.

Martin Arjovsky and Leon Bottou. Towards principled methods for training generative adversarial networks. In International Conference on Learning Representations , 2017.

Arip Asadulaev, Alexander Korotin, Vage Egiazarian, and Evgeny Burnaev. Neural optimal transport with general cost functionals. arXiv preprint arXiv:2205.15403 , 2022.

Julio Backhoff-Veraguas, Mathias Beiglböck, and Gudmun Pammer. Existence, duality, and cyclical monotonicity for weak transport costs. Calculus of Variations and Partial Differential Equations , 58(6):1-28, 2019.

Jean-David Benamou and Yann Brenier. A computational fluid mechanics solution to the monge-kantorovich mass transfer problem. Numerische Mathematik , 84(3):375-393, 2000.

Charlotte Bunne, Ya-Ping Hsieh, Marco Cuturi, and Andreas Krause. The schrödinger bridge between gaussian measures has a closed form. In International Conference on Artificial Intelligence and Statistics , pages 5802-5833. PMLR, 2023.

Tianrong Chen, Guan-Horng Liu, and Evangelos Theodorou. Likelihood training of schrödinger bridge using forward-backward sdes theory. In International Conference on Learning Representations , 2021.

Yongxin Chen, Tryphon T Georgiou, and Michele Pavon. Stochastic control liaisons: Richard sinkhorn meets gaspard monge on a schrodinger bridge. SIAM Review , 63(2):249-313, 2021.

Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems , 26, 2013.

Marco Cuturi and Arnaud Doucet. Fast computation of wasserstein barycenters. In International conference on machine learning , pages 685-693. PMLR, 2014.

Max Daniels, Tyler Maunu, and Paul Hand. Score-based generative neural networks for largescale optimal transport. Advances in neural information processing systems , 34:12955-12965, 2021.

Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion schrödinger bridge with applications to score-based generative modeling. Advances in Neural Information Processing Systems , 34:17695-17709, 2021.

Jiaojiao Fan, Shu Liu, Shaojun Ma, Hao-Min Zhou, and Yongxin Chen. Neural monge map estimation and its applications. Transactions on Machine Learning Research , 2023. Featured Certification.

Jiaojiao Fan, Qinsheng Zhang, Amirhossein Taghvaei, and Yongxin Chen. Variational wasserstein gradient flow. In International Conference on Machine Learning , pages 6185-6215. PMLR, 2022.

Milena Gazdieva, Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. Extremal domain translation with neural optimal transport. In Advances in Neural Information Processing Systems , 2023.

Milena Gazdieva, Litu Rout, Alexander Korotin, Alexander Filippov, and Evgeny Burnaev. Unpaired image super-resolution with optimal transport maps. arXiv preprint arXiv:2202.01116 , 2022.

Nathael Gozlan, Cyril Roberto, Paul-Marie Samson, and Prasad Tetali. Kantorovich duality for general transport costs and applications. Journal of Functional Analysis , 273(11):3327-3405, 2017.

A Hitchhiker's Guide. Infinite dimensional analysis . Springer, 2006.

Pierre Henry-Labordere. (martingale) optimal transport and anomaly detection with neural networks: A primal-dual algorithm. arXiv preprint arXiv:1904.04546 , 2019.

Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. GANs trained by a two time-scale update rule converge to a local nash equilibrium. In Advances in neural information processing systems , pages 6626-6637, 2017.

Xun Huang, Ming-Yu Liu, Serge Belongie, and Jan Kautz. Multimodal unsupervised imageto-image translation. In Proceedings of the European conference on computer vision (ECCV) , pages 172-189, 2018.

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.

Alexander Korotin, Vage Egiazarian, Arip Asadulaev, Alexander Safin, and Evgeny Burnaev. Wasserstein-2 generative networks. In International Conference on Learning Representations , 2021.

Alexander Korotin, Alexander Kolesov, and Evgeny Burnaev. Kantorovich strikes back! wasserstein GANs are not optimal transport? In Thirty-sixth Conference on Neural Information Processing Systems Datasets and Benchmarks Track , 2022.

Alexander Korotin, Lingxiao Li, Aude Genevay, Justin M Solomon, Alexander Filippov, and Evgeny Burnaev. Do neural optimal transport solvers work? a continuous wasserstein-2 benchmark. Advances in Neural Information Processing Systems , 34:14593-14605, 2021.

Alexander Korotin, Lingxiao Li, Justin Solomon, and Evgeny Burnaev. Continuous wasserstein2 barycenter estimation without minimax optimization. In International Conference on Learning Representations , 2021.

Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. Kernel neural optimal transport. In The Eleventh International Conference on Learning Representations , 2023.

Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. Neural optimal transport. In The Eleventh International Conference on Learning Representations , 2023.

Christian Léonard. A survey of the schrödinger problem and some of its connections with optimal transport. arXiv preprint arXiv:1308.0215 , 2013.

Alex Tong Lin, Samy Wu Fung, Wuchen Li, Levon Nurbekyan, and Stanley J Osher. Alternating the population and control neural networks to solve high-dimensional stochastic mean-field games. Proceedings of the National Academy of Sciences , 118(31):e2024713118, 2021.

Guan-Horng Liu, Tianrong Chen, Oswin So, and Evangelos Theodorou. Deep generalized schrödinger bridge. Advances in Neural Information Processing Systems , 35:9374-9388, 2022.

Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of International Conference on Computer Vision (ICCV) , December 2015.

Guansong Lu, Zhiming Zhou, Jian Shen, Cheng Chen, Weinan Zhang, and Yong Yu. Largescale optimal transport via adversarial training with cycle-consistency. arXiv preprint arXiv:2003.06635 , 2020.

Andreas Lugmayr, Martin Danelljan, and Radu Timofte. Ntire 2021 learning the super-resolution space challenge. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 596-612, 2021.

Ashok Makkuva, Amirhossein Taghvaei, Sewoong Oh, and Jason Lee. Optimal transport mapping via input convex neural networks. In International Conference on Machine Learning , pages 6672-6681. PMLR, 2020.

François-Pierre Paty, Alexandre d'Aspremont, and Marco Cuturi. Regularity as regularization: Smooth and strongly convex brenier potentials in optimal transport. In International Conference on Artificial Intelligence and Statistics , pages 1222-1232. PMLR, 2020.

Michele Pavon and Anton Wakolbinger. On free energy, stochastic control, and schrödinger processes. In Modeling, Estimation and Control of Systems with Uncertainty: Proceedings of a Conference held in Sopron, Hungary, September 1990 , pages 334-348. Springer, 1991.

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

Litu Rout, Alexander Korotin, and Evgeny Burnaev. Generative modeling with optimal transport maps. In International Conference on Learning Representations , 2022.

Filippo Santambrogio. Optimal transport for applied mathematicians. Birkäuser, NY , 55(5863):94, 2015.

Erwin Schrödinger. Über die umkehrung der naturgesetze . Verlag der Akademie der Wissenschaften in Kommission bei Walter De Gruyter u. Company., 1931.

Vivien Seguy, Bharath Bhushan Damodaran, Remi Flamary, Nicolas Courty, Antoine Rolet, and Mathieu Blondel. Large scale optimal transport and mapping estimation. In International Conference on Learning Representations , 2018.

Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusion schrödinger bridge matching. In Thirty-seventh Conference on Neural Information Processing Systems , 2023.

Christian Szegedy, Wei Liu, Yangqing Jia, Pierre Sermanet, Scott Reed, Dragomir Anguelov, Dumitru Erhan, Vincent Vanhoucke, and Andrew Rabinovich. Going deeper with convolutions. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 1-9, 2015.

Francisco Vargas, Pierre Thodoroff, Austen Lamacraft, and Neil Lawrence. Solving schrödinger bridges via maximum likelihood. Entropy , 23(9):1134, 2021.

Cédric Villani. Optimal transport: old and new , volume 338. Springer Science & Business Media, 2008.

Gefei Wang, Yuling Jiao, Qian Xu, Yang Wang, and Can Yang. Deep generative learning via schrödinger bridge. In International Conference on Machine Learning , pages 10794-10804. PMLR, 2021.

Yujia Xie, Minshuo Chen, Haoming Jiang, Tuo Zhao, and Hongyuan Zha. On scalable and efficient computation of large scale optimal transport. volume 97 of Proceedings of Machine Learning Research , pages 6882-6892, Long Beach, California, USA, 09-15 Jun 2019. PMLR.

A Extended background: KL divergence with the Wiener process plan

This section illustrates that (7) holds. Consider a process T ∈ F ( P 0 ) , i.e., T is a probability distribution on Ω with the marginal P 0 at t = 0 .

Let W ϵ be the Wiener process with variance ϵ starting at P 0 , i.e., it satisfies dX t = √ ϵdW t with X 0 ∼ P 0 . Hence, π W ϵ ( y | x ) is the normal distribution dπ W ϵ ( y | x ) dy = N ( y | x, ϵI ) . Then KL ( π T || π W ϵ ) between joint distributions at times t = 0 and t = 1 of these processes is given by:

where dπ ( x,y ) d [ x,y ] denotes the joint density of distribution π . We derive

After substituting this result into (18), one obtains:

B Proofs

In this section, we provide the proof for our main theoretical results (Theorems 4.1 and 4.3). The proofs require several auxiliary results which we formulate and prove in M B.1 and M B.2.

In M B.1, we show that entropic OT can be reformulated as a maximin problem. This is a technical intermediate result needed to derive our main maximin reformulation of SB (Theorem 4.1). More precisely, in M B.2, we show that these maximin problems for entropic OT and SB are actually equivalent. By using this observation and related facts, in M B.3, we prove our Theorems 4.1 and 4.3.

B.1 Relaxation of entropic OT

To begin with, we recall some facts regarding EOT and SB. Recall the definition of EOT (2):

Henceforth, we assume that P 0 and P 1 are absolutely continuous . The situation when P 0 or P 1 is not absolutely continuous is not of any practical interest: there is no π ∈ Π( P 0 , P 1 ) for which the differential entropy H ( π ) is finite which means that (20) equals to + ∞ for every π ∈ Π( P 0 , P 1 ) , i.e., every plan is optimal. In turn, when P 0 and P 1 are absolutely continuous, the OT plan is unique thanks to the strict convexity of entropy (on the set of absolutely continuous plans).

Recall equation (19) for KL ( π || π W ϵ ) :

We again note that

i.e., problems (20) and (21) can be viewed as equivalent as their minimizers are the same. For convenience, we proceed with inf π ∈ Π( P 0 , P 1 ) KL ( π || π W ϵ ) and denote its optimal value by L ∗ , i.e.,

For a given β ∈ C b, 2 ( Y ) , we define an auxiliary joint distribution dπ β ( x, y ) = dπ β ( y | x ) d P 0 ( x ) , where dπ β ( y | x ) is given by

where C x β ( x ) def = ∫ Y exp( β ( y )) dπ W ϵ ( y | x ) . Note that C x β < ∞ since β ∈ C b, 2 ( Y ) is upper bounded. Before going further, we need to introduce several technical auxiliary results.

Proposition B.1. For ν ∈ P 2 ( Y ) and x ∈ X it holds that

Proof of Proposition B.1. We derive

Lemma B.2. For π ∈ Π( P 0 ) , i.e., probability distributions π ∈ P 2 ( X × Y ) whose projection to X equals P 0 , we have

Proof of Lemma B.2. For each x ∈ X , we substitute ν = π ( ·| x ) to (22) and integrate over x ∼ P 0 . For the left part, we obtain the following:

For the right part, we obtain:

Hence, the equality (23) holds.

Now we introduce the following auxiliary functional ˜ L :

Recall that π 1 denotes the second marginal distribution of π . We use this functional to derive the saddle point reformulation of EOT.

Lemma B.3 (Relaxation of entropic optimal transport) . It holds that

where sup is taken over potentials β ∈ C b, 2 ( Y ) and inf over π ∈ Π( P 0 ) .

Proof of Lemma B.3. We obtain

where C ( x, ν ) def = KL ( ν || π W ϵ ( y | x ) ) . The last problem in (25) is known as weak OT [7, 20] with a weak OT cost C . For a given β ∈ C b, 2 ( Y ) , consider its weak C -transform given by:

Since C : X × P 2 ( Y ) → R is lower bounded (by zero), convex in the second argument and jointly lower semi-continuous, the following equality holds [7, Theorem 1.3]:

where sup is taken over β ∈ C b, 2 ( Y ) . We use our Proposition B.1 to note that

This allows us to derive

Here in transition to line (29), we use our Lemma B.2. It remains to take sup β in equality between (28) and (29) and then recall (27) to finish the proof and obtain desired (24).

Thus, we can obtain the value L ∗ (8) by solving maximin problem (24) with only one constraint π ∈ Π( P 0 ) . Moreover, our following lemma shows that in all optimal pairs ( β ∗ , π ∗ ) which solve maximin problem (24), π ∗ is necessary the unique entropic OT plan between P 0 and P 1 .

Lemma B.4 (Entropic OT plan solves the relaxed entropic OT problem) . Let π ∗ be the entropic OT plan between P 0 and P 1 . For every optimal β ∗ ∈ argsup β inf π ∈ Π( P 0 ) L ( β, π ) , we have

Proof of Lemma B.4. Since β ∗ is optimal, we know from Lemma B.3 that inf π ∈ Π( P 0 ) L ( β ∗ , π ) = L ∗ . Thanks to π ∗ ∈ Π( P 0 , P 1 ) , we have π ∗ 1 = P 1 . We substitute π ∗ to L ( β ∗ , π ) and obtain

The functional π ↦→L ( β ∗ , π ) is strictly convex (in the convex subset of Π( P 0 ) of distributions π for which KL ( π || π W ϵ ) < ∞ ). Thus, it has a unique minimizer, which is π ∗ .

From our Lemmas B.3 and B.4 it follows that to get the OT plan π ∗ one may solve the maximin problem (24) to obtain an optimal saddle point ( β ∗ , π ∗ ) . Unfortunately, it is challenging to estimate KL ( π || π W ϵ ) from samples, which limits the usage of this objective in practice.

B.2 Equivalence of EOT and DSB relaxed problems

Below we show how to relax SB problem (11) and link its solution to the relaxed entropic OT (24).

For a given β ∈ C b, 2 ( Y ) , we define an auxiliary process T β such that its conditional distributions are T β | x,y = W ϵ | x,y and its joint distribution π T β at t = 0 , 1 is given by π β .

To simplify many of upcoming formulas, we introduce C β def = ∫ X log C x β d P 0 ( x ) . Also, we introduce F ( P 0 ) to denote the set of processes starting at P 0 at time t = 0 .

Lemma B.5 (Inner objectives of relaxed EOT and SB are KL with T β and π T β ) . For π ∈ Π( P 0 ) and T ∈ F ( P 0 ) , the following equations hold:

Note that the last two terms in each line depend only on β but not on π or T .

Proof of Lemma B.5. The first equation (32) directly follows from Lemma B.2. Now we prove (33):

In the transition to line (34), we use the disintegration formula (6). In line (35), we use the definition of T β , i.e., we exploit the fact that T β | x,y = W ϵ | x,y and again use (6).

As a result of Lemma B.5, we obtain the following important corollary.

Corollary B.6 (The solution to the inner problem of relaxed SB is a diffusion) . Consider the problem

Then T β is the unique optimizer of (36) and it holds that T β ∈ D ( P 0 ) , i.e., it is a diffusion process:

Proof. Thanks to (33), we see that T β is the unique minimizer of (36). Now let Q def = π T β 1 . Then

In transition to (38), we use the fact that the process solving the Schrödinger Bridge (this time between P 0 and Q ) with the Wiener Prior is a diffusion process (see Dynamic SB problem in M 2.2 for details). As a result, we obtain T β ∈ D ( P 0 , Q ) ⊂ D ( P 0 ) and finish the proof.

Below we show that for a given β , minimization of the SB relaxed functional L ( β, T f ) over T f is equivalent to the minimization of relaxed EOT functional ˜ L ( β, π ) (24) with the same β .

Lemma B.7 (Equivalence of the inf values of the relaxed functionals) . It holds that

Moreover, the unique minimizers are given by T β ∈ D ( P 0 ) and π T β ∈ Π( P 0 ) , respectively.

Proof of Lemma B.7. Follows from Lemma B.5 and Corollary B.6.

Finally, we see that both the maximin problems are equivalent.

Corollary B.8 (Equivalence of EOT and DSB maximin problems) . It holds that

Proof of Corollary B.8. We take sup β of both parts in equation (39).

Also, it follows that the maximization of inf T f ∈D ( P 0 ) L ( β, T f ) over β allows to solve entropic OT.

B.3 Proofs of main results

Finally, after long preparations, we prove our main Theorem 4.1.

Proof of Theorem 4.1 and Corollary 4.2. From our Lemma B.7 and Corollary B.8 it follows that

i.e., both maximin problems share the same optimal β ∗ . Thanks to our Lemma B.7, we already know that the process T β ∗ ∈ D ( P 0 ) and the plan π T β ∗ ∈ Π( P 0 ) are the unique minimizers of problems

respectively. Therefore, T f ∗ = T β ∗ and, in particular, π T f ∗ = π T β ∗ . Moreover, since ( β ∗ , π T f ∗ ) is an optimal saddle point for ˜ L , from Lemma B.4 we conclude that π T f ∗ = π ∗ , i.e., π T f ∗ is the EOT plan between P 0 and P 1 . In particular, π T f ∗ ∈ Π( P 0 , P 1 ) which also implies that T f ∗ ∈ D ( P 0 , P 1 ) . The last step is to derive

which concludes that T f ∗ is the solution to SB (5).

Proof of Theorem 4.3. Part 1. From Lemma B.5 and Corollary B.6 it follows that that inf T f L ( ˆ β, T f ) has the unique minimizer T ̂ β whose conditional distributions are T ̂ β | x,y = W ϵ | x,y . Therefore,

Part 2. Now we consider ϵ 2 . We know that

From Lemma B.5 and Corollary B.6, we also know that

Therefore:

Thus, we obtain ϵ 2 = KL ( T f ∗ || T ̂ β ) .

Part 3. By summing (41) and (42) and using the Pinsker inequality, we obtain

Here we use the triangle inequality in line (43). Therefore, ρ TV ( T ˆ f , T f ∗ ) ϵ 1 + ϵ 2 . Part 4. By summing (41) and (42) and using the Pinsker inequality, we obtain

Thus, TV 1 2

C Euler-Maruyama

In our Algorithm 1, at both the training and the inference stages, we use the Euler-Maruyama Algorithm 2 to solve SDE.

D Drift Norm Constant Multiplication Invariance

Our Algorithm 1 aims to solve the following optimization problem:

with C = 1 . At the same time, we use C = 1 2 ϵ in our theoretical derivations (12). We emphasize that the actual value of C > 0 does not affect the optimal solution T f ∗ to this problem. Specifically, if

Algorithm 2: Euler-Maruyama algorithm

( β ∗ , T f ∗ ) is the optimal point for the problem with C = 1 , then ( ˜ Cβ ∗ , T f ∗ ) is the optimal point for C = ˜ C . Indeed, for a pair ( β, T f ) it holds that

where we use ˜ β def = ˜ Cβ . Hence problems sup β inf T f L 1 ( β, T f ) and sup ˜ β inf T f L ˜ C ( ˜ β, T f ) can be viewed as equivalent in the sense that one can be derived one from the other via the change of variables and multiplication by ˜ C > 0 . For completeness, we also note that the change of variables β ↔ ˜ β actually preserves the functional class of β , i.e., β ∈ C b, 2 ( Y ) ⇐⇒ ˜ β ∈ C b, 2 ( Y ) .

For convenience, we get rid of dependence on ϵ in the objective (12) and consider L 1 for optimization, i.e., use C = 1 in Algorithm 1. Still the dependence on ϵ remains in sup β inf T f L 1 ( β, T f ) as T f ∈ D ( P 0 ) is a diffusion process with volatility ϵ . Interestingly, this point of view (optimizing L 1 instead of L 1 2 ϵ ) technically allows to consider even ϵ = 0 . In this case, the optimization is performed over deterministic trajectories T f determined by the velocity field f ( X t , t ) . The problem sup β inf T f L 1 ( β, T f ) may be viewed as a saddle point reformulation of the unregularized OT with the quadratic cost in the dynamic form, also known as the Benamou-Brenier formula [43, M 6.1]. This particular case is out of scope of our paper (it is not EOT/SB) and we do not study the properties of L 1 in this case. However, for completeness, we provide experimental results for ϵ = 0 .

E ENOT for Toy Experiments and High-dimensional Gaussians

In 2D toy experiments, we consider 2 tasks: Gaussian → 8 gaussians and Gaussian → Swiss roll . Results for the last one (Figure 5) are qualitatively similar to results of the first one (Figure 2), which we discussed earlier ( M 5.1). For both tasks, we parametrize the SDE drift function in Algorithm 1 by a feedforward neural network f θ with 3 inputs, 3 linear layers (100 hidden neurons and ReLU activations) and 2 outputs. As inputs, we use 2 coordinates and time value t (as is). Analogically, we parametrize the potential by a feedforward neural network β ϕ with 2 inputs, 3 linear layers (100 hidden neural and ReLU activations) and 2 outputs. In all the cases, we use N = 10 discretization steps for solving SDE by Euler-Maruyama Algorithm 2, Adam with lr = 10 -4 , batch size 512. We train the model for 20000 total iterations of β ϕ , and on each of them, we do K f = 10 updates for the SDE drift function f θ .

(a) Input and target samples

(b) ENOT

(ours)

,

ϵ

= 0

(c) ENOT

(ours)

,

ϵ

= 0

.

01

(d) ENOT

(ours)

,

ϵ

= 0

.

1

2

0

CInputdistribution

2

Fitted distribution

2

Fitted distribution

2

Fitted distribution

2

-1

0

-1

0

-1

0

Trajectories

Trajectories

Trajectories

2

Targetdata

Targetdata

Targetdata

Q

2

Targetdistribution

-2

-1

0

1

2

In the experiments with high-dimensional Gaussians, we use exactly the same setup as for toy 2D experiments but chose N = 200 discretization steps for SDE, all hidden sizes in neural networks are 512, and we train our model for 10000 iterations. To illustrate the stability of the algorithm, we provide the plot of BW 2 2 -UVP (%) between the ground truth EOT plan π ∗ and the learned plan π of ENOT during training for DIM = 128 in Figure 6.

ENOT (ours) BW-UVP plan during the training

3.0

2.5

BW-UVPplan

2.0

1.5

1.0

0.5

0.0

0

2000

4000

6000

8000

10000

Outerproblemsteps

F ENOT for Colored MNIST and Unpaired Super-resolution of Celeba Faces

For the image tasks ( M 5.3, M 5.4), we find out that using the following reparametrization of EulerMaruyama Algorithm 2 considerably improves the quality of our Algorithm 1. In the Euler-Maruyama Algorithm 2, instead of using a neural network to parametrize drift function f ( X t , t ) and calculating the next state as X t +1 = X t + f ( X t , t )∆ t + √ ϵ ∆ t , we parametrize g ( X t , t ) = X t + f ( X t , t )∆ t by a neural network g θ , and calculate the next state as X t +1 = g θ ( X t , t ) + √ ϵ ∆ t . In turn, the drift function is given by f ( X t , t ) = 1 ∆ t g ( X t , t ) -X t . Also, we do not add noise at the last step of the Euler-Maruyama simulation because we find out that it provides better empirical performance.

X~Po t=0.1 t=0.2 t=0.3 t=0.4 t=0.5 t=0.6 t=0.7t=0.8 t=0.9 t=1.0

0

2

1233]

333

3

II

3

L

3

=

3

10

3

=

3

We use WGAN-QC discriminator's ResNet architecture 5 for the potential β . We use UNet 6 as g θ ( X t , t ) of SDE in our model. To condition it on t , we first obtain the embedding of t by using the positional embedding 7 . Then we add conditional instance normalization (CondIN) layers after each UNet's upscaling block 8 . We use Adam with lr = 10 -4 , batch size 64 and 10:1 update ratio for f θ /β ϕ . For ϵ = 0 and ϵ = 1 our model converges in ≈ 20000 iterations, while for ϵ = 10 it takes ≈ 70000 iteration to convergence. The last setup takes more iterations to converge because adding noise with higher variance during solving SDE by Euler-Maruyama Algorithm 2 increases the variance of stochastic gradients.

In the unpaired super-resolution of Celeba faces, we use the same experimental setup as for the colored MNIST experiment. It takes ≈ 40000 iterations for ϵ = 0 and ≈ 70000 iterations for ϵ = 1 and ϵ = 10 to converge. In Figures 7, 1 we present trajectories provided by our algorithm for Colored MNIST and Celeba experiments.

Computational complexity. In the most challenging task ( M 5.4), ENOT converges in one week on 2 × A100 GPUs.

G Details of the baseline methods

In this section, we discuss details of the baseline methods with which we compare our method.

G.1 Gaussian case ( M 5.2).

SCONES [14]. We use the code from the authors' repository https://github.com/mdnls/scones-synthetic

for their evaluation in the Gaussian case. We employ their configuration blob/main/config.py .

LSOT [45]. We use the part of the code of SCONES corresponding to learning dual OT potentials blob/main/cpat.py and the barycentric projection blob/main/bproj.py in the Gaussian case with configuration blob/main/config.py .

FB-SDE-J [10]. We utilize the official code from https://github.com/ghliu/SB-FBSDE

with their configuration blob/main/configs/default_checkerboard_config.py for the checkerboard-to-noise toy experiment, changing the number of steps of dynamics from 100 to 200 steps. Since their hyper-parameters are developed for their 2-dimensional experiments, we increase the number of iterations for dimensions 16, 64 and 128 to 15 000.

FB-SDE-A [10]. We also take the code from the same repository as above. We base our configuration on the authors' one ( blob/main/configs/default_moon_to_spiral_config.py ) for the moon-to-spiral experiment. As earlier, we increase the number of steps of dynamics up to 200. Also, we change the number of training epochs for dimensions 16, 64 and 128 to 2,4 and 8 correspondingly.

DiffSB [15] . We utilize the official code from

with their configuration blob/main/conf/dataset/2d.yaml for toy problems. We increase the amount of steps of dynamics to 200 and the number of steps of IPF procedure for dimensions 16, 64 and 128 to 30, 40 and 60, respectively.

MLE-SB [48]. We use the official code from https://github.com/franciscovargas/GP_Sinkhorn

with hyper-parameters from blob/main/notebooks/2D Toy Data/2d_examples.ipynb . We set the number of steps to 200. As earlier, we increase the number of steps of IPF procedure for dimensions 16, 64 and 128 to 1000, 3500 and 5000, respectively.

G.2 Colored MNIST ( M 5.3)

SCONES [14]. In order to prepare a score-based model, we use the code from

https://github.com/ermongroup/ncsnv2

with their configuration blob/master/configs/cifar10.yml . Next, we utilize the code of SCONES from the official repository for their unpaired Celeba super-resolution experiment ( blob/main/scones/configs/superres_KL_0.005.yml ). We adapt it for 32 × 32 ColorMNIST images instead of 64 × 64 celebrity faces.

DiffSB [15]. We use the official code with their configuration blob/main/conf/mnist.yaml adopting it for three-channel ColorMNIST images instead of one-channel MNIST digits.

G.3 CelebA ( M 5.4)

SCONES [14]. For the SCONES, we use their exact code and configuration from blob/main/scones/configs/superres_KL_0.005.yml . As for the score-based model for celebrity faces, we pick the pre-trained model from

https://github.com/ermongroup/ncsnv2

It is the one used by the authors of SCONES in their paper.

Augmented Cycle GAN [2]. We use the official code from

https://github.com/NathanDeMaria/AugmentedCycleGAN

with their default hyper-parameters.

ICNN [38]. We utilize the reworked implementation by https://github.com/iamalexkorotin/Wasserstein2Benchmark .

which is a non-minimax version [26] of ICNN-based approach [38]. That is, we use blob/main/notebooks/W2_test_images_benchmark.ipynb and only change the dataloaders.

H Mean-Field Games

This appendix discusses the relation between the Mean-Field Game problem and Schrödinger Brdiges.

H.1 Intro to the Mean-Field game.

Consider a game with infinitely many small players. At time moment t = 0 , they are distributed according to X 0 ∼ ρ 0 . Every player controls its behavior through drift α of the SDE:

Here ρ t is the distribution of all the players at the time moment t . When we consider a specific player, we consider ρ t as a parameter. Each player aims to minimize the quantity:

Here L ( x, α, ρ ) is similar to the Lagrange function in physics and describes the cost of moving in some direction given the current position and the other players' distribution. The additional function f ( X t , ρ t ) is interpreted as the cost of the player's interaction at coordinate x with all the others. Now we can introduce the value function ϕ ( x, t ) , which for position x and start time t returns the cost in case of the optimal control:

Before considering the Mean-Field game, we need to define an additional function H ( x, p, ρ ) . It is similar to the Hamilton function and is defined as the Legendre transform of Lagrange function L :

Mean-Field game implies finding the Nash equilibrium for all players of such the game. It is known [1] that the Nash equilibrium is the solution of the system of Hamilton-Jacobi-Bellman (HJB) and Fokker-Planck (FP) PDE equations. For two functions H ( x, p, ρ ) and f ( x, ρ ) , Mean-Field game formulates as a system of two PDE with two constraints:

The solution of this system is two functions ρ ( x, t ) and ϕ ( x, t ) , which describe all players' dynamics. Also, in Nash equilibrium, the specific player's behavior is described by the following SDE:

H.2 Relation to our work.

In recent work [34], the authors show that the Schrodinger Bridger problem could be formulated as a Mean-Field game with hard constraints on distribution ρ ( · , T ) = ρ target ( · , T ) via choosing proper function g ( x, ρ ( · , T )) such as:

̸

Also, the authors proposed an extension of DiffSB [34] algorithm for the Mean-Field game problem.

In [33], the authors in their experiments consider only soft constraints on the target density. More precisely, they consider only simple constraints such as g ( x, ρ ) = || x -x target || 2 , where x target is a given shared target point for every player, and every player is penalized for being far from this. Such soft constraint force players to have delta distribution at point x target .

To solve the Mean-Field problem, the authors parameterize value function ϕ ( x, t ) by a neural network and use different neural network N θ to sample from ρ t . The authors penalize the violation of Mean-Field game PDEs for optimizing these networks. After the convergence, one can sample from the distribution ρ t by using neural network N θ . Approach [33] has the advantage that authors do not need to use SDE solvers, which require more steps with growing parameter ν of diffusion operator . However, computation of Laplacian and divergence for high-dimensional spaces (e.g., space 12228-dimensional space of 3x64x64 images) at each iteration of the training step may be computationally hard, restricting the applicability of their method to large-scale setups.

In our approach , we initially work with the SDE:

which describes the player's behavior and use a neural network to parametrize the drift α . We consider only hard constraints on the target distribution, f ( X t , ρ t ) = 0 and L ( X t , α t , ρ t ) = 1 2 || α t || 2 since this variant of Mean-Field game is also the particular case of Schrodinger Bridge problem and is equivalent to the entropic optimal transport. Since we do not need to compute Laplacian or divergence, our approach scales better with the dimension . However, for high values of diffusion parameter ν (which is equal to the 1 2 ϵ in our notation, where ϵ is the entropic regularization strength), our approach needs more steps for accurate solving of the SDE to provide samples, as we mentioned in limitations.

I Extending ENOT to other costs

In the main text, we focus only on EOT with the quadratic cost c ( x, y ) = 1 2 ∥ x -y ∥ 2 which coincides with SB with the Wiener prior W ϵ . However, one could use a different prior Q v instead of W ϵ in (5):

and solve the problem

Here we just use the known expression (6) for KL ( T f || Q v ) between two diffusion processes through their drift functions. Using the same derivation as in the main text M 2.2, it can be shown that this new problem is equivalent to solving the EOT with cost c ( x, y ) = -log π Q v ( y | x ) , where π Q v ( y | x ) is a conditional distribution of the stochastic process Q v at time t = 1 given the starting point x at time t = 0 . For example, for W ϵ (which we consider in the main text) we have

i.e., we get the quadratic cost. Thus, using different priors for the Schrodinger bridge problem makes it possible to solve Entropic OT for other costs. We conjecture that most of our proofs and derivations can be extended to arbitrary prior process Q v just by slightly changing the minimax functional (12):

We conduct a toy experiment to support this claim and consider Q v with ϵ = 0 . 01 and v ( x, t ) = ∇ log p ( x ) , where log p ( x ) is a 2D distribution with a wave shape, see Figure 8. Intuitively, it means that trajectories should be concentrated in the regions with a high density of p . In Figure e 8, there the grey-scale color map represents the density of p , start points ( P 0 ) are green, target points ( P 1 ) are red, obtained trajectories are pink and mapped points are blue.

J ENOT for the unregularized OT ( ϵ = 0 )

Our proposed algorithm is designed to solve entropic OT and the equivalent SB problem. This implies that ϵ > 0 . Nevertheless, our algorithm technically allows using even ϵ = 0 , in which case it presumably computes the unegularized OT map for the quadratic cost. Here we present some empirical evidence supporting this claim as well some theoretical insights.

EMPIRICAL EVIDENCE. We consider the experimental setup with images from the continuous Wasserstein-2 benchmark [28, M 4.4]. The images benchmark provides 3 pairs of distributions (Early, Mid, Late) for which the ground truth unregularized OT map for the quadratic cost is known by the construction. Hence, we may compare the map learned with our method ( ϵ = 0 ) with the true one.

We train our method with ϵ = 0 on each of 3 benchmark pairs and present the quantitative results in Table 6. We use the same L 2 -UVP metric [28, M 4.2] as the authors of the benchmark. As the baselines, we include the results of MM:R method from [28] and the method from [3]. Both methods are minimax and have some similarities with our approach. As we can see, ENOT with ϵ = 0

2.0

startpoints

finish points

1.5

targetpoints

1.0

0.5

0.0-

0.5

1.0

1.5

-2.0

-3.0

2.5

2.0

1.5

1.0

0.5

0.0

0.5

1.0

works better than the MM:R solver but slightly underperforms compared to [3]. This evaluation demonstrates that our method recovers the unregularized OT map for the quadratic cost with the comparable quality to the existing saddle point OT methods.

THEORETICAL INSIGHTS. We see that empirically our method with ϵ = 0 recovers the unregularized OT map. At the same time, this is not supported by our theoretical results as they work exclusively for ϵ > 0 and rely on the properties of the KL divergence.

Overall, it seems like for ϵ = 0 our method yields a saddle point reformulation of the BenamouBrenier (BB) [8] problem which is also known as the dynamic version of the unregularized OT ( ϵ = 0 ) with the quadratic cost. This problem can be formulated as follows:

i.e., the goal is to find the process T f of the minimal energy which moves the probability mass of P 0 to P 1 . BB (45) is very similar to DSB (11) but there is no multiplier 1 ϵ , and the stochastic process T f is restricted to be deterministic ( ϵ = 0 ). It is governed by a vector field f . Just like the DSB (11) is equivalent to EOT (2), it is known that BB (45) is equivalent to unregularized OT with the quadratic cost ( ϵ = 0 ). Namely, the distribution π T f ∗ is the unregularized OT plan between P 0 and P 1 .

In turn, our Algorithm 1 for ϵ = 0 optimizes the following saddle point objective:

where T f : dX t = f ( X t , t ) dt with X 0 ∼ P 0 (the constraint X 1 ∼ P 1 here is lifted) and β ∈ C 2 ,b ( Y ) . Just like in the Entropic case, functional L can be viewed as the Lagrangian for BB (45) with β playing the role of the Lagrange multiplier for the constraint dπ T f 1 ( y ) = d P 1 ( y ) . Naturally, it is expected that the value (45) coincides with (46), and we provide a sketch of the proof of this fact.

Overall, the proof logic is analogous to the Entropic case but the actual proof is much more technical as we can not use the KL -divergence machinery which helps to avoid non-uniqueness, etc.

Step 1 (Auxiliary functional, analog of Lemma B.3). We introduce an auxiliary functional

where β is a potential and H : R D → R D is a measurable map. This functional is nothing but the well-known max-min reformulation of static OT problem (in Monge's form) with the quadratic cost [3, Eq. 4], [28, Eq.9]. Hence,

Step 2 (Solution of the inner problem is always an OT map). An existence of some minimizer H = H β in inf H ˜ L ( β, H ) can be deduced from the measurable argmin selection theorem, e.g., [21, Theorem 18.19]. For this H β we consider P ′ def = H β ♯ P 0 . Recall that

Here we may add the fictive constraint H♯ P 0 = P ′ which is anyway satisfied by H β and get

The last equality holds since ∫ β ( H ( x )) d P 0 ( x ) = ∫ β ( y ) d P ′ ( y ) does not depend on the choice of H due to the constraint H♯ P 0 = P ′ . The latter is the OT problem between P and P ′ and we see that H β is its solution.

Step 3 (Equivalence for inner objective values). Since H β is the OT map between P 0 , P ′ (it is unique as P 0 is absolutely continuous [43]), it can be represented as an ODE solution T f β to the Benamour Brenier problem between P 0 , P ′ , i.e., T f β : dX t = f β ( X t , t ) dt and H β ( X 0 ) = X 0 + ∫ 1 0 f β ( X t , t ) dt . Furthermore, in this case, ∥ X 0 -H β ( X 0 ) ∥ 2 = ∫ 1 0 || f β ( X t , t ) || 2 dt . Hence, it can be derived that

Step 4 (Equivalence of the saddle point objective). Take sup over β ∈ C b, 2 ( Y ) and get the final equivalence:

Step 5 (Dynamic OT solutions are contained in optimal saddle points). Pick any optimal β ∗ ∈ argsup β inf H L ( β, T f ∗ ) and let T f ∗ be any solution to the Benamou-Brenier problem. Checking that T ∗ ∈ inf H L ( β ∗ , T f ) can be done analogously to [31, Lemma 4], [42, Lemma 4.1]. □

The derivation above shows the equivalence of objective values of dynamic unregularized OT (45) and our saddle point reformulation of BB (45). Additionally, it shows that solutions T f ∗ can be recovered from some optimal saddle points ( β ∗ , T f ∗ ) of our problem. At the same time, unlike the EOT case ( ϵ > 0 ), it is not guaranteed that for all the optimal saddle points ( β ∗ , T f ∗ ) it holds that T f ∗ is the solution to the BB problem. This aspect seems to be closely related to the fake solutions issue in the saddle point methods of OT [30] and may require further studies.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes to solve dynamic entropic optimal tansport (EOT), also known as Schrödinger bridge problem, with nerual solver. Specifically, the authors propose a saddle-point, maximin, formulation of EOT, yielding a GAN-resemble algorithm that can be trained in an end-to-end fashion. Experiments are conducted on 2D toy datasets and images translation.\n\nSTRENGTHS:\n- The saddle point reformulation of EOT is interesting. The proposed algorithm can be viewed as a stochastic control problem, where the terminal cost function $\\beta$ is learned so that the resulting policy approaches terminal distribution. \n- Writing is generally clear and easy to follow. Sufficient related works and preliminaries are included.\n- Experiment are extensive and include many related baselines.\n\nWEAKNESSES:\n- The proposed algorithm aims to minimize (12), which is not well-defined when $\\epsilon$=0, as the KL term will blow up. Even though $\\epsilon$=0 is algorithmically applicable, it makes the current algorithm disconnected from the mathematical framework. I'll be more convinced if the authors can provide additional justification (maybe connection to OT).\n\n- The proposed method is closely related to recent maximim OT [1], which consists of the same two networks (potential + policy) and the same training losses. From my understanding, the two methods coincide when $\\epsilon$=0 and $N=1$. Are the authors aware of [1]? Given that the proposed ENOT seems to work best when $\\epsilon$=0, can the author compare with [1]?\n\n- Given that DiffSB [13] was compared throughout most Sec 5, I suggest the authors to compare in Sec 5.4 to [2], which applies DiffSB to unpaired super-resolution image datasets.\n\n- While Sec 2 has introduced sufficient background and comparison between EOT and SB, which I do appreciate, I think their connection to Sec 4, which is the main content, is rather weak. Given that the proposed algorithm closely relates to dual formulation (e.g. (26) in Appendix), I suggest including those parts in Sec 2.\n\n[1] On amortizing convex conjugates for optimal transport  \n[2] Conditional Simulation Using Diffusion Schrödinger Bridges\n\nQUESTIONS:\n- Experiments seem to support that $\\epsilon$=0, where the dynamics reduce to flow, works much better than SDE. This seems to contradict recent insights from diffusion models [1], where SDE usually performs better than ODE on the same tasks. Can the authors comment on that?\n- why is there a \"1/|f|\" term in Alg 1, when computing the KL? And what's $f_{n,m}$? I understand tht $f_n$ is the drift output at time step $n$. \n\n[1] Elucidating the Design Space of Diffusion-Based Generative Models\n\nLIMITATIONS:\nLimitations are included in Sec6, where the authors mentioned the potential computational burden caused by simulation and back-prop through the SDE dynamics during training.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nInspired by how Sinkhorn duals are derived the authors adapt said derivation to the path measure and via the disintegration theorem they derive a novel unconstrained min-max objective for solving the Schrodinger bridge problem, the authors then proceed to showcase their method in eOT based tasks, introducing a new gaussian benchmark and displaying competitive results to previous approaches.  Additionally, the authors quantify errors in sampling and transport of approximate minimisers to their proposed schemes.\n\nSTRENGTHS:\n1. The extension of the Sinkhorn dual to the dynamic setting is rather elegant and certainly novel in the way it is carried out \n2. The paper is excellent on the presentation side in regards to technical ideas, whilst the contributions are novel/creative they are presented in such a way that understanding them was not difficult.\n3. The new formulation allows for a novel duality gap analysis, which is one of the few works analysing learned SBP methods in the approximate setting. \n4. From a purely conceptual viewpoint the work is great and rather complete, just some clarifications/additions on the experimental side and motivations could be enhanced.\n\nWEAKNESSES:\nOutside of a concern in how different methods are compared (detailed in the questions), this paper is overall well written and has mostly sound experiments. \n\nFrom the method standpoint, there are several potential weaknesses and lacking ablations which I will detail in the limitations.\n\nQUESTIONS:\n1. Some of the methods you compare to are closed form / non-iterative in the way they solve the half bridges e.g. [1] using GPs.  A thing to note is iterative approaches such as SGD can keep resampling from P_0 and P_1 in the toy tasks nonstop, effectively making the dataset set size somewhat proportional to the iterations, this strongly benefits generalisation capabilities, especially in higher dimensions in contrast to a the GP approach in [1] which is likely fitted on a small dataset. If we assume that the dataset is fixed across epochs K_f * batchsize would give a proxy as to the size of the dataset used in the DL approaches , from inspecting the code provided in the supplementary zip in particular the high dim Gaussians and toy examples IPYNB this quantity seems to be above 5000, whilst I suspect for [1] it might be lower. Even then I suspect you resample from the toy distributions (Gaussians) every time as indicated in your pseudocode in which case the dataset sizes are simply not comparable at all and the non-iterative approaches like MLE-SB are prone to be affected significantly by generalisation error (which scales slowly in number of samples for kernel methods, which also don't overcome the curse of dim as well as neural networks)  in particular in high dimensions.\n2. Are the same number of timesteps used across the DL and the GP approaches when training ? computationally it feels unlikely the gram matrix would fit in memory for the given timesteps used with the DL approaches. \n3. A suggestion here would be to compare all approaches in an additional table (e.g. in the appendix) under budgets in the number of steps and number of samples (and fairly ensure these are the same across methods), this can be useful to quantify the data efficiency of each approach. Furthermore, the settings under which [1] was run (steps and dataset size) should be reported.\n\nLIMITATIONS:\nFinding saddlepoints / solving min-max problems is typical quite a challenging task and was often a huge challenge in the training stability of GANs (very high variance training results). In contrast to IPF which is much nicer from an objective viewpoint (sequence of regression losses), this approach poses the question of how stable training is.\n\nAs none of the experiments have error bars (on training runs) it is difficult to see if the proposed approach is robust and “easy” to train.  I would suggest to the authors provide such results and in addition maybe comparisons to training loss with Chen 2021 joint or De Bortoli 2021.", "TITLE": "Official Review", "RECOMMENDATION": 9, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper focuses on neural optimal transport  and more particularly through a \"dynamic schrodinger bridge\" approach. I am not an expert in this particular topic, but I must say that the authors manage to make it quite readable and a good introduction to the methodology.\n\nAs far as I can tell, the particularity of the approach is to start from a known key constrained optimization problem in eq (11) that has been presented before and that is progressively explained. Then, departing from previous methods, the authors propose a Lagrangian approach, for which we now have an unconstrained problem and the need to train not only the drift function that allows sampling, but also some \"beta network\" (coined in as the beta network), in a min-max optimization approach, that is similar --but different -- from the usual GAN methodology.\n\nAfter the rigorous theoretical treatment, the authors present some very nice experiment that seem to support their claims and the interesting features of their approach. I think the paper is quite challenging, but that it is also very stimulating and inspiring.\n\n-- \nafter rebuttal, I am still happy with the paper. Maintaining my score.\n\nSTRENGTHS:\nI feel slightly uncomfortable in assessing whether the proposed method is new or is not exactly or if the authors miss some particular recent and/or relevant method in their references, mostly because I cannot be considered an expert on the topic. Still, it looks to me that they are doing their best at providing a very objective account of the relevant literature and giving many pointers, so I am assuming they can be trusted when they present their contributions.\n\nThe proposed method seems quite ok to implement (cf Algorithm 1) and the results are very nice. I am not really aware of the appropriate metrics and usual evaluation criteria that should be used, but I felt compelled by the experiments, which made me want to try things out. I guess this is the most important aspect.\n\nTo summarize what I feel are the strengths of the paper:\n- good theoretical overview, motivations and derivations\n- the resulting method seems simple to implement\n- interesting performances\n\nWEAKNESSES:\nThe paper is a bit difficult to follow, but I don't think it should really be modified to be simpler. I guess it's mostly a matter of myself not being an expert in the topic.\n\nStill, some changes here or clarifications here and there could help. I will mention them below.\n\nQUESTIONS:\n- p3 \"Hence, one may optimize (5) over processes T for which T|x,y=W_xy for every x, y and set the last term in (6) to zero\". How do you actually do that ? By enforcing gaussian dynamics ?\n- the paragraph \"continuous OT\" reads a bit weak to me. All the references are just given in a row, without great care for actually discussing them and their connections with the proposed approach.\n\nIn algorithm 1, I have some questions.\n- do you actually need to store the gradients when computing {X_n, f_n} for the computation of L_\\beta ? It looks to me that Eul-Mar(X_0, T_{f_\\theta}}) can actually run in some `no_grad` environment, or am I mistaken ? This would mean there is some possibility in strongly parallelizing this or even use some external workers to compute that ?\n\n- On the contrary, in the inner loop (over k), you really need to store the gradients, right ? Is it necessary to use many inner iterations K_f ? Maybe I'm mistaken but I don't see that discussed, although this looks like a key computational burden when I have a look at algorithm 1, right ? Is it feasible to just record gradient for the last k steps ? for some of them ? Some thing that could help going faster ?\n\n- I have no clue what this BW2-UVP metric is. Could you at least give us a hint that would help us avoid checking reference [25] ?\n\n- You must check the references. Most of them are badly formated. You have \"schr\\\"odinger\" everywhere and you are missing many uppercases for proper nouns\n\nLIMITATIONS:\nI would say that the limitations are clearly mentioned", "TITLE": "Official Review", "RECOMMENDATION": 9, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe main idea of the paper is to estimate a stochastic map for the entropic optimal transport problem using its connection to the dynamic Schrödinger bridge (SDB) problem. The authors formulate the SDB as a saddle point problem of an associated Lagrangian. Then they recover the transport plan as the joint distribution of the solution to the dynamic Schrödinger bridge problem's initial and final values, while the transport is encoded in the drift term of the learned stochastic process that is the solution of the SDB. Compared to previous methods, the method at hand offers more stability for small entropic regularization coefficients. The errors on the drift term solution and on the transport map are quantified given the corresponding duality gaps of the inner and outer optimization problems of the saddle point objective. Finally, the approach is supported by experimental evaluation.\n\nSTRENGTHS:\nThe paper is extremely well-written and reader-friendly.\n\n- Several remarks are made to facilitate reading and to provide intuition about technical notions.\n- A guarantee on the quality of the saddle point solution is provided.\n- Addressing the small $\\epsilon$ case, which is a source of instability of several other methods.\n\nWEAKNESSES:\n* In line 223, it is mentioned that the negative entropy is not strongly convex. This is false as the function $p\\mapsto x\\ln(x)$ has second derivate $x \\mapsto \\frac{1}{x}$ which is bounded from below by $1$ on the interval $[0,1]$. See for example Section 4.1 of reference [2]. As a result, the comparison to [1]  (reference [5] in the paper) needs to be reconsidered.\n\n* There is no concluding section.\n\n* Some experimental section metrics are not introduced.\n  \n  * The $\\rm{BW_2^2-UVP}$ is not introduced even in the appendix, although a reference for it is given. \n  \n  * FID: the previous remark applies.\n\n* The parametrization $g(X_t,t) = X_t + f(X_t,t)\\Delta_t $ should have been indicated in the main paper rather than in the appendix as although it is mathematically equivalent to the parametrization presented in the main paper, it allowed better results on the CelebA dataset according to the appendix\n  \n  ## Minor remarks\n\n* Problem with links: For some reason, the bibliographic references links along with links to equations and sections etc. are not working.\n\n* $\\pi^{W^\\epsilon}$ is introduced for the first time in Equation (8) without being defined.\n\n* $W_{|x,y}$ is not explicitly defined, although one can infer its meaning from the definition of $T_{|x,y}$.\n\n* Theorem 4.1: I think it should be \"every pair $(\\beta^*,T_{f^*})$ for (13)\" rather than \"for (12)\" as problem (13) is a saddle point problem.\n\n* Reference to Algorithm 2: in line 195, Algorithm 2 is referenced. However, it is not indicated that it is written in the appendix.\n\n* Suggestion: index $m$ can be removed in the \"$\\widehat{KL} \\leftarrow$\" line of Algorithm 1  since the sum terms are already indicated to be the values of $f_n$, or it is possible to indicate $\\sum_{m=1}^{|f_n|}$.\n\n* Line 158: I think it should be added \"that is bounded from above\" to \"a continuous function\".\n\n## References\n\n[1] Asadulaev, A., Korotin, A., Egiazarian, V., & Burnaev, E. (2022). Neural optimal transport with general cost functionals. *arXiv preprint arXiv:2205.15403*.\n\n[2] Peyré, G., & Cuturi, M. (2019). Computational optimal transport: With applications to data science. *Foundations and Trends® in Machine Learning*, *11*(5-6), 355-607.\n\nQUESTIONS:\n* Computation of $\\mathbb{E}\\left[\\int_0^1\\Vert f(X_t,t)\\Vert^2{\\rm d}t\\right]$ : in line 197, it is indicated that the mean of the $f(x,t)$ is used. Is this justified by the Riemann integral discrete approximation? If so, does a trapezoidal rule for approximating the integral improve the result?\n* Is it straightforward to generalize the approach to costs other than the squared Euclidean distance?  Does it fundamentally change the nature of the associated stochastic process \n* Did the authors try to apply the method to the domain adaptation (DA) problem as several DA methods rely on optimal transport ?\n----------\nI have read the authors's rebuttal. They have addressed my concerns.\n\nLIMITATIONS:\nThe impact of the paper and the limitations of the contribution are clearly discussed in Section 6.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}