SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Skoltech\*, AIRI†
Dmitry Vetrov
HSE University; AIRI
vetrovd  $\text{巧}$  yandex.ru
bltech\*, AIRI†
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.
$X\sim \mathbb{P}_0$  t=0.1 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
Figure 1: Trajectories of samples learned by our Algorithm 1 for Celeba deblurring with  $\epsilon = 0, 1, 10$ .
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
$^{\ddagger}$ HSE University
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 (§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 (§4.2). Unlike many predecessors, our method allows handling small entropy coefficients. This enables practical applications to data  $\rightarrow$  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 (§4.3). Empirically, we illustrate the performance of the method on several toy and large-scale EOT tasks (§5).
Optimal Transport (OT) and Schrödinger Bridge (SB) problems imply finding an efficient way to transform some initial distribution  $\mathbb{P}_0$  to target distribution  $\mathbb{P}_1$ . While the solution of OT only gives the information about which part of  $\mathbb{P}_0$  is transformed to which part of  $\mathbb{P}_1$ , SB implies finding a stochastic process that describes the entire evolution from  $\mathbb{P}_0$  to  $\mathbb{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  $\mathcal{X}$ ,  $\mathcal{Y}$  and use  $\mathcal{P}_2(\mathcal{X}) = \mathcal{P}_2(\mathcal{Y})$  to denote the respective sets of Borel probability distributions on them which have finite second moment. For two distributions  $\mathbb{P}_0 \in \mathcal{P}_2(\mathcal{X}), \mathbb{P}_1 \in \mathcal{P}_2(\mathcal{Y})$ , consider the following minimization problem:
$$
\inf  _ {\pi \in \Pi \left(\mathbb {P} _ {0}, \mathbb {P} _ {1}\right)} \int_ {\mathcal {X} \times \mathcal {Y}} \frac {\| x - y \| ^ {2}}{2} d \pi (x, y), \tag {1}
$$
where  $\Pi (\mathbb{P}_0,\mathbb{P}_1)\subset \mathcal{P}_2(\mathcal{X}\times \mathcal{Y})$  is the set of probability distributions on  $\mathcal{X}\times \mathcal{Y}$  with marginals  $\mathbb{P}_0$  and  $\mathbb{P}_1$ . Such distributions  $\pi \in \Pi (\mathbb{P}_0,\mathbb{P}_1)$  are called the transport plans between  $\mathbb{P}_0$  and  $\mathbb{P}_1$ . The set  $\Pi (\mathbb{P}_0,\mathbb{P}_1)$  is non-empty as it always contains the trivial plan  $\mathbb{P}_0\times \mathbb{P}_1$ . A minimizer  $\pi^*$  of (1) always exists and is called an OT plan. If  $\mathbb{P}_0$  is absolutely continuous, then  $\pi^{*}$  is deterministic: its conditional distributions are degenerate, i.e.,  $\pi^{*}(\cdot |x) = \delta_{T^{*}(x)}$  for some  $T^{*}:\mathcal{X}\to \mathcal{Y}$  (OT map).
Entropic OT formulation. We use  $H(\pi)$  to denote the differential entropy of distribution  $\pi$  and  $\mathrm{KL}(\pi ||\pi^{\prime})$  to denote the Kullback-Leibler divergence between distributions  $\pi$  and  $\pi^{\prime}$ . Two most popular entropic OT formulations regularize (1) with the entropy  $H(\pi)$  or KL-divergence between plan  $\pi$  and the trivial plan  $\mathbb{P}_0\times \mathbb{P}_1$ , respectively  $(\epsilon >0)$ :
$$
\inf  _ {\pi \in \Pi \left(\mathbb {P} _ {0}, \mathbb {P} _ {1}\right)} \int_ {\mathcal {X} \times \mathcal {Y}} \frac {\| x - y \| ^ {2}}{2} d \pi (x, y) - \epsilon H (\pi), \tag {2}
$$
$$
\inf  _ {\pi \in \Pi \left(\mathbb {P} _ {0}, \mathbb {P} _ {1}\right)} \int_ {\mathcal {X} \times \mathcal {Y}} \frac {\left\| x - y \right\| ^ {2}}{2} d \pi (x, y) + \epsilon \mathrm {K L} (\pi | | \mathbb {P} _ {0} \times \mathbb {P} _ {1}). \tag {3}
$$
Since  $\pi \in \Pi (\mathbb{P}_0,\mathbb{P}_1)$ , it holds that  $\mathrm{KL}(\pi ||\mathbb{P}_0\times \mathbb{P}_1) = -H(\pi) + H(\mathbb{P}_0) + H(\mathbb{P}_1)$ , i.e., both formulations are equal up to an additive constant when  $\mathbb{P}_0$  and  $\mathbb{P}_1$  are absolutely continuous. The minimizer of (2) is unique since the functional is strictly convex in  $\pi$  thanks to the strict convexity of  $H(\pi)$ . This unique minimizer  $\pi^{*}$  is called the entropic OT plan.
2.2 Schrodinger Bridge (SB)
SB with the Wiener prior. Let  $\Omega$  be the space of  $\mathbb{R}^D$  valued functions of time  $t\in [0,1]$  describing some trajectories in  $\mathbb{R}^D$ , which start at time  $t = 0$  and end at time  $t = 1$ . We use  $\mathcal{P}(\Omega)$  to denote the set of probability distributions on  $\Omega$ . We use  $dW_{t}$  to denote the differential of the standard Wiener process. Let  $W^{\epsilon}\in \mathcal{P}(\Omega)$  be the Wiener process with the variance  $\epsilon$  which starts at  $\mathbb{P}_0$ . This diffusion process can be represented via the following stochastic differential equation (SDE):
$$
W ^ {\epsilon}: d X _ {t} = \sqrt {\epsilon} d W _ {t}, \quad X _ {0} \sim \mathbb {P} _ {0}. \tag {4}
$$
We use  $\mathrm{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]:
$$
\inf  _ {T \in \mathcal {F} \left(\mathbb {P} _ {0}, \mathbb {P} _ {1}\right)} \mathrm {K L} (T \| W ^ {\epsilon}), \tag {5}
$$
where  $\mathcal{F}(\mathbb{P}_0,\mathbb{P}_1)\subset \mathcal{P}(\Omega)$  is the set of probability distributions on  $\Omega$  having marginal distributions  $\mathbb{P}_0$  and  $\mathbb{P}_1$  at  $t = 0$  and  $t = 1$ , respectively. Thus, the Schrödinger Bridge problem implies finding a stochastic process  $T$  with marginal distributions  $\mathbb{P}_0$  and  $\mathbb{P}_1$  at times  $t = 0$  and  $t = 1$ , respectively, which has minimal KL divergence with the prior process  $W^{\epsilon}$ .
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 (§4) is based on it. Let  $\pi^T$  denote the joint distribution of a stochastic process  $T$  at time moments  $t = 0, 1$  and  $\pi_0^T$ ,  $\pi_1^T$  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  $\mathrm{KL}(T||W^\epsilon)$  as [48, Appendix C]:
$$
\mathrm {K L} (T | | W ^ {\epsilon}) = \mathrm {K L} \left(\pi^ {T} | | \pi^ {W ^ {\epsilon}}\right) + \int_ {\mathcal {X} \times \mathcal {Y}} \mathrm {K L} \left(T _ {| x, y} | | W _ {| x, y} ^ {\epsilon}\right) d \pi^ {T} (x, y), \tag {6}
$$
i.e., KL divergence between  $T$  and  $W^{\epsilon}$  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 \in (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]):
$$
\operatorname {K L} \left(\pi^ {T} \mid \mid \pi^ {W ^ {\epsilon}}\right) = \int_ {\mathcal {X} \times \mathcal {Y}} \frac {\left| \left| x - y \right| \right| ^ {2}}{2 \epsilon} d \pi^ {T} (x, y) - H \left(\pi^ {T}\right) + C, \tag {7}
$$
where  $C$  is a constant which depends only on  $\mathbb{P}_0$  and  $\epsilon$ . 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}^{\epsilon}$  for every  $x,y$  and set the last term in (6) to zero. In this case:
$$
\inf  _ {T \in \mathcal {F} (\mathbb {P} _ {0}, \mathbb {P} _ {1})} \mathrm {K L} (T | | W ^ {\epsilon}) = \inf  _ {T \in \mathcal {F} (\mathbb {P} _ {0}, \mathbb {P} _ {1})} \mathrm {K L} \left(\pi^ {T} | | \pi^ {W ^ {\epsilon}}\right) = \inf  _ {\pi^ {T} \in \Pi (\mathbb {P} _ {0}, \mathbb {P} _ {1})} \mathrm {K L} \left(\pi^ {T} | | \pi^ {W ^ {\epsilon}}\right). \tag {8}
$$
i.e., it suffices to optimize only over joint distributions  $\pi^T$  at time moments  $t = 0,1$ . Hence, minimizing (8) is equivalent (up to an additive constant  $C$ ) to solving EOT (2) divided by  $\epsilon$ , and their respective solutions  $\pi^{T^*}$  and  $\pi^*$  coincide. Thus, SB problem (5) can be simplified to the entropic OT problem (2) with entropy coefficient  $\epsilon$ .
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:
$$
\inf _ {\pi \in \Pi (\mathbb {P} _ {0}, \mathbb {P} _ {1})} \mathbf {K L} (\pi | | \pi^ {W ^ {\epsilon}}) = \sup _ {\beta} \left\{\int_ {\mathcal {X}} \beta^ {C} (x) d \mathbb {P} _ {0} (x) + \int_ {\mathcal {Y}} \beta (y) d \mathbb {P} _ {1} (y) \right\},
$$
where  $\beta^C (x)\stackrel {\mathrm{def}}{=}\inf_{\nu \in \mathcal{P}_2(\mathcal{Y})}\bigl \{\mathrm{KL}\bigl (\nu ||\pi^{W^\epsilon}(\cdot |x)\bigr) - \int_{\mathcal{Y}}\beta (y)d\nu (y)\bigr \} .$  The sup here is taken over  $\beta$  belonging to the set of functions
$$
\mathcal {C} _ {b, 2} (\mathcal {Y}) \stackrel {{\text {d e f}}} {{=}} \left\{\beta : \mathcal {Y} \rightarrow \mathbb {R} \text {c o n t i n u o u s s . t .} \exists u, v, w \in \mathbb {R}: u \| \cdot \| ^ {2} + v \leq \beta (\cdot) \leq w \right\},
$$
i.e.,  $\beta$  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  $\mathcal{D}(\mathbb{P}_0)$  of finite-energy diffusions  $T_{f}$  [32, Proposition 4.1] which are given by:
$$
T _ {f}: d X _ {t} = f \left(X _ {t}, t\right) d t + \sqrt {\epsilon} d W _ {t}, \quad X _ {0} \sim \mathbb {P} _ {0}, \quad \mathbb {E} _ {T _ {f}} \left[ \int_ {0} ^ {1} | | f \left(X _ {t}, t\right) | | ^ {2} d t \right] <   \infty , \tag {9}
$$
where  $f: \mathbb{R}^D \times [0,1] \to \mathbb{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  $\mathrm{KL}(T_f||W^{\epsilon})$  between processes  $T_{f}$  and  $W^{\epsilon}$  is [40]:
$$
\operatorname {K L} \left(T _ {f} \| W ^ {\epsilon}\right) = \frac {1}{2 \epsilon} \mathbb {E} _ {T _ {f}} \left[ \int_ {0} ^ {1} | | f \left(X _ {t}, t\right) \| ^ {2} d t \right]. \tag {10}
$$
By substituting (10) in (5), SB problem reduces to the following problem [11, Problem 4.3]:
$$
\inf  _ {T _ {f} \in \mathcal {D} (\mathbb {P} _ {0}, \mathbb {P} _ {1})} \mathrm {K L} \left(T _ {f} | | W ^ {\epsilon}\right) = \inf  _ {T _ {f} \in \mathcal {D} (\mathbb {P} _ {0}, \mathbb {P} _ {1})} \frac {1}{2 \epsilon} \mathbb {E} _ {T _ {f}} [ \int_ {0} ^ {1} | | f (X _ {t}, t) | | ^ {2} d t ], \tag {11}
$$
where  $\mathcal{D}(\mathbb{P}_0,\mathbb{P}_1)\subset \mathcal{P}(\Omega)$  is the set of finite-energy diffusion on  $\Omega$  having marginal distributions  $\mathbb{P}_0$  and  $\mathbb{P}_1$  at  $t = 0$  and  $t = 1$ , respectively. Note that  $\mathcal{D}(\mathbb{P}_0,\mathbb{P}_1)\subset \mathcal{F}(\mathbb{P}_0,\mathbb{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.
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  $\mathbb{P}_0, \mathbb{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  $\mathbb{P}_0, \mathbb{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  $\mathbb{P}_0$  to  $\mathbb{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 \mapsto \int_{\mathcal{V}} y d\pi^{*}(y|x)$ , i.e., the first conditional moment of the entropic OT plan. This strategy may yield a deterministic approximation of  $\pi^{*}$  for small  $\epsilon$  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  $\pi^{*}(y|x)$  of the OT plan  $\pi^{*}$  from the dual potentials. Unfortunately, this is costly. During the training phase, the method requires learning a score-based model for the distribution  $\mathbb{P}_1$ . More importantly, during the inference phase, one has to run the Langevin dynamic to sample from  $\pi^{*}(y|x)$ .
The optimization of the above-mentioned entropic approaches requires evaluating the exponent of large values which are proportional to  $\epsilon^{-1}$  [45, Eq. 7]. Due to this fact, those methods are unstable for small values  $\epsilon$  in (2). In contrast, our proposed method (§4) resolves this issue: technically, it works even for  $\epsilon = 0$ .
3.2 Approaches to Compute Schrödinger Bridges
Existing approaches to solve DSB mainly focus on generative modeling applications (noise  $\rightarrow$  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  $\pi^*$  of the related EOT problem (8). In §4.1, we theoretically derive the proposed saddle point optimization objective. In §4.2 we provide and describe the practical learning procedure to optimize it. In §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  $\mathbb{P}_0$  and  $\mathbb{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  $\pi \in \Pi(\mathbb{P}_0, \mathbb{P}_1)$  and (b) estimate the entropy  $H(\pi)$  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
$$
\mathcal {L} (\beta , T _ {f}) \stackrel {\text {d e f}} {=} \left\{\overbrace {\mathbb {E} _ {T _ {f}} \left[ \frac {1}{2 \epsilon} \int_ {0} ^ {1} | | f (X _ {t} , t) | | ^ {2} d t \right]} ^ {= \mathrm {K L} \left(T _ {f} \mid \mid W ^ {\epsilon}\right)} - \int_ {\mathcal {Y}} \beta (y) d \pi_ {1} ^ {T _ {f}} (y) + \int_ {\mathcal {Y}} \beta (y) d \mathbb {P} _ {1} (y) \right\}. \tag {12}
$$
This functional can be viewed as the Lagrangian for DSB (11) with the relaxed constraint  $d\pi_1^{T_f}(y) = d\mathbb{P}_1(y)$ , and function  $\beta : \mathcal{V} \to \mathbb{R}$  (potential) playing the role of the Langrange multiplier.
Theorem 4.1 (Relaxed DSB formulation). Consider the following saddle point optimization problem:
$$
\sup  _ {\beta} \inf  _ {T _ {f}} \mathcal {L} (\beta , T _ {f}) \tag {13}
$$
where sup is taken over potentials  $\beta \in \mathcal{C}_{b,2}(\mathcal{Y})$  and inf is taken over diffusion processes  $T_{f}\in \mathcal{D}(\mathbb{P}_0)$ . Then for every optimal pair  $(\beta^{*},T_{f^{*}})$  for (13), i.e.,
$$
\beta^ {*} \in \operatorname * {a r g s u p} _ {\beta} \inf  _ {T _ {f}} \mathcal {L} (\beta , T _ {f}) \qquad a n d \qquad T _ {f ^ {*}} \in \operatorname * {a r g m i n} _ {T _ {f}} \mathcal {L} (\beta^ {*}, T _ {f}),
$$
it holds that the process  $T_{f}^{*}$  is the solution to SB (11).
Corollary 4.2 (Entropic OT as relaxed DSB). If  $(\beta^{*},T_{f^{*}})$  solves (13), then  $\pi^{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  $\pi^{*} = \pi^{T_{f}^{*}}$ . The notable benefits of considering (13) instead of (2), (3) and (11) is that (a) it is as an optimization problem over  $(\beta ,T_{f})$  without the constraint  $d\pi_1^{T_f}(y) = d\mathbb{P}_1(y)$ , and (b) objective (12) admits Monte-Carlo estimates by using random samples from  $\mathbb{P}_0,T_f,\mathbb{P}_1$ . In §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(\pi)$
in EOT. This term is replaced by the energy of the process  $\mathbb{E}_{T_f}\left[\int_0^1 ||f(X_t,t)||^2 dt\right]$  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  $\beta (y)^{4}$  by neural nets  $f_{\theta}:\mathbb{R}^{D}\times [0,1]\to \mathbb{R}^{D}$  and  $\beta_{\phi}:\mathbb{R}^{D}\rightarrow \mathbb{R}$ . We consider the following maximin problem:
$$
\sup  _ {\beta} \inf  _ {T _ {f _ {\theta}}} \left\{\frac {1}{2 \epsilon} \mathbb {E} _ {T _ {f _ {\theta}}} \left[ \int_ {0} ^ {1} | | f _ {\theta} (X _ {t}, t) | | ^ {2} d t \right] + \int_ {\mathcal {Y}} \beta_ {\phi} (y) d \mathbb {P} _ {1} (y) - \int_ {\mathcal {Y}} \beta_ {\phi} (y) d \pi_ {1} ^ {T _ {f _ {\theta}}} (y) \right\}. \tag {14}
$$
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  $\mathbb{E}_{T_f}[\int_0^1 ||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_{\theta}$  and  $\beta_{\phi}$  by optimizing (12) with the stochastic gradient ascent-descent by sampling random batches from  $\mathbb{P}_0, \mathbb{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  $\int_0^1 ||f(X_t,t)||^2 dt$  in the training objective.
Remark. For the image tasks (§5.3, §5.4), we find out that using a slightly different parametrization of  $T_{f}$  considerably improves the quality of our Algorithm, see Appendix F.