Expectile Regularization for Fast and Accurate Training of Neural Optimal Transport
Authors: Nazar Buzun
Maksim Bobrin
Dmitry V. Dylov
Venue: NeurIPS 2024 spotlight (2024)
Keywords: Optimal Transport
Neural Optimal Transport
Expectile
Regularization
Kantorovich potentials
Primary area: probabilistic_methods
TL;DR: A 3-fold better and a 10-fold faster way to train neural optimal transport: all thanks to a new expectile regularizing loss function that sets an upper bound over the distribution of learning conjugate potentials.

Abstract: We present a new approach for Neural Optimal Transport (NOT) training procedure, capable of accurately and efficiently estimating optimal transportation plan via specific regularization on dual Kantorovich potentials. The main bottleneck of existing NOT solvers is associated with the procedure of finding a near-exact approximation of the conjugate operator (i.e., the c-transform), which is done either by optimizing over non-convex max-min objectives or by the computationally intensive fine-tuning of the initial approximated prediction. We resolve both issues by proposing a new theoretically justified loss in the form of expectile regularization which enforces binding conditions on the learning process of the dual potentials. Such a regularization provides the upper bound estimation over the distribution of possible conjugate potentials and makes the learning stable, completely eliminating the need for additional extensive fine-tuning. Proposed method, called Expectile-Regularized Neural Optimal Transport (ENOT), outperforms previous state-of-the-art approaches in the established Wasserstein-2 benchmark tasks by a large margin (up to a 3-fold improvement in quality and up to a 10-fold improvement in runtime). Moreover, we showcase performance of ENOT for various cost functions in different tasks, such as image generation, demonstrating generalizability and robustness of the proposed algorithm.

ENOT: Expectile Regularization for Fast and Accurate Training of Neural Optimal Transport
Nazar Buzun" AIRI, MIPT buzun@airi.ner
Maksim Bobrin" Skoltech, AIRI m w..bobrin@skcltech..r
Dmitry V. Dylov
Skoltech, AIRI
d. wwd.ly@sskoteec.hcom
Abstract
We present a new approach for Neural Optimal Transport (NOT) training procedure, capable of accurately and efficiently estimating optimal transportation plan via specific regularization on dual Kantorovich potentials. The main bottleneck existing NOT solvers is associated with the procedure of finding near-exact approximation the conjugate operator (i.e., the c-transform), which is done either by optimizing
non-convex max-min objectives by the computationally intensive fine-tuning of the initial approximated prediction. We resolve both issues by proposing new theoretically justified loss in the form of expectile regularization which enforces binding conditions on the learning process of the dual potentials. the distribution of
Such a regularization provides the upper bound estimation possible conjugate potentials and makes the learning stable, completely eliminating the need for additional extensive fine-tuning. Proposed method, called Expectile- Regularized Neural Optimal Transport (ENOT), outperforms previous state-of- the-art approaches the established Wasserstein-2 benchmark tasks by large margin (up to a 3-fold improvement in quality and up to a 10-fold improvement in runtime). Moreover, we showcase performance of ENOT for various cost functions in different tasks, such as image generation, demonstrating generalizability and robustness of the proposed algorithm.
Project page with code https://skyloog:: github. io/enot/
Introduction
Computational optimal transport (OT) has enriched machine learning (ML) by offering a new view- angle the conventional MI. tasks through the lens of comparison of probability measures (Villani et al. [2009], Ambrosio et al. [2003], Peyré et al. [2019], Santambrogio [2015]). In different works, OT is primarily employed either 1) differentiable proxy, with the OT distance playing the role of a similarity metric between the measures, or 2) as generative model, defined by the plan of optimal transportation. One notable advantage of using OT in the latter setting is that, compared to other generative approaches, such as GANs, Normalizing Flows, or Diffusion Models, there is no assumption for of the measures be defined in closed form (e.g., Gaussian uniform) or to be Buirwise-aligned admitting various applications of the OT theory. Both loss objective and generative formulations of OT proved successful vast range of modern ML areas, including
"Equal contribution.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

generative modelling (Arjovsky et al. [2017], Gulrajani et al. [2017], Korotin et al. [2021], Liu et al. [2019], Leygonie al. [2019]), reinforcement learning (Fickinger et al. [2022], Haldar al. [2023], Papagiannis and Li [2022]. Luo et al. [2023]), domain adaptation (Xie et al. [2019]. Shen et al. (ll122)))). change point detection (Shvetsov et al. [2020], barycenter estimation (Kroshnin et al. [2021], Buzun [2023], Bespalov et al. [2022a,b]), biology and genomics (Bunne et al. [2022]). Low dimensional discrete OT problems are solved via Sinkhorn algorithm (Cuturi [2013]), which employs entropic regularization. This technique makes the entire optimization problem differentiable and efficient computationally. but may require numerous iterations to converge to an optimal solution, whereas the OT problem for the tasks supported on high- -igh-dimensional measure spaces usually intractable, oftentimes solvable only for the distributions which admit closed-fore density formulations. As a result, the need for computationally efficiennt OT solvers has become both evident (Peyré al. [2019]) and pressing (Montesuma et al. [2023], Khamis et al. [2024]).
In this work, we will be concerned with the complexity, the quality, and the runtime speed of the computational estimation of deterministic OT plan T between two probability measures a and [3 supported measurable spaces X, X C _d with Borel sigma-algebra. The OT problem in Monge's formulation (MP) for a cost function e X x y + IR is stated as:
MP 3) T: Ty / c(x, ((:))a(())
inf
where T
= [ (3)[3))))) is the set of measure-preserving maps, defined by a push forward operator Tya B) = aB)) = (B) for any Borel subset B c y. The minimizer of the cost above exists if X' is compact, a is atomless (i.e. Vx X' ax [(x+)) = 0) and the cost function is continuous (ref. Santambrogio [2015] Theorem 1.22 and Theorem 1.33).
However, MP formulation of the OT problem is intractable, since it requires finding the maps T under the coupling constraints (which non-convex optimization problem) and is not general enough to provide a way for some mass-splitting solutions. By relaxing constraints in equation (1), the OT problem becomes convex and this form is known as Kantorovich problem (KP) (ref. Villani et al. [2009]):
inf Ex [c(2, 4)], f (r,y)dd((E,)) TEAT(
KP (3 inf TEilla 3) Jxxy
(2)
where [I 3] { P(X x V) (d y) = dx(x), dT y) dB(r)) is a set of admissible couplings with respective marginals 3. MP and KP problems are equivalent in case
X' = y are compact, the cost function c(x, v) is continuous and a is atomless. Since KP (2) is convex, it admits dual formulation (DP), which constrained concave maximization problem and derived via Lagrange multipliers (Kantorovich potentials) f and g:
DP(,,
sup
Ex [x[f(z)] + +Bgls())]
+ inf 7 Ex f(x) -(v)), (3) #,730"
(L)
where L1 is a set of absolutely integrable functions with respect to underlying measures a. B. The exchange between infimum and supremum is possible by strong duality (Slater's condition). If one decomposes the outer expectation Ex in the last equation as Bu(w)) E (yll))) we can notice that the supremum by fxx) should satisfy to the condition:
((x) (x) inf Ex(ol=)c((, v) e(v)]= inf [o(a, T(x)) - -(((())) T:X+y
otherwise, the infimum by would yield the -c value. Operator g is called c-conjugate transfor- mation. If MP-KP, the solution π (y]z² deterministic and one may set Tty(2) [z)). Finally, DP (3) may be reduced to a single potential optimization task (using inequality (4), ref. Villani et al. [2009] Theorem 5.10):
D) = sup (E) T: T
(x)) Eglo(2)]
(5)
ge Le (3)
sup inf
Ea [c(æ, T(=))] + Eelg(v)] -
(6)
In practice, during optimization process the infimum by T in c-conjugate transformation is approxi- mated by a parametric model Te, such that
g0(x) S [g c(x, To()) - g(Ta(z)).

---

A number of approaches were proposed to model To in equation (6) e.g., with the Input Convex Neural Networks (ICNN) (Amos al. [2017]. Makkuva et al. [2020], Taghvaei and Jalali [2019]) or with arbitrary non-convex neural networks (Rout et al. [2021]. Korotin et al. [2023b]. Most of these approaches make an assumption that the cost is squared Euclidean and utilize Brenier theorem (Brenier [1991]), from which optimal map recovers as gradient of convex function. The main bottleneck of these parametric solvers is their instability in finding the optimal c-conjugate potential g" from equation (7) and the rough estimation of T often results in situation where the sum of potentials g and gf diverges. These instability problems were thoroughly discussed in works Amos [2023], Korotin et al. [2021]. Recently, Amos [2023] showed that it is possible to find near exact conjugate approximation by performing fine-tuning on top of initial guess (Ig) in order to achieve closest lower bound in the inequality (7). Despite being an exact approximation to the true conjugate, such procedure requires extensive hyperparameter tuning and will definitely introduce an additional computational o
In this work, we propose to mitigate above issues by constraining the solution class of conjugate potentials through a novel form of expectile regression regularization Rg. In order to make joint optimization of g and To stable and more balanced (optimize g and To in the OT problem (6) synchronously and with the same frequency), we argue that it is possible to measure proximity of
potential to )² without an explicit estimation of the infimum in c-conjugate transform (4) and instead optimize the following objective:
(x)) + Egla(v)] - Ea, (Rg
(8)
The regularizer will have to constrain the differences g (() and (x)) that should match at the end of training. We show that such natural regularization outperforms the of the-art NOT approaches in all of the tasks of the established benchmark for the computational OT problems (the Wasserstein-2 benchmark, presented in Korotin et al. [2021]), with remarkable 5 to 10-fold acceleration of training compared to previous works and achieving faster convergence on synthetic datasets with desirable properties posed on OT map. Moreover, we show that proposed method obtains state-of-the-an results on generative image-to-imas -ingge tasks in terms of FID and MSE.
2 Related Work
In the essence, the main challenge of finding the optimal Kantorovich potentials in equation (6) lies in alternating computation of the exact c-conjugate operators (4). Recent approaches consider the dual OT problem from the perspective of optimization over the parametrized family of potentials. Namely, parametrizing potential 9, either as a non-convex Multi-Layer Perceptron (MLP) (Dam et al. [2019]) or an Input-Convex Neural Network (ICNN) (Amos et al. [2017]). Different strategies for finding the solution to the conjugate operator can be investigated under a more general formulation of the following optimization (Makkuva et al. [2020], Amos (2023]): min Ex (T T())]
max Eal9,, + [vl,,
(9)
with ((x) being the fine-tuned argmin of e-conjugate transform (4) with initial value (zr). Loss objective Lamor can be one of three types of amortization losses which makes (x(x) converge to ((x).. This max-min problem is similar to adversarial learning. where In acts discriminator and Te finds a deterministic mapping from the measure a to B. The first objective in equation (9) is well-defined under certain assumptions and the optimal parameters can be found by differentiating w.r.t. 1), according to the Danskin's 's envelope theorem (ref. Danskin [1966]). We briefly overview main design choices of the amortized models To (a) in the form of continuous dual solvers and the corresponding amortization objective options for Lamor in the Appendix C.
Another method considers the solution to the optimal map in (1) from different perspective by introducing a regularization term named Monge Gap (Uscidda and Cuturi [2023]) and learns optimal T map from Monge formulation directly without any dependence on conjugate potentials. More explicitly, by finding the reference measure fa with Support(a)) C Support(μ4), the following regularizer quantifies deviation of T from being optimal transport map:
M f
Ex [c(, T(x)] KP (H, Tw
(10)
"Tii intuition is supported by a direct evaluation in Section below.
3

---

with KP being cetropy-regularized Kantorovich problem (2). However, despite its elegance, we still need some method to compute KP and the underlying measure μ should be chosen thoughtfully. considering that its choice impacts the resulting optimal transport map and the case when l Oz does not always provide expected outcomes.
3 BBckground
Bidirectional transport mapping. We employ notations with hats (#. T. F. f (0)9 to indicate the correspondence to the solution (argmins) of the OT problem (6). Optimality in equation (6) is obtained whenever complementary slackness is satisfied, namely: v(x, y) E Support( (x)x)
f(y) c(r,,)) Consider specific setting when the optimal transport plan (1x, is deterministic and MP=KP. Let the domains of a, 3 be equal and compact, i.e. X = y, for some strictly convex function h the cost c(z, y) = h(x v). Denote by h" the convex conjugate of h. implying that (0h)- Vh.. If a is absolutely continuous, then 7 (zz, y) is unique and concentrated on graph (z,, (z)). Moreover,
may link it with Kantorovich potential f as follows (Santambrogio [2015] Theorem 1.17): Vf(x) E Ozc(a, (x)) and particularly for c(x, y) = h(x - y) ((x) = a Vh (V]()))
(11)
the same conditions are met for measure /3 we can express the inverse mapping T through the potential g:
()) VA Vyl)).
(12)
Max-min optimization in problem (6) by means of parametric models Ja and In, is unstable due non-convex nature of the problem. One way improve robustness is to simultaneously train bidirectional mappings ((x) and (y) expressed by formulas (11) and (12), thus yielding self- improving iterative procedure. During the training, we also may use the equations (11) and (12) with non-optimal functions fe and Iy, because there are no restrictions on T in problem (6) and we can use any representation for the transport mapping function. Under weaker constraints (for example X == = Rd),, the e-concavity of the potentials I and may be required. local e-concavity in the data concentration region or, if the conditions for equations (11) and (12) are not satisfied, we can use an arbitrary function for Ta and do not express through the potential fe. Under Brenier's theorem conditions (Brenier [1991]] 1]) in domain H for the squared Euclidean cost, holds that ((x) x Vf(x), where f is some l2-concave function. It follows that the optimal potentials f and ğ are l2- lo-concave, even if one uses not 12-concave potentials f. g in the training process.
this case, we can rely on
Expectile regression. The idea behind the proposed regularization approach to minimize the least asymmetrically weighted squares. It is a popular option for estimating conditional maximum of a distribution through neural networks. Recently, expectile regression was used in some offline Reinforcement Learning algorithms and representation learning approaches (Ghosh et al. [2023]). Let fe IRd + IR be some parametric model from L2 (IRd) space and x, y be dependent random variables in Rd x R. where y has finite second moment. By definition (Newey and Powell [1987]). the expectile regression problem is:
min E Cv B[ 0
fa(x)) = min E T Ily I fa(x)) (y fo(x))) 2 0.5.
(13)
The expectation is taken over the {æ, 3) pairs. The asymmetric loss C, reduces the contribution of those values of that are smaller than (z). while the larger values are weighted more heavily (ref. Figure 5). The expectile model fo (x: strictly monotonic parameter т. Particullary, the important property for us is when + 1, it approximates the conditional (on x) maximum operator over the corresponding values of y (Bellini al. [2014]). Below we compute c-conjugate transformation by means of expectile.
4 Proposed Method
The main motivation behind our method to regularize optimization objective in DP (6) with non-exact approximation of c-conjugate potential (z), defined in (7). The regularisation term

---

[z, y) should "pull" ( towards (g)) and (z) towards g²(x:). Instead of finding explicit
c-conjugate transform, we compute 7-expectile of random variables g (x) c(z, y), treating as condition. From the properties of expectile regression described above and equation (4) follows that when r + 1, the expectile converges to
max
[g" (x) c(z, y)] = -9 ZEX
(14)
Let the parametric models of Kantorovich potentials be represented as fo())) and g,(y). The transport mapping (1( has the same parameters as fo(œ) can be expressed through fo (ref. 11), or otherwise, when fe not used dinectictti training), it is its own parameters. Let approximate the maximum of eq. (14) by r-expectile of g (æ) - czz y) conditioning on So the target (term in eq. 13) of the expectile regression here is 9 (x) c(x, y). The model in this case is -y) It has negative sign because we approximate c-transform of 9 (x). which equals inf, c(z, y) 9 (x)) The corresponding regression loss C, (o (x) c(x, y) + 9, Accounting the definition of g" (7) obtain the regularization loss for potential 94,*
(n x. v) = L. (c(x, To(x)) Gn (To())) c(., y) + 9n(1))-
(15)
The proposed expectile regularisation is incorporated into alternating step of learning the Kantorovich potentials by implicitly estimating e-conjugate transformation, additionally encouraging model g to satisfy the e-concavity criterion (Villani et al. [2009] Proposition 5.8). We minimize Rg(n) = Ea, Fa.aR, (R, 2, y) by 77 and simultancously do training of the dual OT problem (6), splitting into two losses
(n) - E3ll9, (v)) + B. (T))((()))))
(16)
( = - (To (To Ex (c2, Te(x)
(17)
Algorithm ENOT Training
Input: samples from unknown distributions and Bi cost function c(z, y):
Parameters: parametric potential model f or vector field f = T, parametric potential model g, optimizers opt_f f and opt_g, batch size n, train steps N, expectile т, expectile loss weight X. bidirectional training flag i
function train_step (f, fæ1...., ..). [yi,...,Wy. ...))
[Assign OT mapping T(x)]
2: if is wu.bidrreeiionnl is true T(x) = r - Vh (VE()) else T(x) f(x)
3: (Compute dual OT losses and expectile regularisation Rg)
4: Lg 1 E ((((x))) 1 - g(w)
5: L = - xi=1 [o(x,, T(x.)) -(())))(((xxi)))
-
6: Rg = E L (c(o,, T(x:) c(xi, vt) + g(y;) - g(T(x:)))
7: (Apply gradient updates for parameters of models and g)
8: opt_f. mi i mi ze (f, loss = L₂)
9: opt_g. opt_g.minimme minimine ze (g, loss Ly ABg)
end function
[Main train loop)
11: for 1, do
sample 21...., a, V1. y1..... Un 13: if is bidirectional is false or t mod 2 = then
train train_step(f, g, (f1,..., 2n), {y vn.)
15: else
(Update inverse mapping B by swapping and g) 17: train train_step(q, f, (U }n), (I1,..... )
18: end if
19: end for
20: sample x ,,,..., In ~ a. "I...., Un ~ B
21: (Approximate OT distance by sum of conjugate potentials, get T from step 2)
22: dist En (g()) c(x, T(æ,)) - (T(()))) f(x) = c(z, T(x)) - <(T(x))
23: return f. g, dist
5

---

Algorithm I describes complete training loop with full objective expression (0) and hyperparameters 7 (expectile) and includes two training options: direciicati with models By and To; and bidirectional for strictly convex cost functions in form of hn( y) with models So: In and Te, 2T11 (the latter are represented in terms of fe. In by formulas (11), (12)). The bidirectional training procedure updates In, To in one optimization step and then switches to fo, T update in the next step. This option includes analogical regularisation term for the potential fo:
R2 x, v) C. (c(T7 (v). 3) (T (w)) c(z, y) fa()))
(w))
(18)
In the end of training approximate the correspondent Wasserstein distance by expression We ax 3) = Eglg,(u) Bal9;
Ea (c]c[[((,,] Tx(x))
(19)
with optimized parameters 0, f7. We include formal convergence amalysis for 7 regularized functions being a tight bound on the exact solution to the c-conjugate transform in Appendix D.
Experiments
In this section, we provide thorough validation of ENOT on popular W2 benchmark to test the quality of recovered OT maps. We compare ENOT with the state-of-the-ae approaches and also showcase its performance in generative tasks. Additional results and visualizations on 2D synthetic tasks are provided in Appendix F.
5.1 Results on W Benchmark
While evaluating ENOT, we also measured the wall-clock runtime all Wasserstein-2 benchmark data (Korotin et al. [2021]). The tasks the benchmark consist of finding the optimal map under the squared Euclidean norm e(2:, y) = || y/l between either: (1)
(HD) pairs (a, 3) of Gaussian mixtures, where the target measure is constructed as an average of gradients of learned ICNN models via WKoooon et al. [2019]) (2) samples from pretrained generative model W2GN (Korotin et al. [2021]) on CelebA dataset (Liu al. [2015]). The quality of the map T from a to is evaluated against the ground truth optimal transport plan T via unexplained variance percentage metric (CIVV) (Korotin et al. [2019, 2021]. Makkuva et al. [2020]), which quantifies deviation from the optimal alignment T", normalized by the variance of B:
Ex EallT(x) -<)))) (x))
C (T, az 100
(20)
Vargly]
The results of the experiments are provided in Table for CelebA64 (64 x 64 image size) and in Table 2 for the mixture of Gaussian distributions with varying number of dimensions D. Overall, ENOT manages to approximate optimal plan T* accurately and without any computational overhead compared to the baseline methods which require an inner conjugate optimization loop solution. To be consistent with the bascline approaches, we averaged our results across 3-5 different seeds. All the hyperparameters are listed in Appendix E.2 (Table 6).
Despite the fact that we compute at each train step a non-exact c-transform, the expectile regularization enables the method to outperform all exact methods in all our extensive tests. In actuality, the regularization does not introduce an additional bias, neither in theory, nor in practice. At the end of training (or upon convergence), we obtain the exact estimate of the c-conjugate transformation. Other methods demand near-exact estimation at each optimization step, requiring additional inner optimization and introducing significant overhead. We assume that introduces imbalance in the simultaneous optimization by g and T in equation (6), underestimating the OT distance as a result.
5.2 Different Cost Functionals
We further investigate how ENOT performs for different cost functions and compare Monge gap regularization (Uscidda and Cuturi [2023]) and ENOT between the measures defined on 2D synthetic datasets. In Figure 1, we observe that despite recovering Monge-like transport maps To. ENOT achieves convergence up to 2x faster and produces more desirable OT-likk optimal maps. To test other specific use cases, conducted experiments 2D spheres data (Figure 2), where parametrize

---

Method
Conjugate
Early Generator
Mid Generator
Late Generator
W2-Cycle
None
1.7
0.5
0,25
MM-Ohjective
None
2.2
0.9
0.53
MM-R-Obe MM-R-Objective
None
1.4
0.4
0.22
W2OT-Cycle
None
>
26,50= 60.14
0.29 0.59
W2O1-Objective
None
> 100
0.291 0.15
0.69 1 0.9
W2OT-Cycle
L-BFGS
0.62 1 0.01
0.20 1 0.00
0.09 0.00
W20T-Objective
L-BFGS
0.61 = 0.01
0.20 = 0.00
0.09 = 0.00
W20T-Regression
L-BFGS
0.62 + 0.01
0.20 = 0.00
0.09 + 0.00
W2OT-Cycle
Adam
0.65 + 0.02
0.21 = 0.00
0.11 0.05
W2OT-Objective Adam
0,65 : 0.02
0,21 + 0.00
0.11 + 0.05
W20T-Regnession
Adam
0.66 = 0.01
0.21 = 0.00
0.12 = 0.00
ENOT (Ours)
None
0.32 + 0.011
0.08 + 0.004
0.04 + 0.002
Table 1: CVV comparison of ENOT CelebA64 tasks from the Wasserstein-2 benchmark. The attributes after the method names ('Cycle', "Objective", "Regression') correspond to the type of amortisation loss. Column 'Conjugate' indicates the selected optimizer for the internal fine-tuning of c-conjugate transform. The results of our method include the mean and the standard deviation across 3 different seeds. The best scores are highlighted.
Fitted map Sréhcmm Divergnnce
Ritted map Ta I. Monge cap
Fitted with ENCT
sorre
Figure Fitting of three different transport maps Ta between source and target measures in R with Fuclidean cost function c(t, y) | y/l- We use the same number of iterations and MLP architecture for each method. Left: Sinkhorn divergence; Middle: Monge gap; Right: ENOT.
Fitted Map Tg. Sinkhom divergence
Filted Map Tp. Monge gap
Filled Map T. ENOT regularization
Suer
Recovered OT maps To between synthetic measures on 2-sphere with geodesic cost c(x, y) arccos(<" y). All models are MLPs with outputs normalized to be on a unit sphere. Blue dots the empirical source measure, red crosses the empirical target measure and the orange crosses are the result of the found transport map. Left: Sinkhorn; Middle: Monge; Right: ENOT.
Figure 2:
the map To as a MLP and the algorithms with the geodesic cost c(r,y) arccos(x y) with 1000 iterations. We set here flag is wwwdiiiiiiiiiii...cce wwwrreccinaaaaallee.. (meaning the training mode one-directional this example). Remarkably, the time required for convergence is minimal for ENOT, while the Monge gap takes up to three times longer. Moreover, in our experiments, Monge gap solver diverged for n > 1300 iterations. ENOT consistently estimates accurate and continuous OT maps.
7

---

Method
Conjugate
D
D= 4
D ==8
D 16
D 32
W2-Cycle
None
0.1
0.7
2.6
3.3
6.0
MM-Objective
None
0.2
1.0
1.8
1.4
6.9
MM-R-Objective
None
0.1
0.68
2.2
3.1
Monge Gap
None
0.1 1.00
0.57 T 0.0
2.05 1 0.06
4.22 T 0.1
7.24
W2OT-Cycle
None
0.05 0.0
0.35 0.01
> 100
100
> 100
W2OT-Objective
None
100
> 100
> 100
100
> 100
W2OT-Cyccl
L-BFGS
100
100
> 100
100
100
W2OT-Objective
L-BFGS
0.03 0,03 0.0
0,22 : 0.01
0.11
2.09
0.31
W2OT-Ree Regg--sgr
L-BFGS
0.03 + 0.0
0.22 + 0.01
0.61 0.04
0.77 0.1
1.97
0.38
W2OT- W2OT-CCycle
Adam
0.18 0.03
0.69 + 0,56
1.62 = 2.82
100
> 100
W20TT-Objective
Adam
0.06 0.01
0.26 + 0.22
0.63 + 0.07
0.81 0.10
1.99
- 0.32
W2OT-
Adam
0.22 0.01
0.28 0.02
0.61 + 0.07
0.8 0.10
2.07
ENOT (Ours)
None
0.02 + 0.0
0.03 1 0.001
0.14 1 0.01
0.24 1 0.03
0.67
E 0.02
Method
Conjugate
D 64
D 128
D 256
W2-Cycle MM-Objective
None
7.2
2.0
None
8.1
2.2
2.6
MM-R-OUS OOM-R-Obeective
None
10.1
3.2
2.7
Monge Gap
None
7.99 1 0.19
9.1 + 0.29
9.41 = 0.21
W2OT-Cycle
None
> 100
> 100
> 100
W2OT-Objective
None
> 100
> 100
> 100
W2OT-Cycll
L-BFGS
> 100
> 100
> 100
W2OT-Objective
L- L-BFGS
2.08 0.40
0.67 0.05
0.59 0.04
W2OT-Regression
L-BFGS
2.08
0.39
0.67 + 0.05
0.65 = 0.07
W2OT-Cyele
Adam
>> 100
> 100
> 100
W2OT-Ohjective
Adam
2.21 + 0.32
0.77 + 0.05
0.66 + 0.07
W20T-Regression
Adam
2.37 + 0.46
0.77 + 0.06
0.75 0.09
ENOT (Ours)
None
0.56 + 0.03
0.3 + 0.01
0.51 0.02
Table LUV comparison of ENOT with baseline methods on the high-dimensional (HD) tasks from Wasserstein-2 benchmark. The suffixes ('Cycle', 'Objective', 'Regression') correspond to the type of amortisation loss. Column 'Conjugate' indicates the selected optimizer for the internal fine-tuning of c-conjugate transform. D is the dimension of the measures domain. The mean and the standard deviations of our method are computed across different seeds. The best scores are highlighted.
5.3 Unpaired Image-to-Inee Image Translation
To showcase the power of expectile regularization beyond the W2 benchmarks, we apply our method to unpaired image-to-image translation task. The corresponding image datasets are: female subset of Celebrity faces (CelebA(f)) (Liu et al. [2015]), Anime Faces (Anime)², FlickrFaaces-HQ (FFHQ) (Karras al. [2019]), comic faces v2 (Comics)", Handbags and Shoes'. The datasets pre pre-processed in the conventional way as described in (Gazdieva et al. [2023]). The trained transport maps include: Handbags to Shoes, FFHQ Comics, CelebA(f) to Anime. We employ squared Euclidean cost function divided by the image size (64 or 128), basic U-Net architecture (Ronneberger al. [2015]) for the transport map (1x)), and ResNet from WGAN-QC (Liu et al. [2019]) potential 9,, (y). ENOT trains one-directional mode with total steps count = 120k. Appendix Table 7 contains complete list of hyperparameters (conventionally, we have used FID (Heusel et al. [2017]) metric for the hyperparameters tuning). We report the learned transport maps in Figure 3 as well the widely used FID and MSE metrics Table 3. Appendix 10 includes additional evaluation on test images.
kaaggl e wttp///tsstseeetttaaaaaaaaaaannn http:///aaaaaeeeaaaaaaoo///oooaa wwwpaao-ooiicaccappapaa--ee-e-pp red- mad-syytheti c-v2

---

Image-to to- tange-to-mmage translation baselines include popular GAN-based approaches: CycleGAN (Zhu et al. (2017]) and StarGAN-v2 (Choi al. [2020]), and two other recent neural OT methods: Extremal OT (Gazdieva et al. [2023]) and Kernel OT (Korotin et al. [2023a]). ENOT outperforms the baselines in all tasks therms of FID score and, as all other experiments, significantly speeds up the computation. It takes about hours to train the transport model on one GPU RTX 3090 Ti with image size 64 64 and about 16 hours when the image size is 128 x 128, while an approximate training time of the other OT algorithms and GANs is about 3 days on the same GPU.
Figure 3: Left: Handbags to Shoes; Middle: FFHQ Comies; Right: CelebA(f) to Anime; all images sizes are 128x128, the 1st row contains the source images, the 2 row contains predicted generative mapping by ENOT; Cost function: L divided by the image size.
Task and image size
CycleGAN
StarGAN
Extr. OT
Ker. OT
ENOT
Handbags + Shoes 128
23.4
22.36
27.10
26.7
19.19
FFHQ + Comics 128
20.95
20.81
17.11
CelebA(f) + Anime 64
20.8
22.40
14.65
18.28
13.12
CflebA(f) + Anime 128
19.44
21.96
18.85
FID Metric
Task and image size
CycleGAN
StarGAN
Extr. OT
Ker. OT
ENOT
Handbags + Shoes 128
0.43
0.24
0.37
0.37
0.34
FFIIQ + Comics 128
0.22
0.21
0.20
ColebA(f) > Anime 64
0.32
0.21
0.30
0.34
0.26
CelebA(f) + Anime 128
0,31
0.36
0.28
MSE Metric
Table 3: Comparison of ENOT baseline methods for image-to-image translation. We evaluate generation task between two different datasets: Source = Target. And compare resulting images based on Frechet Inception Distance (FID) and Mean Squared Error (MSE). Empty cells indicate that original authors of particular method did not include results for those tasks.
5.4 Ablation Study: Varying hy hyperparameters expectile and regularization weight
Figure presents the study of the impact of the proposed expectile regularization on the ΣV metric. This is done by varying the values of the expectile hyperparameter 7 and the scaling of the expectile loss coefficient in Algorithm 1. Colored contour plots show the areas of the lowest and the highest values of C.V. The grey areas depict the cases when the OT solver diverged. For example, in high-dimensions, 2 64, is the case for À = 0, pointing out that the expectile regularization with is necessary prevent the instability during training.
The ablation study shows that, even when the parameter choices of I and À are not optimal, ENOT still outperform the other baseline solvers in Table 2, making ENOT approach robust to extensive hyperparameter tuning, compared to amortized optimization approach (Amos [2023]) sensitive to the hyperparameters of the conjugate solver. All ablation study experiments were conducted using the network structure and the learning rates in Appendix E.2 (Table 4) (they coincide with those in Table 2).

---

08
weight
weight
030
0.5
4.55
des
Expoctile Values 7
Expectile Values
abrs
Expectile Values
225 0275
Figure 4:
Contour plots of L dependence the values of À and in Algorithm 1 for the dimensions of = 256 (Left, NaN values are greyed out), D = 128 (Middle), and D = 64 (Right).
Conclusion, Limitations and Future Work
Our paper introduces neural optimal transport with the help of expectile regularisation. We show that a solution to such a regularization objective is indeed close upproximation to the true e-conjugate potential. Remarkably, ENOT surpasses the current state-of- of-thh--rttrr upproaches, yielding an up to a 10-fold improvement in terms of the computation speed both synthetic 2D tasks and well-recognized Wasserstein-2 benchmark.
method, ENOT, for efficient computation of the conjugate potentials in
The proposed regularized objective on the conjugate potentials relies two additional hyperparam- eters, namely: the expectile coefficient т and the expectile loss trade-off sealer X, thus requiring re-evaluation for new data. However, given the outcome of our ablation studies, the optimal parameters found on the Wasserstein-2 benchmark are optimal enough or least provide a good starting point.
We believe that ENOT will become new baseline to compare against for the future NOT solvers and will accelerate research in the applications of optimal transport in high-dimensional tasks, such as generative modelling. As for future directions, ENOT can be tested with the other types of cost functions, such as Lagrangian costs, defined on non-Euclidean spaces and in the dynamical optimal transport settings, such as flow matching.
References
L. Ambrosio, K. Deckelnick, G. Dziuk, M. Mimura, V. A. Solonnikov, H. M. Soner, and L, Ambrosio. Lecture notes optimal transport problems. Springer, 2003.
B. Amos. Tutorial on amortized optimization for learning to optimize over continuous domains. MXXi:22002.00665. 2022.
B. Amos. amortizing convex conjugates for optimal transport. The Eleventh International Conference on Learning Representations (ICLR), Kigali, Rwanda, 2023.
B. Amos, Xu, and Z. Kolter. Input convex neural networks. In D. Precup and Y. W. Teh, editors, Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 146-155. PMLR, 06-1 Aug 2017.
M. Arjovsky, S. Chintala, and L. Bottou. Wasserstein generative adversarial networks. In International conference on machine learning, pages 214-223. PMLR, 2017.
Bellini, B. Klar, A. Müller, and E. Rosazza Gianin. Generalized quantiles as risk measures. Insurance: Mathematics and Economics, 54(C):41-48, 2014.
I. Bespalov, N. Buzun, and D. Dylov. Brulé: Barycenter-regullriize unsupervised landmark extraction. Pattern Recognition, 131:108816, 2022a.
Bespalov, N. Buzun, O. Kachan, and D. V. Dylov. Lambo: Landmarks augmentation with manifold-barycentiee oversampling. IEEE Access, 10:117757-117766 2022b.
Y. Brenier. Polar factorization and monotone rearrangement of vector-valued functions. Communica- tions pure and applied mathematics, 44(4):375-41 1991.
10

---

C. Bunne, L. Papaxanthos, A. Krause, and M. Cuturi. Proximal optimal transport modeling of population dynamics. In International Conference on Artificial Intelligence and Statisties, pages 6511-6524 PMLR, 2022.
N. Buzun. Gaussian Approximation for Penalized Wasserstein Barycenters. Math. Meth. Stat., 32(3):
1-26, 2023. doi: 01310//$1066330233100000
Y. Chen, M. Telgarsky, C. Zhang, B. Bailey, D. Hsu, and J. Peng. A gradual, semi-discrete approach to generative network training via explicit wasserstein minimization. 36th International Conference on Machine Learning. ICML, pages 1845-1858, 2019.
Y. Choi, Y. Uh, J. Yoo, and J.-W. Ha. Stargan v2: Diverse image synthesis for multiple domains. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), June 2020.
D.-A. Clevert, T. Unterthiner, and S. Hochreiter. Fast and accurate deep network learning by exponential linear units (elus). arXiv:1511.07288 2015.
M. Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, 26, 2013.
M. Cuturi, L. Peng--aapaxnntos, Y. Tian, C. Bunne, G. Davis, and O. Teboul. Optimal transport tools (ott); A jax toolbox for all things wasserstein. arXiv preprint arXiv:2201.12322 2022.
N. Dam, Q Hoang, T. Le. T. D. Nguyen, H. Bui, and D. Phung. Three-player wasserstein gan via amortised duality. Proceedings of the Twenty-Eighth International Joint Conference on Artificial Intelligence, IJCAI-19, pages 2202-2208. International Joint Conferences on Artificial Intelligence Organization, 2019. doi: 10. 00.22963/jccai2211/3333
J. M. Danskin. The theory of max-min, with applications. SIAM Journal on Applied Mathematics, 14 (4):641-664, 1966.
A. Fan Jiaojiao, Taghvaei and Y. Chen. Scalable computations of wasserstein barycenter via input convex neural networks. In M. Meila and T. Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 1571-1581. PMLR, 18-24 Jul 2021.
A. Fickinger, S. Cohen, Russell, and B. Amos. Cross-domain imitation learning via optimal transport. International Conference on Learning Representations, 2022.
M. Gazdieva, A. Korotin, D. Selikhanovych, and E. Burnaev. Extremal domain translation with neural optimal transport. In sevethh Conference on Neural Information Processing Systems, 2023.
D. Ghosh, C. A. Bhateja, and Levine. Reinforcement learning from passive data via latent intentions. Proceedings of the 40th International Conference Machine Learning, 202:1 1321-11333 23-29 Jul 2023.
I. Guirajani, F. Ahmed, M. Arjovsky, V. Dumoulin, and A. C. Courville. Improved training of wasserstein gans. Advances in neural information processing systems, 30, 2017.
S. Haldar, V. Mathur, D. Yarats, and L. Pinto. Watch and match: Supercharging imitation with regularized optimal transport. In Conference on Robot Learning, pages 32-43. PMLR, 2023.
M. Heusel, H. Ramsauer, T. Unterthiner, B. Nessler, and S. Hochreiter. Gans trained by a two time- scale update rule converge a local nash equilibrium. In I. Guyon, U. von Luxburg, S. Bengio, H. M. Wallach, R. Fergus, V. N. Vishwanathan, and R. Gamett, editors, Advances in Neural Information Processing Systems 30: Annual Conference Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pages 6626-6637, 2017.
T. Karras, S. Laine, and T. Aila. A style-based generator architecture for generative adversarial networks. In 2019 IEEF/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 1396-4408 2019. doi: 10. 10.1109/CVPR.220900002

---

A. Khamis, R. Tsuchida, M. Tarek, V. Rolland, and L. Petersson. Scalable optimal transport methods in machine learning: A contemporary survey. IEEE Transactions on Pattern Analysis and Machine Intelligence, page 1-20, 2024. ISSN 1939-3539. doi: 10.1
A. Korotin, V. Egiazarian, A. Asadulaev, A. Safin, and E. Burnaev. Wasserstein-2 generative networks. arXiv preprint arXiv: 1 30882, 2019.
A. Korotin, L. Li, A. Genevay, J. M. Solomon, A. Filippov, and E. Burnaev. Do neural optimal transport solvers work? continuous Wasserstein-2 benchmark. Advances in Neural Information Processing Systems, 34, 2021.
A. Korotin, D. Selikhanovych, and E. Burnaev. Kernel neural optimal transport. In International Conference on Learning Representations, 2023a.
A. Korotin, D. Selikhanovych, and Burnaev. Neural optimal transport. The 11th International Conference on Learning Representations, 2023b.
A. Kroshnin, V. Spokoiny, and A. Suvorikova. Statistical inference for B barycenters. The Annals of Applied Probability, 31(3):1264 1298, 2021. doi: 10.1214/20- AAP1618.
J. Leygonie, She, A. Almahairi, S. Rajeswar, and A. Courville. Adversarial computation of optimal transport maps, arXiv preprint arXiv: 1906.9961, 2019.
H. Liu, X. Gu, and D. Samaras. Wasserstein gan with quadratic transport cost. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), October 2019.
Z. Liu, P. Luo, X. Wang, and X. Tang. Deep learning face attributes in the wild. In Proceedings of International Conference on Computer Vision (ICCV), December 2015.
Y. Luo, Z. Jiang, S. Cohen, Grefenstette, and M. Deisenroth. Optimal transport for offline imitation learning. arXiv preprint arXiv:2303. 2023.
A. Makkuva, A. Taghvaei, S. Oh, and J. Lee. Optimal transport mapping via input convex neural networks. In International Conference on Machine Learning, pages 6672-6681. PMLR, 2020.
A. Mallasto, J. Frellsen, W. Boomsma, and A. Feragen. (q, w)-wasscrstein gans: Comparing ground metrics for wasserstein gans. ArXiv, abs/1992..3642, 2019.
E. F. Montesuma, F. N. Mboula, and A. Souloumiac. Recent advances in optimal transport for machine learning. arXiv:2306.1615 2023.
W. Newey and J. Powell. Asymmetric least squares estimation and testing. Econometrica, 55(4): 819-47, 1987.
G. Papagiannis and Y. Li. Imitation learning with sinkhorn distances. In Joint European Conference on Machine Learning and Knowledge Discovery in Databases, pages 116-131. Springer, 2022.
G. Peyré, M. Cuturi, et al. Computational optimal transport: With applications to data science. Foundations and Trends in Machine Learning, 11(5-6):355-6000 607, 2019.
O. Ronneberger, P. Fischer, and T. Brox. U-net: Convolutional networks for biomedical image segmentation. In N. Navab, J. Hornegger, W. M. Wells, and A. F. Frangi. editors, Medical Image Computing and Computer-Assinn Assisted Intervention MICCAI 2015, pages 234-241, Cham, 2015. Springer International Publishing.
L. Rout, A. Korotin, and E. Burnaev. Generative modeling with optimal transport maps. 10.00:2110.02999 2021.
W. Rudin. Principles of Mathematical Analysis. International series in pure and applied mathematics. McGraw-Hill 1976. ISBN 9780070856134,
Santambrogio. Optimal transport for applied mathematicians. Birkäuser, NY, (5((5-53))5544 2015.
Shen, Qu, W. Zhang, and Y. Yu. Wasserstein distance guided representation learning for domain adaptation. In Proceedings of the AAAI Conference Artificial Intelligence, volume 32, 2018.
12

---

N. Shvetsov, N. Buzun, and D. V. Dylov. Unsupervised non parametric change point detection in clectrocardiography. Proceedings of the 32nd International Conference on Scientific and Statistical Database Management, 2020.
A. Taghvaei and A. Jalali. 2- 2-waserrstein approximation via restricted convex potentials with applica- tion to improved training for gans. arXiv:190.00117,,20 2019.
Uscidda and M. Cuturi. The monge gap: A regularizer to learn all transport maps. arXiv preprint 0et::22022049533, 2023.
C. Villani et al. Optimal transport: old and new, volume 338. Springer, 2009.
Y. Xie, M. Chen, H. Jiang, T. Zhao, and H. Zha. On scalable and efficient computation of large scale optimal transport. In International Conference on Machine Learning, pages 6882-6892. PMLR, 2019.
J.-Y. Zhu, T. Park, Isola, and A. Efros. Unpaired image-to-image translation using cycle- consistent adversarial networks. In 2017 IEEE International Conference Computer Vision (ICCV), pages 2242-2251, 2017. doi: 10.1109/1CCV2201722244.
13

---

Appendix
A Expectile visualisations
Expectile loss function
Expectiles of cond. distribation
0.03
42
34
Figure 5: Expectile regression. Left: the asymmetric squared loss Ly The value T = 0.5 corresponds to the standard MSE loss, while 0.9 and 0.99 give more weight to the positive differences. Right: expectile models (x). The value r = 0.5 corresponds to the conditional statistical mean of the distribution, and when 1 approximates the maximum operator over the corresponding values of y.
B W case.
For squared Euclidean cost c(x. v) = 2 || v11 one may use ordinary conjugation and replace the vector norms outside the supremum in (6). Let Kantorovich potential g equals 11ly u(y), then
(x)) inf (211== w/12 2 ||vl12 u(v) ||0:|| (1)
(21)
and consequently from (6) we derive that
W2(a, 3) = - Ex |
sup
Bal-u" (x)] Eal Sssluuu) u(]]
(22)
2
ue uEL, (g)
By equation (12) the corresponding optimal transport map ((x) equals to the gradient of i (argmax- imum from the last formula):
f(x) - r- -
(23)
C Expanded Review of Related Work
In this section of the Appendix briefly outline categorization different approaches for estimating dual Kantorovich problem as proposed by (Amos [2023, 2022]), where the following differentiable amortization loss design choices are highlighted:
A. Objective-based learning: (Camor = Lobi) methods utilize local information (4) to establish optimal descent direction for model's parameters 0. (Dam al. [2019]) predicts approximate amortized solution To (z) from equation (7) by minimizing the next expression over mini- batch of samples from ac
(Tx()) (o)))) - (T)
(24)
Methods max-min [MM] (Dam et al. [2019]). max-min batch- wisc [MM-B] (Mallasto al. [2019], Chen et [2019]), max-min + ICNN [MMv1] (Taghvaei and Jalali [2019]), Max-min + 2 ICNNs [MMv2] (Makkuva et al. [2020], Fan Jiaojiao and Chen [2021]). [W20T-Object ve] (Amos (2023]) use such objective-based amortization in order to learn optimal prediction. However, objective-base methods are limited by computational costs and predictions made by amortized models can be overestimated, resulting in sub- optimal solution.
14

---

B. Regression-based Amortization (Amos [2022]) (Caner = Leg)) is an instance of regression- based learning, which can be done by fitting model's prediction Ta(x) into ground-truth solution fxx), taking Euclidean distance as proximity measure:
(To() [T -x))]"
(25)
Such choice for learning (x) is computationally efficient and works best when ground- truth solutions T(x) are provided. However, there guarantees for obtaining optimal solution when Txx) not unique.
C. Cycle-based Amortization (Camor Leyele) based the first order optimality criteria for equation (4), i.e Vyc(a, y) (y) If c(1,, y) |a yll then Vyc(x, v) and one may use the following expression in the loss
min E. Cyylle (T)) min E., || To( (To(
(26)
It is called cycle-consisteng, regularization. Method [W2] Korotin et al. [2019] uses this choice and substitutes it from the dual loss (9) to avoid solving max-min problem.
D Conjugate Function Approximation by Expectile
Lemma D.1 (Rudin [1976])- Let random vector & have a compact support Ω and Vo E Ω: fa+1(x)) (a(r) be sequence of continuous functions. Then from functional convergence f -> follows convergence of fa(E) to f(E) with probability
Theorem D.2. Denote by fr C a non-parametric solution of expectile regression in class of continuous functions that approximates the T-th (r > 0.5) conditional expectile of c-conjugate transform g(n) - ct, 1)). Let function be support 52, then with probability 1
and random vectors &, 1j have compact
lim lim (+(() = gg((() = sup (g(n) c(E,n)} T-1
(27)
NEA
Proof. First note that (E( < (5) with probability 1, otherwise it would be possible to reduce the average value of the loss function L, by taking (g) By the monotonicity property of the expectile (ref. Bellini al. [2014]) Va E S2 T1: frz(2)) > fr. (z). When + 1 for each x l2 it holds that (z) converges to (x) as monotone and bounded sequence. By means of Lemma B. we also derive that (E) converges to (E) with probability 1.
E
Implementation Details
E.1
Environment and Libraries
We implement ENOT in JAX framework, making fully compatible and easily integrable with the OTT-JAX library (Cuturi et al. [2022]). Moreover, since ENOT introduced expectile regularization, there no additional overhead and whole procedure easily jilcompiled, which a drastic difference with previous approaches. To find the optimal hyperparameters in Appendix (E.2), we used Weights & Biases sweeps for hyperparameter grid search and Hydra for managing different setup configurations. ENOT implementation consists of only single file, which is casy to reproduce and can be tested on other datasets of interest. We provide step-by-step tutorial of benchmarking ENOT by the following link at OTT-JAX on the website https: sskylooop gi github. e/ontt/.
E.2 Hyperparameters for Wasserstein-2 Benchmark Tasks
Tables and 6 provide detailed hyperparameter values used the experiments in Section 5. To find the values of these parameters, we performed an extensive grid search sweep across different seeds, yielding the best results among the seeds, on average. We tried be as close as possible in terms of hyperparameters previous works. Likewise, we tested different choices of hidden layers
15

---

and found that the most stable training occurs at 2 512, but we found that for low www-dimensionll tasks (i.e D < 64), 128 neurons are enough to achieve lowest L compared results reported in (Amos [2023], Makkuva et al. [2020], Korotin et al. [2023b]). Since ENOT does not introduce any additional computational overhead, increasing number of neurons will not slow down overall training time. Runtime comparison with (Amos [2023]) for W-2 benchmark presented in Table 9.
Hyperparameter
Value
potential model fo conjugate model go fa hidden layers go hidden layers training iterations activation function oottiiier
non-convex MLP
non-convex MLP
[512, 512, 512] is D 64, else [128, 128, 128)
200 000
ELU (Clevert al. [2015])
Adam with cosine annealing (a = 1c-4)
g optimizer Adam f 3 Adam 3
[0.9, 0.9]
[0.9, 0.7]
initial learning rate
3e-4
expectile coef. À expectile
0.3
batch size
1024
Table 4: Hyperparameters for D-dimensional Gaussian Mixture Wasserstein-2 benchmark tasks.
Hyperparameter
Value
potential model fo conjugate model ge fo hidden layers ge hidden layers training iterations activation function f optiiiierr g optimizer Adam f Adam 3 initial learning rate expectile coef. À expectile batch size
non-convex MLP
non-convex MLP
[64, 64, 64, 64]
100 000
ELU (Clevert et al. [2015])
Adam with cosine annealing (a = Le-4)
[0.9, 0.9991
5e-4
0.3
0.99
1024
Table Hyperparameters for Synthetic 2D datasets from (Rout et al. [2021])
Hyperparameter
Value
potential model fo conjugate model go hidden layers training iterations activation function optimizer g optimizer Adam f Ø Adam g 3 initial leamning rate expectile coef. A expectile batch size
ConvPotential (Amos [2023])
6 Conv Layers
80 000
ELU (Clevert al. [2015])
Adam with cosine annealing 1c-4)
[0.5, 0.5]
3e-4
1.0
0.99
64
Table 6: Hyperparameters for CelebA64 Wasserstein-2 benchmark tasks.
16

---

Hyperparameter
Value
potential model fo conjugate model de training iterations activation function f Toptimizer g optimizer Adam /
UNet (Ronneberger et al. (2015))
ResNet (Liu (2019])
120 000
ReLU and LeakyReLLU(0.2)
Adam with cosine anncaling (a = 1c-2)
[0.5, 0.5]
Adam 9
initial learning rate
le-4
initial learning rate
5e-5
expectile onef. À expectile
1.0
0.98
batch size
64
Table
Hyperparameters for image to image translation tasks.
MLP Hidden layers
Method
Runtime
W20T (L-BFGS)
60 min
[64, 64, 64, 64]
ENOT
1.3 min
W20T (L-BFGS) ENOT
120 min
[128, 128, 128, 128]
1.3 min
W20T (L-BFGS)
300 min
256,, 256, 256, 256]
ENOT
1.3 min
Table 8: Runtime comparison for different layers sizes between W2OT (Amos [2023]) with default hyperparameters and ENOT on synthetic 2D data tasks from Rout al. [2021].
F Results on Synthetic 2D Datasets
Additionally, we evaluate the performance of ENOT on synthetic datasets, introduced in Makkuva al. [2020] and Rout et [2021]. Here, all neural networks are initialized as non-convex MLPs. and for each optimal plan found by ENOT, we demonstrate difference between ground truth Sinkhorn W2(a, (3) distance and optimal plan found by ENOT, which recovered from learned potentials by equation (23). Figures 7 and 8 show the estimated optimal transport plans (in blue) both in forward and backward directions recovered by To R 3 and T. B R a and the countour plots of the learned potential functions respectively. Also, in Table we compare the runtime to complete 20k iterations using amortized method from W2OT (Amos [2023]). Table depicts how runtime changes for varying number of hidden layers in non-convex MLP, while keeping other hyperparameters for amortized model to those recommended from original paper with LBFGS solver. Additional details on the full list of hyperparameters is included in Appendix E.2 (Table 4).
F1 Synthetic 2D Tasks Details
Table 5 lists optimal parameters for ENOT for synthetic-2D tasks from Rout al. [2021]. Amos [2023] pointed out that LeakyReLU activation works better compared to ELU used for Wasserstein-2 benchmark. However, for expectile regularisation we found out that ELU works as well for synthetic 2d tasks. We keep Adam 3 parameters as default [0.9, 0.999] and observe that 25k training iterations are enough to converge for tasks from Rout et [2021]. Moreover, we tried different neurons per layer and Table 8 shows runtime in minutes in comparison to previous state-of-ther -ftth-r-- approach. Such speedups are made possible due to efficient utilization of jit compilation since ENOT does not use any inner optimizations.
F.2 Additional results with varying expectile hyperparameter
To characterize ENOT performance as a function of expectile т, we performed evaluation with ranging it from 0.5 to 0.999 in several tasks:
17

---

Method
D =0
2
D = 4
D
D
16
32
D
64
D
128
D
256
W2OT
157
108
91
140
246
397
571
1028
ENOT (Ours)
14
14
15
15
15
16
21
21
Table Comparison of runtimes (in minutes) against the baseline (W2OT-Objective L-BFGS) the high-dimenne diah--iimnsionll (HD) tasks from the Wasserstein-22 benchmark with same networks architecture.
7
2
(D 256)
W2, Synth. 2D
FID (CelebA + Anime)
MSE (CelebA + Anime)
0.5
0.55
33.47
16.43
0,264
0.6
0.52
12.63
16.28
0.260
0.7
0.51
9.59
15.95
0.265
0.8
0.49
1.4
15.19
0.262
0.9
0.5
0.06
13.87
0.266
0.95
0.54
0.03
14.27
0.267
0.999
0.55
0.02
13.91
0.288
Table 10: Performance of ENOT with varying levels of expectile hyperparameter 7 on W2 benchmark (Ist column), showcasing intuition on convergence as 7 - 1; Synthetic 2D data (2nd column); Image-to-lmage translation FID (3rd column), and MSE (4th column).
Image-to-image translation dataset (CelebA(f) to Anime with image size 64, Table 3);
Wasserstein-2 benchmark with D = 256 (Table 2);
Synthetic 2D dataset from Figure 7.
In these experiments (Table 10), we observe significant drop in performance when approaches 0.5 on the Synthetic 2D dataset and (CelebA(f) to Anime (in terms of FID). On Wasserstein-2 benchmark, the tendency is less evident. At the same time, values of in the range [0.9, 1.0) always demonstrate convergence of ENOT, giving good results in all experiments. Setting an instability. This can be the case because, under certain conditions, the overall contribution of proposed regularization term will be zero, which means that the potentials can become unbounded. However, in our experiments, such an instability occurred extremely rarely (mostly due to bad optimizer parameters), resulting only in a slight drop performance.
may cause
2.4
Loos wright A
0.0
16
0.75
0.6
0.4
0.8
05
0.2
08
3.07
Dimeusiion
Dimeusion
Figure 6: Ablating varying Left: loss weight and Right: expectile T coefficients in ENOT based dimension of task from W2 benchmark. Ly is shown.
F.3 Details on Generative Tasks
Table provides details on hyperparameters used for unpaired image-to-image translation from Section 5.3. We observe that ENOT outperforms GAN based approaches, such as CycleGAN and StarGAN-v2. ENOT also outperforms the closest similiar recent approachs for generative modelling based on NOt such as Extremal OT (Gazdieva et al. [2023]) and Kernel OT (Korotin et al. [2023a]).
18

---

Potential
Potential g
Source
brge:
Gloend
surce
Sager
Vyllage)
Sall" 8.38
Sinkhorn WilValaadnn o) 0.01
Potentials est.: Wafe,,
Ma - 0.06
Sinkhom est.: WollVolad.o) 0.06
Potentials est: (1(0.)) 21.22
Figure 7: Recovered optimal transport plans (f()) and (y) from (23)) and learned potentials contour plots obtained from solving OT dual problem (22) with squared Euclidean cost via ENOT regularisation synthetic datasets from Makkuva et al. [2020]. Evaluation metric is Sinkhorn distance between the measures, i.c. W2 (3, (2 T B). The estimated distance (22) from learned potentials compared with the reference value W2(a, (3).
B).
SCLETP
Figure 8: Recovered optimal transport push forward map (23) visualization for squared Euclidean cost using ENOT algorithm on synthetic datasets from Rout et al. [2021].
F.4 W2- Wy2-Benchmark Tasks Details
High- linee-iimnsional measures (HD) task from Korotin et al. [2019] tests whether OT solvers redistribute mass among modes varying measures. Different instantiations of Gaussian mixtures in dimensions D=2, 4, 16, ... 256 are compared between each other via OT. In the benchmark. Mix3t oMi x10 1 used, where source measure a can consist of random mixture of Gaussians and target measure consist of two random mixtures 91 32 of 10 Gaussians. Afterwards, pretrained OT potentials # = 3 are used to form the final pair as (ax, ] (V + Vx)2)
Images task produces pair candidates for OT solvers in the form of high--dimensional images from CelebA64 faces dataset (Liu al. [2015]). Different pretrained checkpoints (Early, Mid, Late, Final) from WGAN-QC model (Liu et al. [2019]) are used to pretrain potential models. Target measure for final checkpoint constructed as average between learned potentials via [W2] solver and forms a pair input for OT algorithm (CKCelebA. B BCkpt) (OeFinal, [ ( V (66))+ #CFmma))) ))). Figure 9 shows an example pair of two such measures (a. B))
19

---

Figure 9: Example pair from W2-Benchmark of CelebA faces. Top row: Images, which were produced as final checkpoint from WGAN-QC model. Bottom row: Images, obtained from early checkpoint of WGAC-QC model.
E5 Unpaired Image to Image Additional Results
Figure 10: Optimal transportation mapping found by ENOT for Top: Handbag (top row) = Shoes
(bottom row); Middle: FFHQ (top row) Comies (bottom row); Bottom: CelebA(f) (top row) Anime (bottom row) image-to-image translation tasks.
20

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We made a thorough experiments on established benchmarks and compare proposed method with the other NOT solvers. In all extensive tests, ENOT outperforms the competition both in terms of time and accuracy.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
We provided discussions on limitations in Section main paper (last two paragraphs).
Justification:
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full of assumptions and complete (and correct) proof?
21

---

Answer: [Yes]
Justification: We provided theoretical proof that expectile regularization approximates upper bound the exact c-transform.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: Along with detailed hyperparameter specifications in the Appendix, in- cluded easy to follow Jupyter notebook which can be found in supplementary materials, enabling the others fully reproduce the results in the paper.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g.. in the of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
22

---

5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: We provided fully reproducible code, which could be easily integrated into OTT-JAX framework. Moreover, we provided step-by-step jupyter notebook, showcasing the performance of the proposed algorithm in all discussed tasks.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https: //tp://ippppp.. cc/
public /guidessssseeeeeee ProdeSubmi ss wesoooooiioo for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (ht
//nips cc/ http:c/pppcccccccccccccccccc.cct cubbic//ge ides (codeSubmt i LonPol icy) for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We provided detailed hyperparameters specifications in the Appendix for each the tested benchmarks.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: All reported results are statistically significant. We include evaluation error (StDev) for each model and each dataset in the study across different runs and seeds.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by enror bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
23

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative еπτог rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We include the exact computer configuration in model in the main text.
Appendix and mention GPU
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: /neurips...... wccppublic/tthicscuii..o
Answer: [Yes]
Justification:
We read it and adhered to the ethical guidelines.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [No]
Justification: Our paper does not address the societal impact as we operate with common datasets and benchmarks for testing Neural Optimal transport solvers.
Guidelines:
The answer NA means that there is no societal impact of the work performed.

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), faimess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Not applicable this work.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but encourage authors to take this into account and make best. faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We properly refer to the original papers and use the open source codes from official repositories, providing the direct URLs to them.
Guidelines:
The answer NA means that the paper does
existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
25

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, papersw sw thoode com/ datasets s has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Question:
Answer: [Yes]
Justification: We include all of the details corresponding to train procedures, datasets used, and citations. Moreover, we provide a readme file for repository details. The released code is legally approved for the publication; no special documentation is needed.
Guidelines:
The answer NA means that the paper does release assets.
Researchers should communicate the details the dataset/code/mooell as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: No crowdsourcing was used this study.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
No human studies/IRB was needed for this study.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
26

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
27


Reviews:
[{"review_id": "RkkWkz8wup", "reviewer": "Reviewer_2fef", "summary": "This work tackles the problem of estimating the dual Kantorovich OT potentials \nparametrized via neural networks.  They propose to approximate conjugate operator \nusing a well motivated loss called expectile regularization which approximates a conditional\nmaximum operator, well suited for estimating the c-transform. The authors illustrate the \nperformance of their method against several benchmarks showcasing a significant improvement \nin both performance and runtime.", "questions": "1. The expectile loss is well suited to approximate the c-transform only if tau \nis large enough (close to 1). Therefore, it would make sense to see a significant \ndrop in performance when tau is close to 0.5. It would be have been nice to validate \nthis intuition in figure 4 for instance.\n\n2. I would expect this to lead to some form of tradeoff. Letting tau->1 should lead to\n some form of instability and/or slower convergence ? Is that the case ? How difficult \nis this tradeoff to settle in practice ? \n\n3. The fitted maps of Sinkhorn in fig1 and fig2 are odd, I suspect the entropic regularization \nwas very high, leading to a very very smooth transportation plan from which the fitted map\n was obtained by taking the maximum over values very close to each other.\n\n- L74: is *squared\n- L137: the*  proposed regularization approach\n- L160: the* proposed", "limitations": "none", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Very simple, novel and intuitive idea with important and significant impact\n- Thoroughly carried out experiments and evaluation", "weaknesses": "- Little analysis of the drawbacks of the proposed method as well as the effect of setting \nthe threhshold parameter tau (see Q1-2 below)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Q1: The expectile loss is well suited to approximate the c-transform only if tau is large enough (close to $1$). Therefore, it would make sense to see a significant drop in performance when tau is close to $0.5$. It would be have been nice to validate this intuition in figure 4 for instance.**\n\nThank you for raising this question. To characterize ENOT performance as a function of expectile $\\tau$, we performed additional evaluation with $\\tau$ ranging from $0.5$ to $0.999$ on **1)** image-to-image translation dataset (CelebA(female) $\\Rightarrow$ Anime with dim=$64$, Table $3$), **2)** $W_2$ benchmark with dim=$256$ (Table $2$), and **3)** 2D dataset from Figure $7$ (second row). In these experiments, we indeed observe a significant drop in performance when $\\tau$ approaches $0.5$ ($0.7, 0.6, 0.5$) on the 2D dataset and CelebA (female) $\\Rightarrow$ Anime with dim=$64$ (in terms of FID). On $W_2$ benchmark, the tendency is less evident. At the same time, values of $\\tau$ in the range $\\left[0.90, 1.0\\right]$ always demonstrate convergence of ENOT, giving good results in all experiments. We will reflect on these new experiments in the revised manuscript. \n\n| Expectile ($\\tau$) |  *$\\mathcal{L}_2 ^\\text{UV}$* (dim=$256$) | $W_2$ - 2D | FID (CelebA(f) -> Anime) | MSE (CelebA(f) -> Anime) |\n| :---------: | :----------------: | :--------------: | :------------------------:| :------------------------: |\n| 0.5       | 0.55             | 33.47          | 16.43                    | 0.264                    |\n| 0.6       | 0.52             | 12.63          | 16.28                    | 0.26                     |\n| 0.7       | 0.51             | 9.59           | 15.95                    | 0.265                    |\n| 0.8       | 0.49             | 1.4            | 15.19                    | 0.262                    |\n| 0.9       | 0.5              | 0.06           | 13.87                    | 0.266                    |\n| 0.95      | 0.54             | 0.03           | 14.27                    | 0.267                    |\n| 0.999     | 0.55             | 0.02           | 13.91                    | 0.288                    |\n\n**Q2: I would expect this to lead to some form of tradeoff. Letting tau->1 should lead to some form of instability and/or slower convergence ? Is that the case ? How difficult is this tradeoff to settle in practice ?**\n\nSetting $\\tau \\simeq 1$ may indeed cause an instability. This can be the case because, under certain conditions, the overall contribution of proposed regularization term will be zero, which means that the potentials can become unbounded. However, in our experiments, such an instability occurred extremely rarely (mostly due to bad optimizer parameters), resulting only in a slight drop in performance.  Yet, we carefully show that the optimal $\\tau$ is usually in the range $\\left[0.9, 1.0\\right]$, with different values in this range yielding approximately the same result; so, selecting the optimal value of $\\tau$ becomes simple.\n\n**Q3: The fitted maps of Sinkhorn in fig1 and fig2 are odd, I suspect the entropic regularization was very high, leading to a very very smooth transportation plan from which the fitted map was obtained by taking the maximum over values very close to each other.**\n\nThank you for pointing this out. In order to be consistent with the experiments from the Monge gap paper (Uscidda and Cuturi [2023]), we used their exact default parameters of the entropic regularization, namely $\\epsilon = 0.01$ $\\cdot$ mean(C), with C being the cost matrix (in Figure 1, $\\epsilon \\sim 0.12)$, replicating the 'odd' look from the original paper. Moreover, we observe that lowering the values of $\\epsilon$ produces plots of similar transportation smoothness.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FTQWcyDK9Z", "reviewer": "Reviewer_E19z", "summary": "This paper introduced a new framework for the training of Neural Optimal Transport, a recently emergent paradigm to enable the training of optimal transport plan in larger scale settings. The key contribution lies in the new and novel formulation of a regularized training loss that takes into account an expectile regularization term, which makes the max-min training of the generative NOT more stable. Empirical evaluations demonstrated improvments on both reducing the evaluated metrics and shortening the training runtime.", "questions": "See weaknesses.", "limitations": "See weaknesses.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The authors seem to have done a thorough literature review.", "weaknesses": "- Theoretical motivation unclear: lack of references on the instability of finding optimal c-conjugate (i.e. the motivation of this work). Where are the citations for the sentence from line 76 to 78? What does it mean to make the optimization problem more \"balanced\" with the expectile regression regularization on (line 86)? The same applies when the authors claim expectile regression is a popular option for estimatiing maximum of  distributions through neural networks (line 130).\n\n- I suggest the authors rewrite their methodological section to be more compact and to the point (less verbose with the derivations, these can be put in the Appendix). Overall I was getting lost in this section, with the full main training objective never fully written, but have to resort to the pseudocode in algorithm 1 to see what was going on, and what the final loss was being used.\n\n\n- **Major:** for eq (13) and eq (14) to hold you need both $f$ and $g$ to be convex functions (a la Brenier's theorem) with respect to the input; the only option is to use parameterization with ICNN. However, in table 4-5 and 7 the authors used non-convex MLP as parameterization; this is contradicting with the write up in methodology part. Could the authors elaborate on this?\n\n- Ambiguous report of the empirical evaluations: most of the image task should use standard metric such as the FID. Many of the FIDs for the baselines are missing. \n\nWriting and notation nitpick: \n\n- I don't get the the derivation of the main loss function eq (17). The term in eq (15) is wrt to $\\theta$ and a term $y - f_\\theta$. What is the equivalence of $y$ and $f_\\theta$ for $\\mathcal{L}_\\tau$ in equation (17) then? \n\n- There should be explanation that we replace the quantity in Eq (7) to (5).\n\n- Optimization and optimisation are used interchangeably. Stick to one of these. \n\n- What is $\\hat{\\pi}$ right in the beginning of section 3? How are they different from $\\pi$?. The same with $\\hat{f}$ and $\\hat{g}$. In the background section, all of the notations are added with hats $\\hat{T}, \\hat{\\pi}, \\hat{f}, \\hat{g}$, then these when back to normal in Section 4; are they different problems from the introduction?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Q1:  Lack of references on the instability of finding optimal c-conjugate (the motivation of this work). What does it mean to make the optimization problem more \"balanced\" with the expectile regression regularization on (line 86)? The same applies when the authors claim expectile regression is a popular option for estimating maximum of distributions through neural networks (line 130).**\n\nWe are thankful for questions and for suggesting extra citations. We would like to stress that the main motivation behind our work was in finding fast, simple and accurate solver. The instability of finding accurate c-conjugate transform is the main bottleneck in the existing neural OT solvers. Those problems were thoroughly discussed in the following works (refer to Refs in our paper, e.g.,  Amos 2022a and Korotin et al 2021). For example, Korotin et al 2021 (Section 4.3, page 7) state that \"The exact conjugate solver is slow since each optimization step solves a hard subproblem for computing the conjugate. Solvers that approximate the conjugate are also hard to optimize: they either diverge from the start or diverge after converging to nearly-optimal saddle point\".\nMoreover this claim is empirically supported by the results in Tables 1 and 2 (i.e., W2OT-Objective, W2OT-Cycle diverged in all W2 tasks). Remarkably, our proposed regularization method stabilizes the overall training procedure.\n\nThe term \"more balanced\" here means that we optimize optimal map $T$ and potential $g$ in the OT problem (Equation 6) synchronously and with the same frequency.\n\nRecently, expectile regression was used in some offline Reinforcement Learning algorithms and representation learning approaches. Among them, Implicit Q-Learning (IQL) methods [R1, R2].\n\n[R1] Ilya Kostrikov, Ashvin Nair and Sergey Levine. Offline Reinforcement Learning with Implicit Q-Learning, 2022. \n\n[R2] Dibya Ghosh, Chethan Anand Bhateja, Sergey Levine. Reinforcement Learning from Passive Data via Latent Intentions, 2023.\n\n**Q2 Major: for eq (13) and eq (14) to hold you need both f and g to be convex functions (a la Brenier's theorem) with respect to the input; the only option is to use parameterization with ICNN. However, in table 4-5 and 7 the authors used non-convex MLP as parameterization; this is contradicting with the write up in methodology part. Could the authors elaborate on this?**\n\nIn Section (3), the variables with hats ($\\hat{\\pi}$, $\\hat{T}$, $\\hat{f}$, $\\hat{g}$) correspond to the solution (argmins) of the OT problem (6).  The equations (13) and (14) hold under the specified conditions (source and target domains $X$, $Y$ are compact and the measures are absolutely continuous). So there is no strict requirement on $c$-concavity (or convexity); the potentials $f$ and $g$ must merely be the solution to the OT problem (ref. Santambrogio Theorem 1.17). During the training, we also may use the equations (13) and (14) with non-$c$-concave functions $f$ and $g$, because there are no restrictions on $T(x)$ in problem (6) and we can use any representation for the transport mapping function.  \n \nUnder weaker constraints (for example $X=Y=R^n$), the $c$-concavity of the potentials $f$ and $g$ may be required. In this case, we can rely on a local $c$-concavity in the data concentration region. In addition, if the conditions in equations (13) and (14) are not satisfied, we can use the mode (is\\_bidirectional = False), in which we use an arbitrary function for $T(x)$ and do not express it through the potential.  \n\nIn practice, the is no benefit in using ICNN, because it makes sense only for the squared Euclidean cost, and even in this case the optimization problem deals with additional constraints, becoming more difficult. Similar empirical observations of under-performance of ICNNs were reported in (Amos 2022a). \n\n**Q3: Ambiguous report of the empirical evaluations: most of the image task should use standard metric such as the FID. Many of the FIDs for the baselines are missing.**\n\nRespectfully, the missing values in Table 3 mean that the corresponding experiments were not performed in the baseline publications. Given the suggestion, we performed additional experiments using the official implementation of the OT methods reported by the other authors to fill in the gaps, which we include in a table below and propose to add to our revision. \n\n| Task (FID)              | Extr. OT | Ker. OT | ENOT  |\n| :---------------------: | :--------: | :-------: | :-----: |\n| Handbags ⇒ Shoes 128  | 27.1     | 26.7    | 19.19 |\n| FFHQ ⇒ Comics 128     | 20.95    | 20.81\\* | 17.11 |\n| CelebA(f) ⇒ Anime 64  | 14.65    | 18.28\\* | 13.12 |\n| CelebA(f) ⇒ Anime 128 | 19.44\\*  | 21.96   | 18.85 |\n\n**Q4: What are $\\hat{\\pi}, \\hat{T}, \\hat{f}, \\hat{g}$?**\n\nWe are sorry for the ambiguity. In Section 3, the hats ($\\hat{\\pi}, \\hat{T}, \\hat{f}, \\hat{g}$) correspond to the solution of the OT problem in equation 6. As written in the beginning of this section, $\\hat{\\pi}$ is the **optimal** transport plan, $\\hat{T}$ - corresponds to the **optimal** transport mapping,  $\\hat{f}, \\hat{g}$ are the **optimal** Kantorovich potentials. The notation without $\\hat{\\cdot}$ merely means that the optimization is done over the given variable.\n\n**Q5: Derivation of Equation 17:**\n\nWe approximate the maximum of eq. (16) by $\\tau$-expectile of $g^T(x) - c(x, y)$ conditioning on $y$. So the target (term $y$ in eq. 15) of the expectile regression here is  $g^T(x) - c(x, y)$. The model is then $-g_{\\eta}(y)$ with a negative sign as we approximate c-transform of $g^T(x)$ as $\\inf_x c(x, y) - g^T(x)$. The corresponding regression loss is: $\nL_{\\tau} \\big(\n     g^T(x) - c(x, y)  + g_\\eta(y)\n    \\big).$\nTaking definition of $g^T(x)$ from eq. (7), we derive the expression (17): \n$L_{\\tau} \\big(\n     c(x, T_\\theta(x)) - g_\\eta(T_\\theta(x)) - c(x, y)  + g_\\eta(y)\n    \\big).$\n\n**Q6: Full Objective:** is defined in lines 168-169. The expression is $L_g(\\eta) + L_f(\\theta)  + \\lambda R_g(\\eta)$.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "IYzhtw34DC", "reviewer": "Reviewer_ZPj5", "summary": "Authors provided a new, theoretically justified loss in the form of expectile regularisation which stabilize the learning of Neural Optimal Transport.  Importantly proposed method outperforms previous state-of-the-art approaches on the established Wasserstein-2 benchmark tasks and image-to-image by a large margin.", "questions": "- The authors consider two training modes depending on the ‘is_bidirectional’ parameter. In what cases the bidirectional training mode works better and what it depends on?\n- The algorithm (Algorithm 1 line 2) states that in case 'is_bidirectional' is false $T(x)=\\nabla f(x)$. From what it follows?\n- Other algorithms that use expectile regression tend to use 0.9 expectile, why in your method you take the expectile parameter closer to 1?\n- By what metric can we choose the hyperparameters of $\\lambda$ and $\\tau$ ?\n- Partial OT is often required in practical tasks. It is possible to generalize this method to solve the problem of partial OT?\n- L2 cost is probably not the best choice in the image-to-image translation task, can we use any cost, e.g. parameterised by the neural network, in coupling with your method?", "limitations": "Limitations are discussed.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "Originality: \n- Expectile regularisation appears to be novel and not previously considered in OT literature.\n\nClarity:\n- Paper is well written and easy to follow.\n\nSignificance:\n- Paper providing an efficient non max-min solver for computing neural optimal transport maps, which is important.", "weaknesses": "Quality: \n- Comparison with other methods is focusing mainly on Wasserstein-2 benchmark and image-to-image translation. Comparison with some popular in OT biology problems (see 1, 2) would improve the contribution. \n\nReferences:\n1) TrajectoryNet: A Dynamic Optimal Transport Network for Modeling Cellular Dynamics Alexander Tong, Jessie Huang, Guy Wolf, David van Dijk, Smita Krishnaswamy\n\n2) Light Schrödinger Bridge: Alexander Korotin, Nikita Gushchin, Evgeny Burnaev", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Q1: Comparison with other methods is focusing mainly on Wasserstein-2 benchmark and image-to-image translation. Comparison with some popular in OT biology problems (see 1, 2) would improve the contribution.**\n\nWe thank the reviewer for the links to these interesting articles and the application of OT in biology. We analyzed the cited works and concluded that our method is not explicitly suitable for solving the problems from these papers, because they require finding highly non-linear maps. But we have made a rough linear estimation on the proposed biological dataset (taken from TrajectoryNet paper) with our method, given in the table below. Despite learning linear maps, ENOT outperforms baseline methods based on score-matching (DSBM) and flow matching (SB-CFM).\n| Solver  | W1 metric    |\n| ------- | ------------ |\n| ENOT    | 0.88 +- 0.03   |\n| LightSB | 0.82 +- 0.01   |\n| SB-CFM  | 1.22 +- 0.4   |\n| DSBM    | 1.78 +- 0.42 |\n\n**Q2  The authors consider two training modes depending on the ‘is_bidirectional’ parameter. In what cases the bidirectional training mode works better and what it depends on?**\n\nThe bidirectional mode works only with strongly convex functions $h(x-y)$. It significantly improves the solution when $T(x)$ is a discontinuous function (as in Figure 7). We propose to mention this in the Discussion Section.\n\n**Q3 The algorithm (Algorithm 1 line 2) states that in case 'is_bidirectional' is false . From what it follows?**\n\nWhen the parameter  is\\_bidirectional = False, we can use any parameterization for $T(x)$ by a neural network. The expression $T= \\nabla f$ is used to avoid introducing additional variables, but here, we assume **any** function.  \n\n**Q4 Other algorithms that use expectile regression tend to use 0.9 expectile, why in your method you take the expectile parameter closer to 1?**\n\nThe point is that with $\\tau$ converging to 1, we get a more accurate estimate of the c-conjugate transformation. But from the experiments in Figure 4, we can conclude that the values in the range 0.9-1.0 also give a good result. \n\n**Q5 By what metric can we choose the hyperparameters**\n\nConventionally, one could use L2-UVP metric for the experiments on W2 benchmark (Tables 1,2) and FID metric for the Image-to-Image tasks (Table 3). \n\n**Q6  Partial OT is often required in practical tasks. It is possible to generalize this method to solve the problem of partial OT?**\n\nOur method can be easily adapted for a special case of partial OT, when the whole source measure $\\alpha$ and only a part of the target measure $\\beta$ are used. Such problem is considered in Gazdieva et al. 2023 (citation from the paper).  For the general case, we are currently unable to suggest a proper method of adaptation. Still, if a new method for a neural partial OT is proposed elsewhere, it will also positively benefit from the use of our expectile regularization.  \n\n**Q7 L2 cost is probably not the best choice in the image-to-image translation task, can we use any cost, e.g. parameterised by the neural network, in coupling with your method?**\n\nIndeed, the L2 distance between the activations from VGG16 layers can be used as a cost function. We also conducted experiments with it, but it turned out that the ordinary L2 cost gives more accurate results. Perhaps, this stems from the fact that the faces are centered in all images, enabling a direct comparison.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "TCiYDK6x6C", "reviewer": "Reviewer_QVdS", "summary": "The paper introduces ENOT (Expectile-Regularized Neural Optimal Transport), a new method for training Neural Optimal Transport (NOT) models. It improves the efficiency of NOT solvers by incorporating a novel expectile regularization on dual Kantorovich potentials. Empirically, the authors use a Wasserstein-2 benchmark to demonstrate their method's improvement in accuracy and runtime.", "questions": "The quantile regularization term biases the estimation of the conjugate operator, making it non-exact. Despite this, the authors note that the regularization enables the method to outperform all exact methods. Could the authors speculate on the reasons for this? Is it solely due to the numerical stability of exact methods?\n\nIn Section 5.4, the authors mentioned\n\n>Figure 8 presents the study of the impact of the proposed expectile regularization on the L_2^{\\text{UV}} metric. \n\nI believe this is a typo. The figure that the authors refer to is Figure 6 in the Appendix.", "limitations": "See \"Weakness\".\n\nNitpicking: the authors may want to consider differentiating the commands \\citep and \\citet in latex citation and use them accordingly.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well-organized and clearly written.\n\n- The proposed expectile regularization provides a new way to stabilize the learning of dual potentials. Empirically, this works really well, compared to existing NOT formulations.\n\n- The experimental results are comprehensive.\n\n- The authors were very upfront about the limitations of their approach, e.g., the requirement of extra hyperparameters $\\tau$ and $\\lambda$.", "weaknesses": "To my understanding, while the empirical performance of ENOT is strong, the technical novelty of the approach is limited. Expectile regularization combines the c-conjugate transform (7) and the quantile regression loss (15). ENOT merges two established ideas to produce strong empirical results.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are thankful for raising important questions and for the positive feedback. \n\nFirst, we would like to address the concerns about novelty:\n\nTo the best of our knowledge, ENOT is the first approach that ventures into the approximation of the c-transform by means of expectile regularization. Moreover, ENOT is first to apply this approximation for efficient and accurate solution of Neural OT by completely mitigating additional $c$-conjugate optimization as done in prior works. These features provide at least a 3-fold better and a 10-fold faster training performance against the competitors.\n\n**Q1: The quantile regularization term biases the estimation of the conjugate operator, making it non-exact. Despite this, the authors note that the regularization enables the method to outperform all exact methods. Could the authors speculate on the reasons for this? Is it solely due to the numerical stability of exact methods?**\n\nWe acknowledge the need to discuss this and would like to add the following text to the paper. In actuality, the expectile regularization does not introduce an additional bias, neither in theory, nor in practice.  The formal convergence for the c-conjugate transform, when $\\tau$ converges to 1, is discussed in Appendix D.  At the end of training (or upon convergence), we obtain the exact estimate of the $c$-conjugate transformation. Other methods demand near-exact estimation at each optimization step, requiring additional inner optimization and introducing significant overhead. We assume that introduces an imbalance in the simultaneous optimization by $T$ and $g$ in Equation 6, underestimating the OT distance as a result.\n\n**Q2: Section 5.4 typo**\n\nThank you for careful reading. We have made the correction.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "zVXiy5iW3l", "reviewer": "Reviewer_qjtB", "summary": "-\tThis paper introduces a new regularizer for learning Neural Optimal Transport. The proposed method, called Expectile-Regularized Neural Optimal Transport (ENOT), is based on the expectile regression. ENOT demonstrates competitive results on the Wasserstein-2 benchmark and Unparied Image-to-Image Translation.", "questions": "- The ENOT and the Monge gap are similar in that they both serve as regularizers for inducing Neural Optimal Transport. Could you provide an additional quantitative comparison between ENOT and Monge gap? Currently, there are only qualitative examples in Fig 1 and 2.\n- Table 3 presents the FID results on the Unpaired Image-to-Image Translation task. On the other hand, Table 10 in the appendix includes additional MSE results. For the Unpaired Image-to-Image Translation task, the goal of Neural Optimal Transport is to accurately approximate the target distribution (FID) while minimizing the cost function (MSE). Hence, the MSE results should be included in the main part of the manuscript.\n\n**Typo:**\n- Line 220: Table 10 -> 3\n- Line 230: Figure 8 -> 4", "limitations": "-\tThe authors addressed the limitations of their work.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- The proposed method is easy to implement as it only requires an additional regularizer.\n- The ENOT achieves competitive results on the Wasserstein-2 benchmark.\n- The ENOT proposes a new regularizer that softly guides the potential $g$ towards $g^{cc}$ using expectile regression.", "weaknesses": "- The derivation of the ENOT regularizer (Eq 17) requires further clarification.\n- Please see the Questions Section below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Q1: The ENOT and the Monge gap are similar in that they both serve as regularizers for inducing Neural Optimal Transport. Could you provide an additional quantitative comparison between ENOT and Monge gap? Currently, there are only qualitative examples in Fig 1 and 2.**\n\nWe are thankful for this suggestion. We performed an additional comparison with the Monge Gap on $W_2$ benchmark (dataset from Table 2), which can be found below. Given the potential impact of these results on the ultimate algorithm selection for NOT, we propose to add the table to the main text. \n| Monge Gap    | ENOT         | DIM (W2 bench) |\n| :------------: | :------------: | :--------------: |\n| 0.1 +- 0.0   | 0.02 +- 0.0  | 2              |\n| 0.57 +- 0.0  | 0.03 +- 0.0  | 4              |\n| 2.05 +- 0.06 | 0.14 +- 0.01 | 8              |\n| 4.22 +- 0.1  | 0.24 +- 0.03 | 16             |\n| 7.24 +- 0.17 | 0.67 +- 0.02 | 32             |\n| 7.99 +- 0.19 | 0.56 +- 0.03 | 64             |\n| 9.1 +- 0.29  | 0.3 +- 0.01  | 128            |\n| 9.41 +- 0.21 | 0.51 +- 0.02 | 256            |\n\n**Q2: Table 3 presents the FID results on the Unpaired Image-to-Image Translation task. On the other hand, Table 10 in the appendix includes additional MSE results. For the Unpaired Image-to-Image Translation task, the goal of Neural Optimal Transport is to accurately approximate the target distribution (FID) while minimizing the cost function (MSE). Hence, the MSE results should be included in the main part of the manuscript.**\n\nWe put the results with the MSE metric in the Appendix, because the majority of baselines we compare against did not include the outcome for those tasks. For convenience, we performed the experiments with the baseline methods and can include them into the main part of the paper:\n\n| Tasks (MSE)              | Extr. OT | Ker. OT | ENOT |\n| :---------------------: | :--------: | :-------: | :----: |\n| Handbags ⇒ Shoes 128  | 0.37     | 0.37    | 0.34 |\n| FFHQ ⇒ Comics 128     | 0.22     | 0.21\\*  | 0.2  |\n| CelebA(f) ⇒ Anime 64  | 0.3      | 0.34\\*  | 0.26 |\n| CelebA(f) ⇒ Anime 128 | 0.31\\*   | 0.36    | 0.28 |\n\nOther issues:\nEquation (17): please kindly refer to our response to Reviewer E19z above.\n\nWe are also thankful for spotting the typos, which we have corrected.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank our reviewers for their constructive feedback. In the attached PDF, we provide additional experiments requested by the reviewers, showcasing the performance of ENOT. Corresponding tables are duplicated within individual responses to each reviewer.

Meta-review: The paper introduces a new method for estimating the dual Kantorovich OT potentials, parameterized via neural networks, along with an interesting new loss function. It includes extensive benchmarking that demonstrates the practical benefits of the method, as well as a thorough literature review on the topic. Additionally, the paper proposes a new, efficient solver and is well-organized and clearly written. It also explores the limitations of the method and the impact of extra hyperparameters. While the motivation could be strengthened, and the reviewers have provided several suggestions for improvement in this direction, I recommend acceptance.
Decision: Accept (spotlight)