Expectile Regularization for Fast and Accurate Training of Neural Optimal Transport

Authors: nazar  buzun
  airi
maksim  bobrin
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-13 13:26:36

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.

Main paper: 
 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:EQUATIONwhere {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] ):EQUATIONwhere Π[α, β] = {π ∈ 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:EQUATIONwhere 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:EQUATIONotherwise, 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):EQUATION= supEQUATIONIn practice, during optimization process the infimum by T in c-conjugate transformation is approximated by a parametric model T θ , such thatEQUATIONA 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:EQUATIONThe 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. 
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] ):EQUATIONwith 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:EQUATIONwith 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. 
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)EQUATIONIf the same conditions are met for measure β we can express the inverse mapping T -1 (y) through the potential g: T -1 (y) = y -∇h * (∇ g(y)).(12) 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:EQUATION)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. 
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 toEQUATIONLet 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 toinf 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 η :EQUATION)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 EQUATIONEQUATION(x)} 2: if is_bidirectional is true T(x) = x -∇h * (∇f(x)) else T(x) = f(x) 3: {Compute dual OT losses and expectile regularisation R g } 4: L g = 1 n n i=1 g(T(x i )) -1 n n i=1 g(y i ) 5: L f = 1 n n i=1 c(x i , T(x i )) -g(T(x i )) 6: R g = 1 n n i=1 L τ c(x i , T(x i )) -c(x i , y i ) + g(y i ) -g(T(x i ))7: {Apply gradient updates for parameters of models f and g} 8: opt_f.minimize(f, loss = L f ) 9: opt_g.minimize(g, loss = L g + λR g ) end function 10: {Main train loop} 11: for t ∈ 1, . . . , N do 12: sample x 1 , . . . , x n ∼ α, y 1 , . . . , y n ∼ β 13: if is_bidirectional is false or t mod 2 = 0 then 14: train_step(f, g, {x 1 , . . . , x n },: dist = 1 n n i=1 g(y i ) + c(x i , T(x i )) -g(T(x i )) ; f(x) = c(x, T(x)) -g(T(x)) 23: return f, g, distAlgorithm 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 θ :EQUATIONIn the end of training we approximate the correspondent Wasserstein distance by expressionEQUATIONwith 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. 
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.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 (Korotin et al. [ , 2021]] , Makkuva et al. [2020] ), which quantifies deviation from the optimal alignment T * , normalized by the variance of β:EQUATION)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.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 Table 2: L UV 2 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 5 different seeds. The best scores are highlighted.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 Shoes5 . 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 = 120k. 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. 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 ).  
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. The value τ = 0.5 corresponds to the standard MSE loss, while τ = 0.9 and τ = 0.99 give more weight to the positive differences. 
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. 
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. 
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. 
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. 
 
 Appendix: Right: expectile models f τ (x). The value τ = 0.5 corresponds to the conditional statistical mean of the distribution, and when τ → 1 it approximates the maximum operator over the corresponding values of y.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), thenEQUATIONand consequently from (6) we derive thatEQUATIONBy equation ( 12) the corresponding optimal transport map T (x) equals to the gradient of u * (argmaximum from the last formula):EQUATIONIn this section of the Appendix we briefly outline categorization of different approaches for estimating dual Kantorovich problem as proposed by (Amos [2023 (Amos [ , 2022]] ), where the following differentiable amortization loss design choices are highlighted:A. 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 α:L obj (T θ (x)) = c(x, T θ (x)) -g η (T θ (x)).(24) 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. B. 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:L reg (T θ (x), T (x)) = ∥T θ (x) -T (x)∥ 2 (25)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.C. 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 lossEQUATION)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.Lemma D.1 (Rudin [1976] ). Let random vector ξ have a compact support Ω and ∀x ∈ Ω: 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.f n+1 (x) ≥ f n (x)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/.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 . 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 ).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.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. 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] ). 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 (α, β). • 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. 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).This intuition is supported by a direct evaluation in Section 5 below.kaggle.com/datasets/reitanaka/alignedanimefaceskaggle.com/datasets/defileroff/comic-faces-paired-synthetic-v2github.com/junyanz/iGAN/blob/master/train_dcgan 


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

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: 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: limitations: none
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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).

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


- **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?

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

Writing and notation nitpick: 

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

- There should be explanation that we replace the quantity in Eq (7) to (5).

- Optimization and optimisation are used interchangeably. Stick to one of these. 

- 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?
questions: questions: See weaknesses.
limitations: limitations: See weaknesses.
paper_summary: paper_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.
strength_weakness: strength_weakness: Quality: 
- 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. 

References:
1) TrajectoryNet: A Dynamic Optimal Transport Network for Modeling Cellular Dynamics Alexander Tong, Jessie Huang, Guy Wolf, David van Dijk, Smita Krishnaswamy

2) Light Schrödinger Bridge: Alexander Korotin, Nikita Gushchin, Evgeny Burnaev
questions: 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?
- The algorithm (Algorithm 1 line 2) states that in case 'is_bidirectional' is false $T(x)=\nabla f(x)$. From what it follows?
- Other algorithms that use expectile regression tend to use 0.9 expectile, why in your method you take the expectile parameter closer to 1?
- By what metric can we choose the hyperparameters of $\lambda$ and $\tau$ ?
- Partial OT is often required in practical tasks. It is possible to generalize this method to solve the problem of partial OT?
- 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: Limitations are discussed.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.
questions: 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?

In Section 5.4, the authors mentioned

>Figure 8 presents the study of the impact of the proposed expectile regularization on the L_2^{\text{UV}} metric. 

I believe this is a typo. The figure that the authors refer to is Figure 6 in the Appendix.
limitations: limitations: See "Weakness".

Nitpicking: the authors may want to consider differentiating the commands \citep and \citet in latex citation and use them accordingly.
paper_summary: paper_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.
strength_weakness: strength_weakness: - The derivation of the ENOT regularizer (Eq 17) requires further clarification.
- Please see the Questions Section below.
questions: 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: limitations: -	The authors addressed the limitations of their work.

Comments: title: Reponse to rebuttal
comment: Thank you for raising the score and for helping us improve the paper.
Thank you for your feedback and for the final score. We are pleased to see you appreciate the empirical performance and are available for further conceptual explanations if needed during the next stage of the discussion.
Thank you. We appreciate your suggestions and feedback.
Thank you for your input and for raising the score. We commit to implement all the changes in the revised version of the paper.
Thank you for your response. I especially appreciate your discussion on Q1. It would be great to see a similar discussion included in the paper. 

I've increased my score to 7. However, I must note that although I've used optimal transport in previous work, I don't consider myself an expert in the field. My support for the paper is mainly because of its strong empirical performance. For this reason, I recommend that the area chair give more weight to the opinions of OT experts when making the final decision.
I appreciate the author for their clarifications and additional experiments. These have been helpful in addressing my concerns. Hence, I will raise my rating to 6.
Thank you for the reply. Since my major concern is resolved, I'm raising the score towards acceptance. I would however expect the authors to better phrase their implementation of the experiments as they have promised in the rebuttal and response to my review.
I appreciate the detailed answers and clarifications. I find this paper novel, easy to follow, and of interest to the community. The strengths of the paper significantly outweigh the weaknesses, and I increase my score.
**Q1**: Thank you for requesting this clarification. Yes, as indicated in Sections 5.1 and 5.3, and in most of our experiments, we used squared Euclidean cost. Under Brenier's theorem conditions for the squared Euclidean cost, it holds that $\hat{T}(x) = \nabla \hat{f} (x)$, where $\hat{f}$ is some convex function. It follows that the optimal potential $\hat{f}$ is convex, even if one uses non-convex potentials $f$ in the training process. 

Also, to reiterate our initial answer, there are no restrictions on $T(x)$ in Eq. (6) during the training and we can use any representation for this function, including $T(x) = \nabla f(x)$. Therefore, there is no need for $c$-concavity through the neural network parameterization whatsoever, and we do not use it, as it makes the solution worse in practice. 

The beauty is that our expectile regularization actually forces functions to be more $c$-concave according to the $c$-concavity criterion (refer to Villani et al. [2009] Proposition 5.8). 

&nbsp;

**Q2**: We apologize for the ambiguity, which could have been avoided had we explicitly mentioned that we used flag `is_bidirectional=False` (meaning the training mode is one-directional in this example). Thank you for spotting the mismatch. So, in Section 5.2, we do not use (13) and (14) and we do not express $T(x)$ through the potentials, and the convexity of the cost function $h$ is not required. As such, we believe the issue should be resolved and there is no contradiction with Remark 1.18 in Santambrogio (2015). We will definitely add this additional clarification to the text. 

Please kindly notice that we do mention it in another part of the main paper (citations below). So, by adding this clarification to Section 5.2 also, we avoid any potential for the same confusion of the future readers of the paper. Thank you. 

> Lines 155-157 (Method section):  “The transport mapping $T_θ(x)$ has the same parameters as $f_θ(x)$ if it can be expressed through $f_θ$ (ref. 13), or otherwise, when $f_θ$ is not used (one-directional training), it is its own parameters.”

> Lines 204-205 (Section 5.2): “we parametrize the map Tθ as a MLP”
I thank the authors for the detailed rebuttal. 

However, my major concern (Q2) regarding the soundness of this work remains: 

- I am fully understand Brenier's theorem only deals with squared Euclidean cost (W2 distance), but isn't this cost function was what you used in a major part of the experimental section anyway(Section 5.1 and Section 5.3)? This should mean c-concavity of $f$ and $g$ have to hold by enforcing it through the neural network parameterization? 

- Moreover, the L2 norm used in Section 5.2 (different cost functionals) violates the condition that h(x - y) must be strictly convex -- you can also verify this by yourself, or check Remark 1.18 in Santambrogio (2015) -- so I am not sure why your experimental results can hold all the case.