Neural Hamilton--Jacobi Characteristic Flows for Optimal Transport
Authors: Yesom Park
Shu Liu
Mo Zhou
Stanley Osher
Venue: ICLR 2026 Poster (2026)
Keywords: Optimal Transport
Hamilton--Jacobi Equations
Method of Characteristics
Class-Conditional Optimal Transport
Primary area: optimization

Abstract: We present a novel framework for solving optimal transport (OT) problems based on the Hamilton--Jacobi (HJ) equation, whose viscosity solution uniquely characterizes the OT map. By leveraging the method of characteristics, we derive closed-form, bidirectional transport maps, thereby eliminating the need for numerical integration. The proposed method adopts a pure minimization framework: a single neural network is trained with a loss function derived from the method of characteristics of the HJ equation. This design guarantees convergence to the optimal map while eliminating adversarial training stages, thereby substantially reducing computational complexity. Furthermore, the framework naturally extends to a wide class of cost functions and supports class-conditional transport. Extensive experiments on diverse datasets demonstrate the accuracy, scalability, and efficiency of the proposed method, establishing it as a principled and versatile tool for OT applications with provable optimality.

Published as a conference paper at ICLR 2026
NEURAL HAMILTON-JACOBA CHARACTERISTIC FLOWS FOR OPTIMAL TRANSPORT
Yesom Park', Shu Liu', Mo Zhou', Stanley Osher
Departmeet of Mathematics, University of California, Los Angeles Department of Mathematics, Florida State University yei i somêmath ucla. edu, s125bn@ @ fsu scu.edde .edu, mozhou3 66 660mmt6.....mm ucla..l....o..o edu
ABSTRACT
We present a novel framework for solving optimal transport (OT) problems based the Hamilton- Jacobi (HJ) equation, whose viscosity solution uniquely char- acterizes the OT map. By leveraging the method of characteristics, we derive closed-form, bidirectional transport maps, thereby eliminating the need for nu- merical integration. The proposed method adopts a pure minimization frame- work: a single neural network is trained with a loss function derived from the method of characteristics of the HJ equation. This design guarantees convergence to the optimal map while eliminating adversarial training stages, thereby substan- tially reducing computational complexity. Furthermore, the framework naturally extends to a wide class of cost functions and supports class-conditional transport. Extensive experiments on diverse datasets demonstrate the accuracy, scalability, and efficiency of the proposed method, establishing as a principled and versatile tool for OT applications with provable optimality.
INTRODUCTION
Optimal transport (OT) is a fundamental problem that seeks the most cost-efficient transform from one probability distribution into another by minimizing a transportation cost function, which quan- tifies the effort to move mass. With its strong theoretical foundation and broad practical relevance, OT has been widely applied diverse areas, including traffic control (Carlier et al. 20088] Danila cet al., 2006 Barthélemy & Flammini 2006), biomedical data analysis (Schiebinger et al., 2019 Koshizuka & Sato [022] | Bunne et al. 2023). generative modeling (Wang et al.] adl][222]]Onke
et al,
2021
Zhang & Katsoulakis lakis 2023 Liu et al. 2019), and domain adaptation (Courty et al. 2016
Damodaran et al. 2018; Balaji et al.] 2020). In recent years, there has been growing interest deep learning techniques solve OT problems, leading to the development of methods grounded in various mathematical formulations. Early approaches were primarily built upon the classical Monge formulation (Lu et al. 2020 Xie al. 2019) and its relaxation into the Kantorovich frame- work (Makkuva et al.) 2020). While theoretically rigorous, these methods often suffer from high computational complexity. The primal-dual formulation, which recasts the OT problem as a saddle- point optimization over the generative map and the Kantorovich potential function, has inspired scalable algorithms (Liu et al., 2019 Taghvaai & Jalali 2019, Korotin et al., [2211a] [Lu et al. [02T]] Choi et al. 2024). Similar approaches have also been proposed for the Monge problem with gen- cral costs (Asadulacv et al.| 2024 Fan et al. (002)). However, these approaches typically rely on adversarial training of two neural networks, which is challenging manage and often introduces instability and inefficiency into the optimization process. Alternative approaches have investigated dynamical formulations using ordinary differential equations (ODEs) (Yang & Karniadakis 2020 Onken et al.) [022 Tong et al., 2020: Huguet et al.] 2022) and entropie-regularized models involy- ing stochastic differential equations (SDEs) (Genevay al. 2016 Seguy et al. 2017 Daniels al.; 2021 Gushchin et al. 2023, Zhou et al. 2024). Machine learning algorithms that unify Lagrangian and Eulerian perspectives of Mean Field Control problems Ruthotto et al. (2020) ;
2017
Lin et al. (2021) )
Zhao et al. (2025) likewise provide a computational framework for OT. Nevertheless, these methods typically require solving systems of differential equations, resulting in substantial computational
overhead during both training and inference. Moreover, many existing methods yield bias maps that deviate from the OT solution due to the incorporation regularization terms into the formulation.
"Corresponding author. Email: 8 sjo@math, ucl a.l..

---

Published as conference paper at ICLR 2026
Method (representative reference)
Optimization
# Networks
OT direction
Sampling
Optimality
Dual Formulation (Asadulaev 2024 Dynumical Models (Onken el u. 2021
Min-Max
Two
One-way
Direct
No
Min Min
Single
Bidirectional
Rerative
No
HJ--used (Proposed)
Single
Bidirectional
Direct
Yes
Table 1: Comparison of key features across different OT model approaches. Optimality indicates
whether an approach is guaranteed, in principle, to recover the true OT map under sufficient network capacity.
Contributions. We propose a novel and efficient framework, termed neural characteristic flow (NCF), for solving OT problems via the Hamilton-Jacobi (HJJ equation, whose viscosity solution characterizes the OT map- Despite its strong theoretical foundation for OT, the HJ formulation poses two major challenges: non-uniqueness of solutions and the need to solve ODEs in dynamical formulations. We overcome both by leveraging the method of characteristics and an implicit solution formula (Park & Osher 2025) to obtain closed-form, bidirectional transport maps without numerical integration of ODEs. NCF uses single neural network and avoids adversarial training dual- network architectures, reducing complexity while improving efficiency. Our framework guarantees theoretical consistency with OT optimality conditions and supports a broad class of cost functions, including class-conditional transport. We also provide convergence analysis for Gaussian settings and demonstrate strong empirical performance across datasets of varying dimensions. comparison of key features across different OT model approaches is summarized in Table
2 PRELIMINARY
2.1 MONGE'S S OPTIMAL TRANSPORT PROBLEM
For domain Ω C Rd. we denote (I2) as the space of probability measures on Q. Let ΩxΩ - [0, oo] be a cost function that measures the cost of transporting one unit of mass. For μ4, v E P (9), the classical Monge problem formulates OT as finding measurable map T $2 > St that pushes forward ja to V. i.e., Tyj v/ while minimizing the transportation cost:
We (n, v) inf [ c(x, T (x)) du (x). Typ=v
(I)
We call a solution T ① an OT map between μ and v. the case where the cost is expressed as function of the difference between the two variables, T is characterized as follows:
Theorem 2.1 (Santambrogio )2155). When c yx, = L (x -y) for lower semi-contimuous
(l.s.c.), sub-differentiable, and strictly convex function f Q + R. the optimal map is expressed
in terms of the Kantorovich dual potential function 4² Ω2 + R as
T (x) = x Vh (x)), where h (z) SupyER {zT l (y)} is the Legendre transform of l.
(2)
2.2 DYNAMICAL FORMULATION Benamou & Brenier (2000) formulate the OT in continuous-time dynamical formulation:
inf E,, 1/" El (u (x((())
(3)
s.t. x v, x(0 μ. x(ty) ~
(4)
where the terminal time fy > is typically set to 1. Within this dynamical framework, the associated optimality condition is governed by the Hamilton-Jacol (HJ) equation (Evans) 2022 chapter 10):
+ h (Vu) = on Ω x {1 = 0},
Ω x (0,t,)
(5)
2 ー 9
coupled with the continuity equation 0,p + Vh (Vu) 0 that governs the evolution of the probability distribution. Here, Vu denotes the gradient of u with respect to the spatial variable and g represents the initial condition, whose explicit analytic form is typically intractable. The optimal velocity field is then determined by v" = Vh (Vu). where u is the viscosity solution to equation (5).
2

---

Published as a conference paper at ICLR 2026
RELATED WORKS
Deep learning methods for OT have gained traction following the development of scalable OT solvers (Genevay
2016 Seguy et al.] 2017) and WGANs (Arjovsky al. (017). Many approaches utilize GAN-based models to approximate OT plans, although they often suffer from training instability and extensive hyperparameter tuning. Another major line of work is based on the Comullation(Kantoooo (Kantorovich) (006). where the OT map is recovered via optimization of dual potentials, typically parameterized by input convex neural networks (ICNNs) (Amos et al.) 2017). While theoretically sound, these methods involve unstable min-max optimization. To ad- dress these issues, natural gradient methods have been proposed to improve computational efficacy (Shen et al. 2020; Liu ct al. 2024). Regularization techniques such L penalties (Genevay al., 2016 Sanjabi et al. 2018) and yyclscconsistency constraints (Korotin et al.) 2019] 202[b) have been proposed, though unconstrained alternatives have shown stronger empirical performance (Korotin et al. 2021 la Fan et al. 20022..
Kantorovich dual formulation
To address the settings where deterministic OT maps may not exist, recent work has considered weak OT formulations (Backhoff- Veraguas et al. 2019). Neural approaches for weak OT and class-
conditional transport have been proposed (Korotin et al. 2023 Asadulaev et al. 2024). but may
yield spurious solutions under weak quadratic costs. Kernalized costs (Korotin et al. 2022) have been introduced to mitigate this.
OT has also been modeled as a dynamical system via continuous flows (Yang & Karniadakis 2020
Tong et al. 2020; Onken al. 2021 Huguet al. 2022). While expressive, these methods require solving ODEs during training and inference, making them computationally expensive. Entropic and f-divergence regularized stochastic models (Daniels al., 202T Gushchin et al. 2023) improve smoothness but often rely on Langevin dynamics, which can be biased in high dimensions (Korotin et al., 2019). The HJ equation has been used to improve OT models, with physics-informed neural network (PINN) (Raissi al. 2019) approaches applying L penalties on HJ residuals to improve continuous normalizing flows, ODE-based formulations (Yang & Karniadakis 2020; Onken et al.)
2021), and stochastic variants (Zhang & Katsoulakis 2023). However, due to the ill-posed nature of the HJ equation, this approach lacks guarantees for recovering the viscosity solution.
HJ CHARACTERISTIC FLOWS FOR OT
In this section, we represent the OT map through the characteristics of the HJ equation, offering a principled and efficient framework for OT. Note that solving the HJ equation directly challenging due to its inherent ill ill-posedness, non-smoothness of solutions, and gradient discontinuities, all of which complicate both theoretical analysis and numerical approximation.
Method of Characteristics. The viscosity solution to 5 is theoretically characterized by the following system of characteristic ordinary differential equations (CODEs) Evans (2022): Park &
Osher 2025):
x = (p) ú h(()) + P Vh(p)
(6a)
(6b)
p= 0
(6c)
where denotes the shorthand for Vu. CODE for p 6cp implies that p remains constant along each characteristic trajectory. Consequently, the characteristics are straight lines of the form x(t) = (V/(p) + x(0). which coincide with the OT map in (2) at terminal time t = ty. From dynamical perspective, the ODE 4 can be interpreted as the characteristic equations (6a) of the HJ equation that determine the OT map (2). In other words, the transported point 7* (x) of sample x ~ J corresponds to the terminal position of the characteristic line that originates from x.
Our CODE formulation not only provides principled construction of the forward transport map but also naturally characterizes the backward map. We denote by T the forward OT map transporting μ v. and by TH the backward map transporting to μ.
Proposition 4.1 (Bidirectional OT Map). There exists viscosity solution u² to the HJ equation 55) that characterizes both the forward and backward OT maps through its forward and backward characteristic flows:
3

---

Published as a conference paper at ICLR 2026
T (x) = x =x x (x,'(,))), X- To (y)= y V, (Vu" (y,4)), you.
Accordingly, the viscosity solution of the HJ equation enables a bidirectional characterization of the OT map via forward and backward characteristic flows. Notably, since the characteristics straight lines, both the forward and inverse transport maps admit explicit closed-form eoormexxpeeee expressions. This obviates the need for numerical integration of ODEs typically required in conventional dy- namical formulations. Consequently, the CODE-based formulation addresses a key computational bottleneck, enabling efficient and direct computation bidirectional transport maps.
Implicit Solution Formula. Recently, novel mathematical formulation for the viscosity solution of HJ equations has been developed using the system of CODEs (Park & Osher (2025). Within this formulation, the viscosity solution admits the following implicit formula:
(x, = -th (Vu) + tVu Vh (Vu) + g gg(x tVh (Vu)).
Proposition 4.2. For OT problems ① where & satisfies the conditions in Theorem 2.1 the implicit solution formula ( characterizes the viscosity solution of the HJ equation () almost everywhere.
Proof. Detailed proof is provided in Appendix A.1
METHODS
5.1 OT WITH GENERAL COSTS
We propose novel deep learning method, termed neural characteristic flow (NCF), for learning bidirectional OT maps under general cost & by solving the HJ equation (5) vis its implicit solution formula (9). The HJ equation characterizes the OT map as the gradient of the viscosity solution, en- suring that the resulting map minimizes the given cost functional. When coupled with the continuity equation, it also describes the evolution of probability distributions, thus guaranteeing correct mass transport from source to target. However, jointly solving this coupled system of PDEs is computa- tionally expensive. To address this, the proposed NCF computes the OT map solely through the HJ equation, avoiding the need to solve the continuity equation explicitly.
Implicit Neural Representation. We represent the solution u of the HJ equation using an implicit neural representation (INR) ug IRd x IR + E parameterized by The network takes the spatial variable x and temporal variable t as input. By the universal approximation theorem (Hornik et al. 1989 Leshno al. 1993), the INR can approximate the viscosity solution the HJ equation. We denote by T [ug] as the transport map that aims to map μ to v defined by (7) through ug:
Ty [oo] (x) x + f Vh ( u4g (x,0) 0))
(10)
The backward map Th [ug] analogously defined according to via uy evaluated f. = ly-
HJ-based Training Loss. While the HJ equation does not directly encode distributional informa- tion, it can recover the desired OT map, provided that an appropriate initial function reflects the relationship between the source and target distributions. However, practice, where only finite sam- ples from these distributions are available, deriving an analytic form for is generally intractable. To address this challenge, we introduce loss term ensure that the initial condition is appropri- ately learned during training, thereby steering the HJ solution toward accurately solving the desired OT problem. Specifically, this term enforces alignment between the generated samples obtained via T [uo] and the given target data. This alignment can be effectively quantified using discrepancy mea- sures such as the maximum mean discrepancy (MMD) (Smola et al., 20066.. whose value between two distributions μ and v are defined as follows:
// k(x, y) d(u(x) v(x)) - -v(y))
MMD v)²
(11)
where k(.,)): :2x Ω R is kernel function. The population loss for the MMD is
LuMD(uo)) MMD [uolet,
(12)
4

---

Published as a conference paper at ICLR 2026
We adopt the negative distance kernel k (x, x
y yll2. which has proved to handle high- With this kernel, the MMD loss becomes
dimensional problems efficiently (Hertrich et al. 2024).
the squared energy distance (Rizzo Székely 22016).
In our implementation of the implicit solution formula, we replace the initial function g with us evaluated at t = O, and train the model using the following g-weighted loss function
Luu (uo) - 11==INaxa /aatuss, (no + th (Vuo) NV Vugg Vh (Vua) uo (x 1V (Vua). 0)) de(x) dt,
(13)
probability measure o used to weight the residual so that the implicit solution formula (9) is enforced across the entire spatial domain while allowing Lu be efficiently approximated via Monte Carlo sampling. When Ω is bounded, a natural choice is the uniform distribution on Ω. In practice, e may also be chosen as a uniform distribution supported on the region covered by the available samples, concentrating computational effort on the portions of the domain most relevant to the transport dynamics.
where i a
The overall loss combines the implicit HJ loss and the MMD loss with a weight À >> 0:
Lu (uo) + >CMMD(us).
(14)
We refer to Appendix B for practical choices of and the Monte Carlo estimation of the loss.
Advantages of the Proposed Approach. Our method offers several key advantages over existing
OT frameworks, as summarized in Table 1 First, it jointly learns both forward and backward OT maps using a single neural network in one training phase. This contrasts with prior methods that require multiple networks, either due to the lack of invertibility or the use of adversarial dual for- mulations-leadingg to increased model complexity and training cost. Our method also avoids the instability of min-max optimization common in dual approaches, resulting in more stable training. Second, unlike dynamical OT models that require solving ODEs or SDEs, we use the method of characteristics to obtain OT maps in closed form. This removes the need for iterative solvers and improves sampling efficiency at both training and inference time. Third, our model directly incor- porates the HJ equation via an implicit solution formula that reliably recovers the viscosity solution, as supported by the numerical results in Section This not only aligns with the theoretical optimal- ity conditions of OT but also helps identify and correct deviations from the target solution during training. Finally, our framework supports broad class of cost functions beyond the quadratic case, offering greater flexibility and wider applicability across OT tasks.
5.2 THEORETICAL ANALYSES
In this section, we present theoretical analyses of our method. focusing on the OT problem with Ω = Rd and the quadratic cost E(.) = ğll 1| for which the corresponding Hamiltonian is given by ((-) E -| as well. We prove that the minimizer of the loss (14) exactly recovers the true OT maps. Moreover, in the Gaussian setting, we establish stability analysis by showing that a small loss guarantees convergence to the true solution.
Consistency Analysis
With some mild convexity assumption, we establish that the minimizer of (4 leads precisely to the optimal transport map.
Theorem 5.1 (Consistency of loss), Suppose the probubility distributions u.v have finite second moments and @ E (RF) is strictly positive. Assume u E Cl. x f1). and define u1(+) == u((,tt) E Ohe (IE)) with Vuj E L2ddd, (Rd, If minimizes the loss functional (14), i.e.
Lur(u) + ALsuu(()) =0,
and the map Ty (u] is bijective with its Jacobian D₂ (u(((x)) is positive definite for any x E RRd, then Ty (u) and Th [u] are the optimal transport maps from μt, and vice versa.
The proof is provided in Appendix A.2 See also Remark A.5 for further discussion on the mono- tonicity condition for [_TH(uu.
Remark 5.2 (On regularity assumption of u). It is worth noting that the transport curves associated with the W OT problem do not intersect for t [0, (cf. Chap. 8 of (Villani et al., 2008)). Since these curves constitute the characteristics of the HJ equation associated with the OT problem, can expect classical solutions to the HJ equation, provided that μ and v admit
5

---

Published as a conference paper at ICLR 2026
sufficiently regular density functions. This observation motivates the regularity assumption on in Theorem [.. Moreover, is parametrized with neural networks in practice, which naturally preserve the regularity.
Stability Analysis the Gaussian setting. Let μ Mb, m,). (b,,, E.), then the OT map
The loss 14 also exhibits favorable stability properties, which we illustrate in
Tx (x) A(x b,,) by
(15)
where A := En (E (E x For analytical tractability, we consider a simplified quadratic parameterization us uo(x, 1) = -2)x)02(()x 01 x + (o(()), where = (02(-), (1).), 0o(-) :
[0.ty] + Rand x R x R. Although this represents a restricted subclass of neural networks, it
permits rigorous analysis and yields insights relevant to more general architectures.
Assumption 5.3. O(t) is bounded by K and K- K-Lipschitz. ||b,l]| ||ball || IEElll K. A
is strictly positive definite with smallest eigenvalue X > 0.
Theorem 5,4 (Stability of loss). Under Assumption 5.3 the errors for ug and TV [uo] satisfy
||1g- u² (111111--) + [[TT] -T
(1 -
SC sC(Ciu +
(16)
where u² and Ty* are the true solution and OT map. C only depends on d, K and XA-
The theorem implies that sufficiently small loss guarantees convergence of the approximate solution ug-and consequently the resulting transport map Ty [ug]-to their true counterparts. Furthermore, the proof shows that while multiple transport maps may minimize the MMD loss, the implicit HJ loss ensures that the OT map is uniquely recovered. The detailed description and proof for the theorem are deferred to Appendix A.3
5.3 CLASS- CONDITIONAL OT
We extend our HJ-based framework to class- class-conditional OT, transporting source to target indepen- dently within each of the K labeled classes so as to preserve label consistency and class-specific structure. This formulation is particularly well-suited for domain adaptation and class-conditional generative modeling, where preserving class-specific features is crucial.
The OT map between samples of the k-th class must satisfy the HJ equation within the support of the corresponding class-specific distribution, as dictated by the optimality condition. Consequently, the global transport map TH satisfies the HJ equation (5) across the entire domain. Although non- differentiable regions may arise due to intersections between transport maps of different classes, such discontinuities occur primarily in the boundaries between class supports. Since the gradient of the HJ solution is computed only within the support of each class-specific distribution, the transport map remains expressible in these regions. Accordingly, we retain the implicit HJ loss function and modify the MMD loss to account for class conditioning as follows:
(13)
1
Ectass (TH (uol); p,
E((T) [uol); Me. vx).
(17)
K
A similar approach was proposed by Asadulaev al. (2044
EXPERIMENTAL RESULTS
We evaluate the effectiveness of the proposed neural characteristic flow (NCF) across diverse OT tasks. All experiments in this section employ the quadratic cost function & = 11-112. which is the canonical cost associated with the Wasserstein-2 distance. Computations were performed on single NVIDIA GV100 (TITAN V) GPU. Further implementation details are provided in Appendix B
6.1 UNCONDITIONAL OT
6.1.1 2D TOY EXAMPLES
We test the proposed NCF on a 2D toy dataset. We also compare our model with the neural opti- mal transport (NOT) framework (Korotin al. 2023 including both the strong (deterministic) and

---

Published as a conference paper at ICLR 2026
(a) Source y
(b) Strong NOT
(c) Weak NOT
(d) HJ-PINN
(c) NCF
(f) Target v
g Strong NOT
(h) Weak NOT
(1) HU-PINN
(i) NCF
Figure 1: Swiss roll (μ) = Double moons (v): The top row shows transport in the direction - ,4, and the bottom row shows μ + v. The leftmost column displays μ and v for reference.
weak (stochastic) variants. Since NOT directly parameterizes the transport map, it requires sepa- rate training for each transport direction. Additionally, we include an ablation study replacing our implicit solution formula loss (13) with PINN loss on the HJ equation, referred to as HJ-PINN.
Figure [ shows bidirectional transport results on 2D distributions. In addition to visualizing the transported distributions, we overlay the learned transport maps as black solid lines assess whether each model has captured an OT plan. For weak NOT, the map is the average over noise inputs, as in the original work. Compared tions more accurately and learns transport maps closely aligned with the optimal solution. Strong NOT produces noisy, incoherent transport. Weak NOT performs better but still shows overlapping trajectories, indicating an incomplete OT representation. HJ-PINN yields noisy, intersecting trans- port paths, suggesting failure to learn OT dynamics. In contrast, our model learns accurate OT maps without trajectory crossings. Moreover, unlike NOT, which requires four separate networks for bidirectional training, our method achieves more accurate bidirectional transport with a single network. These results highlight the superior accuracy and efficiency of our approach. For further experimental results on the 2D example, please refer to Appendix C.T
all baselines, our method captures source and target distribu-
6.1.2 EVALUATION ON HIGH-DIMEN DIGG-DIMENSIONNL GAUSSIANS
For general distributions, the ground truth OT solution is unknown, making quantitative evalua- tion challenging. To enable precise assessment, we consider the Gaussian case: / N E,,) and V N (0, ),,), where a closed-form solution is available via (5). Following Korotin et al. 2021a). we vary the dimension from to 64, with Ny and En generated using random eigen- vectors uniformly sampled the unit sphere and logarithms of eigenvalues drawn uniformly from [_..2]. In addition to strong NOT and HJ-PINN, we evaluate several established OT methods: MM-v1 (Taghvaei & Jalali 2019 Korotin et al.) (021a). which solves min-max dual problem using input-convex neural networks (ICNNs), alternating between optimizing the potential and its convex conjugate; MM:R (Korotin al. 2021a) also employs min-max framework but does not enforce convexity, instead learning separate networks for forward and backward maps via a negative Wasserstein loss combined with a conjugacy loss: LS (Seguy al. (017). which addresses the dual problem via entropic regularization; and WGAN-QC (Liu et al.) 2019), which employs a WGAN ar- chitecture with quadratic cost. Except for NOT which directly parameterizes transport maps all models use shared architecture for potential functions.
Performance is measured using the unexplained variance percentage (UVP) (Korotin et al., 2019). which quantifies the 7,2 error of the estimated transport map, normalized by Var(p). Computational efficiency is also evaluated in terms training and inference time, peak memory usage, and memory required to store bidirectional OT maps. Table reports UVP across models and dimensions, while Figure summarizes computational metrics.
Our method consistently yields accurate OT maps with favorable scaling behavior, outperforming NOT, WGAN-QC, and LS, which exhibit greater deviation from the ground-truth transport. In higher dimensions, the performance of NCF is slightly reduced, which attribute to the fact that, unlike baseline methods that parameterize separate networks for forward and backward maps, our approach represents both directions using single network. Specifically, NCF approximates both

---

Published as a conference paper at ICLR 2026
Table 2: Quantitative evaluation on Gaussian distributions. UVP (V) is measured across different OT methods as the data dimension increases.
Method d 2
d
d 8
d
16
d 32
d 64
NOT
77.248
125.419
114.056
176.086
182.287
196.831
WGAN-QC
1.596
5.897
31.0367
59.314
113.237
141.407
LS
5.806
9.781
15,963
25.232
41.445
55.360
MM-v1
0.161
0.172
0.173
0.210
0.374
0.415
MM:R
0.012
0.048
0.0117
0.202
0.354
0.604
HJ-PINN
0.080
0.069
0.163
0.458
0.576
1.683
NCF
0.010
0.021
0.086
0.146
0.436
0.858
NCF-Adaptive
0.010
0.022
0.090
0.155
0.307
0,407
Training Time
Evaluztion Time
Max Memory Allocation
Didirectional OT Map Storage.
15
3.5
Mexuury (MB)
120
Mezuoy (MCES)
Timno (5)
36
Dimension
Dimension
Dimension
Dimension
NOT
WGAN-OC
1.
MM-1
MMR
EIJPINN
NCF
XCF-Adaptive
Figure 2: Computational comparison. Training time (s/epoch), evaluation time (s/epoch), peak memory (MB) during training, and memory (MB) for storing bidirectional OT maps are reported.
maps using exactly half the network capacity required by methods such as MM-v1. This parameter reduction is advantageous in low dimensions; however, in higher dimensions (e.g.. d 32, 64), the limited capacity constrains the simultaneous learning of both directional maps. To address this, conducted experiments with slightly enlarged network for nigh-dimensional cases- wisthot denoted modifying any training hyperparameters- denoted as NCF-Adaptive in the results. With this ad- justment, method recovers performance in high dimensions while still using fewer parameters than MM-v1 or MM:R.
Moreover, baseline methods such MM-v1 and MM:R incur substantially longer training times and significantly higher memory usage. In addition, MM:R, which relies on min-max optimization, exhibits noticeable training instability (see Appendix C.2 In contrast, our approach avoids expen- sive nested min-max max optimization and leverages a single network, resulting in faster, stable, and more memory-efficient training. At inference, NOT achieves the lowest latency due to direct map parameterization, whereas other methods, including ours, require gradient-basee based evaluation, intro- ducing additional overhead. This overhead, however, decreases with increasing dimension. Finally, comparison with HJ-PINN highlights the superior effectiveness of our implicit loss in approximating the viscosity solution to the underlying HJ equation.
6.1.3 APPLICATION TO COLOR TRANSFER
We employ the dataset provided by CycleGAN (Zhu et al. 2017) for image color transfer experi- ments. From each of the three available groups of image pairs, we selected 10 representative pairs. For each pair, we perform both forward and backward color transfer. To evaluate the effectiveness of our model, include comparisons with two widely used classical color transfer methods: stan- dard per-channel histogram matching technique and the approach of Reinhard et al. (2001). which aligns the mean and standard deviation of color channels. These baselines represent statistical meth- ods that do not rely on OT, providing a complementary perspective on performance. We include NOT and MM-vl as deep learning OT baselines.
To quantitatively evaluate color fidelity and distributional consistency, we employ two widely used histogram- based metrics: Earth Mover's distance (EMD) and histogram intersection (HI), summa- rized in Table 3 Across all three domains, our method consistently achieves superior performance compared to all baselines in both metrics. In particular, our proposed method exhibits superior robustness in handling more complex and multimodal color distributions compared to MM-v1, pecially in contrast to the simpler Gaussian settings examined in the previous section. Qualitative results provided in Appendix
8

---

Published as a conference paper at ICLR 2026
Table 3: Quantitative evaluation of color transfer. Earth mover distance (EMD) and histogram
intersection (HI) between color distributions of target and transported images are reported.
Winter-Summed
Monet-Photogree
Gogh-Photogrant
Method
EMD (L)
Hi (*)
EMD (4)
HI (?)
EMD (4)
HI (1)
HisMatching
0.0012
0.7296
0.0013
0.7532
0,0010
0.7668
Reinhard
0.0013
0.6255
0.0012
0.7255
0,0009
0.7406
NOT
0.0008
0.8002
0.0008
0.8210
0.0008
0.8247
MM-v1
0.0014
0.7295
0.0011
0.7722
0.0007
0.8265
MM:R
0.0015
0.6404
0.0013
0.6810
0.0018
0.6260
NCF
0.0005
0.8914
0.0004
0.9174
0.0003
0.9117
Vertical Gaussians
Horizontal Gaussians
(a) Distributions
(b) Unconditional NCF
(c) GNOT
(d) Class-conditionall NCF
Figure 3: 2D class-conditional OT. The leftmost column displays μ (red) and (bluc), with class la- bels indicated by distinct markers. In the remaining columns, blue dots denote transported samples, while solid black and dotted gray lines represent the learned transport maps for each class.
6.2 CLASS-CONDITIONAL OT
6.2.1 2D TOY EXAMPLES
We present experimental results on a 2D synthetic dataset consisting of class-labeled samples. de- signed to evaluate class-conditional OT. To assess the ability of the proposed class-conditional NCF variant model class-guided transport, we compare it against an unconditional NCF, which does not utilize label information. Furthermore, to benchmark our method against existing approaches, we include NOT with general cost functionals (GNOT) (Asadulaev et al. 2024). recent model designed to perform class-conditional OT.
Figure 3 presents results on a Gaussian mixture dataset, where each data point is associated with class label. The unconditional NCF, lacking access to label information, learns global transport map that ignores class structure, aligning source and target points purely based on W2 distance. In contrast, both GNOT and the proposed class-conditional NCF learn separate transport maps per class. However, GNOT exhibits intersecting transport paths between classes, suggesting suboptimal- ity with respect to the transport cost. The caas-connitional NCF effectively disentangles transport across classes and yields maps that closely approximate the optimal solutions. These results high- light the accuracy and effectiveness of our approach, grounded in CODE-based formulation of the HJ equation, for learning class- conassiooniii transport in structured settings.
6.2.2 MNIST & FASHION MNIST
We apply model to the MNIST (LeCun 1998) and Fashion MNIST (Xiao et al. (017)) datasets, each comprising 10 classes. Given their substantially lower intrinsic dimensionality relative to the ambient space (Pope et al. 2021). solve class-conditional OT problems in latent spaces obtained via 3-,33-V VAEs (Higgins et al. 2017) ) see Appendix B.4 for details.
We consider transport from each Fashion MNIST class to its corresponding MNIST class; additional class-conae conditional OT tasks on MNIST are provided in Appendix C.4] We compare against basclines from Asadulaae et al, 2024), including NOT and GNOT, well domain adaptation OT method

---

Published as a conference paper at ICLR 2026
2
3
4
5
6
7
8
9
7917/00 7d
1 2
3
4
5
6 7
8
9
N
7
TORTHOND
7
2
3
4
5
6
7
8
9
Figure 4: Class-conditional OT between MNIST and Fashion MNIST. Left: Transport Fashion
MNIST data z forward T [x)((r) then backward T (T
(x)). (T][o]]]]) Right: Transport MNIST
data backward Т (z) then forward T (Tb (x)] The first row shows the source data,
the second & third rows present the data being transported back and forth using the learned map.
Table 4: Comparison of the accuracy and FID scores for the forward class-conditioned maps (Fash- ion MNIST + MNIST) learned using different methods. The accuracy and FID scores for the baseline methods are adopted from (Asadulaev et al. 2024).
NOT
GNOT
Discrete OT
Metric
AugCycleGAN
MUNIT
NCF [Ours]
1,2 cost
Stochastic map SinkhomL.pI.1
Accuracy(S))
10.96
83.22
10.67
12.03
8.93
83.42
FID
7.51
5.26
>100
26.35
7.91
18.27
Courty et al. 2016 Flamary et al. 2021) that uses discrete OT with label-supervised regularization. Additionally, we evaluate unsupervised image translation methods AugCycleGAN (Almahairi al. 2018) and MUNIT (Huang et al. 2018).
Figure 4 shows bidirectional transported samples by NCF; uncurated results are in Appendix C.4 These results qualitatively demonstrate NCF's ability to perform bidirectional, class-conditional OT on real images. For quantitative evaluation, report Fréchet Inception Distance (FID) (Heusel et al. 2017) and class-wise accuracy, which measures how well the class identity is preserved dur- ing transport, in T Our method achieves the highest accuracy. indicating its strong class- transport performance. Although the FID score is relatively high, this is largely due to the discrep- ancy introduced by the VAE decoder. To isolate this effect, we compute the FID between the NCF outputs and the VAE-decoded images. The resulting low score 2.73 indicates that the transport map in the latent space faithfully reproduces the target distribution. This is further supported by the KDE plots in Figure 13 showing close alignment between the transported and target latent distributions along principal components.
7 CONCLUSION
We introduced theoretically grounded OT framework that recovers forward and backward maps in closed form via HJ characteristics. The resulting single-network, integration-free algorithm gives accurate, bidirectional maps, supports broad class of costs, and extends cans-coondiional trans- port with pairwise MMD alignment. We establish consistency and stability. Several tasks including synthetic, color-transfer and MNIST demonstrate accuracy and efficiency of our algorithm.
Future directions include improving high-dimensional performance beyond latent-space implemen- tations by developing more efficient gradient evaluations and scalable network designs. Extending the stability analysis to general neural architectures would provide a deeper theoretical understand- ing of our method and its convergence behavior. Moreover, as demonstrated by our numerical experiments, the proposed NCF framework offers accurate and computationally efficient ap- proach for estimating candicoondi transport maps. Exploring real-world applications in this direction- such domain adaptation Nguyen al. (2024). cross-domain retrieval Chuang et al. (203), and biomedical conditional modeling Manupriya et al. (2024) with a particular emphasis on pursuing improved semantic correctness, represents an important avenue for future research.
ACKNOWLEDGMENTS
Y. Park and S. Osher were supported in part by DARPA under grant 12200112590074 S. Liu, M. Zhou, and S. Osher were supported in part by NSF under grant 2208272. The work S. Liu and M. Zhou was supported in part by the AFOSR YIP award No. FA9550-23-1-0082 M. Zhou and S. Osher were additionally supported by AFOSR under MURI grant N00014-20-1-2788 S. Osher was further supported by ARO under grant W91 W911NFF24-1-0155 and by NSF under grant 1554564.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Amjad Almahairi, Sai Rajeshwar, Alessandro Sordoni, Philip Bachman, and Aaron Courville. Aug- mented cyclegan: Learning many-to-many mappings from unpaired data. In International con- ference machine learning, pp. -95-204. PMLR, 2018.
Brandon Amos, Lei Xu, and J Zico Kolter. Input convex neural networks. In International confer- ence on machine learning, pp. 146-155. PMLR, 2017.
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International conference on machine learning, pp. 214-223 PMLR, 2017.
Arip Asadulaev, Alexander Korotin, Vage Egiazarian, Petr Mokrov, and Evgeny Burnaev. Neural optimal transport with general cost functionals. International conference on learning representa- tions, 2024.
Julio Backhoff- Veraguas, Mathias Beiglböck, and Gudmun Pammer. Existence, duality, and cyclical monotonicity for weak transport costs. Calculus of Variations and Partial Differential Equations, 58(6):203, 2019.
Yogesh Balaji, Rama Chellappa, and Soheil Feizi. Robust optimal transport with applications in generative modeling and domain adaptation. Advances in Neural Information Processing Systems, 2020.
Mare Barthélemy and Alessandro Flammini. Optimal traffic networks. Journal of Statistical Me- chanics: Theory and Experiment, 2006(07):L07002, 2006.
Jean-David Benamou and Yann Brenier. A computational fluid mechanics solution the monge- kantorovich mass transfer problem. Numerische Mathematik, 84(3):375-393, 2000.
David Berthelot, Colin Raffel, Aurko Roy, and Ian Goodfellow. Understanding and improving interpolation autoencoders via an adversarial regularizer. arXiv preprint arXiv:1807.07543 2018.
Rajendra Bhatia. Matrix analysis, volume 169, Springer Science & Business Media, 2013.
Gary Bradski and Adrian Kaehler. Learning OpenCV: Computer vision with the OpenCV library. O'Reilly Media, Inc.", 2008.
Charlotte Bunne, Stefan G Stark, Gabriele Gut, Jacobo Sarabia Del Castillo, Mitch Levesque, Kjong-Van Van Lehmann, Lucas Pelkmans, Andreas Krause, and Gunnar Rätsch. Learning single- cell perturbation responses using neural optimal transport. Nature methods, 20511111599-1768, 2023.
Guillaume Carlier, Cristian Jimenez, and Filippo Santambrogio. Optimal transportation with traffic congestion and wardrop equilibria. SIAM Journal on Control and Optimization, 47(3):1330-1333 2008.
Jaemoo Choi, Yongxin Chen, and Jaewoong Choi. Improving neural optimal transport via displace- ment interpolation. arXiv preprint arXiv:2410.03783, 2024.
Ching-Yas Yao Chuang, Stefanie Jegelka, and David Alvarez-Melis. Infoot: Information maximizing optimal transport. In International Conference on Machine Learning, pp. 6228-6242. PMLR, 2023.
Nicolas Courty, Rémi Flamary, Devis Tuia, and Alain Rakotomamonjy. Optimal transport for do- main adaptation. IEEE transactions on pattern analysis and machine intelligence, 39(9):1853- 1865, 2016.
Nicolas Courty, Rémi Flamary, Amaury Habrard, and Alain Rakotomamonjy. Joint distribution optimal transportation for domain adaptation. Advances in neural information processing systems, 30, 2017.
Bharath Bhushan Damodaran, Benjamin Kellenberger, Rémi Flamary, Devis Tuia, and Nicolas Courty. Deepjdot: Deep joint distribution optimal transport for unsupervised domain adaptation. In Proceedings of the European conference on computer vision (ECCV), pp. 447-463, 2018

---

Published as a conference paper at ICLR 2026
Max Daniels, Tyler Maunu, and Paul Hand. Score-based generative neural networks for large-scale optimal transport. Advances in neural information processing systems, 14:12955-12995. 2021.
Bogdan Danila, Yong Yu, John A Marsh, and Kevin Bassler. Optimal transport on complex networks. Physical Review E Stattisttiaal, Nonlinear, and Soft Matter Physics, 74(4):046106, 2006.
Lawrence C Evans. Partial differential equations, volume 19. American Mathematical Society, 2022.
Jiaojiao Fan, Qinsheng Zhang, Amirhossein Taghvaei, and Yongxin Chen. Variational wasserstein gradient flow. International Conference Machine Learning, 2022.
Jiaojiao Fan, Shu Liu, Shaojun Ma, Hao-Min Zhou, and Yongxin Chen. Neural monge map estima- tion and its applications. Transactions Machine Learning Research, 2023.
Xue Feng and Thomas Strohmer. Improving autoencoder image interpolation via dynamic optimal transport. arXiv preprint arXiv:2444.088000 2024.
Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon, Stanis- las Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier Léo Gautheron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamoniy, levgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78): 1-8, 2021. URL http:/// org/ wttpp//pppppppp/pppppaeee http://ppppes/////2/0/// 451 html
Torres Fremlin. Measure theory, volume 4. Torres Fremlin, 2000.
Aude Genevay, Marco Cuturi, Gabriel Peyré, and Francis Bach. Stochastic optimization for large- scale optimal transport. Advances neural information processing systems, 29, 2016.
Nikita Gushchin, Alexander Kolesov, Alexander Korotin, Dmitry P Vetrov, and Evgeny Burnaev. Entropic neural optimal transport via diffusion processes. Advances in Neural Information Pro- cessing Systems, 36:75517 2023.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), June 2016.
Johannes Hertrich, Christian Wald, Fabian Altekrüger, and Paul Hagemann. Generative sliced mnn flows with riesz kernels. International conference learning representations, 2024.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-vae: Learning basic visual concepts with constrained variational framework. In International conference on learning representations, 2017.
Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks are uni-
versal approximators. Neural networks, 2(5):359-366, 1989.
Xun Huang, Ming-Yu Liu, Serge Belongie, and Jan Kautz. Multimodal unsupervised image-to- image translation. In Proceedings of the European conference on computer vision (ECCV), pp- 172-189, 2018.
Guillaume Huguet, Daniel Sumner Magruder, Alexander Tong, Oluwadamilola Fasina, Manik Kuchroo, Guy Wolf, and Smita Krishnaswamy. Manifold interpolating optimal-transport flows for trajectory inference. Advances in neural information processing systems, 35:29705-29718. 2022.
Leonid Kantorovich. On the translocation masses. Journal of mathematical sciences, 133(4), 2006.
12

---

Published as a conference paper at ICLR 2026
Diederik Kinga, Jimmy Ba Adam, et A method for stochastic optimization. arXiv preprint arXiv: 1412(6), 2014.
Diederik Kingma. Max Welling, al. Auto-encoding variational bayes, 2013.
Alexander Korotin, Vage Egiazarian, Arip Asadulaev, Alexander Safin, and Evgeny Burnaev. Wasserstein-2 generative networks. arXiv preprint arXiv: /3002, 2019.
Alexander Korotin, Lingxiao Li, Aude Genevay, Justin M Solomon, Alexander Filippov, and Evgeny Burnaev. Do neural optimal transport solvers work? continuous wasserstein-2 benchmark. Advances in neural information processing systems, 34:14593-14605 2021a.
Alexander Korotin, Lingxiao Li, Justin Solomon, and Evgeny Burnaev. Continuous wasserstein-2 barycenter estimation without minimax optimization. arXiv preprini arXiv:2102.01752, 2021b.
Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. Kernel neural optimal transport. arXiv preprint S.MXi::2205./5266 2022.
Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. Neural optimal transport. Interna- tional conference on learning representations, 2023.
Takeshi Koshizuka and Issei Sato. Neural lagrangian schr\"" odinger bridge: Diffusion modeling for population dynamics. arXiv preprint wrXiv:2244.0485: 2022.
Yann LeCun. The mnist database of handwritten digits. http://yann. lecun. com/exdb/mmistt, 1998.
Moshe Leshno, Vladimir Ya Lin, Allan Pinkus, and Shimon Schocken. Multilayer feedforward net- works with nonpolynomial activation function can approximate any function. Neural networks, 6(6):861-86) 1993.
Alex Tong Lin, Samy Wu Fung, Wuchen Li, Levon Nurbekyan, and Stanley Osher. Alternating the population and control neural networks to solve high-dimensional stochastic mean- field games. Proceedings of the National Academy of Sciences, (18831)e2002413313 2021.
Huidong Liu, Xianfeng Gu, and Dimitris Samaras. Wasserstein gan with quadratic transport cost. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 4832 4841, 2019.
Shu Liu, Shaojun Ma. Yongxin Chen, Hongyuan Zha, and Haomin Zhou. Learning high dimensional wasserstein geodesics. arXiv preprint arXiv:2102.02992 2021.
Shu Liu, Stanley Osher, and Wuchen Li. A natural primal-dual hybrid gradient method for adversarial neural network training on solving partial differential equations. arXiv preprint arXiv:2411.0627 2024.
Guansong Lu, Zhiming Zhou, Jian Shen, Cheng Chen, Weinan Zhang, and Yong Yu. Large-scale op- timal transport via adversarial training with cycle-consistency. arXiv preprint arXiv:2003.066. 2020.
Ashok Makkuva, Amirhossein Taghvaei, Sewoong Oh, and Jason Lee. Optimal transport mapping via input convex neural networks. In International Conference on Machine Learning, pp. 6672- 6681. PMLR, 2020.
Piyushi Manupriya, Rachit K Das, Sayantan Biswas, and SakethaNath N Jagarlapudi. Consistent optimal transport with empirical conditional measures. In International Conference on Artificial Intelligence and Statistics, pp. (656-3654. PMLR, 2024.
Tuan Nguyen, Van Nguyen, Trung Le, He Zhao, Quan Hung Tran, and Dinh Phung. A class- aware optimal transport approach with higher-order moment matching for unsupervised domain adaptation. arXiv preprint arXiv:2401. 15952, 2024.
Derek Onken, Samy Wu Fung, Xingjian Li, and Lars Ruthotto. Ot-flow: Fast and accurate continu- ous normalizing flows via optimal transport. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 9223-9232, 2021.
13

---

Published as a conference paper at ICLR 2026
Yesom Park and Stanley Osher. Neural implicit solution formula for efficiently solving hamilton- jacobi equations. 2025. URL
19351
Phillip Pope, Chen Zhu, Ahmed Abdelkader, Micah Goldblum, and Tom Goldstein. The intrinsic dimension of images and its impact on learning, arXiv preprint arXiv:2104.08882 2021.
Maziar Raissi, Paris Perdikaris, and George Karniadakis. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational physics, 378:686-707, 2019.
Erik Reinhard, Michael Adhikhmin, Bruce Gooch, and Peter Shirley. Color transfer between images. IEEE Computer graphics and applications, 21(5):34-41, 2001.
Maria L Rizzo and Gábor J Székely. Energy distance. wiley intendisciplinary reviews: Computa- fional statistics, 8(1):27-38, 2016.
Lars Ruthotto, Stanley Osher, Wuchen Li, Levon Nurbekyan, and Samy Wu Fung. A machine learning framework for solving high-dimensional mean field game and mean field control prob- lems. Proceedings of the National Academy of Sciences, 117 117(17:91133-1193, 2020.
Maziar Sanjabi, Jimmy Ba, Meisam Razaviyayn, and Jason D Lee. On the convergence and bustness of training gans with regularized optimal transport. Advances in Neural Information Processing Systems, 31, 2018.
Filippo Santambrogio. Optimal transport for applied mathematicians. Birkäuser, NY, 55(58-63):94, 2015.
Geoffrey Schicbinger, Jian Shu, Marcin Tabaka, Brian Cleary, Vidya Subramanian, Aryeh Solomon, Joshua Gould, Siyan Liu, Stacie Lin, Peter Berube, et al. Optimal-tramspe analysis of single-cell gene expression identifies developmental trajectories reprogramming. Cell, 176(4):928-99)) 943, 2019.
Vivien Seguy, Bharath Bhushan Damodaran, Rémi Flamary, Nicolas Courty, Antoine Rolet, and Mathieu Blondel. Large-scale optimal transport and mapping estimation. arXiv preprint 1rXiv:1711.02283, 2017.
Zebang Shen, Zhenfu Wang, Alejandro Ribeiro, and Hamed Hassani. Sinkhorn natural gradient for generative models. Advances in Neural Information Processing Systems, 33:1646-1656, 2020.
Alexander J Smola, A Gretton, and K Borgwardt. Maximum mean discrepancy. In 13th international conference, ICONIP, volume 6, 2006.
Gábor J Szěkely and Maria I. Rizzo. Energy statistics: A class of statistics based on distances. Journal of statistical planning and inference, 1249 2013.
Amirhossein Taghvaei and Amin Jalali. 2-wasserstein approximation via restricted convex potentials with application improved training for gans. arXiv preprint 1rXiv:/902.07/999 2019.
Alexander Tong, Jessie Huang, Guy Wolf, David Van Dijk, and Smita Krishnaswamy. Trajectorynet: A dynamic optimal transport network for modeling cellular dynamics. In International conference on machine learning, pp. 9526-9536 9536. PMIR, 2020.
Cédric Villani al. Optimal transport: old and new, volume 338, Springer, 2008
Gefei Wang, Yuling Jiao, Qian Xu, Yang Wang, and Can Yang. Deep generative learning via schrödinger bridge. In International conference on machine learning, pp. 10794-10804, PMLR, 2021.
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: novel image dataset for benchmark- ing machine learning algorithms. arXiv preprint arXiv: 778.007747 2017.
Yujia Xie, Minshuo Chen, Haoming Jiang, Tuo Zhao, and Hongyuan Zha. On scalable and efficient computation of large scale optimal transport. In International Conference on Machine Learning, pp. 6882-6892, PMLR, 2019.
14

---

Published as a conference paper at ICLR 2026
Liu Yang and George Em Karniadakis. Potential flow generator with 12 optimal transport regularity for generative models. IEFE Transactions on Neural Networks and Learning Systems, 33(2): 528-538, 2020.
Benjamin Zhang and Markos Katsoulakis. A mean-field games laboratory for generative mod- eling. arXiv preprint arX44 2023.
Jiaxi Zhao, Mo Zhou, Xinzhe Zuo, and Wuchen Li. Variational conditional normalizing flows for computing second-order mean field control problems. arXiv preprint arXiv:2503. 19580, 2025.
Mo Zhou, Stanley Osher, and Wuchen Li. Score-based neural ordinary differential equations for computing mean field control problems. arXiv preprint arXiv:2409. 16471, 2024.
Jun-Yan Zhu, Tacsung Park, Phillip Isola, and Alexei A Efros. Unpaired maage-to-image translation using cycle-consistent adversarial networks. In Proceedings of the IEEE international conference on computer vision, pp. 2223-2223 2017.
15

---

Published as a conference paper at ICLR 2026
A PROOF
A.1 PROOF OF PROPOSITION 4.2
Since is I.s.c,, sub-differentiable, and strictly convex, its Legendre transform h is also l.s.c., sub- differentiable, and strictly convex. For such Hamiltonian h, has been proven (Park & Osher 2025) that the viscosity solution satisfies the implicit solution formula (9) almost everywhere. Con- sequently, the optimal solution to these OT problem is characterized by the implicit solution formula.
A.2 PROOF OF THROREM 5.1
Lemma A.1. Suppose that μ. are probability distributions on IR: Suppose has finite second order moment, i.e.. Jaa [[(x) < co. Assume that sb E + [U++ is convex and differ- entiable p-a.e. Set T Vń and suppose Jgs [T(()))Pp((z) < -00. Then T optimal for the transport cost 211= - yll² between and v.
This lemma proved in Theorem 1.48 Santambrogio (015).
Lemma A.2. Suppose T Rdd -> R&d a bijective map on RR. Assume that T is strictly monotone, that is, (T(x) -T(y), x y) > 0 for arbitrary I, E Rd + Here, we denote (:,-) as the l² inner product on IRd. Suppose : Rd + IRd satisfies S SoT = Id, then is also bijective, and strictly monotone.
Proof. It is straightforward to verify that S is surjective and injective, and it is thus bijective. Now, for arbitrary x; y Rdd # y, there exists unique x', y E Rd, x' # y', such that = TTx'), T(y)). Thus, we have S(x) = S(T(=')) x', Sly) = S(T() = n/. and (S(x) - S(y), x y) (x' T(x') T(y') ) O.
For brevity, in the following discussion, we denote Ct (TR:) by C* (RR) for any & N. We denote Ou as the d x d zero matrix. For symmetric matrices A, B E Rdxd, we denote A x B if A B is positive definite.
Theorem A.3. Civen the probability distributions μ. E P (Rd) with fo= ||0|| du, Szx ||~ dv +o0, suppose un E Cl (IR",, u1 E (I2 Vui E LRR [Rd satisfy
4 _ (x)(1))
un 14 Vui (x)) = un (1))
V
(18)
Assume further that (Id + I, Vuo and Ia tj 22 (x) Os. then Id + t, Vuo(-) is the optimal transport from μ to and Id iy Vun (.) is the optimal transport map from to J..
v. the mapping Id Vu Rd Rd is bijective
Proof. We split the proof into several steps:
Step 1. We first prove the fact that Vun(æ Vui (u, E)) = Vu (x) for arbitrary x E Rd. This can be shown by taking gradient with respect on both sides of (88
(La L t, (u) -t (x) (u) V
Re-arrange this equation yields
(L. -
vu - (x)) =
As I, -,, (2 (I) > Og, we deduce that Vuo( (x)) = (z) for arbitrary E Rdd
Step For the sake of brevity, we denote To(-) == Id + v, and T.(:) Id - (t) We prove that To o Ti- Id. This can be derived by straightforward calculation
(z) (x) Vuo (T) t Vx
x + (V - (x)) - Vur(z) = x,
for any Rd.
16

---

Published as a conference paper at ICLR 2026
Step Before we prove the assertion, we show that SR* [[To(z))] du < +00, Sad [T)()) < +00. The latter inequality can be shown using
/ [T) 2
(Le [TT z112 dv / | dv
V + #/5e ||x||/ du) < +00.
For the former inequality, we have
J Rd
[To(2)]) s 217 / [uo( +2 S ||=]| du.
Using the fact that TiF μ. The first term above equals (eoo (x)()) dv Jas V (x)]12 du [[ (1 < +00, where use the fact that Vuo(Ti (x)) (u) established in step 1. This accomplishes the proof for fqa [[T[(=)]12 dp < +00.
Step We now prove the conclusion. Firstly, recall that DT} L. - ty V (z) Os, this leads to the fact that (.) is strictly monotone. We now apply Lemma A.2 to show that To is also bijective, and strictly monotone.
To is bijective suggests that Ti (-). is the inverse mapping of Tu(-), this leads to Tig = p. As (-) = (-) with yp(-) = ||-I"" Lemma A. proves that Ti is the optimal transport map from μ..
2 - (,) being convex, combinging the fact established in step 3,
Furthermore, To is strictly monotone yields that DTo(x) = Ia + un (x) > Od, this indicates that To(:) = [ + tyua(-) is the gradient a convex function. Combining with the fact that Togm = v and finite 1,2(N) cost for To. we deduce that To is the optimal transport map form j to
Recall (IR)), and the implicit HJ loss Cur(u) defined in (13 The following Theorem is natural corollary of Theorem A.33
Theorem A.4 (Consistency result), Suppose the probability distributions μ., v possess finite second- order moments. Assume u E Cla x [0, fy)) and define u1(-) ;= u(-,ty) E C²(R)) with Vui E (RR Rd;v). v). Denote hyperparameter measure Rd. Suppose that u minimizes the loss functional, i.e..
> 0. Assume that is strictly positive probability
Cur(u) + >LumD(u)= =.
Assume further that the map [H[u] Rdd - RRd. is bijective and its Jucobian D₂ D2 [u](a)) > Oa for any x E Rdd, Then the maps TV ud) and [ylu] are the optimal transport maps from v to j and from Ju to v, respectively.
Proof. Denote E(x,t) = (u(x - - u(x,t,)+ + t/)()) we have E E C(R x [0, ty)), and E(x,t) 2 0 on Rd x [0.ty]. Denote m as the Lebesgue measure on ty].
Then we have
Cml(u) /
E(x, f) dedm 0.
Tonelli's Theorem Fremlin (2000) leads to
[ (1." E(z, t) dt do(x)= = 0.
Now, Lemma A.6 indicates that ((x)= == 1'' E(x, t) dt is continuous on Rdd. And we have
Sau (xx) do(x) 0. Lemma A.7 suggests that d(x) 0, which f'y E(x, t) dt 0, Ya æd
Using similar argument presented in the proof of Lemma A.7 shows E(x,t) 0 for all & E IRd, I (0,
Now applying Theorem A.3 with up u(-,0) 0) and u(((4) proves the assertion.
17

---

Published as a conference paper at ICLR 2026
Remark A.5 (On the Monotonicity Condition). The monotonicity condition D, D [u] O2 closely related to the e-concavity of which provides sufficient condition in OT theory (Villani al., 2008; Santambrogio 2015). In practice, our proposed method successfully computes the OT map for the benchmark problems even without explicitly enforcing this condition. Nonetheless, it remains an interesting direction for future research to investigate efficient strategies for enforcing monotonicity and to assess the potential benefits of doing so.
Lemma A.6. Suppose E' C R is compact. Assume that f E C(R x E), then v(x) = SN fax, continuous with respect to variable x.
Proof. Fix arbitrary E Rd pick continuous on the compact set B x E, we know that [s'] is bounded from above. We can then apply the dominated convergence theorem to show that lime-+o v(2) lim, SE f(2, dt SF. (x, 1) dt y(r). Thus, y is continuous on Rd.
O, we denote B = {y lly x|| S r). Since f is
Recall that a Borel measure a on R is strictly positive if o(B) for any open set B R Lemma A.7. Suppose a P (Rd) is strictly positive, assume g
Rdd + [0, +00) is continuous
function. If Sea g(x) (x) = O, one has g(x) = for any x E R
Proof. Suppose not true, we have specific no E Rd such that g(x)) > 0. Since is continuous, one can find 0, such that [g(z) g(zo)l < [/g(xo) for arbitrary x B : { || Fo || < 60)-
E B one has 9(.) 2 g( X (-) on
Consider the indicator function XBŞ, (x) = otherwise,
Rd. This yields JR+ g(x) darx) 2 J&g g(
XBS (x) dorz) = a(zo)a 0, where the last 2
inequality is due to is strictly positive and the ball B& is an open set. This is a contradiction.
O
A.3 PROOF OF THEOREM 5.4
Without loss of generality, we we prove the theorem with I = 1. The proof can be generalized to arbitrary time horizon with no difficulty. We first state and prove two auxiliary lemmas. Lemma A.8. Let (() be bounded and Lipschitz. Lipschitt.uunciion [0, 1). Then
|| SCII
This C only depends the bound and Lipschitz constant of
Proof. Let = |I z and L be the Lipsehitz constant of f. We split into two cases.
Case 1. If K 2 L. then by the Lipschitz condition of f, we have
[ f(t)2 dt 2 [ (K (K- - LO) dt = K2 LK + 1 K2.
1 3
Case 2. If K < L,, since f is I--Lipschitz, f(t) non-zero for an interval of length at least K/I, and
S f())
K/L (K Lt)² dt
K3
3L
In both cases, conclude that
11 S SC111
Next, present lemma that shows the stability of the MMD with respect the mean and covari- ance of Gaussian distributions. The MMD with kernel k(x, y) x y/ is known as the energy
distance, which is extensively studied in Szěkely & Rizzo 2013 Rizzo & Székely (2016).
18

---

Published as a conference paper at ICLR 2026
Lemma A.9 (Stability of MMD on Gaussian). Let μ N(bps mn) and u N(by, Eny) be two
Gaussian distributions Rdd with ||b ||by | |Evll2 S K. Then, there exists a constant C' that only depends on d and K s.f.
b b | - Ex 12 S CMMD(W,
(19)
Remark A.10. The boundedness assumption is necessary. As a counterexample, let μ (0, o²) and v/ ~ (1, )²). Let o² + co, 1b,, - byl = 00 1 = 1 remains unchanged, while MMD(H, v)
converges to 0 (see (0)).
Proof. Throughout the proof, we will use the notations e and C' to denote positive constants that only depends on d and K. These constants may change from line line. We start by recalling important result. Let (t), (t) R2 7 R be the characteristic functions of jr, V
(t(t) exp [b) -
1 ET,t)
( G ,,
Then (Székely & Rizzo
2013
Proposition establishes that
MMD(U, v)² = d + 1
(tm(t) - (t)
dt.
(20)
2
Rd
||| d+1
Next, we split into two steps and give bounds for the mean and covariance separately.
Step 1. We give bounds for the mean. For any given t E Rd the four complex numbers exp (b] ftEtt) (i, k v) forms a isosceles trapezoid in the complex domain. an isosce- les trapezoid, the length of each diagonal is greater than equal the arithmetic mean of the lengths of the two parallel sides. Therefore,
|exp exp
(tt(t) - (t) =
1 [xTt) - exp AbJ 2t tt_ TTT,tt Ext
2
exp
1 t T
1ntt/ + exp
- [ 2i,t )] |1- - exp (i(b,, - -b,) Tt)l
22itt)
2
exp( - C
[|t||3) 1
cos (b,- by) T) isin((b,. bt)
Substitute this estimation into 20
obtain
exp( C' [|t|]]] cos((b, - bb) isin((b,, - bb,) Tt) dt
-
MMD(U, U)2 2
||t||ut2
3dd
exp(-C11t11) ||t||
sin² (b,- bT) dt.
4c
Since ||b, bull S 2K is bounded, can find r > that only depends on K such that sin(w) >
r for all [0. Y/2 | byll r]. Denote B, the ball in Rd centered at the origin with radius have
exp(-C|||) sin (5(b, bb) 8) dt
M)3 20
||t||ut1
13
exp(-
(7 - -b) ) dt
2 C
|t||d+T
B,
[bb - -bo) t/² dt.
2
B. ||t|
we further restrict the 3D angle of in the set B, (t E B, 1(b - Tt) 2
| bull ||t|] }. ie., the angle between and b, by is close to or a. Then,
MMD(U, v)² 20 B. ||t|l4i (b,, bb) b Tt|² dt
4. B. ||t||=11 | byll dt ellb, byll
19

---

Published as a conference paper at ICLR 2026
Step 2. We give bounds for the covariance. This time, we use the fact that the length of each diagonal in isosceles trapezoid is greater than or equal to the length of either of the non-parallel (equal) sides. This gives
1 I Ext exp (bbtt 21
(t)(t) - (b]l
exp
ab
2t
2
Exnt) )
2 exp
Ext
exp
1-
(-2
(
1 (En T(E - En)t -En)t)
exp
(-2
exp
2
T (Ee, En)t
where we used - 2 2 | in the last inequality. Next, we diagonalize the symmetric matrix Σ, En as Q(E, Σn] A diag(),) Without loss of generality, we assume that
|^l |Σ Enll2 I12 We denote gd-1 c Rd the unit sphere in R Substituting the estimation
above into 20 we get
D)²
exp(-C|lt|) -E,))t dt En)t)
3z ||t||ut1
2 c / B1 ||t||ctt (t (2iu - -m,)t)² dt
/ -d-1
R R Rd (sT "EE - -E))s) 2 Rd-1 dsdR
/
SE ds c S As)) ds.
= 0
88--
We further pick subset g--1 {s Egd-1
| 2 3) (i.e., points the unit sphere with the
first coordinate 2 3). For all s E 5d-1
1s TAS
2 | 2 -
Σ
1Ag
83 N 2
7A1l.
1
Mnl.
j=2
Therefore,
( (s As)²
1
f
MMD(H, D)² 2 J0--
ds
2 c11Ep
Ex
Finally, combining Step 1 and Step 2, we reach the conclusion that
| 1y + EE Σ, S CMMD(u.
Before proving Theorem 5.4 we clarify the result we need to show. The HJ equation is
0, t) + V = 0.
The optimal push forward map is (x) = x + ux 0), which implies Vzu(x, -r* x.
The optimal trajectory for OT has constant velocity, given by
x; = x + (xx 0) = x t(x T (x)).
Therefore, the optimal push forward map
1(x, (1 x)x+ ((1 1)T + AL) + (b, Abb,).
20

---

Published as a conference paper at ICLR 2026
Taking derivative in 1, we get the optimal velocity Lagrange coordinate
a.f(x.t) = -x (x)): = 1 ((((x, - -x
= [(xx - -1 A1) (f(,1)) -1-
(b, Abb,)()
= ((+U(A-))+ ((A-)(( +b.- Ab,).
Therefore, the optimal velocity field in Eulerian coordinate is
Vxu(x, t) = xxx t) = (I t(A- ((A 1)x + b, Abb)).
(21)
The most important challenge for convergence analysis is the term uø(x tV ux, t), 0) in the HJ loss, which contains the composition of the u. In order to address this issue, we consider the quadratic parametrization
us (x,, t) -- ( (2())x + +tt)) x 0o(t)
(22)
2
where 0 = (02(-), [.(-)] [0,ty] + Roym* x Rd x R is bounded and Lipschitz. According to (21) the optimal 0% (f) and (t) are uniquely determined, and the optimal (t) is uniquely determine up to an additive constant.
(t) = ((1 - --t At) (1 - A)
(t) ((1 - -1)7++ At) (Abb - -by)
(23)
(1) (0) + (b, - Abp) T. ((1 - t)T AC) (b, - Abp)
Now we are ready to prove Theorem 5.4 We denote D = 1-1, Throughout the proof, we will e = 2-11p n and the time domain 0,4y) [0, 1]. which coincide with our numerical imple- mentation. The proof can be extended to general domain without essential difficulty. Throughout the proof, when we say a function is bounded and Lipschitz continuous, we mean the bound and Lipschitz constant only depends on d, AA, and K. We will use C' to denote an absolute constant that only depends on AA, and K. The value of C may change from line to line.
Proof for theorem 5.4 We only need show 16 when CHs and LAMD sufficiently small. The proof consists four steps.
Step 1. We analyze the MMD loss in this step. Recall the MMD loss is
/ k(x, y) d((Id + Veue(, 0)) (x) d((ld v=uo(:0)) # -
LMMD
Under with the parametrization (22). the MMD loss is between (Id + V 2u 01) N ((I (z)) (0), (I (2(0)) (2()))) N(bb, Enr
and (b,, ). By Lemma A.9 we have
| - by + Ex S CLMMD
(24)
Multiplying Σ3 both sides for the covariance, we get
(x =0))))) - M < C
By the 1 --Höldrr continuity of matrix square root in operator norm (Bhatia 2013 Theorem X.I.I)
( [ Bå < I|A B|B for any symmetric positive definite matrix A, B), we have
( (1- (==00= - -(EXE,))
< sc
(25)
21

---

Published as a conference paper at ICLR 2026
Next, we diagonalize 02(1). Since (2(() E Raxd is symmetric, we can find unitary matrix Q(1) and diagonal matrix A(t) = diag( (A:(1) s.t.
(2) = Q(t)A(()(())*
and 2.. (1(.. The column vectors of (t)) are the orthonormal eigenvectors of @2 Since (t) is bounded and Lipschitz continuous in t. its eigenvalues and eigenvectors are also bounded and Lipschitz continuous in / by Weyl's inequality. Q(1) is uniquely defined up to sign shift for each column. Then / (2( = Q(0)(T (())
Next, we define notation of "absolute value" for symmetric matrices. Given a symmetric matrix, we diagonalize it through unitary transform, take absolute value of the diagonal element, and then apply the inverse unitary transform back. As result, - - A(t) diag(1) ([]l (x(()) and - (20)) = Q(0) / A()) Q(0) then we have
( (1-000))))
E 02))
Here, we remark that while ((t) is not uniquely deñned, the "absolute value" |T (02(0)) is uuiquely defined given (2))). Plugging this expression into (25). we get
$
E T 02 (2)
E
MMD
Multiplying = both sides, get
- 02(1) - All2 S CCAMD
(26)
which implies
||0) ((0))220)) S &CRAND
(27)
Therefore, the off diagonal elements of (Q(0)AQ(0)])) )is (i + j) and diagonal elements (Q(())AQ())))) s satisfy
[(Q(0)AQ(O)))) )ayl ((Q(0)A(())))) - 11- (o)) s C.
Here, we add a subscript in the constant C1 in order to keep track of this constant. Later, whenever we use C1, means this fixed constant that does not change from line line.
We remark that there are 2d choices of (0) such that (1 (0)))) - 02())) = Σ. through letting (0) = = (i = l...., d), where X is the i-th eigenvalue of A. All these choices gives a push forward map that transport j to y (if we set 0₃ (0) = (I Og (0)) by). However, only (0) the HJ loss Lıu is necessary.
A gives the OT map. The MMD loss LMDD cannot distinguish these choices, so
Step 2. We analyze the implicit HJ loss in this step. Under the parametrization 22). the HJ loss is
/
/
1
CHU
x (t(() (0.(t)x + (02(t)x (t))
5
2 (x (() - (t))) (x ((())) (t)))
- (0) (x + +00g(()
- (0())] -0()))))))))))))))))))))))))))))
o(x) dt.
Reorganizing the terms, we have
C 1/' /
1 5 (02(t) + (t) (1))
x
+x (I (t)(t)) (01 (1) - (1)))
(o(() + t (1)) (1()
t2 ( T (0) - 6o(0) ) e(xx dx dt.
+
2
2
/''' / x +2(1)++xx x TT(() To(1)
1
e(x) dx dt.
9 2
22

---

Published as a conference paper at ICLR 2026
We observe that (2) is symmetric. The integration x for the loss can be computed directly. The zero-th to third order integration in x can be easily obtained by symmetry (recall e(x) and = [11
J (1,
2d0u3 200
dx - 2
In order compute the fourth order integration x, we temporally denote (2 by I' for notational simplicity and compute the integration Jo (x Tx)2 dx. Expanding everything, the integration is
( Tx dx
E 6.R.l:11
TiJTRl
/ Z.0,TK01 dx.
All the non-zero terms in the form ITIT'M lt ID 2,zzZ221 dx can be categorized into 4 cases.
1. j k L. The integration is 1 Jn o. dx 1 20.
2. i = j + k = I. The integration is T ak J xΩ dx 9 TRI
3. i = k + j = 1. The integration is r f x x3 dx = T
4. 1 + j k. The integration is T,yT, JD 2 dx = ; 2r
Summing them together, we have
+ + (33) )
1 (x (xTx dx==xx=> 1xx=2" 2d
Σ
Σ in
+
Therefore, after integration in x, the implicit HJ loss becomes
CHI
/
7
d Σ Ta
+ 7 ン
Σ
+
+
(T)
(t)] + +o(t)} dt
d
1 |1:(001 (fmaat0)+ +ro(e) dt
1
(2
Σ
T2(1) +
45
18
1=1
/
1
In (f dt
|T2(t))
+
J
45
Therefore,
L (IT2(0)) + ([T(()) +T(())))++))))))))))) at s CLu-
Therefore, by Lemma A.8 we have
max (t)ll max
max [[T (t)ll + max [T(()] S C2
(28)
Here, C2 also does not change from line to line.
Step 3. In this step, show that (2) must be close to 07( -- A, provided that ||T2(t)]] p is
sufficiently small. Since A has minimum eigenvalue a > (Q())AQ(O)T) " 2 Aa 2 C CiND where the last inequality is because LAMD is sufficiently small. We recall that in Step 1, we showed for any i ......d
(((())A(0))) )a - xx(0]|| C. C CMD
We want to show that,
(((())4((0)) )) - -11-((0)) C C.
(29)
23

---

Published as a conference paper at ICLR 2026
for all i. Lee, we want show xx(0) 2 and
(0) E [ (Q()))QQ(0))): C CM (Q(0))Q(()) )s C2 CLCDD
for all i. In order to show this, we assume to the contrary that ( > 1 (recall (1) 2 Ax(O)) and
2
(0) (2(0)Q(())) - 1 (Q(())AQ(0)) )) C.CiND)
(30)
We will derive contradiction. We denote (z(()) Q()) (2)))(()). Since unitary transform does not change Frobenius norm,
Q() T())((()
+ (1 (0000000000000000000000000000000
(A(I))
(31)
A(t) + (t)) (I +A(t))
shares the same estimation as 28 Q(1) Q(1)T02l (t)) Since Q(t) is bounded and Lipschitz continuous, ū2(() and (t) are also bounded and Lipschitz continuous. We will also use K to denote their bound and Lipschitz constant. Since (0) > 1 an eigenvalue of (0) and hence also an eigenvalue of (t). we have
Let {T.(t)]( be the diagonal element of (t
max (1) > 1
AA [0, 1].
Next, we focus on the diagonal elements of T((Q(())) Since [Q())==()00)) s C2c
(32)
its i-th diagonal element (recall (11
(x(() + (Ax(1)² -1
= =(1+(v(0) (1 -00+4000000)))))))))0)))0)))
=11+ (m)))))))()))))))(())))))))))))) (()
also satisfies
1(1 () - -00] S C2
-
(33)
for all t E [0, 1], where recall that (t) is the i-th diagonal element for A(t) = Q(t) (t)) and (t) is the i-th diagonal element for uiz()) (tt) (0)
The rest of the proof for deriving contradiction to (30) is technical, so we explain the main idea first. In order that [(1 + (1)(1) [[1 (1) (1)] is small for all t E [0, 1]. either of the following must hold
1. 1+ (t, 0, which implies (t()= -
(t(
1 1 1
2. - (t) - (t) se 0, which implies Ax(t) 2 1- (T, 1
(T((1) (At
0 the function is (t).)
When 4- -+ blows up and we cannot have 1 + (f= R O.
Since max (t) > 1, we know from intermediate value theorem that there exists at least one index
(t) blows up as L - 44-
i and 4 E (0, 1) s.t. - (4,(4)) = 0. This implies that the function
T (t,
As result, in order that [1 + (A)( [[1 (t))))))) - (r)) is small for all E (0, 1). there has
to be some "shift" between two functions: (t) is sometimes close - | and sometimes close to
(t)
1 1
However, note that the difference between the two functions
(t((t)
8l (t)
and 1
1
has a positive lower bound
t1 - (t((t)
> 1 (t)
(34)
1 (t)
(tt(t)
Tt(t)
||.4 + 111, + C2 C.
24

---

Published as a conference paper at ICLR 2026
Therefore, for some t in the middle, both [A:(t) + il and [2:(1)
r.(0) I-17((1)
are larger than
+| This gives contradiction to 33
Next, we give rigorous proof for this contradiction. Let
t" inf {t E [0. 1] (t) = 0 for some i}.
Note that the set above is non-empty because max T/( for all E [0, 1). If we do not have the assumption A (0) 1, then (1) < 1 may not well-defined. By definition of 1", we can find an index j such that 1 (1*) 0. Therefore, Ty("") 2 1. Let ln = 3K then for E [0, fo). we have
|1 [1) 2 t
(35)
3
Let 1 = tt
At, where At = 2k(1+2K) then for all t E [tt.t"] we have
1
2 [ry(t") - Klt (z)) - (t)(())))))
(T)-))(t))A
2 1 K&t K (1++'r((+))++ (t" (t))
(36)
2 1 KA K (0 2K (*1) 2 (/K + +1= = 2
If to 2 fi we pick E [81, tol and then multiply
35 and 36 we reach contradiction with
33
to < f1, then we consider the behavior of
+ thy (t)) -
tT (y(+)
(t)
When E t0, (t) is close
because
35 and 33
implies
1 - (t)
3 [(1 - (T)-)((()))(()) Ty(1 s C2Ci
2
which gives
20(1 (()
3C2Cin 2(1 (t)
9C24i
(37)
tTy
When E (,,tt))Ay(0) is close to 7 because (36) and 33 implies
| + (A₂(1) < C0
This implies
(,(1) + 1 < 2C2C# < 2C2Ci S 6KC
(38)
where the last inequality is because t1 to Therefore, explained before, there has to be shift between the two approximations 377 and (38) in the middle when E [to, t1] because (y is I Lipschitz continuous. However, the difference between the two functions has positive lower bound
- (t) 1 - (0)
1
1(1 (T,(1))
|| A + 1l12 - C, CIADD
Therefore, there exists t2 E [t.. s.t.
(2((22) toT
(12)
2
(39)
2 (11A + /ll2 + Ci Ci
and
(3((2)) + t2
1
2
(40)
(IA + Ill CL
Finally, we split into two cases.
25

---

Published as a conference paper at ICLR 2026
Case 1. If 12Tz((2)) S 3K then
1
1
T5(12)
2
1
2
1
12
3K
3K
and
(b)))
2 (12) - [1 - 12T,
(12))
1 1 K
1
2 3K
3K
3
This implies
111 +1+12,10 (22) 12Ty (2Tg((2))A)((1)) (12)]
Σ 3
1 11 (2
-3
(t22)
1/
2 1 1
1
= 0(1)
3 3K 2
(IA+ /112 + C)CAMD))
which contradicts to (33
Case 2. If 1 (t))l >
3k then
1(1 (t2)
[[1 (2T (t2)ll
1
(1(22)
|2||
{g(12)
|1 - (1))
(1(12)
12
(-(27((12)
O1),
3K
[|A+ + [|A+ C
3K 2 (IIA+ Ill2 + C1
which also contradicts to 333
Combining Case / and Case 2, we conclude that the assumption A1 (0) 1 cannot hold. Therefore,
29 hold. This further implies | A(0) I A(0). Plugging back into 27 and (26). we get
||T - A() Q(0)AQ(0)) s C
and
[|0((0) (I A)ll2 I CC
Therefore, we obtain
[[0((0) (I (T-A)|le S CC
(41)
Step We show that (2(t), (t). and (t( satisfies the error estimations (16).
Step 4.1. We estimate (t) first.
We first show that 1- (t)(t) has positive lower bound. Recall that (t(1) is the diagonal element of ūz(t) Q(t) (()) We first observe that any diagonal element for
Q() t(1 ()) Q(t) ((1 t)1 + tA))(t)
must be larger than or equal min(1,A4) By (4T). Lr, (T,( a diagonal element of
- 1Q(1 (0(())T((0)00))))) = Q T (1 4 02 (()))((()
must satisfies
1 - 1-tr(1)) 1 - iC > min(1,AA} CLi. > 2
min [1, AA
Therefore, 1 1- (T)(() has positive lower bound.
26

---

Published as a conference paper at ICLR 2026
Next, we claim that, for any (fixed) + (1r(1) has positive lower bound. Similar step 3, 33 can be rewritten as
Tt(t)
11 + -1,0(()111 (T,( 21(() (T(((
S CC
Therefore, the lower bound for 1 - (t) implies
|| t A1(t)) - :i(()
S CL
(42)
(t,
for all E [0, 1] and i. If we further restrict ourself to t E 1 1]. we have
(:(0)) -(0)-(-t)) --)0((-t)/ A:(e) - (7)
T((t)
- CC for tE
3K 11
(43)
Tt(t)
implying that (t) must be close to either -7 or 1- (T:
Next, we show the lower bound for + (t,)).. For E [0, 3k
1 1 tt, 2 tk
Therefore, we must have V E [0, 3k]
Ti(t)
(1(()
S CC
(44)
T- (T)
33 1]. similar to the argument in step 3, by (-)ry(A)* but cannot be close to both because the difference between the two functions has a positive lower bound of O(1):
For t E 1 (.lt)
43 Ax(t) must be close to either or
1
(t)
1
2
1
1 - (()
11
(T.(1)
(T:(() max (1)1)(()))
(45)
max 1, || A Cdiff-
max(1, [1, 111 - All2 + CCANND
where the second last inequality is because between and t (t) during E 3K 1]. Since we already have at t
and 455 imply that (x cannot "shift"
1
1
43
3K*
7.()) implies that (t) is close to [-) for all E latk 1] and hence
Ai(t) (- t
2 Caill CL 2 Caul O(1)
Therefore, for all & [ske1] 1]
11 + (X((1)2
Caiff
0(1)
GK
Combining the lower bound for t E [0, J. we finish proving the claim that 1 1+ t+;((tt) has a positive lower bound of (1) that is independent of This positive lower bound also implies that I 10 (t) invertible and has positive lower bound (recall (x(t) are eigenvalues of (1)). Therefore, by
definition of (t) and 28
||(T 102(0))0(()) 02(0) [[T (02( CLH
(46)
Next, give positive lower bound for 1-10t (0). Note that
10) = (1 1)T LAA + ((T - A 0z())
By 41 we know that the smallest eigenvalue of I 102(0) is larger than equal to
(1- + - tcct 2 2
min (1, AA) = O(1),
27

---

Published as a conference paper at ICLR 2026
which gives a lower bound for 102(0). Applying this bound 46).. we obtain
02(() - -1 ( (0(0)-((-10:0)000111111 S SCC
(47)
We further notice that by (41
1(T - (T
(T- - (T-)))(0))00))
(I- - -(I--))) (I (I- 02(0))]
+ -1- - (()--))) tII- A (2())) (I (02(()))1 F
Smin [1,Aa]= CC + min [1,2a] -t cD
2min K
= =CCAMD-
Therefore, for any E [0, 1],
[82(t) - F [2(() (I t(I A[) -))
'(I (I- (I-
< |02(() (I 1 t02) -1 (0) F + ||0 101--8p (()(())))) 02(0) - (1- -(1 -()- A))
(1 l (
I C ( + ckMMD)
(48)
Step 4.2. We verify that (t) has small error. We first give an error estimation for (0). Recall the true value is
(0) Ab, b, = -I--0(0))D, -b.
Therefore
(0) (0))
< (0) (II- - (2(O))b, b,)ll + |1(T A A- (20))b,b
(49)
CC MMD + CC MMD CC
where we used (4 and 41 in the second inequality. Next, we give error estimate of (1) for E [0, 1]. Recall that
(t) (1 + t0)(t)) (1) (0)())) (0))
satisfies the estimation 28 have
Since 1 + (0) has positive lower bound (shown in step 4. 4. we
(1) - 100((0)) (1) - 0, (0) = +
S CC
(50)
Therefore, for any E [0, 1].
(1) - 01 (()1 = (1(()) (t) ((1 1) + At) (Ab,, - bb)
-1|0000 (I (0)0(0)) (0) < C CC (0) (1(() - 0T(1)]
(51)
5 < (0)) (t) 0 (0)ll + [|t((0(0) (0) (1)l + (1)) 0 (0)1))
sc (ct C C LMDD
C ++ -
In the third inequality, we used (0 (4). and 49
Step 4.3. Finally, we verify that ft( has small error. Recall that 0 (t) is uniquely defined up to an additive constant and
t (b, Ab₂) ((1 ((1-t) T 1)-+t)++A4 At) (b, Abb,)
((t)) (0(
2
00() T (I (1-13(00) (0 (0)()00000
t
2
28

---

The code for Published as a conference paper at ICLR 2026
Also recall that
t 02000))0200))1 [o(() 0o(t) 0o(0) (1))))))1))))))))00))))))))))))) 2000==00))+00)))211000000000 --000))))000))0)010)000000)) 02(0)), 10 (0)
= (t()) (0( + -0 (t) (1- - (0) - 2 (0)
Therefore,
(t) (0)) (1)(t) ()))
[o(() _ 0. (1) T [(-- 102(0)) (1) 0,(0)) -:00 (9100)000
01)
(52)
C L
t K + t (10 (1) (1) (0) (0) 0 (0) e2(0n))
cc (c + CDMD
where 28 and 50p are used the first inequality.
51 and 49
are used in the second inequality.
Finally, combining
48
and 52 conclude
(l)2(t) (t)l + |00 (t) (t) + (tu(t) (t)))) s (c cuMD):
max
(E[0,1]
which implies 16
Finally we make two remarks to this stability analysis. First, we omit the discretization emror and generalization error in this analysis in order to obtain a clear environment for studying the loss function. Second, while we prove the stability result in Gaussian setting, we believe the stability result hold for general distributions, as long as they belongs to some class with sufficient regularity condition.
B IMPLEMENTATION DETAILS
The implementation and training code is available at https:/giitbb..om/ebb///V section 6.2.2 is provided in httpss//gttub.com//SSSiiuuuuuuuu Class conditional OT
Empirical Loss via Monte Carlo Approximation. In practice, the training loss function is ap- proximated via Monte Carlo estimation. For Lu (13), we set @ as the uniform distribution on com- pact computation domain D Ω. We uniformly sample batch of collocation points {(x(0) (t)) (xx)), from the space-time computational domain D x [0, ty] to obtain the empirical loss
NN
+ 1.h (Vugg) (1
V. ( Vh (Vug)
(x(*) (x(0) 47h (Vuj")
Ew =
up
N
where ug ((x(1), tt). Similarly, the MMD term (2) is estimated empirically through samples ((x(0) y 1( from the initial and target distributions
EDDD
1 NZ
Σ
(x(00) x00) k(y - 2k (x( (0))))))))
where (0) (+) + t, (x( 0). To better learn bidirectional OT, we employ the MMD loss in both forward and backward directions. The total loss is
min Eµr (uo) + XENMD (uol); H, 2) AEMMD (TH [uol]/ v).
(53)
B. IMPLEMENTATION DETAILS FOR 2D EXPERIMENTS
Training. For the experiments for 2D toy distributions in Sections 6.1. 1 and 6.2.1 we use a sim-
ple 5-layer MLP with hidden dimension 64 and Softplus activation (with = 100). The model trained using the Adam optimizer with learning rate of 10-3, We sampled 50,000 points from each distribution to create the corresponding sample datasets. At each training epoch, we uniformly sample 1,000 collocation points from the computational domain D [-1. 1 to compute the im- plicit solution formula loss dataset each epoch.
L3 For the MMD loss, we randomly select 750 samples from the given
29

---

Published as a conference paper at ICLR 2026
Baselines. For the NOT baseline, we follow the official implementation provided in the public repository without any modification. The HJ-PINN ablation model was trained under the exact same experimental settings as our proposed NCF across experiments. For GNOT in the class- conditional setting, we use the official code released by the authors without modification.
[1]
B.2 IMPLEMENTATION DETAILS FOR GAUSSIAN EXPERIMENTS
Training. For the high-dimensional Gaussian experiments in Section 6.1.2 we employ the Den- seICNN architecture, which is fully connected neural network with additional input-quadratic skip connections, to ensure a fair comparison with the baseline models provided in (Korotin et al. [202 Since our method does not require input convexity, we omit the commonly imposed constraints that enforce positivity of certain neural network weights, which are typically used to guarantee convex- ity. Following Korotin et al. 2021a), we adopt the network architecture DenselCNN[1; max(2d.64) max(2d,64), max(d,32) for a d-dimensional problem. The model is optimized using Adam with fixed learning rate 10 regardless of the input dimension. To construct dataset, we randomly sample 10s points from each of the source and target distributions. We set 1) as the bounding box (i.e.. axis-aligned minimum and maximum values) of these samples and define as the compu- tational domain for solving the HJ equation. At each training epoch, we uniformly sample 1,000 collocation points from D to compute the implicit solution formula loss (13). For the MMD loss, randomly select 2,000 points from the given source and target datasets every epoch.
The baselines LS, WGAN-QC. MM-vl and MM:R are all used via the official imple- mentations from the public repository of Korotin al. (2021a The implementations of NOT and HJ-PINNs follow the same settings described in Appendix B.1
Baselines.
Evaluation Metric.
Unexplained Variance Percentage (UVP): Given the predicted transport map T from μ to v, UVP is defined by 2 UVP T
100
f - T
L2(M)) Var (v) (%). A UVP
value approaching 0% indicates that 7" provides close approximation to the OT map T** whereas values substantially exceeding 100% imply that the estimated map fails to capture the underlying structure of the OT. We use 10 random samples drawn from μ to compute UVP
Memory and Time Metries: Memory consumption is reported as the peak memory usage during training. Training time is measured as the average runtime per epoch over 100 epochs. Inference time refers to the time required to transport 10 test samples using the learned map. Additionally, measure the memory required to store the trained networks for the bidirectional OT maps. For our method, this corresponds to the storage size of a single spatio-temporal solution function for the HJ equation. For dual-based baselines, this reflects the memory needed to store both the primal and dual potential functions. For the NOT baseline, which leams the forward and backward OT maps separately, we report the total memory required to store both learned transport maps.
B.3 IMPLEMENTATION DETAILS FOR COLOR TRANSFER
Training. The color transfer experiments in Section 1.3 are trained using exactly the same ex- perimental setup as in the high-dimensional Gaussian case described in Appendix B.2 to ensure fair comparison with the baseline models.
Baselines. For the classical methods, implemented Reinhard color transfer using OpenCV's Bradski & Kachler (2008) color space conversion and channel-wise mean-std matching. Histogram matching was implemented by computing per-channel histograms and CDFs, then applying the resulting pixel value mapping directly. Since these methods only support one-way transfer, we con- ducted forward and backward transfers separately. Both methods serve as standard, straightforward baselines.
30

---

Published as a conference paper at ICLR 2026
Evaluation Metries.
Earth-Mover Distance (EMD); For both the target and transported images, we compute normalized color histograms separately for each BGR channel. The EMD quantifies the minimal cost required to transform one histogram into another, offering perceptually meaningful measure of distributional difference. We compute the EMD independently for each channel and report the average across all three. Lower EMD values indicate greater similarity.
Histogram Intersection (HII): HI measures the overlap between the normalized color his- tograms of the target and transported images. For each BGR channel, we compute the intersection as the sum of the minimum values across corresponding bins. The final score is obtained by averaging over all three channels. Higher values (closer to 1) indicate greater similarity.
B.4 IMPLEMENTATION DETAILS FOR SECTION 6.2.2
This series of experiments focuses on the MNIST dataset (LeCun 1998), which comprises 10 classes of 28 28 grayscale images of handwritten digits ranging from 9; And the Fashion MNIST dataset consisting of 10 classes of 28 x 28 grayscale images of clothing items, labeled from 0 to 9.
For both MNIST & Fashion MNIST datasets, the value of each pixel of the grayscale images takes integer value from 1 to 255. We always normalized the pixel values of each data point to [0, 1] by dividing by 255 before calculation.
VAE Pretraining. In our study, we employ pretrained 3- VAE models (Kingma et al. 2013 Hig- gins et al.) 2017), which offer satisfactory generative quality and faithful manifold representations for image encoding. Advanced auto-encoder architectures (Berthelot et al. 2018 Feng & Strohmer, 2024) that better preserve the interpolation quality of decoded images will be considered in future work. Although the ambient dimension of the data is 784, prior work has shown that the dataset exhibits a moderately low intrinsic dimension (Pope et al. (021). Thus, in our implementation, we set the latent dimension to di = 10 for in-domain transport tasks on MNIST data set, and to di = 35 for cross- dooss- transport between Fashion MNIST and MNIST data sets.
To train the VAE, we consider the encoder (.). RRd + R4, and encoding variance Sa(+) : Rd R4: which share the same parameter 6, together with the decoder (w Rd + Rd, with o, w being the tunable parameters. For arbitrary x; from the dataset and the latent variable z E Rd the ELBO-type loss L3(ó, x;), == (x/x)( Pw (z) 3DKI (qo(lx)) (lp)))) is considered, where we set the conditional probability (-Iz) = NDD (Z), La.) the prior px(-) (0, and the posterior (-x))x,) =
(x.)). Here of is predetermined variance, and E (x.) =
(S We optimize the following obtain Eg, Dw:
M
M
Σ (p(x,wxx))
E Be - -D( Eo (x)
max 0,W
M
2M
σ²
1=1
+ +3( 1-3(-Sop (x,)
1+ [x)())) exp((S(xx))
Here we denote 1 = (1....,1) E Rd,, In our experiment, we pick o² = TOU and set 3 = 0.1 to ensure reconstruction fidelity over regularization.
We train the VAE pairs (E)(), (-) and (E()),)()) on MNIST dataset (x(2) and Fashion MNIST dataset (x(2) respectively. We set batch size as 32, and apply the Adam algorithm (Kinga et al. 2014) with learning rate 10- 4 for 150 epochs. In practice, the trained VAE reproduces MNIST images with an accuracy 98.2%, and reproduces Fashion MNIST images with an accuracy 87.0%
Encoding & Normalization. Denote y = (x. (k)), k = 1,2, we normalize the latent samples hessy by (k) = (y )(k) for < < N, k 1,2. Here we denote y(h) =
N 1 (k) the mean of the dataset, and o (k) diag (}) the entrywise variance, where
31

---

Published as a conference paper at ICLR 2026
diag( (() denotes diagonal matrix with its diagonal entries taken from the empirical covariance (K) x(()) (x. xT
matrix (k) 2==1
Loss function & Training. We denote μr, v as the distribution of the normalized latent samples y!") where k 1 or 2. To compute the OT map between μ., v, we set y 1, and introduce neural network ug [Rd: [0. f + R. In practice, we incorporate the loss functional for backward OT into the original loss (14) that is, we consider
min
{ (wo) (uo) è
(uo) + (Eclaas((T)) [uo]) 1 11,1) Ectuss (1a, (TH (uol) w))}
where we denote (ua) =- LHi(ue) defined in 13 and define the corresponding backward implicit loss as
Cij (wo) = I (us th (Vus) + Vug VA (Vug) - ug (x iVh (Vua), cl((x) (((()(ul) dt.
Since the latent samples are normalized as described above, we set = No La, and independently draw I; and I uniformly from [0, Ly] to form the collocation points {(z,, (3) for approximating fi and fiu In implementation, we set À = 500 in order to balance the scales of the implicit HJ loss and the MMD loss. We denote NHj, NMMD and MMD between distributions certain classes. In our experiments, we choose Nu = 4000 and NMMD 400. We apply the Adam method with learning rate 10-4 for optimizing 0. The algorithm is conducted for 1000000 iterations.
the batch size for evaluating the implicit loss
Neural Net Architectures. The architectures for the VAE encoder and decoder are summa- Regarding the OT map, we parameterize ug Rd++1 IR using ResNet architecture He al. (016)) with depth I. and width (hidden dimension) d 128. Specifi- cally, we define
rized in Table 5 and Table
ux(x,t) fL f2-10.
f2(2,t).
where each layer Su is given by
A + bb
k: == A Eddx(d+t), by E ER
fv(
3 + (A) + bb.
2 k < L 1,
A e pixd by E R
Auy + bu,
k = L, Ar. Rixd by ER.
We use the hyperbolic tangent activation (-) tanh(-) and set the residual scaling parameter = 1. We set L = 5 for in-domain transports on MNIST, and use L = 6 for cross-domain trasport task between Fashion MNIST and MNIST.
Table 5: Encoder architecture for 3 VAE for image size (H,WOC C) (28, 28, 1), latent dimension
di = 10.
Layer
Parameters
Output Shape
Input
(H, W)
Conv2D
128 filters, 5, stride 1, ReLU
(H, W, 128)
Conv2D
128 filters, x 3, stride 1, ReLU
(H, W, 128)
Conv2D
64 filters,
stride ReL.
(H/2, 64)
Conv2D
64 filters, 3 x 3, stride 1, ReLU
(1/2, W/2,
Conv2D
64 filters, x stride ReLU
(1//2, 64)
Cony2D
64 filters, 3 x stride ReLU.
Conv2D
64 filters, 3 x 3, stride ReLU
(11/2, 64)
Flatten
(H/2- WW/2-64)
Dense
(16 Н W. 64). ReLU (64, de)
(16- 11 W)
Dense (mean (x))
(d)
Dense (log variance S,(x))
(64, de)
(d)
Output (reparam.)
y (x(x) + exp( (x))
(d))
Evaluation Metrics.
All methods evaluated the festing portions of the MNIST datasets.
32

---

Published as a conference paper at ICLR 2026
Table Decoder architecture for 3 VAE for image size (H,W,C) (28, 28, 1). latent dimension d: 10.
Layer
Parameters
Output Shape
Input (y)
(de
(de, 16-16666 H W), ReLU
(16 H W)
Reshape
(H/2, W/2, 64)
Cora2DTaaseesee
64 filters, stride 1, ReLUU
(H/2, W/2, 64)
Cons2D/Transpose
64 filters, 3 x 3, stride 1, ReLU
(H/2, W/2, 64)
Coraapora
filters, stride 1, ReLU
(H/2, W/2, 64)
Con>2DTraassose
64 filters, 3 x 3, stride 1, ReLU
(H/2, W/2, 64)
DDranporeee
filters, stride 2, ReLU
(H, W.64)
Conv Drramnsoe
128 filters, x 3, stride 1, ReLU
(H, W, 128)
128 filters, x stride ReLU
(H, W, 128)
Con>2DThransooe
filter, x 5, stride 1. ReLU x D.íy)
(H. W.1)
Output
(H, W)
Source j
(b) Strong NOT
(c) Weak NOT
(d) HJ-PINN
(e) NCF
Target
(g) Strong NOT
(h) Weak NOT
(i) HJ -INP
( NCF
Figure 5: Checkerboard (μ)) = Eight Gaussians (v)): The top row shows transport in the direction μr, and the bottom shows μ + v, with y and the leftmost column.
Classification Accuracy: We evaluate the class-wise accuracy of the generated data. Fol- lowing (Asadulaev et al. 2024), we train ResNet-18 classifiers achieving 98.85% accuracy
the MNIST test set.
FID score: The FID score evaluated on the entire test set, which consists of approxi- mately 1,000 samples per class.
C FURTHER RESULTS
C.I ADDITIONAL RESULTS FOR 2D TOY EXAMPLES
"Results on 2D distributions with multiple modes are presented in Figure 3 As in Section 6.1.1 the proposed NCF successfully learns bidirectional OT even in network. Compared to baselines, not only transports the distributions more accurately but also produces transport maps with less overlap, indicating that it learns more optimal transport paths.
settings with a single
Ablation Study on the Regularization Parameter We further present an ablation study inves- tigate the effect of the regularization parameters Xy and A in the proposed loss function (14). These parameters control the balance between the implicit HJ loss and the MMD loss. Specifically, the implicit HJ loss promotes the optimality of the transport map by encouraging alignment with the HJ equation, while the MMD loss measures how well the transported distribution matches the target distribution. Since all experiments in the paper are conducted under the setting 'y >s. we vary yy to examine how this trade-off influences the learned transport map. We conduct experiments on the owo-dimensional examples introduced in Section 6.1.1 and above. The results summarized in Figure ( The case Xy = oo corresponds to training without the implicit loss, using only the MMD loss.
33

---

Published as a conference paper at ICLR 2026
(a) Xy 00
(b) 1
(c) A = 0.1
(d) X = 0.05
(c) Ay = 0.01
Figure 6: Effect of the regularization parameter ' in the loss function. Each figure visualizes the transport from the source distribution (blue) to the target distribution (red) for varying values y. Results shown for two examples introduced in Section 6.1.1 and Appendix C.1
As shown in the figure, when Xy is small (i.e., the MMD loss dominates), the transported distribution aligns well with the target, but the resulting transport map becomes highly entangled. This indicates that the model learns map far from the optimal one, due to the lack of guidance from the HJ constraint. contrast, increasing Xy enforces stronger adherence to the HJ equation, resulting in a smoother, more structured transport map that closely resembles the optimal solution. However, when A becomes too small, the influence of the MMD term diminishes, leading to inaccurate matching of the distributions. These results highlight the complementary roles of the implicit HJ loss and the MMD loss, as also supported by Theorem [5.4] and underscore the importance of appropriately tuning Xy. Additionally, the relatively small difference in performance between Xy 0.1 and A5 = 0.05 suggests that the model is not overly sensitive to the choice of this parameter.
C.2 FURTHER RESULTS ON TRAINING STABILITY
Figure illustrates the training stability of our method compared to the representative baseline MM:R in high-dimensional Gaussian experiments in Section 6.1.2 MM:R optimizes a min-max problem, which exhibits severe training instability: although the UVP enror initially decreases in the early epochs, it tends to grow substantially training continues. Following the original paper (Korotin et al. 2021a). the smallest UVP errors observed during training are reported Table 2 However, in practical OT scenarios the true optimal transport map unknown, making it unclear when training should be stopped or which solution is preferable. Therefore, the ability to maintain stable training many epochs, demonstrated by method, is crucial advantage.
In the original MM:R implementation (Korotin et al., (021a), the model relies on a special identity- potential pretraining stabilize training. Without this initialization, training becomes highly unsta- ble and can completely fail to approximate the OT map. In contrast, our method does not require any special pretraining and maintains stable UVP reduction throughout training, highlighting key practical strength.
In contrast, our method demonstrates robust and stable learning, shown in Figure 7 does not require any special pretraining and maintains stable UVP reduction throughout training, steadily decreasing over epochs. This stability represents significant practical advantage, ensuring reliable convergence even over long training schedules.
Our method
C.3 ADDITIONAL QUALITATIVE RESULTS FOR COLOR TRANSFER
Figures and present qualitative results for bidirectional color transfer across three distinct categories of image pairs. In each figure, the leftmost columns display the source and target images, while the remaining columns show the results of various methods applied in both the forward (source target) and backward (target source) directions.
34

---

Published as a conference paper at ICLR 2026
MMER wi Inttiiization
MXCK EnIiATuiniinn
NMMR wiy tritiitrontoon
10
MM. Irittaliintion
10
NEP
10
NCC
107
100
10
100
16
10
100
10000
20:000
30000
69030
50003
10000
200000
30509
40000
50030
Ipodh
Ippoc
Figure 7: Training stability of our method compared to MM.R high-dimensional Gaussian periments.
C.4 ADDITIONAL RESULTS FOR CAAS--CONDIIIONAL OT FOR MNIST & FASHION MNIST DATASETS
In this section, we provide additional results from all the tasks conducted in our experiments. All methods are evaluated on the testing portions of the MNIST and Fashion MNIST datasets. Following (Asadulaev et al., 2024), we train ResNet- 18 classifiers achieving 98. 95% accuracy for evaluation. A generated sample is deemed accurate if the trained classifier assigns it to the corresponding target class. Table 7 reports the accuracy and Fréchet Inception Distance (FID) the generated images across the first two experimental tasks.
Table 7: Performance metrics of NCF across all tasks. For the FID row, each bracketed value repre- sents the distance between the OT-generated distribution (Dw, o Th [Fe])e (resp. (Dw o ("(u)))))) and the decoded distribution DWE (resp. Dwsj), where D., denotes the appropriate VAE decoder corresponding to f or v.
Task
Task
Metric
Forward
Backward
Forward
Backward
Accuracy(%) t FID L
95.58 19.93 (2.65)
95.08 18.91 (2.45)
92.51
92.73
18.98 (2.25) 19.01 (2.26)
Task (In-class transfer: Map each MNIST class i to the class + 5, for O, .5...,4) 4.) We present in Figure 11 the uncurated MNIST images generated using the forward (resp. backward) mapping, D1 TW [uo] (f(1)), (()) ~ (resp. D (To)))()())) yr v))
Task (In-class shift: Map each MNIST class the class (i + 1) mod 10, for i 0, ........ 9.) Similar to Task 1, Figure [2] shows the uncurated MNIST images generated by the computed map- pings. this experiment, the forward (resp. backward) maps are also trained without incorporating the implicit HJ loss LI,, and are therefore not guaranteed to be optimal. By contrast, the mappings produced by our proposed method explicitly designed to account for optimality- preservation of MNIST digit styles (e.g.. thickness, orientation, etc.), further illustrated in Figure
superior
13
Task 3 (Inter-class transport: Map each class in Fashion MNIST its corresponding class in MNIST.) The uncurated MNIST and Fashion MNIST images generated by the computed maps are shown in Figure [4] While our method effectively recovers the overall profiles of Fashion MNIST images, the decoderr scheme faces difficulties in capturing fine texture details. As a future research direction, we aim to enhance our approach by incorporating U-net architectures and directly performing OT in the pixel space. Furthermore, in Figure [5] we present the KDE plots of the pushforward distributions TV [ua]= (resp. 7 ("(l)) together with their targets (resp. 11), which demonstrate the satisfactory generative quality the computed OT map Ty [uo]: Figure presents the classification accuracy (%) of the generated images on the test dataset over training iterations. We display only the first 70000 iterations, since the accuracy no longer improves as training progresses.
35

---

Published as a conference paper at ICLR 2026
Source
HisMatching (Fed)
Reinhard (Fwd)
NOT (Fwd)
MM--1 (Fwd)
MMER (Fwd)
NCF (Fwd)
Target
HisMatching (Bwd)
Reinhard (Bwd)
NOT (Bwd)
MM (Bwdd
MM.R (Bwd)
NCF (Bwd)
HisMatching (Fwd)
Reinhard (Fwd)
(Fwd)
Target
NisMetching (Bwd)
Reinhard (Bwd)
NOT (Bwd)
MM-v1 (Bwdd
MM. (Bwd)
NCT (Owd)
Source
Hisskttching (fidd
Reinhard (Fw
NOT (Fwd)
MM4-1 (1 wd)
MM:R IT ('dd
NCT (Fwd)
Thrgat
HisMatching
Reinhard (Bwd)
MM-11
MM.R (Bwd)
Source
HisMatching (Fwd)
Reinhard (Fwd)
NOT MM-v1
MM:R (Fwd)
(Fwdd
Target
HisMatching (Bwd)
Reinhard (Bwd)
NOT (Bwd)
MM (Bwdd
MM.R (Bwd)
NCF (Bwd)
Source
HisMatching (Fwd)
Reinhard (Fwd)
NOT
MM
MM.R (Fwd)
(Fwd)
Target
HisMatching (Bwd)
(Bwdd)
(Bwi))
(Bwd)
(Bwi))
Figure 8: Winter + Summer: Qualitative results for bidirectional color transfer between seasonal image pairs.
36

---

Published as a conference paper at ICLR 2026
Source
HisMatthing (Fwd)
Reinhard (Fwd)
NOT (Fwd)
MM-v1 (Fwd)
MMBR (Fwd)
NCF (Fwd)
Target
HisMatching (Bwd)
Reinhard (Bwd)
NOT (Bwd)
MM (Bwid
MM.R (Bwd)
NCF (Bwd)
HisMatching (Fad)
Reinhhard (Fwd)
(Fwd)
Target
NisMetching (Bwd)
Reinhard (Bwd)
NOT (Bwd)
MM-v1 (Bwdd
MM.R (Bwd)
NCT (0wd)
Source
HisMatching (fwd)
Reinhard wdd
NOT (Fwd)
MM4-1 (f vd)
MM:R (Fidi
NCT (Fwd)
Throet
HisMatching
Reinhard (Bwd)
NOT
NM.R (Bwd)
Source
HisMatching (Fwd)
Reinhard (Fwd)
MM-v1 (F (Fwi)
MM:R (Fwd)
(Fwd)
Target
HisMatching (Bwd)
Reinhard (Bwd)
NOT (Bwd)
MMvvl (Bwdd
MM.R (Bwd)
NCF (Bwd)
Source
HisMatching (Fwd)
Reinhard (Fwd)
NOT
MM-v1 (Fwd)
MM.R (Fwd)
(Fwd)
Target
HisMatching (Rwd)
(Bwd)
(Bwd)
(Bwi))
MMR (Bwd)
(Bwit)
Figure 9; Gogh painting ++ Photograph: Qualitative results for bidirectional color transfer between Gogh paintings and real-world photographs.
37

---

Published as a conference paper at ICLR 2026
Source
Hislaatching (Fwd)
Reinhard (Fwd)
NOT (Fwd)
MM-v1 (F Fwdd
MM.R (Fwd)
MCF (Fwd)
Target
HisMatching (Bwd)
Reinhard (Bwd)
NOT (Bwd)
MM (Bwd)
MM.R (Bwd)
NCF (Bwd)
Relnnhard (Fwd)
(Fwd)
Target
MisMetching (Bwd)
Reinhard (Bwd)
NOT (Bwd)
MM-v1 (Bwdd
MM-R (Bwd)
NCT (Bwd)
Source
HisMatching (fidd)
Reinhard (Fwd
NOT (Fwd)
MM-v1 (f wd)
MM:R (fidd)
NCT (Fwd)
Tbrost
HisMatching
Reinhard (Bwd)
MM.11
NM.R (Bwd)
Source
HisMatching (Fwd)
(Fwd)
MM-vll
MM:R (Fwd)
(Fwd)
Target
HisMatching (Bwd)
Reinhard (Bwd)
NOT (Bwd)
MM (Bwd)
MM.R (Bwd)
NCF (Bwd)
Source
HisMatching (Fwd)
Reinhard (Fwd)
NOT
MM-11
MM.R (Fwd)
(Fwd)
Target
HisMatching (Bwd)
(Bidd)
(Bed)
(Bwi))
(Bwit)
Figure 10: Monet painting + Photograph: Qualitative results for bidirectional color transfer be- tween Monet paintings and real-world photographs.
38

---

Published as a conference paper at ICLR 2026
6
6
7
8
9
0
/
2
3
4
6
7
8
9
0
2
4
5
6
7
9
0
2
3
4
6
7
8
9
0
2
4
5
6
7
8
9
0
/
2
3
4
5
6
7
8
9
0
2
4
5
6
7
8
9
0
/
2
3
4
5
6
7
9
0
2
3
4
5
6
2
9
0
2 3
4
<
6
7
8
9
0
2
4
(a) Images generated @ ty-
(b) Images generated @ 0.
Figure 11: Task 1: Uncurated images generated using the computed forward (Left) backward
(Right) OT maps.
2
3
4
5
6
7
9
0
9
2
3
4
5
0
/
2
3
4
5
6
7
8
9
0
9
0
1
2
3
4
5
6
8
2
3
4
5
6
7
8
9
0
9
7
3
4
5
6
7
8
2
3
4
5
6
2
9
0
9
0
2
3
4
5
6
7
8
/
2
3
4
5
6
7
8
9
O
9
O
/
2
3
3
6
7
8
2
6
4
6
7
8
0
9
O
Z
3
4
5
6
7
8
2
3
4
5
6
7
8
9
0
9
O
2
3
4
5
6
7
8
2
3
4
5
6
5
8
9
O
9
0
2
3
4
5
6
8
9
2
3
4
5
6
7
8
@
6
9
0
2
3
4
6
6
7
8
2
3
4
6
7
8
9
0
9
O
2
3
5
6
7
8
(a) Images generated @ ly.
(b) Images generated @ 0.
Figure 12: Task 2: Uncurated images generated using the computed forward (Left) & backward (Right) OT maps.
3
4
g
9 9
4
3
3
4
4
g
g
9
9 9
3
3
8
8
8
9
9
9
9
9 9
3
3
3
4
4
U
8
8
8
9
9 9
9
999
3
3
3
3
3
4
4
4
4
4
4
8
8
&
4
9
9 99991
(a)
Manifold of digit (b)
Generated by OT (c) Generated by
(d) Manifold of digit
(e)
Genersse by or (f) Generated by noe.
optimal
optimal
Figure 13:
The style of each MNIST digit is better preserved by the computed optimal transport map. The triangular table is produced using linear interpolation in VAE latent space.
39

---

Published as a conference paper at ICLR 2026
2
3
4
5
6
5
8
7
0
/
2
3
4
5
6
7
8
9
0
|
Z
3
4
5
6
9
8
9
0
/
2
3
4
5
6
7
8
9
0
/
2
3
4
5
6
8
9
P
2
3
4
6
6
7
8
2
3
4
5
6
7
8
9
0
2
3
4
0
7
9
0
1
2
3
4
5
í
7
8
9
2
3
4
5
6
1
4
9
(a) MNIST images generated @ ty.
(b) Fashion MNIST images generated @ 0.
Figure 14: Task 3: Uncurated images generated using the computed forward (Left) & backward (Right) OT maps.
Figure 15: KDE contours the MNIST latent samples generated using the computed map (blue) and the target samples (orange), conditioned on each MNIST class (0-9, arranged left to right and top to bottom). The samples are projected onto the first two PCA dimensions. The heat maps illustrate the discrepancies between the two distributions.
40

---

Published as a conference paper at ICLR 2026
38
10002 20002
30303
40003
30000
63002
70092
Figure 16: Accuracy(%) of trained forward class-conditional transport map versus training itera- tions. Results are displayed for the first 70000 iterations, beyond which the accuracy exhibits no significant improvement.
41


Reviews:
[{"review_id": "D7c7k4ZMFN", "reviewer": "Reviewer_9wfo", "summary": "Single-network and pure minimization training for conditional OT maps. The proposed method eliminates the need for adversarial training or saddle-point optimization that many \"OT-GAN\" or dual potential methods use. NCF uses a single neural network and a single loss function to compute bidirectional maps. This is a significant simplification: prior approaches often involve two networks (e.g. a generator and a discriminator or dual potentials) and tricky min–max optimization, which can be unstable and require extensive hyper-parameter tuning.", "questions": "**Q1**: The Theorem 5.1 assumes regularity conditions. How does the use of an infinitely-smooth neural network reconcile with approximating a potentially non-smooth solution? Does this implicit regularization bias the solution away from the true, non-smooth solution?\n\n**Q2**: Could you please clarify, is \"Optimality\" in Table 1 means guaranteed convergence to the optimum? Why Dual Formulation and Dynamical Models are mentioned as \"No\"? \n\n**Q3**: The GNOT paper considered an image-to-image paired cost setup. Is your method applicable to high-dimensional paired data, forward and backward? Regarding the class-conditional results in Table 4: Given the poor FID score (18.27) compared to baselines like GNOT (5.26), can the method truly be considered competitive for high-fidelity generative tasks without using any decoders? \n\n\n**Minor**: \n\n*Q4*: For Figure 4, it would be interesting to see how the model maps the generated forward samples backward.\n\n*Q5*: The 'integration-free' nature of the model is a key feature. By deriving a closed-form map from the HJ characteristics, the method avoids the need to solve ODEs. However, are  ODEs a real bottleneck for dynamical OT models? There is a well-known trilemma between coverage, sampling and 'accuracy' of generative models (https://arxiv.org/pdf/2112.07804). It would be interested to see a dynamic formulation of the proposed method. Can be obtain bidirectional maps in dynamical OT formulation? Maybe the authors have conducted some experiments using an ODE-like formulation, as done by (https://arxiv.org/pdf/2507.17513) using the HJB framework? The comparison would be very interesting.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The paper is very well written and easy to follow. The theoretical motivation, from the dynamical formulation of OT and the Hamilton-Jacobi (HJ) equation to the method of characteristics, is presented clearly. The core idea of leveraging this to derive closed-form, bidirectional maps is elegant.\n\nThe proposed Neural Characteristic Flow (NCF) framework is compelling. Its main advantages using a single neural network and, crucially, avoiding adversarial min-max optimization  are significant practical contributions over common dual-formulation approaches.\n\nThe algorithm is supported by theoretical arguments. The authors provide a consistency analysis (Theorem 5.1) showing that a zero loss recovers the true OT map and a stability analysis (Theorem 5.4), but demonstrating that a small loss guarantees convergence in the Gaussian setting", "weaknesses": "The HJ formulation for OT is well-established. The \"implicit solution formula\" (Equation 9) , which is the foundation for the proposed loss , is explicitly cited from a very recent work (Park & Osher, 2025). The main contribution appears to be the application of this new formula to the OT problem, combined with a standard MMD loss, which was already considered in (Asadulaev, et. al. 2024).\n\nSimilar to other papers in this area, the primary class-conditional experiment is transporting Fashion MNIST to MNIST. This setup (e.g., mapping a \"Trouser\" to a \"1\") feels artificial and does not provide a compelling, practical use case for the class-conditional framework. The colour transfer setup is also very simple and is often easier than mapping between mixtures of Gaussians. \n\nThe experimental validation, while broad, has a significant weakness in the class-conditional setting. In Table 4 (Fashion MNIST → MNIST) , while NCF achieves the highest accuracy (83.42%), its FID score of 18.27 is substantially worse than GNOT (5.26), NOT (7.51), and MUNIT (7.91). The paper attempts to explain the poor FID by blaming the VAE decoder. The authors then provide a \"re-calculated\" FID of 2.73, computed between NCF outputs and VAE-decoded images, but no comparison with other methods are done.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7L3KwtqyA9", "reviewer": "Reviewer_nTXZ", "summary": "This paper proposes a new algorithm for computing optimal transport maps between the distributions based on the usage of Hamilton-Jacobi (HJ) equation. For this purpose, the authors, first, represent the OT maps using the characteristics of the HJ equations; second, consider an implicit formula for the viscosity solution of HJ equation; third, parametrize the viscosity solution using neural networks and learn it by optimizing the  loss derived from the implicit solution formula and additional MMD regularizer. As a result, this allows for deriving the approximations of forward and backward OT maps.", "questions": "- Does your method have connections with the unbalanced optimal transport? This question appears since your loss (14) uses a regularization which motivates the learned OT maps to push input distribution to the target one. Such a strategy of softening the marginal constraint of the OT problem is usually exploited in unbalanced OT field. And the related question is – is it necessary to use MMD loss in this  regularization? Or maybe it can be replaced by some f-divergences?\n- Is it possible to generalize your theoretical results in section 5.1 for general types of costs (other than quadratic)?\n- Does the result of your algorithm depend on the choice of the parameter $\\lambda$?\n- Please explain the intuition behind the weighting function $\\rho$ introduced in equation (13).", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "In general, the paper is quite well-written and suggests some theoretical results supporting the constructed algorithm. The proposed algorithm allows for computation of OT maps without the need of min-max optimization or numerical integration of ODEs.", "weaknesses": "Some of the mathematical concepts introduced in the paper lack necessary details. First, the definition of Hamilton-Jacobi (HJ) equation given in section 2.2 lacks reference to the literature where it was introduced. The continuity equation mentioned in line 102 is not written directly in the text. Besides, the characterization of the viscosity solution by the system of characteristic ordinary differential equations (CODE) should be supported by the relevant reference. \n\nMeanwhile, the intuition behind some of the ideas which appear in the derivation of the main objective is not sufficiently clear. In lines 204-209, the authors write that since the initial condition in the HJ equation is not known analytically, they introduce an additional loss term to ensure that this condition is appropriate. However, this loss term does not deal with the initial condition directly but rather enforce the solutions of the OT problem derived from the HJ (viscosity) solution to satisfy the marginal condition, i.e., to map the source distribution to the target. Thus, it seems that the addition of this loss term to the main objective serves another goal than the initially stated one, see questions section for my additional concerns on this topic. Moreover, the intuition behind the weighting function $\\rho$ introduced  in the equation (13) is not clarified.\n\nStill, while the authors provide some theoretical results justifying their approach, my major concerns correspond to the **error in the experimental comparisons** provided by the authors in section 6.1. The authors perform testing of their approach on the benchmark providing the ground-truth OT maps between the specific constructed pairs of distributions (Korotin, 2021a) and claim that their approach leads to the best performance. However, one of the competitive approaches (NOT) seems to be not tested correctly. Indeed, NOT approach actually corresponds to [MM:R] which was tested in the benchmark paper (Korotin, 2021a) and outperformed all other methods considered there. However, according to Table 2 in the paper under review, the NOT ([MM:R]) approach leads to awful results in comparison to another [MM-v1] method from the benchmark. This, actually, contradicts the results from the benchmark where the NOT method leads to the best results as I already said. Thus, I kindly suggest the authors check their results and implementation of the NOT approach which probably contains errors. These errors might also have affected the experimental results on the considered color transfer task. Meanwhile, the experimental results in section 6.2.2 show that the proposed approach did not beat the competitor GNOT according to FID metric. Thus, overall practical validity of the approach is unclear until the stated errors are resolved.\n\n**In summary**, I am mostly confused by the revealed errors in the experimental comparisons provided in the paper. These errors raise doubts regarding the results reported for the experiments with high-dimensional Gaussians and the color transfer task. Thus, I could not assign a positive score to the paper until the issues are resolved.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Vwz7OhbQFS", "reviewer": "Reviewer_r22u", "summary": "The paper proposes a Neural Characteristic Flow (NCF) method for static optimal transport (OT) with translation-invariant costs. Despite solving static OT, the method’s theory begins with the equivalent dynamical formulation, whose optimality condition is governed by the Hamilton–Jacobi (HJ) equation.\nThe core idea is to enforce an implicit solution formula (obtained via the method of characteristics) for the HJ PDE, whose viscosity solution uniquely determines the OT map.\nTo tie the PDE solution to the data (i.e., to satisfy boundary conditions), the authors add an MMD loss between the transported source and the target, combined with a data-driven HJ implicit-solution loss.\nThis design has several advantages: (1) it is a pure minimization approach that avoids the min–max optimization typical of dual methods and trains a single network; (2) although grounded in a dynamical viewpoint, it does not require ODE/SDE simulation; and (3) it jointly learns forward and backward maps from the same potential.\nThe authors also show that the framework extends to class-conditional OT by applying class-wise MMD within the same HJ formulation.\nEmpirically, with the standard quadratic cost, the method is shown to work in low-dimensional settings.", "questions": "1. Why use a VAE latent space for MNIST experiments? Did you try solving the task in the original pixel space? If so, how do performance, stability, and runtime compare? If not, what specifically fails in pixel space?\n2. Why does your method perform best in dimensions d≤16 but lose to MM-v1 at d=32 and d=64? Could you, please, discuss scalabilty of your method taking into account concerns from the weaknesses?\n3. Could you provide an experimental comparison between NCF and DIOTM [1] and ENOT [2] ?\n4. What are some real-world applications for the class-conditional OT method?\n5. What is the advantage of using the Implicit Solution Formula (9) over the original HJB equation (5) (like in DIOTM [1] or HOTA [3])? Were any ablation studies conducted on this point?\n\n[3] Buzun et al. HOTA: Hamiltonian framework for Optimal Transport Advection. https://arxiv.org/abs/2507.17513", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. Implementation of a novel approach via an implicit HJ formulation (from paper Park & Osher, 2025).\n2. Efficient training (no minimax optimization; no ODE/SDE simulation).\n3. Theoretical analyses and guarantees (primarily for the quadratic-cost case).", "weaknesses": "1. **Lack of scalability evidence or discussion.**\nAuthors evaluate the proposed NCF only on low-dimensional datasets; even for MNIST they work in a low-dimensional VAE latent space. The method raises subtle potential scalability issues. First, it requires backpropagation through Jacobian–vector product, which could lead to instability and high memory consumption in higher dimensions. Second, the assumptions required by the theory—global bijectivity, a positive-definite Jacobian everywhere, and a strictly positive sampling distribution—may not hold in practical settings. \n\n2. **Missed state-of-the-art baselines.** \nThe experimental evaluation would be strengthened by including comparisons with recent and relevant static Optimal Transport methods, specifically DIOTM [1] and ENOT [2]. These approaches have demonstrated superior performance and rapid convergence, and a direct comparison is crucial for a comprehensive assessment of the proposed method's contributions.\n\n[1] Choi et al. Improving Neural Optimal Transport via Displacement Interpolation, ICLR 2025\n\n[2] Buzun et al. ENOT: Expectile Regularization for Fast and Accurate Training of Neural Optimal Transport, NeurIPS 2024", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)