SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \operatorname {K L} \left(\pi^ {T _ {f ^ {*}}} | | \pi^ {W _ {\epsilon}}\right) + \int_ {\mathcal {X}} \log C _ {\beta} ^ {x} d \mathbb {P} _ {0} (x) - \int \hat {\beta} (y) d \mathbb {P} _ {1} (y) = \\ \int_ {\mathcal {X}} \operatorname {K L} \left(\pi^ {T _ {f ^ {*}}} (\cdot | x) | | \pi^ {W ^ {\epsilon}} (\cdot | x)\right) d \mathbb {P} _ {0} (x) + \int_ {\mathcal {X}} \log C _ {\hat {\beta}} ^ {x} d \mathbb {P} _ {0} (x) - \int \hat {\beta} (y) d \mathbb {P} _ {1} (y) = \\ \int_ {\mathcal {X}} \operatorname {K L} \left(\pi^ {T _ {f ^ {*}}} (\cdot | x) | | \pi^ {W ^ {\epsilon}} (\cdot | x)\right) d \mathbb {P} _ {0} (x) + \int_ {\mathcal {X}} \log C _ {\hat {\beta}} ^ {x} d \mathbb {P} _ {0} (x) - \int \hat {\beta} (y) d \pi^ {T _ {f ^ {*}}} (y | x) d \mathbb {P} _ {0} (x) = \\ \int_ {\mathcal {X}} \left\{\mathrm {K L} \left(\pi^ {T _ {f ^ {*}}} (\cdot | x) | | \pi^ {W ^ {\epsilon}} (\cdot | x)\right) + \log C _ {\hat {\beta}} ^ {x} - \int_ {\mathcal {Y}} \hat {\beta} (y) d \pi^ {T _ {f ^ {*}}} (y | x) \right\} d \mathbb {P} _ {0} (x) = \\ \int_ {\mathcal {X}} \left\{\mathrm {K L} \left(\pi^ {T _ {f ^ {*}}} (\cdot | x) | | \pi^ {T ^ {\hat {\beta}}} (\cdot | x)\right) - \log C _ {\hat {\beta}} ^ {x} + \log C _ {\hat {\beta}} ^ {x} \right\} d \mathbb {P} _ {0} (x) = \\ \int_ {\mathcal {X}} \operatorname {K L} \left(\pi^ {T _ {f ^ {*}}} (\cdot | x) | | \pi^ {T ^ {\hat {\beta}}} (\cdot | x)\right) d \mathbb {P} _ {0} (x) = \operatorname {K L} \left(\pi^ {T _ {f ^ {*}}} | | \pi^ {T ^ {\hat {\beta}}}\right) = \\ \mathrm {K L} \left(\pi^ {T _ {f ^ {*}}} \| \pi^ {T ^ {\hat {\beta}}}\right) + \underbrace {\int_ {\mathcal {X} \times \mathcal {Y}} \mathrm {K L} \left(T _ {f ^ {*} | x , y} \| T _ {| x , y} ^ {\hat {\beta}}\right) d \pi^ {T _ {f ^ {*}}} (x , y)} _ {= 0, \text {s i n c e} T _ {f ^ {*}} | x, y = T _ {| x, y} ^ {\hat {\beta}} = W _ {| x, y} ^ {e}} = \mathrm {K L} \left(T _ {f ^ {*}} \| T ^ {\hat {\beta}}\right). \tag {42} \\ \end{array}
$$
Thus, we obtain  $\epsilon_{2} = \mathrm{KL}(T_{f^{*}}||T^{\widehat{\beta}})$
Part 3. By summing (41) and (42) and using the Pinsker inequality, we obtain
$$
\begin{array}{l} \epsilon_ {1} + \epsilon_ {2} = \mathrm {K L} (T _ {\hat {f}} | | T ^ {\hat {\beta}}) + \mathrm {K L} (T _ {f ^ {*}} | | T ^ {\hat {\beta}}) \geq 2 \rho_ {\mathrm {T V}} ^ {2} (T _ {\hat {f}}, T ^ {\hat {\beta}}) + 2 \rho_ {\mathrm {T V}} ^ {2} (T _ {f ^ {*}}, T ^ {\hat {\beta}}) \geq \\ \left[ \rho_ {\mathrm {T V}} \left(T _ {\hat {f}}, T ^ {\hat {\beta}}\right) + \rho_ {\mathrm {T V}} \left(T _ {f ^ {*}}, T ^ {\hat {\beta}}\right) \right] ^ {2} \geq \rho_ {\mathrm {T V}} ^ {2} \left(T _ {\hat {f}}, T _ {f ^ {*}}\right). \tag {43} \\ \end{array}
$$
Here we use the triangle inequality in line (43). Therefore,  $\rho_{\mathrm{TV}}(T_{\hat{f}}, T_{f^*}) \leq \sqrt{\epsilon_1 + \epsilon_2}$ .
Part 4. By summing (41) and (42) and using the Pinsker inequality, we obtain
$$
\epsilon_ {1} + \epsilon_ {2} = \mathrm {K L} (T _ {\hat {f}} | | T ^ {\hat {\beta}}) + \mathrm {K L} (T _ {f *} | | T ^ {\hat {\beta}}) =
$$
$$
\mathrm {K L} \left(\pi^ {T _ {\hat {f}}} | | \pi^ {T ^ {\hat {\beta}}}\right) + \int_ {\mathcal {X} \times \mathcal {Y}} \mathrm {K L} \left(T _ {\hat {f} | x, y} | | T _ {| x, y} ^ {\hat {\beta}}\right) d \pi^ {T _ {\hat {f}}} (x, y) +
$$
$$
\mathrm {K L} \left(\pi^ {T _ {f ^ {*}}} | | \pi^ {T ^ {\widehat {\beta}}}\right) + \int_ {\mathcal {X} \times \mathcal {Y}} \mathrm {K L} \left(T _ {f ^ {*} | x, y} | | T _ {| x, y} ^ {\widehat {\beta}}\right) d \pi^ {T _ {f ^ {*}}} (x, y) \geq
$$
$$
\begin{array}{l} \mathbf {K L} \left(\pi^ {T _ {\hat {f}}} | | \pi^ {T ^ {\hat {\beta}}}\right) + \mathbf {K L} \left(\pi^ {T _ {f ^ {*}}} | | \pi^ {T ^ {\hat {\beta}}}\right) \geq 2 \rho_ {\mathrm {T V}} ^ {2} \left(\pi^ {T _ {\hat {f}}}, \pi^ {T ^ {\hat {\beta}}}\right) + 2 \rho_ {\mathrm {T V}} ^ {2} \left(\pi^ {T _ {f ^ {*}}}, \pi^ {T ^ {\hat {\beta}}}\right) \geq \\ \left[ \rho_ {\mathrm {T V}} \left(\pi^ {T _ {\hat {f}}}, \pi^ {T ^ {\hat {\beta}}}\right) + \rho_ {\mathrm {T V}} \left(\pi^ {T _ {f ^ {*}}}, \pi^ {T ^ {\hat {\beta}}}\right) \right] ^ {2} \geq \rho_ {\mathrm {T V}} ^ {2} \left(\pi^ {T _ {\hat {f}}}, \pi^ {T _ {f ^ {*}}}\right). \\ \end{array}
$$
Thus,  $\rho_{\mathrm{TV}}(\pi^{T_f*},\pi^{T^{\hat{\beta}}})\leq \sqrt{\epsilon_1 + \epsilon_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:
$$
\sup_{\beta}\inf_{T_{f}\in \mathcal{D}(\mathbb{P}_{0})}\underbrace{\left\{\mathbb{E}_{T_{f}}[\int_{0}^{1}C||f(X_{t},t)||^{2}dt] + \int_{\mathcal{Y}}\beta(y)d\mathbb{P}_{1}(y) - \int_{\mathcal{Y}}\beta(y)d\mathbb{P}_{1}^{T_{f}}(y)\right\}}_{\stackrel {\mathrm{def}}{=}\mathcal{L}^{C}(\beta ,T_{f})},
$$
with  $C = 1$ . At the same time, we use  $C = \frac{1}{2\epsilon}$  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
Input: batch of initial states  $X_0$  at time moment  $t = 0$ ;  
SDE drift network  $f_{\theta}:\mathbb{R}^{D}\times [0,1]\to \mathbb{R}^{D}$ ;  
number of steps for the SDE solver  $N\geq 1$ ;  
noise variance  $\epsilon \geq 0$ .  
Output: batches  $\{X_{n}\}_{n = 0}^{N}$  of intermediate states at  $t = \frac{n}{N}$  simulating the process  
 $dX_{t} = f(X_{t},t)dt + \sqrt{\epsilon} dW_{t}$ ;  
batches  $\{f_n\}_{n = 0}^N$  of drift values  $f(X_{n},t_{n})$  at  $t = \frac{n - 1}{N}$  simulating the process;  
 $\Delta t\gets \frac{1}{N}$ ;  
for  $t = 1,2,\ldots ,N$  do  
for  $i = 1,2,\ldots ,|X_0|$  do  
Sample noise  $W$  from  $\mathcal{N}(0,I)$ ;  
 $f_{t - 1,i}\gets f(X_{t - 1},t - 1)$ ;  
 $X_{t,i}\gets X_{t - 1,i} + f_{t - 1,i}\Delta t + \sqrt{\epsilon\Delta t} W$ ;
$(\beta^{*},T_{f^{*}})$  is the optimal point for the problem with  $C = 1$ , then  $(\widetilde{C}\beta^{*},T_{f^{*}})$  is the optimal point for  $C = \widetilde{C}$ . Indeed, for a pair  $(\beta ,T_{f})$  it holds that
$$
\mathcal {L} ^ {1} (\beta , T _ {f}) = \mathbb {E} _ {T _ {f}} [ \int_ {0} ^ {1} | | f (X _ {t}, t) | | ^ {2} d t ] + \int_ {\mathcal {Y}} \beta (y) d \mathbb {P} _ {1} (y) - \int_ {\mathcal {Y}} \beta (y) d \mathbb {P} _ {1} ^ {T _ {f}} (y) =
$$
$$
\frac {1}{\widetilde {C}} \left\{\mathbb {E} _ {T _ {f}} \left[ \int_ {0} ^ {1} \widetilde {C} | | f (X _ {t}, t) | | ^ {2} d t \right] + \int_ {y} \widetilde {C} \beta (y) d \mathbb {P} _ {1} (y) - \int_ {y} \widetilde {C} \beta (y) d \mathbb {P} _ {1} ^ {T _ {f}} (y) \right\} =
$$
$$
\frac {1}{\widetilde {C}} \left\{\mathbb {E} _ {T _ {f}} \left[ \int_ {0} ^ {1} \widetilde {C} | | f (X _ {t}, t) | | ^ {2} d t \right] + \int_ {\mathcal {Y}} \widetilde {\beta} (y) d \mathbb {P} _ {1} (y) - \int_ {\mathcal {Y}} \widetilde {\beta} (y) d \mathbb {P} _ {1} ^ {T _ {f}} (y) \right\} = \frac {1}{\widetilde {C}} \mathcal {L} ^ {\widetilde {C}} (\widetilde {\beta}, T _ {f}), \tag {44}
$$
where we use  $\widetilde{\beta} \stackrel{\mathrm{def}}{=} \widetilde{C}\beta$ . Hence problems  $\sup_{\beta} \inf_{T_f} \mathcal{L}^1(\beta, T_f)$  and  $\sup_{\widetilde{\beta}} \inf_{T_f} \mathcal{L}^{\widetilde{C}}(\widetilde{\beta}, 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  $\widetilde{C} > 0$ . For completeness, we also note that the change of variables  $\beta \leftrightarrow \widetilde{\beta}$  actually preserves the functional class of  $\beta$ , i.e.,  $\beta \in \mathcal{C}_{b,2}(\mathcal{Y}) \Longleftrightarrow \widetilde{\beta} \in \mathcal{C}_{b,2}(\mathcal{Y})$ .
For convenience, we get rid of dependence on  $\epsilon$  in the objective (12) and consider  $\mathcal{L}^1$  for optimization, i.e., use  $C = 1$  in Algorithm 1. Still the dependence on  $\epsilon$  remains in  $\sup_{\beta}\inf_{T_f}\mathcal{L}^1 (\beta ,T_f)$  as  $T_{f}\in \mathcal{D}(\mathbb{P}_{0})$  is a diffusion process with volatility  $\epsilon$ . Interestingly, this point of view (optimizing  $\mathcal{L}^1$  instead of  $\mathcal{L}^{\frac{1}{2\epsilon}}$ ) technically allows to consider even  $\epsilon = 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_{\beta}\inf_{T_f}\mathcal{L}^1 (\beta ,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, §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  $\mathcal{L}^1$  in this case. However, for completeness, we provide experimental results for  $\epsilon = 0$ .
E ENOT for Toy Experiments and High-dimensional Gaussians
In 2D toy experiments, we consider 2 tasks: Gaussian  $\rightarrow$  8 gaussians and Gaussian  $\rightarrow$  Swiss roll. Results for the last one (Figure 5) are qualitatively similar to results of the first one (Figure 2), which we discussed earlier (§5.1). For both tasks, we parametrize the SDE drift function in Algorithm 1 by a feedforward neural network  $f_{\theta}$  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  $\beta_{\phi}$  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  $\mathrm{lr} = 10^{-4}$ , batch size 512. We train the model for 20000 total iterations of  $\beta_{\phi}$ , and on each of them, we do  $K_{f} = 10$  updates for the SDE drift function  $f_{\theta}$ .
(a) Input and target samples
(c) ENOT (ours),  $\epsilon = 0.01$
(d) ENOT (ours),  $\epsilon = 0.1$
Figure 5: Gaussian  $\rightarrow$  Swiss roll, learned stochastic process with ENOT (ours).
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  $\mathrm{BW}_2^2$ -UVP (\%) between the ground truth EOT plan  $\pi^{*}$  and the learned plan  $\pi$  of ENOT during training for  $DIM = 128$  in Figure 6.
Figure 6:  $\mathrm{BW}_2^2$ -UVP  $\downarrow$  (\%) between the EOT plan  $\pi^{*}$  and the learned plan  $\pi$  of ENOT and MLE-SB during the training (DIM = 128).
F ENOT for Colored MNIST and Unpaired Super-resolution of Celeba Faces
For the image tasks (§5.3, §5.4), we find out that using the following reparametrization of Euler-Maruyama 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)\Delta t + \sqrt{\epsilon\Delta t}$ , we parametrize  $g(X_{t}, t) = X_{t} + f(X_{t}, t)\Delta t$  by a neural network  $g_{\theta}$ , and calculate the next state as  $X_{t+1} = g_{\theta}(X_{t}, t) + \sqrt{\epsilon\Delta t}$ . In turn, the drift function is given by  $f(X_{t}, t) = \frac{1}{\Delta 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.
Figure 7: Trajectories from our learned ENOT (ours) for colored MNIST for different  $\epsilon$ .
We use WGAN-QC discriminator's ResNet architecture  $^5$  for the potential  $\beta$ . We use UNet  $^6$  as  $g_{\theta}(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  $\mathrm{lr} = 10^{-4}$ , batch size 64 and 10:1 update ratio for  $f_{\theta} / \beta_{\phi}$ . For  $\epsilon = 0$  and  $\epsilon = 1$  our model converges in  $\approx 20000$  iterations, while for  $\epsilon = 10$  it takes  $\approx 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  $\approx 40000$  iterations for  $\epsilon = 0$  and  $\approx 70000$  iterations for  $\epsilon = 1$  and  $\epsilon = 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 (§5.4), ENOT converges in one week on  $2 \times \mathrm{A}100$  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 (§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
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/defaultmoon_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.
$^{5}$ github.com/harryliew/WGAN-QC
$^{6}$ github.com/milesial/Pytorch-UNet
$^{7}$ github.com/rosinality/denoising-diffusion-pytorch
$^{8}$ github.com/kgkgzrtk/cUNet-Pytorch
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
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 (§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 \times 32$  ColorMNIST images instead of  $64 \times 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 (§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
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 Brediges.
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 \sim \rho_0$ . Every player controls its behavior through drift  $\alpha$  of the SDE:
$$
d X _ {t} = \alpha \left(X _ {t}, t, \rho_ {t}\right) d t + \sqrt {2 \nu} d W _ {t}
$$
Here  $\rho_{t}$  is the distribution of all the players at the time moment  $t$ . When we consider a specific player, we consider  $\rho_{t}$  as a parameter. Each player aims to minimize the quantity:
$$
\mathbb {E} \left[ \int_ {0} ^ {T} \left(L \left(X _ {t}, \alpha_ {t}, \rho_ {t}\right) + f \left(X _ {t}, \rho_ {t}\right)\right) d t + g \left(X _ {T}, \rho_ {T}\right) \right].
$$
Here  $L(x, \alpha, \rho)$  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}, \rho_{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  $\phi(x, t)$ , which for position  $x$  and start time  $t$  returns the cost in case of the optimal control:
$$
\phi (x, t) \stackrel {\text {d e f}} {=} \inf  _ {\alpha} \mathbb {E} \left[ \int_ {t} ^ {T} \left(L \left(X _ {t}, \alpha_ {t}, \rho_ {t}\right) + f \left(X _ {t}, \rho_ {t}\right)\right) d t + g \left(X _ {T}, \rho_ {T}\right) \right].
$$
Before considering the Mean-Field game, we need to define an additional function  $H(x,p,\rho)$ . It is similar to the Hamilton function and is defined as the Legendre transform of Lagrange function  $L$ :
$$
H (x, p, \rho) \stackrel {\mathrm {d e f}} {=} \sup  _ {\alpha} [ - \alpha p - L (x, \alpha , \rho) ].
$$
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,\rho)$  and  $f(x,\rho)$ , Mean-Field game formulates as a system of two PDE with two constraints:
$$
\begin{array}{l} - \partial_ {t} \phi - \nu \Delta \phi + H (x, \nabla \phi , \rho) = f (x, \rho) (\mathrm {H J B}) \\ - \partial_ {t} \rho - \nu \Delta \rho - \mathbf {d i v} (\rho \nabla_ {p} H (x, \nabla \phi)) = 0 (\mathrm {F P}) \\ \text {s . t .} \rho (x, 0) = \rho_ {0}, \phi (x, T) = g (x, \rho (\cdot , T)) \\ \end{array}
$$
The solution of this system is two functions  $\rho(x, t)$  and  $\phi(x, t)$ , which describe all players' dynamics. Also, in Nash equilibrium, the specific player's behavior is described by the following SDE:
$$
d X _ {t} = - \nabla_ {p} H (X _ {t}, \nabla \phi (X _ {t}, t), \rho) d t + \sqrt {2 \nu} d W _ {t}.
$$
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  $\rho (\cdot ,T) = \rho_{target}(\cdot ,T)$  via choosing proper function  $g(x,\rho (\cdot ,T))$  such as:
$$
g (x, \rho (\cdot , T)) = \left\{ \begin{array}{l l} \infty , & \text {i f} \rho (\cdot , T) \neq \rho_ {t a r g e t} (\cdot , T) \\ 0, & \rho (\cdot , T) = \rho_ {t a r g e t} (\cdot , T) \end{array} \right.
$$
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,\rho) = ||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  $\phi (x,t)$  by a neural network and use different neural network  $N_{\theta}$  to sample from  $\rho_{t}$ . The authors penalize the violation of Mean-Field game PDEs for optimizing these networks. After the convergence, one can sample from the distribution  $\rho_{t}$  by using neural network  $N_{\theta}$ . Approach [33] has the advantage that authors do not need to use SDE solvers, which require more steps with growing parameter  $\nu$  of diffusion operator. However, computation of Laplacian and divergence for high-dimensional spaces (e.g., space 12228-dimensional space of  $3\times 64\times 64$  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:
$$
d X _ {t} = \alpha \left(X _ {t}, t, \rho_ {t}\right) d t + \sqrt {2 \nu} d W _ {t},
$$
which describes the player's behavior and use a neural network to parametrize the drift  $\alpha$ . We consider only hard constraints on the target distribution,  $f(X_{t},\rho_{t}) = 0$  and  $L(X_{t},\alpha_{t},\rho_{t}) = \frac{1}{2} ||\alpha_{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  $\nu$  (which is equal to the  $\frac{1}{2}\epsilon$  in our notation, where  $\epsilon$  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) = \frac{1}{2}\| x - y\| ^2$  which coincides with SB with the Wiener prior  $W^{\epsilon}$ . However, one could use a different prior  $Q_{v}$  instead of  $W^{\epsilon}$  in (5):
$$
Q _ {v}: d X _ {t} = v \left(X _ {t}, t\right) d t + \sqrt {\epsilon} d W _ {t},
$$
and solve the problem
$$
\inf  _ {T _ {f} \in \mathcal {D} (\mathbb {P} _ {0}, \mathbb {P} _ {1})} \mathrm {K L} (T _ {f} | | Q _ {v}) = \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) - v (X _ {t}, t) | | ^ {2} d t ].
$$
Here we just use the known expression (6) for  $\mathrm{KL}(T_f||Q_v)$  between two diffusion processes through their drift functions. Using the same derivation as in the main text §2.2, it can be shown that this new problem is equivalent to solving the EOT with cost  $c(x,y) = -\log \pi^{Q_v}(y|x)$ , where  $\pi^{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^{\epsilon}$  (which we consider in the main text) we have
$$
c (x, y) = - \log \pi^ {W ^ {\epsilon}} (y | x) = \frac {1}{2 \epsilon} (y - x) ^ {T} (y - x) + \operatorname {C o n s t},
$$
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):
$$
\sup  _ {\beta} \inf  _ {T _ {f}} \left(\frac {1}{2 \epsilon} \mathbb {E} _ {T _ {f}} \left[ \int_ {0} ^ {1} | | f (X _ {t}, t) - v (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}} (y)\right).
$$