Title: Expectile Regularization for Fast and Accurate Training of Neural Optimal Transport
Venue: NeurIPS 2024
Accepted: yes
Keywords: Optimal Transport

Neural Optimal Transport

Expectile

Regularization

Kantorovich 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

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 ExpectileRegularized Neural Optimal Transport (ENOT), outperforms previous state-ofthe-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.

Project page with code https://skylooop.github.io/enot/

1 Introduction

Computational optimal transport (OT) has enriched machine learning (ML) by offering a new viewangle on the conventional ML 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) as a differentiable proxy, with the OT distance playing the role of a similarity metric between the measures, or 2) as a 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 one of the measures to be defined in a closed form ( e.g. , Gaussian or uniform) or to be pairwise-aligned, admitting various applications of the OT theory. Both loss objective and generative formulations of OT proved successful in a vast range of modern ML areas, including

generative modelling (Arjovsky et al. [2017], Gulrajani et al. [2017], Korotin et al. [2021], Liu et al. [2019], Leygonie et al. [2019]), reinforcement learning (Fickinger et al. [2022], Haldar et al. [2023], Papagiannis and Li [2022], Luo et al. [2023]), domain adaptation (Xie et al. [2019], Shen et al. [2018]), 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-dimensional measure spaces are usually intractable, oftentimes solvable only for the distributions which admit a closed-form density formulations. As a result, the need for computationally-efficient OT solvers has become both evident (Peyré et 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 a deterministic OT plan T between two probability measures α and β supported on measurable spaces X , Y ⊂ R d with Borel sigma-algebra. The OT problem in Monge's formulation (MP) for a cost function c : X × Y → R is stated as:

where { T : T # α = β } is the set of measure-preserving maps, defined by a push forward operator T # α ( B ) = α ( T -1 ( B )) = β ( B ) for any Borel subset B ⊂ Y . The minimizer of the cost above exists if X is compact, α is atomless (i.e. ∀ x ∈ X : α ( { 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 is 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]):

where Π [ α , β ] = { π ∈ P ( X × Y ) : ∫ Y d π ( x, y ) = d α ( x ) , ∫ X d π ( x, y ) = d β ( y ) } is a set of admissible couplings with respective marginals α , β . MP and KP problems are equivalent in case X = Y are compact, the cost function c ( x, y ) is continuous and α is atomless. Since KP (2) is convex, it admits dual formulation (DP), which is constrained concave maximization problem and is derived via Lagrange multipliers (Kantorovich potentials) f and g :

where L 1 is a set of absolutely integrable functions with respect to underlying measures α , β . The exchange between infimum and supremum is possible by strong duality ( Slater's condition). If one decomposes the outer expectation E π in the last equation as E π ( x ) E π ( y | x ) , we can notice that the supremum by f ( x ) should satisfy to the condition:

otherwise, the infimum by γ would yield the -∞ value. Operator g c is called c-conjugate transformation. If MP = KP, the solution π ( y | x ) is deterministic and one may set π ( y | x ) = T ( x ) . Finally, DP (3) may be reduced to a single potential optimization task (using inequality (4), ref. Villani et al. [2009] Theorem 5.10):

In practice, during optimization process the infimum by T in c-conjugate transformation is approximated by a parametric model T θ , such that

A number of approaches were proposed to model T θ in equation (6) e.g. , with the Input Convex Neural Networks ( ICNN ) (Amos et 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 a convex function. The main bottleneck of these parametric solvers is their instability in finding the optimal c-conjugate potential g T from equation (7) and the rough estimation of T often results in a situation where the sum of potentials g and g T 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 a near exact conjugate approximation by performing fine-tuning on top of initial guess ( T θ ) 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 overhead. 2

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 R g . In order to make joint optimization of g and T θ stable and more balanced (optimize g and T θ in the OT problem (6) synchronously and with the same frequency), we argue that it is possible to measure proximity of potential g to ( g T ) c without an explicit estimation of the infimum in c -conjugate transform (4) and instead optimize the following objective:

The regularizer will have to constrain the differences g ( y ) -( g T ) c ( y ) and g T ( x ) -g c ( x ) that should match at the end of training. We show that such a natural regularization outperforms the state-ofthe-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 a 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-art results on generative image-to-image 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 g η either as a non-convex Multi-Layer Perceptron ( MLP ) (Dam et al. [2019]) or as 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]):

with ̂ T ( x ) being the fine-tuned argmin of c -conjugate transform (4) with initial value T θ ( x ) . Loss objective L amor can be one of three types of amortization losses which makes T θ ( x ) converge to ̂ T ( x ) . This max-min problem is similar to adversarial learning, where g η acts as a discriminator and T θ finds a deterministic mapping from the measure α to β . The first objective in equation (9) is well-defined under certain assumptions and the optimal parameters can be found by differentiating w.r.t. η , according to the Danskin's envelope theorem (ref. Danskin [1966]). We briefly overview main design choices of the amortized models T θ ( x ) in the form of continuous dual solvers and the corresponding amortization objective options for L amor in the Appendix C.

Another method considers the solution to the optimal map in (1) from a 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 µ with Support( α ) ⊂ Support( µ ), the following regularizer quantifies deviation of T from being optimal transport map:

with KP ε being entropy-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 µ = α does not always provide expected outcomes.

3 Background

Bidirectional transport mapping. We employ notations with hats ( ̂ π , ̂ T , ̂ g , ̂ f = ( ̂ g ) c ) 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: ∀ ( x, y ) ∈ Support ( ̂ π ) : ̂ g c ( x ) + ̂ g ( y ) = c ( x, y ) . Consider a specific setting when the optimal transport plan ̂ π ( x, y ) is deterministic and MP=KP. Let the domains of α , β be equal and compact, i.e. X = Y , for some strictly convex function h the cost c ( x, y ) = h ( x -y ) . Denote by h ∗ the convex conjugate of h , implying that ( ∂h ) -1 = ∇ h ∗ . If α is absolutely continuous, then ̂ π ( x, y ) is unique and concentrated on graph ( x, ̂ T ( x )) . Moreover, one may link it with Kantorovich potential ̂ f as follows (Santambrogio [2015] Theorem 1.17): ∇ ̂ f ( x ) ∈ ∂ x c ( x, ̂ T ( x )) and particularly for c ( x, y ) = h ( x -y )

If the same conditions are met for measure β we can express the inverse mapping ̂ T -1 ( y ) through the potential ̂ g :

Max-min optimization in problem (6) by means of parametric models f θ and g η is unstable due to non-convex nature of the problem. One way to improve robustness is to simultaneously train bidirectional mappings ̂ T ( x ) and ̂ T -1 ( y ) expressed by formulas (11) and (12), thus yielding selfimproving iterative procedure. During the training, we also may use the equations (11) and (12) with non-optimal functions f θ and g η , 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 = Y = R d ), 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 or, if the conditions for equations (11) and (12) are not satisfied, we can use an arbitrary function for T θ and do not express it through the potential f θ . Under Brenier's theorem conditions (Brenier [1991]) in domain R n for the squared Euclidean cost, it holds that ̂ T ( x ) = x -∇ ̂ f ( x ) , where ̂ f is some l 2 -concave function. It follows that the optimal potentials ̂ f and ̂ g are l 2 -concave, even if one uses not l 2 -concave potentials f , g in the training process.

Expectile regression. The idea behind the proposed regularization approach is 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 f θ : R d → R be some parametric model from L 2 ( R d ) space and x, y be dependent random variables in R d × R , where y has finite second moment. By definition (Newey and Powell [1987]), the expectile regression problem is:

The expectation is taken over the { x, y } pairs. The asymmetric loss L τ reduces the contribution of those values of y that are smaller than f θ ( x ) , while the larger values are weighted more heavily (ref. Figure 5). The expectile model f θ ( x ) is strictly monotonic in 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 et al. [2014]). Below we compute c-conjugate transformation by means of expectile.

4 Proposed Method

The main motivation behind our method is to regularize optimization objective in DP (6) with non-exact approximation of c -conjugate potential g T ( x ) , defined in (7). The regularisation term

R g ( x, y ) should 'pull' g ( y ) towards ( g T ) c ( y ) and g T ( x ) towards g c ( x ) . Instead of finding explicit c -conjugate transform, we compute τ -expectile of random variables g T ( x ) -c ( x, y ) , treating y as a condition. From the properties of expectile regression described above and equation (4) follows that when τ → 1 , the expectile converges to

Let the parametric models of Kantorovich potentials be represented as f θ ( x ) and g η ( y ) . The transport mapping T θ ( x ) has the same parameters as f θ ( x ) if it can be expressed through f θ (ref. 11), or otherwise, when f θ is not used (one-directional training), it is its own parameters. Let approximate the maximum of eq. (14) by τ -expectile of g T η ( x ) -c ( x, y ) conditioning on y . So the target (term y in eq. 13) of the expectile regression here is g T η ( x ) -c ( x, y ) . The model in this case is -g η ( y ) . It has a negative sign because we approximate c-transform of g T η ( x ) , which equals to inf x c ( x, y ) -g T η ( x ) . The corresponding regression loss is L τ ( g T η ( x ) -c ( x, y ) + g η ( y ) ) . Accounting the definition of g T (7) we obtain the regularization loss for potential g η :

The proposed expectile regularisation is incorporated into alternating step of learning the Kantorovich potentials by implicitly estimating c -conjugate transformation, additionally encouraging model g to satisfy the c -concavity criterion (Villani et al. [2009] Proposition 5.8). We minimize R g ( η ) = E α , β R g ( η, x, y ) by η and simultaneously do training of the dual OT problem (6), splitting it into two losses

Algorithm 1 ENOT Training

Input : samples from unknown distributions x ∼ α and y ∼ β ; cost function c ( x, y ) ;

Parameters : parametric potential model f or vector field f = T , parametric potential model g , optimizers opt_f and opt_g , batch size n , train steps N , expectile τ , expectile loss weight λ , bidirectional training flag is_bidirectional ;

Algorithm 1 describes a complete training loop with full objective expression L f ( θ )+ L g ( η )+ λR g ( η ) and hyperparameters τ (expectile) and λ . It includes two training options: one-directional with models g η and T θ ; and bidirectional for strictly convex cost functions in form of h ( x -y ) with models f θ , g η and T θ , T -1 η (the latter are represented in terms of f θ , g η by formulas (11), (12)). The bidirectional training procedure updates g η , T θ in one optimization step and then switches to f θ , T -1 η update in the next step. This option includes analogical regularisation term for the potential f θ :

In the end of training we approximate the correspondent Wasserstein distance by expression

with optimized parameters ̂ θ, ̂ η. We include a formal convergence analysis for τ regularized functions being a tight bound on the exact solution to the c-conjugate transform in Appendix D.

5 Experiments

In this section, we provide a thorough validation of ENOT on a popular W2 benchmark to test the quality of recovered OT maps. We compare ENOT with the state-of-the-art 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 Wasserstein-2 Benchmark

While evaluating ENOT, we also measured the wall-clock runtime on all Wasserstein-2 benchmark data (Korotin et al. [2021]). The tasks in the benchmark consist of finding the optimal map under the squared Euclidean norm c ( x, y ) = ∥ x -y ∥ 2 between either: ( 1 ) high-dimensional (HD) pairs ( α , β ) of Gaussian mixtures, where the target measure is constructed as an average of gradients of learned ICNN models via W 2 (Korotin et al. [2019]) or ( 2 ) samples from pretrained generative model W2GN (Korotin et al. [2021]) on CelebA dataset (Liu et al. [2015]). The quality of the map ̂ T from α to β is evaluated against the ground truth optimal transport plan T ∗ via unexplained variance percentage metric ( L UV 2 ) (Korotin et al. [2019, 2021], Makkuva et al. [2020]), which quantifies deviation from the optimal alignment T ∗ , normalized by the variance of β :

The results of the experiments are provided in Table 1 for CelebA64 ( 64 × 64 image size) and in Table 2 for the mixture of Gaussian distributions with a 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 baseline 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 an 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 T θ , ENOT achieves convergence up to 2 × faster and produces more desirable OT-like optimal maps. To test other specific use cases, we conducted experiments on 2D spheres data (Figure 2), where we parametrize

Fitted map Te, with f Monge gap

Fitted map T, with f ENOT

+

source

source

source

target

xx

target

xx

target

push-forward

push-forward

push-forward

Fitted Map      , Sinkhorn divergence

Source

Target

Push-forward

Fitted Map      , Monge gap

Source

Target

Push-forward

Source

Target

Push-forward

Fitted Map      , ENOT regularization

the map T θ as a MLP and test the algorithms with the geodesic cost c ( x, y ) = arccos( x T y ) with n = 1000 iterations. We set here flag is_bidirectional=False (meaning the training mode is one-directional in 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.

5.3 Unpaired Image-to-Image Translation

To showcase the power of expectile regularization beyond the W 2 benchmarks, we apply our method to an 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) 3 , Flickr-Faces-HQ (FFHQ) (Karras et al. [2019]), comic faces v2 (Comics) 4 , Handbags and Shoes 5 . The datasets are pre-processed in the conventional way as described in (Gazdieva et al. [2023]). The trained transport maps include: Handbags to Shoes, FFHQ to Comics, CelebA(f) to Anime. We employ squared Euclidean cost function divided by the image size ( 64 or 128 ), basic U-Net architecture (Ronneberger et al. [2015]) for the transport map T θ ( x ) , and ResNet from WGAN-QC (Liu et al. [2019]) as a potential g η ( y ) . ENOT trains in one-directional mode with total steps count N = 120 k . Appendix Table 7 contains a 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 as the widely used FID and MSE metrics in Table 3. Appendix 10 includes additional evaluation on test images.

Image-to-image translation baselines include popular GAN-based approaches: CycleGAN (Zhu et al. [2017]) and StarGAN-v2 (Choi et 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 in therms of FID score and, as in all other experiments, significantly speeds up the computation. It takes about 5 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 × 128 , while an approximate training time of the other OT algorithms and GANs is about 3 days on the same GPU.

MSE Metric

5.4 Ablation Study: Varying hyperparameters expectile and regularization weight

Figure 4 presents the study of the impact of the proposed expectile regularization on the L UV 2 metric. This is done by varying the values of the expectile hyperparameter τ 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 L UV 2 . The grey areas depict the cases when the OT solver diverged. For example, in high-dimensions, D ≥ 64 , it is the case for λ = 0 , pointing out that the expectile regularization with τ is necessary to prevent the instability during training.

The ablation study shows that, even when the parameter choices of τ and λ are not optimal, ENOT still outperforms 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).

0.875

0.9

0.925

0.95

0.975

0.99

0

0.25

0.5

0.75

1

1.25

1.5

0.7

0.8

0.9

1

1.1

1.2

0.875

0.9

0.925

0.95

0.975

0.99

0

0.25

0.5

0.75

1

1.25

1.5

0.4

0.8

1.2

1.6

2

2.4

0.875

0.9

0.925

0.95

0.975

0.99

0

0.25

0.5

0.75

1

1.25

1.5

0.55

0.56

0.57

0.58

0.59

0.6

6 Conclusion, Limitations and Future Work

Our paper introduces a new method, ENOT, for efficient computation of the conjugate potentials in neural optimal transport with the help of expectile regularisation. We show that a solution to such a regularization objective is indeed a close approximation to the true c -conjugate potential. Remarkably, ENOT surpasses the current state-of-the-art approaches, yielding an up to a 10-fold improvement in terms of the computation speed both on synthetic 2D tasks and on well-recognized Wasserstein-2 benchmark.

The proposed regularized objective on the conjugate potentials relies on two additional hyperparameters, namely: the expectile coefficient τ and the expectile loss trade-off scaler λ , thus requiring a 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 at least provide a good starting point.

We believe that ENOT will become a 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

Ambrosio, K. Deckelnick, G. Dziuk, M. Mimura, V. A. Solonnikov, H. M. Soner, and L. Ambrosio. Lecture notes on optimal transport problems . Springer, 2003.

Amos. Tutorial on amortized optimization for learning to optimize over continuous domains. arXiv:2202.00665 , 2022.

Amos. On amortizing convex conjugates for optimal transport. The Eleventh International Conference on Learning Representations (ICLR), Kigali, Rwanda , 2023.

Amos, L. Xu, and J. 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-11 Aug 2017.

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.

Bespalov, N. Buzun, and D. V. Dylov. Brulé: Barycenter-regularized unsupervised landmark extraction. Pattern Recognition , 131:108816, 2022a.

Bespalov, N. Buzun, O. Kachan, and D. V. Dylov. Lambo: Landmarks augmentation with manifold-barycentric oversampling. IEEE Access , 10:117757-117769, 2022b.

Brenier. Polar factorization and monotone rearrangement of vector-valued functions. Communications on pure and applied mathematics , 44(4):375-417, 1991.

Bunne, L. Papaxanthos, A. Krause, and M. Cuturi. Proximal optimal transport modeling of population dynamics. In International Conference on Artificial Intelligence and Statistics , pages 6511-6528. PMLR, 2022.

Buzun. Gaussian Approximation for Penalized Wasserstein Barycenters. Math. Meth. Stat. , 32(3): 1-26, 2023. doi: 10.3103/S1066530723010039.

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.

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.07289 , 2015.

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

Cuturi, L. Meng-Papaxanthos, Y. Tian, C. Bunne, G. Davis, and O. Teboul. Optimal transport tools (ott): A jax toolbox for all things wasserstein. arXiv preprint arXiv:2201.12324 , 2022.

Dam, Q. Hoang, T. Le, T. D. Nguyen, H. Bui, and D. Phung. Three-player wasserstein gan via amortised duality. In Proceedings of the Twenty-Eighth International Joint Conference on Artificial Intelligence, IJCAI-19 , pages 2202-2208. International Joint Conferences on Artificial Intelligence Organization, 7 2019. doi: 10.24963/ijcai.2019/305.

M. Danskin. The theory of max-min, with applications. SIAM Journal on Applied Mathematics , 14 (4):641-664, 1966.

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.

Fickinger, S. Cohen, S. Russell, and B. Amos. Cross-domain imitation learning via optimal transport. International Conference on Learning Representations , 2022.

Gazdieva, A. Korotin, D. Selikhanovych, and E. Burnaev. Extremal domain translation with neural optimal transport. In Thirty-seventh Conference on Neural Information Processing Systems , 2023.

Ghosh, C. A. Bhateja, and S. Levine. Reinforcement learning from passive data via latent intentions. Proceedings of the 40th International Conference on Machine Learning , 202:11321-11339, 23-29 Jul 2023.

Gulrajani, F. Ahmed, M. Arjovsky, V. Dumoulin, and A. C. Courville. Improved training of wasserstein gans. Advances in neural information processing systems , 30, 2017.

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.

Heusel, H. Ramsauer, T. Unterthiner, B. Nessler, and S. Hochreiter. Gans trained by a two timescale update rule converge to a local nash equilibrium. In I. Guyon, U. von Luxburg, S. Bengio, H. M. Wallach, R. Fergus, S. V. N. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA , pages 6626-6637, 2017.

Karras, S. Laine, and T. Aila. A style-based generator architecture for generative adversarial networks. In 2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) , pages 4396-4405, 2019. doi: 10.1109/CVPR.2019.00453.

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.1109/tpami.2024.3379571.

Korotin, V. Egiazarian, A. Asadulaev, A. Safin, and E. Burnaev. Wasserstein-2 generative networks. arXiv preprint arXiv:1909.13082 , 2019.

Korotin, L. Li, A. Genevay, J. M. Solomon, A. Filippov, and E. Burnaev. Do neural optimal transport solvers work? a continuous wasserstein-2 benchmark. Advances in Neural Information Processing Systems , 34, 2021.

Korotin, D. Selikhanovych, and E. Burnaev. Kernel neural optimal transport. In International Conference on Learning Representations , 2023a.

Korotin, D. Selikhanovych, and E. Burnaev. Neural optimal transport. The 11th International Conference on Learning Representations , 2023b.

Kroshnin, V. Spokoiny, and A. Suvorikova. Statistical inference for Bures-Wasserstein barycenters. The Annals of Applied Probability , 31(3):1264 - 1298, 2021. doi: 10.1214/20-AAP1618.

Leygonie, J. She, A. Almahairi, S. Rajeswar, and A. Courville. Adversarial computation of optimal transport maps. arXiv preprint arXiv:1906.09691 , 2019.

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.

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.

Luo, Z. Jiang, S. Cohen, E. Grefenstette, and M. P. Deisenroth. Optimal transport for offline imitation learning. arXiv preprint arXiv:2303.13971 , 2023.

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.

Mallasto, J. Frellsen, W. Boomsma, and A. Feragen. (q, p)-wasserstein gans: Comparing ground metrics for wasserstein gans. ArXiv , abs/1902.03642, 2019.

F. Montesuma, F. N. Mboula, and A. Souloumiac. Recent advances in optimal transport for machine learning. arXiv:2306.16156 , 2023.

Newey and J. Powell. Asymmetric least squares estimation and testing. Econometrica , 55(4): 819-47, 1987.

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.

Peyré, M. Cuturi, et al. Computational optimal transport: With applications to data science. Foundations and Trends® in Machine Learning , 11(5-6):355-607, 2019.

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-Assisted Intervention - MICCAI 2015 , pages 234-241, Cham, 2015. Springer International Publishing.

Rout, A. Korotin, and E. Burnaev. Generative modeling with optimal transport maps. arXiv:2110.02999 , 2021.

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 , 55(58-63):94, 2015.

Shen, Y. Qu, W. Zhang, and Y. Yu. Wasserstein distance guided representation learning for domain adaptation. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 32, 2018.

Shvetsov, N. Buzun, and D. V. Dylov. Unsupervised non-parametric change point detection in electrocardiography. Proceedings of the 32nd International Conference on Scientific and Statistical Database Management , 2020.

Taghvaei and A. Jalali. 2-wasserstein approximation via restricted convex potentials with application to improved training for gans. arXiv:1902.07197 , 2019.

Uscidda and M. Cuturi. The monge gap: A regularizer to learn all transport maps. arXiv preprint arXiv:2302.04953 , 2023.

Villani et al. Optimal transport: old and new , volume 338. Springer, 2009.

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, P. Isola, and A. A. Efros. Unpaired image-to-image translation using cycleconsistent adversarial networks. In 2017 IEEE International Conference on Computer Vision (ICCV) , pages 2242-2251, 2017. doi: 10.1109/ICCV.2017.244.

Appendix

A Expectile visualisations

-1

-0.5

0

0.5

1

0

0.2

0.4

0.6

0.8

1

1.2

1.4

0.01

0.50

0.80

0.90

0.99

Expectile loss function

-1

-0.5

0

0.5

1

-3

-2

-1

0

1

2

Expectiles of cond. distribution

y

x

x

B Wasserstein-2 case.

For squared Euclidean cost c ( x, y ) = 1 2 ∥ x -y ∥ 2 , one may use ordinary conjugation and replace the vector norms outside the supremum in (6). Let Kantorovich potential g ( y ) equals 1 2 ∥ y ∥ 2 -u ( y ) , then

and consequently from (6) we derive that

By equation (12) the corresponding optimal transport map ̂ T ( x ) equals to the gradient of ̂ u ∗ (argmaximum from the last formula):

C Expanded Review of Related Work

In this section of the Appendix we briefly outline categorization of different approaches for estimating dual Kantorovich problem as proposed by (Amos [2023, 2022]), where the following differentiable amortization loss design choices are highlighted:

Objective-based learning: ( L amor = L obj) methods utilize local information (4) to establish optimal descent direction for model's parameters θ . (Dam et al. [2019]) predicts approximate amortized solution T θ ( x ) from equation (7) by minimizing the next expression over minibatch of samples from α :

Methods max-min [MM] (Dam et al. [2019]), max-min batch-wise [MM-B] (Mallasto et al. [2019], Chen et al. [2019]), max-min + ICNN [MMv1] (Taghvaei and Jalali [2019]), Max-min + 2 ICNNs [MMv2] (Makkuva et al. [2020], Fan Jiaojiao and Chen [2021]), [W2OT-Objective] (Amos [2023]) use such objective-based amortization in order to learn optimal prediction. However, objective-based methods are limited by computational costs and predictions made by amortized models can be overestimated, resulting in suboptimal solution.

Regression-based Amortization (Amos [2022]) ( L amor = L reg) is an instance of regressionbased learning, which can be done by fitting model's prediction T θ ( x ) into ground-truth solution ̂ T ( x ) , taking Euclidean distance as proximity measure:

Such choice for learning T θ ( x ) is computationally efficient and works best when groundtruth solutions ̂ T ( x ) are provided. However, there are no guarantees for obtaining optimal solution when ̂ T ( x ) is not unique.

Cycle-based Amortization ( L amor = L cycle) is based on the first order optimality criteria for equation (4), i.e ∇ y c ( x, y ) = ∇ y g η ( y ) . If c ( x, y ) = 1 2 ∥ x -y ∥ 2 then ∇ y c ( x, y ) = x -y and one may use the following expression in the loss

It is called cycle-consistency 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 ∀ x ∈ Ω : f n +1 ( x ) ≥ f n ( x ) be a sequence of continuous functions. Then from functional convergence f n → f follows convergence of f n ( ξ ) to f ( ξ ) with probability 1 .

Theorem D.2. Denote by f τ ∈ C 0 a non-parametric solution of expectile regression in class of continuous functions that approximates the τ -th ( τ > 0 . 5 ) conditional expectile of c-conjugate transform g ( η ) -c ( ξ, η ) . Let function g be upper-bounded and random vectors ξ , η have compact support Ω , then with probability 1

Proof. First note that f τ ( ξ ) ≤ -g c ( ξ ) with probability 1 , otherwise it would be possible to reduce the average value of the loss function L τ by taking -g c ( ξ ) . By the monotonicity property of the expectile (ref. Bellini et al. [2014]) ∀ x ∈ Ω,τ 2 > τ 1 : f τ 2 ( x ) > f τ 1 ( x ) . When τ → 1 for each x ∈ Ω it holds that f τ ( x ) converges to -g c ( x ) as monotone and bounded sequence. By means of Lemma B.1 we also derive that f τ ( ξ ) converges to -g c ( ξ ) with probability 1 .

E Implementation Details

E.1 Environment and Libraries

We implement ENOT in JAX framework, making it fully compatible and easily integrable with the OTT-JAX library (Cuturi et al. [2022]). Moreover, since ENOT introduced expectile regularization, there is no additional overhead and whole procedure is easily jit -compiled, which is 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 a single file, which is easy 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 or on the website https://skylooop.github.io/enot/ .

E.2 Hyperparameters for Wasserstein-2 Benchmark Tasks

Tables 4 and 6 provide detailed hyperparameter values used in 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 to be as close as possible in terms of hyperparameters to previous works. Likewise, we tested different choices of hidden layers

and found that the most stable training occurs at n ≥ 512 , but we found that for low-dimensional tasks (i.e D ≤ 64 ), 128 neurons are enough to achieve lowest L UV 2 compared to 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.

F Results on Synthetic 2D Datasets

Additionally, we evaluate the performance of ENOT on synthetic datasets, introduced in Makkuva et al. [2020] and Rout et al. [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 W 2 ( α , β ) distance and optimal plan found by ENOT, which is 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 T θ # α ≈ β and T -1 η # β ≈ α and the countour plots of the learned potential functions respectively. Also, in Table 8 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).

F.1 Synthetic 2D Tasks Details

Table 5 lists optimal parameters for ENOT for synthetic-2D tasks from Rout et 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 β parameters as default [0 . 9 , 0 . 999] and observe that 25k training iterations are enough to converge for tasks from Rout et al. [2021]. Moreover, we tried different neurons per layer and Table 8 shows runtime in minutes in comparison to previous state-of-the-art 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:

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 a 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 τ = 1 may 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.

4

8

16

32

64

128

256

0

0.25

0.5

0.75

1

1.25

1.5

0.2

0.4

0.6

0.8

1

4

8

16

32

64

128

256

0.875

0.9

0.925

0.95

0.975

0.99

0.4

0.8

1.2

1.6

2

2.4

F.3 Details on Generative Tasks

Table 7 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]).

Potential f

Potential g

source

target

Vf(source)

source

target

Vg(target)

W2((Vf)#α,3)=0.08

W2((Vg)#β,α)=0.01

Sinkhorn est.:W2(α,β)=21.97

Potentials est.:W2(α,β)=21.95

W2((Vf)#α,β)=0.08

W2((V9)#3,α)=0.06

Sinkhorn est.:W2(a,③)=21.23

Potentials est.:W2(α,β)=21.22

source

target

Vg(target)

F.4 W2-Benchmark Tasks Details

High-dimensional measures (HD) task from Korotin et al. [2019] tests whether OT solvers can redistribute mass among modes of varying measures. Different instantiations of Gaussian mixtures in dimensions D =2, 4, 16, . . . , 256 are compared between each other via OT. In the benchmark, Mix3toMix10 is used, where source measure α can consist of random mixture of 3 Gaussians and target measure consist of two random mixtures β 1 , β 2 of 10 Gaussians. Afterwards, pretrained OT potentials ∇ ψ i # α = β are used to form the final pair as ( α , 1 2 ( ∇ ψ 1 + ∇ ψ 2 )# α ) .

Images task produces pair candidates for OT solvers in the form of high-dimensional images from CelebA64 faces dataset (Liu et 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 is constructed as average between learned potentials via [W2] solver and forms a pair input for OT algorithm as ( α CelebA, β Ckpt ) = ( α Final , [ 1 2 ( ∇ ψ 1 + ∇ ψ 2 )# α Final )] ). Figure 9 shows an example of pair of two such measures ( α , β ).

F.5 Unpaired Image to Image Additional Results

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 in 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 be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it 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]

Justification: We provided discussions on limitations in Section 6 of main paper (last two paragraphs).

Guidelines:

The answer NA means that the paper has no 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 a 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 on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a 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 of 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 be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important 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 set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: We provided a theoretical proof that expectile regularization approximates upper bound on 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 crossreferenced.

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 experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: Along with detailed hyperparameter specifications in the Appendix, we included easy to follow Jupyter notebook which can be found in supplementary materials, enabling the others to fully reproduce the results in the paper.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to 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 a 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 a 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 to 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 to replicate the results, access to a hosted model (e.g., in the case of a 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 all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a 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 to 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 be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

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

Answer: [Yes]

Justification: We provided a 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 NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not 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 a 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 ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on 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. If only a 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) is 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, hyperparameters, 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 of the tested benchmarks.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a 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 error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

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.

It 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 a 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 of range (e.g. negative error 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 in the text.

8. Experiments Compute Resources

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

Answer: [Yes]

Justification: We include the exact computer configuration in Appendix and mention GPU model in the main text.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or 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 as well as 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 it 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.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We read it and adhered to the ethical guidelines.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or 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), fairness 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 not tied to particular applications, let alone deployments. However, if there is a 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 as 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 or 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 a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

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

Answer: [NA]

Justification: Not applicable to this work.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a 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 to 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 is challenging, and many papers do not require this, but we encourage authors to take this into account and make a 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 not use 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 is used and, if 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.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

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

Answer: [Yes]

Justification: We include all of the details corresponding to train procedures, datasets used, and citations. Moreover, we provide a readme file for the 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 not release new assets.

Researchers should communicate the details of the dataset/code/model as part of 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 is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an 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 as details about compensation (if any)?

Answer: [NA]

Justification: No crowdsourcing was used in this study.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in 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 crowdsourcing nor research with human subjects.

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.

Peer reviews:
comments: SUMMARY:
-	This 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.

STRENGTHS:
- The proposed method is easy to implement as it only requires an additional regularizer.
- The ENOT achieves competitive results on the Wasserstein-2 benchmark.
- 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.
- Please see the Questions Section below.

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.
- 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.

**Typo:**
- Line 220: Table 10 -> 3
- Line 230: Figure 8 -> 4

LIMITATIONS:
-	The authors addressed the limitations of their work.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2024-07-07

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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$.\n\nWEAKNESSES:\nTo 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.\n\nQUESTIONS:\nThe 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.\n\nLIMITATIONS:\nSee \"Weakness\".\n\nNitpicking: the authors may want to consider differentiating the commands \\citep and \\citet in latex citation and use them accordingly.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nAuthors 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.\n\nSTRENGTHS:\nOriginality: \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.\n\nWEAKNESSES:\nQuality: \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\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nLimitations are discussed.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- The authors seem to have done a thorough literature review.\n\nWEAKNESSES:\n- 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?\n\nQUESTIONS:\nSee weaknesses.\n\nLIMITATIONS:\nSee weaknesses.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}

comments: SUMMARY:
This work tackles the problem of estimating the dual Kantorovich OT potentials 
parametrized via neural networks.  They propose to approximate conjugate operator 
using a well motivated loss called expectile regularization which approximates a conditional
maximum operator, well suited for estimating the c-transform. The authors illustrate the 
performance of their method against several benchmarks showcasing a significant improvement 
in both performance and runtime.

STRENGTHS:
- Very simple, novel and intuitive idea with important and significant impact
- Thoroughly carried out experiments and evaluation

WEAKNESSES:
- Little analysis of the drawbacks of the proposed method as well as the effect of setting 
the threhshold parameter tau (see Q1-2 below)

QUESTIONS:
1. 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.

2. 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 ? 

3. 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.

- L74: is *squared
- L137: the*  proposed regularization approach
- L160: the* proposed

LIMITATIONS:
none
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-21