Universal Neural Optimal Transport
Authors: Jonathan Geuter
Gregor Kornhardt
Ingimar Tomasson
Vaios Laschos
Venue: ICML 2025 poster (2025)
Keywords: Optimal Transport
Neural Operators
Meta Learning
Adversarial Training
Primary area: general_machine_learning->unsupervised_and_semisupervised_learning
TL;DR: We train a neural network to accurately predict optimal transport distances across datasets and dimensions with neural operators.

Abstract: Optimal Transport (OT) problems are a cornerstone of many applications, but solving them is computationally expensive. To address this problem, we propose UNOT (Universal Neural Optimal Transport), a novel framework capable of accurately predicting (entropic) OT distances and plans between discrete measures of variable resolution for a given cost function. UNOT builds on Fourier Neural Operators, a universal class of neural networks that map between function spaces and that are discretization-invariant, which enables our network to process measures of varying sizes. The network is trained adversarially using a second, generating network and a self-supervised bootstrapping loss. We theoretically justify the use of FNOs, prove that our generator is universal, and that minimizing the bootstrapping loss provably minimizes the ground truth loss. Through extensive experiments, we show that our network not only accurately predicts optimal transport distances and plans across a wide range of datasets, but also captures the geometry of the Wasserstein space correctly. Furthermore, we show that our network can be used as a state-of-the-art initialization for the Sinkhorn algorithm, significantly outperforming existing approaches.

Universal Neural Optimal Transport
Jonathan Geuter 12*
Gregor Kornhardt
Ingimar Tomasson 3
Vaios Laschos 4
Abstract
0.90
Ones
0.85
Optimal Transport (OT) problems are a comer- stone of many applications, but solving them is computationally expensive. To address this problem, we propose UNOT (Universal Neural Optimal Transport), novel framework capable of accurately predicting (entropic) OT distances and plans between discrete measures for a given cost function. UNOT builds on Fourier Neural Operators, universal class of neural networks that map between function spaces and that are iitiizattooninvarrann which enables our net- work process measures of variable resolutions. The network is trained adversarially using ond, generating network and a self-supervised bootstrapping loss. We ground UNOT in an ex- tensive theoretical framework. Through experi- ments on Fuclidean and non-Fuclidean domains, we show that our network not only accurately predicts OT distances and plans across wide range of datasets, but also captures the geometry of the Wasserstein space correctly. Furthermore, we show that our network can be used as state- of-the-art initialization for the Sinkhorn algorithm with speedups of up to 7.4x, significntly outper- forming existing approaches.
0.80
Gauss
0.75
UNOT
0.70
0.65
0.60
0.55
Relative
0.50
0.45
0.40
0.35
0.30
0.25
0,20
0.15
0.10
0.05
0.00
MNIST
CIFAR
MNIST LFW-BEAR CIFAR
LFW
BEAR
28x28
28x28
28x28
64x64
64x64
64x64
Figure 1. Errors the OT distance after a single Sinkhorn itera- tion for the default initialization (Ones), the Gaussian one (Thorn- ton Cuturi, 2022), and ours (UNOT), for c( y) || yl
regularizer is added, this allows for efficient computation of the solution via the Sinkhorn algorithm (Cuturi, 2013). The entropic OT problem between probability measures μ E P(X), v E P(Y) on Polish spaces X, Y, given a cost
function e X x y + R U {oo], defined as
(T((H,V)= inf / =Ell(u,)) xxxx cd - ( (1)
1. Introduction
where I[ is the set of all transport plans (i.e. mea- sures on X' x ) that admit resp. their marginals), LL @ v) = f log((T/ / @ v)dT is the KL divergence of π from f @ v. and > is a regularizing coefficient..
Optimal Transport (Villani, 2009; Peyré & Cuturi, 2019) plays an increasing role in various areas in machine learn- ing, such as domain adaptation (Courty et al., 2017), single- cell genomics (Schiebinger et al., 2019), imitation learn- ing (Dadashi et al., 2020), imaging (Schmitz et al., 2018), dataset adaptation (Alvarez-Meel & Fusi, 2021), and signal processing (Kolouri et al., 2017). Oftentimes, an entropic
Many of these applications require solving problem (1) re- peatedly, such as in single-cell perturbations (Bunne et al., 2022: 2023b;a), Natural Language Processing (Xu al., 2018), flow matching with OT couplings (Tong et al., 2024; Pooladian et al., 2023), or even seismology (Engquist & Froese, 2013). However, solving OT problems is computa- tionally expensive, and fast approximation methods are an active area of research. Variations of the transport problem, such as (generalized) sliced Wasserstein distances (Kolouri et al., 2015; 2019), reduce computational complexity
'Equal contribution 'Havardd John A. Paulson School of En- gineering and Applied Sciences "Kmmpner Institute at Harvard University 'Department of Mathematics, Technische Universität Berlin, Germany 'Weierstrass Institute, Berlin, Germany. Corre- spondence to: Jonathan Geuter <jonathan, geuther@neeeerreeeeeee gmx.de>.
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
'For background on OT, see Appendix A.

---

Universal Neural Optimal Transport
the cost of accuracy via random projections. Two previous works are aimed predicting good initializations for the Sinkhorn algorithm, which can iteratively solve problem (I). In (Thornton & Cuturi, 2022), initializations are computed from OT problems between Gaussians. (Amos et al., 2023) train a neural network to predict transport plans and costs via the entropic dual OT problem (Section 2). While their framework shares some similarities with ours (see Section 5), it is inherently limited to measures of fixed dimension from the training dataset. Instead, we present Universal Neural OT (UNOT) solver which, given discrete measures E P(()) and V E P())) of variable resolution (viewed as discretizations of continuous measures; see Section 3.1), can accurately predict the OT cost and plan associated with prob- lem (I). To this end, we leverage Fourier Neural Operators (FNOs) (Kovachki et al., 2024), a discretization-invariin class of neural networks that can process inputs of variable sizes. An FNO S&p is trained to predict solution to the dual OT problem (Section 2) given two measures (J,, v), from which the primal problem (I) can be solved. Training is self-supervised with an adversarial generator network Go which creates training distributions (z = (p,, from P== (0, I) (see Section 3.4). We want to high- light that in contrast to most neural OT frameworks, such as (Bunne et al., 2023a; Uscidda & Cuturi, 2023; Korotin et al., 2023), we generalize across OT problems (given fixed cost). GeONet (Gracyk & Chen, 2024) also uses Neu- ral Operators to learn Wasserstein geodesics. In Section 4, we show that UNOT significantly outperforms GeONet on approximating geodesics, despite not being trained on them.
2. Background
We give a brief overview of optimal transport, and how relates to UNOT. For more thorough introduction, see Appendix A.l or (Peyré & Cuturi, 2019; Villani, 2009).
Notation. We write vectors in bold (æ, u, etc.) and matrices in capitals (X, U, etc.). 1, E IR" denotes the all-ones vector, A the simplex in IR", and all elements in A"-1 with positive entries are denoted by A PP())) denotes the space of probability measures on X, and (x( the set of probability measures with finite p-th moments. For i E P(X) and a map T, we denote by Ta f the pushforward of μ under T, i.c. the measure μ T 5 = {I E R |||||| = 1} is the unit sphere in TR3
2.1. Optimal Transport
Unregularized Optimal Transport. The unregularized problem takes the form inl in(, fxxx c(x, ()dr(zz, v), akin to (I) without the regularization term. In the case where X' y, c(2, y) = d(x, y)" for p 2 1 and a metric d on X, the Wasserstein-p distance is defined as
VXec d(z, ())dm( (2)
W v)
inf *EN((a))
which is indeed distance on the space (2)) of Borel measures with finite p-th moments (Villani, 2009).
Dual Optimal Transport Problem. The regularized Kan- torovich problem (1) admits a dual formulation:
Our contributions are as follows:
f f(z)dp(x)) + I -
present UNOT, the first neural OT solver capable of generalizing across datasets and input dimensions
sup (w)
- (ff.g).
(3)
introduce a generator Go (Section 3.3) which can prov- ably generate any discrete distribution (of fixed dimen- sion) during training (in fact, prove this result for very general class of residual networks, see Theorem 3 and Corollary 4)
where
x (f.9) = Jxxy
((1))) ldz(z)dr(y).
propose self-supervised bootstrapping loss which provably minimizes the loss against the ground truth dual potentials (Proposition 5)
It can be shown that if c E L' @ v), the values of the primal (1) and the dual (3) coincide (Nutz, 2022).
Disrete Optimal Transport We want to apply UNOT to discretizations of measures i E P(X), v E P()) (see Section 3.1). To this end, consider measures l and v that are supported on finitely many points I1, I' X, Un ..., Un E y resp, i.e. == Σ a:62,0 By abusing notation, we can write μ E R7 and v E R% Note that the dual potentials in problem (3) are elements in L' (H) resp. L' (z)), hence we can abuse notation again and consider the potentials as well. This point of view gives rise to discrete optimal
show that UNOT can accurately predict OT distances across various datasets, costs, and domains of different dimensions up to few percent error, and that it accu- rately captures the geometry of the Wasserstein space by approximating barycenters (Section 4)
approximate Wasserstein geodesics through barycen- ters and OT plans predicted by UNOT (Section 4)
IR" and g E " to be vectors
demonstrate how UNOT sets a new state-of-the-art for initializing the Sinkhorn algorithm while maintaining its desirable properties, such as parallelizability and differentiability (Section 4)
Whenevvr we view a discrete measure function a vector, we will use bold characters.
2

---

Universal Neural Optimal Transport
Iransport. We set C' RR via Ciyy = c( Mj), and view transport plans as matrices IT Rmxx (see Appendix A. for a more thorough introduction to discrete OT). The fol- lowing proposition shows that the plan II can be recovered from the dual vectors f and g (Peyré & Cuturi, 2019).
Cuturi, 2019), we will instead let UNOT predict the dual potential g = c clog(v), i.e.
5 = g, i PEP(X),
Proposition 1. Define the Gibbs kernel K = exp(-C/e)) The unique solution 11 of the discrete OT problem is given by
The prediction g can then be used to solve the entropic OT problem via the relationship (5) and Proposition 1, or to initialize the Sinkhorn algorithm via v - exp(g/e).
IT diag(u) (u) K Kdiag(r)
(4)
Note that the solution to the entropic dual is not unique (see Proposition 1). How we account for this non-uniqueness is explained in Section 3.4. However, when endowing IR" x R" with the equivalence relation (u1, ~ (u2, 02) e YE > 0: (u,, 01) = (Au2. f22) (i.c. accounting for the non-unigue un-uuiqueeess of the dual solution), the map (pt, v) + v, mapping two measures the associated dual potential in the quotient space, is Lipschitz continuous (Carlier et al., 2022), which supports its learnability by neural network.
for two positive scaling vectors u and v unique up to a scaling constant (i.e. Xu, X for x >> 0). Furthermore, (u, v) are linked to a solution (f 9) of the dual problem via
(u, v) (exp(f/c),
In Section 3.1, we show how the solution to the entropic dual between discrete (Ht,, pa) converges to the solution of the continuous dual (3) as (Mn pn) converge to continuous measures (μ,, P) in some way, which will be crucial for the design of our network Sz.
3. Universal Neural Optimal Transport
Consider the OT problem between two (grayscale) images, encoded as vectors in pl,, Vn E RR". These can be viewed as discrete measures P((0, 1)2). which discretize contin- uous measures μl, E P((0, 1J²). where the discretization depends on the resolution of the image, and the continu- ous measures correspond to the images at "infinite" resolu- tion. UNOT should predict the corresponding dual potential In. E R." solving (3) independent of the resolution n.. In Section 3.1, we establish a convergence result for the dual potentials as n + oo, which justifies the use of Neural Oper- ators (Kovachki et al., 2024) as a paramctrization of Sq: also see Section 3.2. Furthermore, as we want UNOT to work across datasets, we require a generator Gg that can provably generate any pair of distributions during training (Section 3.3). In Section 3.4, we construct an adversarial training ob- jective for So and Gg. Further details about hyperparameter and architecture choices can be found in Appendix C. The implementation and model weights are available at https: /ightb... com /rr///rrrrrroorreeee wow.nardt/nauuraaaaaaaaae
2.2. The Sinkhorn Algorithm
The Sinkhorn Algorithm can iteratively solve the discrete dual problem and was introduced in (Cuturi, 2013). It re- quires an initialization v" E R", which is typically set to 1,, and μ and v to be positive everywhere.
Algorithm Sinkhorn(j4, v > 0, exp(-C/o),c.0)
1: for [ = O.... N do
2: u t- M./KKt
3: v441 + v. K u+1
4: end for
5: II 5 diag( (w') ) (idiag(") (T, < (C)
6: return u, v, II, OT& OTe (H,
In the algorithm, // is to be understood as element-wiss division. Sinkhorn and Knopp (Sinkhorn & Knopp. 1967) showed that the iterates u' and of from the algorithm con- verge to the vectors and e from Proposition 1.
3.1. Convergence of Dual Potentials
2.3. Predicting Dual Potentials
In this section, we prove convergence of the discrete dual potentials 9n as n goes to infinity. For brevity, this section is kept informal; see Appendix B for formal treatment. Assume now that X' = ) C R is compact, and c(r, y) is Lipschitz continuous in both its arguments. For absolutely continuous p.v E P(x)) denote by (Hu))EN, (v, (V.)nEN C P(X') discretizing sequences of μ and (formally defined in
Given discrete measures fa and v, UNOT should ultimately be used to approximate the associated transport plan and cost. However, given an optimal dual potential v, the corre- sponding potential u can be computed as
u = k.
"Note that while we consider images as an example, the learn- ing task is the same for any setting where discrete measures of varying resolution share underlying continuous cost function, which arises in settings such single-cell genomics, fluid dynam- ics, point cloud processing, or economies.
which also holds at convergence of the Sinkhorn algorithm. Thus, solving for the m x n-dimensional plan II can be reduced to a n--dimensional problem over v. Since com- putations in the log space tend to be more stable (Peyré &
3

---

Universal Neural Optimal Transport
Appendix B). While a solution (fuz In) of the discrete dual problem between jr, and Un is only defined Y4n resp. Vn- a.e., it can be canonically extended to all of X (Feydy et al., 2018) (see Appendix B for details). The following proposi- tion shows that the extended potentials (nn,gn) converge to the solution (f g) of the continuous entropic problem.
where Ps = =V(0, Ia) is Gaussian prior, NN₂ is a trainable neural network (in practice, we use a 5-layer fully connected MLP, see Appendix C), La.æ' is an interpolation operator matching the generator's 's output dimension d' and acting as a skip connection reminiscent of ResNets (He et al., 2016), and À > 0 is a constant for the skip connection. Ó >> O is a small constant needed to generate our targets with the Sinkhorn algorithm, as outlined in Section 3.4. R denotes renormalizing two probability measures and downsam- pling them to random dimensions in a set range, such that So trains on measures of varying resolutions, which is known to improve NO training (Li et al., 2024a). More specifically, if we write [21, T2| = ReLU (NN, (z) + Ad (z)) + 5 for two vectors æ1 and 22 of equal size (say both with n sam- ples), R. first maps them to /x1/ Σ, (21) x2/ (12)) and then uses 2D bilinear interpolation to downsample them to samples each. The generator is universal in the follow- ing sense:
Proposition 2. (Informal) Let (Un)nEN, (Vn)nEN be dis- cretizing sequences for absolutely continuous μu, E P(X). Let (Jr. In) be the (unique) extended dual potentials of (pm, Vn) such that (zo) = O for some E X and all n. Let (f g) be the (unique) dual potentials of (1,, v) such that (1o) = 0. Then In and In converge uniformly to f and g on all X.
A formal version and its proof can be found in Appendix B. This proposition is crucial in designing our network Sz. as we discuss in the following section.
3.2. Fourier Neural Operators
Theorem 3. Let 0 < X S 1 and Go Rd -> Ted be defined via
Fourier Neural Operators (FNOs) (Kovachki et al., 2024) are neural networks mapping between function spaces. More precisely, a neural operator is a map F:A- 4 between Banach spaces A and U of functions E :: Da - I., and u Du + Rdd respectively, for bounded domains D. C IRd and D, C Rdd. An input function A evaluated points 2,, .... F. E Rd can be encoded as vector [a(x1), ..... axxn)] E gnxd';; the same applies to the output function u E U4 which can be written as u = [m((1), "(ym)) )] E mmxd"., correspond- ing to the values at //, ..... ym E Rdd. At its core, an FNO applies a sequence of L "kernel layers" to the input vector a. In each of these layers, fixed number of Fourier features of the discrete Fourier transform of the input is computed, the features are transformed by (C-valued) linear layer (we use two-layer network in practice instead, as we found it to improve performance), and then mapped back by the inverse Fourier transform. Importantly, neural operators are by construction discretizatioon-nnaaaant when inputs and outputs correspond to discretizations of underlying func- tions. This is exactly what Proposition 2 guarantees: the dual potentials corresponding to measures 14n and v/, con- verge uniformly to the continuous potentials corresponding to the limiting distributions μ and // as the resolution of He, and in increases. Hence, FNOs are natural choice of architecture in our setting. More details on FNOs, and how we implemented Sæ. can be found Appendix A.5.
Go(z) = (NNo(2) + Az).
where 2 ~ P= -N T), and where NNo Rd - Rdd is Lipschitz continuous with Lip(NNo) = L X. Then Go is Lipschitz continuous with Lip(q) < <L X, and G(2) == NZ + Az is invertible on Rdd Furthermore, for any x E R it holds
PGonPa (x) (L + NN (o))o.))
In other words, Go Pz has positive density at any non- negative x E Rd
This shows that any pair of discrete probability measures (1s, v) of joint dimension can be generated by Go. A di- rect consequence of the theorem is extension to functions that are compositions of functions Go as above, which cov- ers a wide class of ResNets. Both proofs can be found in Appendix B.
Corollary 4, Ler Gg = Go, o Go o... o Gon be a compo- sition of functions Go,, each of which is of the form as in Theorem 3. Let z ~ Pz ==== =N(0,) I). Then
PGoaPa (0) 2 (I. XJRa (Gn
1
for any x Rdd As in Theorem 3, this also holds for any x: E Rd. G is followed by ReLU activation.
3.3. Generating Measures for Training
Although the more general Corollary 4 is not needed for purposes, it might be of independent interest to the research community. Note that the generator in Theorem 3 does not exactly match our generator's 's architecture. A discussion of how the theorem relates to our setting, as well as further details on the generator, can be found in Appendix C.
UNOT is trained on pairs of distributions generated by a generator network Go of the following form:
Go Rd + P(() x P(X)
z ~ P= + R |ReLU (# + (l), +5], (6)
4

---

Universal Neural Optimal Transport
Algorithm 2 UNOT Training Algorithm
1: in cost reg parameter prior Pz. learning rates {as}i.
{ Sinkhorn target generator т
for 1,2,.... T do
- sample(p=)
4: (,, Go(z))
5:
for mini-batch (H", w²) in (H, v) do
Figure 2. Generated pair of training samples (lighter=more mass).
6:
IO t Sa(pe, (1te),, 20)
7:
Ir, t (u 11, ub .9))
8:
Φ + Φ a, V
go)
Figure 2 shows a pair of samples generated by Gg. The generator seems to layer highly structured shapes with more blurry ones. More examples, as well as an analysis of the performance of S& on samples generated by Gg over the course of training, can be found in Appendix D.6.
end for
10: for mini-batch x² in do
11:
(no, vo) + Go(z²)
12:
go So vo)
13:
Ir. + Th Th Vo go)
14: O 0 + B, V (g 90)
3.4. UNOT Training Algorithm
15: end for
16:
end for
Given pair of distributions (H, = (z) (in this section, we will remove the subseript n for clarity), So( =: 96 should predict the true dual potential g associated with l and v. Hence, we could simply compute the true potential g with the Sinkhorn algorithm and use L2(g,.9) However, it would be prohibitively expensive to run the Sinkhorn algorithm until convergence. Hence, we instead employ a bootstrapping loss on the prediction 9 Let The (H. v. g+) + 9r. denote running the Sinkhorn al- gorithm on (Ju, v) with initialization o = exp(gø/e) for very small number of iterations k, i.c. warmstarting the Sinkhorn algorithm with the current prediction Iá. and re- turning clog 9 To ensure uniqueness and improve training, we shift 9-. to have zero sum; this corresponds to the non-uniqueness of the dual potentials, see Proposition 1. Minimizing L2 (g 97.) implies minimizing the ground truth loss L2 (gp. g) against the true potential g.
as the target generation procedure, the training objective for Sg and Gg consists of S& trying to minimize the loss L2 (gggg.gl) g). while Go attempts to maximize it, similar to the training objective in GANs (Goodfellow al., 2014). Putting everything together, our adversarial training objec- tive for Sy and Go reads
ll gli
training loss.
max min Ez-p, [T2 (Th (G(2), So (Go(z)))) 0 o
(7)
where S and G without subscripts denote no gradient track- ing, as the target is not backpropagated through. The train- ing algorithm can be seen Algorithm 2. In practice, train- ing will be batched, which we omitted for clarity. Note that vectors g with subscripts O or & are backpropagated through with respect to these parameters, whereas target vectors (with subscript Tk) are not.
Proposition 5. For two discrete measures (u.v) with n particles, let g be an optimal dual potential, 96 = = So(p. v). and 9T. Tx. (1,, v, v.) Without loss of generality, assume
4. Experiments
that Σ, 9: E.9rn 0. Then
Training Details. We test UNOT in three different settings: a) with c(ar, y) ||= y on the unit square X' [0, 1
L2 9) S <( k. (nnnnT.. (2(
b) with c(x, 3) | y1l on [0, 1|2; e) with the spher- ical distance c(æ, y) arccos( (x. y) on the unit sphere S = fz R3 ||||| = =13. For each of these settings, we train a separate model on 200M samples z, where training samples v)) are between 10x 10 and 64 x 64 dimensional (randomly downsampled in Gø). Training takes around 35h on an H100 GPU. S. is an FNO with 26M parameters op- timized with AdamW (Loshchilov & Hutter, 2019); Go is fully connected MLP with 272k parameters optimized with Adam (Kingma & Ba, 2017). In the spherical setting c(z, y) arccos((2: (xx, y¹) we parametrize S, with a Spher- ical FNO (SFNO) (Boney et al., 2023) instead, which is essentially an FNO adapted to the sphere; for more details
for some constant c(K, k. n) > depending only on the Gibbs kernel K, I and n.
The proposition shows that minimizing L2(gø- 9r.) implies minimizing (2(ge: g). i.e. the loss between the prediction and the ground truth potential. The proof is based on the Hilbert projective metric (Peyré & Cuturi, 2019) and can be found in Appendix B.
Training objective. Having defined the loss for Sg, as well
"The Sinkhorn algorithm requires input measures be positive;
this the reason we add the constant 8 > in the generator.
5

---

Universal Neural Optimal Transport
MNIST22228
CIFAR22.28)
MNIST-CIFAR R
Table Mean number of iterations needed to achieve 0.01 relative error on the OT distance for c(zz, y) = [ vll?.
0.3
0.2
UNOT (OURS)
ONES
GAUSS
0.1
MNIST
3 + 5
16 9
10
+ 7
Relative Error
CIFAR
3 6
80
52
19
0.0
CIFAR-MNIST
4
32:15
13 : 9
LFW(04464
BEAR)64004)
LFW MON-BEARIcom (64x64)
0.3
LPW
7
78 20
35
- 14
BEAR
4
41
16
25
- 13
0.2
LFW-BEAR
4 6
53 +
18
29 +
13
0.1
Table 2. Relative speedup of Sinkhorn with UNOT and cost (.. 3) || yl? Time in to achieve 0.01 relative error
0.0
Iteration
Iteration
Iteration
on the OT distance.
UNOT
Gauss
Ones
UNOT (OURS)
ONES
SPEEDUP
MNIST
1.2 10-3
1.5 10-3
1.25
Figure 3. Relative error on the OT distance for Sinkhorn with our initialization (UNOT), compared to the default (Ones) and Gaus- sian initialization (Gauss) (Thornton & Cuturi, 2022).
CIFAR
9.
7.1 10-3
7.4
CIFAR-MNIST
1.3-10-
2.7-10-3
2.07
LFW
3.0-10-3
1.5-10
5
BEAR
2.6-10
1.0-10-2
3.8
LFW-BEAR
2.7-10-3
1.2-10
4.4
the default initialization 1, and the Gaussian initializa- tion from (Thornton & Cuturi, 2022), which is based on closed-form solutions for Gaussian distributions. Note that the Gaussian initialization is only valid for c(æ, y)
on FNOs and SFNOs see Appendix A.5. We highlight that So is relatively small, such that its runtime vanishes com- pared to the runtime of even just a few Sinkhorn iterations, making it much cheaper to run than Sinkhorn (see Section 4. We set & (the number Sinkhorn iterations Tk) to 5, and e 0.01. Additional training details can be found in Appendix C.
| y hence we omit it when c(ær, y) = || y/ or c(x, y) = arccos (a, y)).² We do not compare to Meta OT (Amos et al., 2023) here, as their approach is inher- ently dataset dependent and breaks down when testing on cittr--disttibttion data. 6 For completeness, we include a detailed comparison in Appendix D.2, which shows that UNOT significantly outperforms Meta OT on all datasets except MNIST, the training dataset of Meta OT. Surprisingly, UNOT also almost matches Meta OT on MNIST, despite not having seen any MNIST samples during training, while Meta OT was explicitly trained on them.
We demonstrate the performance of the three models vari- ous tasks, such as predicting transport distances, initializing the Sinkhorn algorithm, computing Sinkhorn divergence barycenters, and approximating Wasserstein geodesics. For the Euclidean settings a) and b) (from above), we view im- ages as discrete distributions on the unit square, and test on MNIST (28x28), grayscale CIFAR1O 10 28), the teddy bear class from the Google Quick, Draw! dataset (64x 64), and Labeled Faces in the Wild (LFWW, 64 64x,4), well as cross-datasets CIFAR-MNIST and LFW-Bear (where fe comes from one dataset and v from the other). For the spherical setting c), we project these images onto the unit sphere in R (for details, see Appendix D.I). Unless oth- erwise noted, we perform a single Sinkhorn iteration on g v)) in all experiments in order to compute the second potential f. Errors are averaged over 500 samples. Additional experiments, including a sweep over input sizes 10 x 10 to 64 x 64, as well as variants of UNOT for fixed input dimension or variable e, can be found in Appendix D.
Figure 1 (Section 1) shows the relative error on OT, (p. v) after a single Sinkhorn iteration for c(x, y) = || - y and Figure 4 shows the same plot for c(x: y) = ||x y on the square, and c(x, y) = arccos (æx, y on the sphere. In Figure 3, we plot the relative error on the OT distance over the number of Sinkhorn iterations for c(z, y) = || y/ (for the equivalent plots for the other cost functions, please see Appendix D.7), demonstrating that UNOT can be used as a state-of- the-art initialization. Table 1 shows the average number of Sinkhorn iterations needed to achieve 0.01 rela- tive error on O v) for c(x, y) = || y||² In Table we show the relative speedup achieved by initializing the Sinkhorn algorithm with UNOT implemented in JAX over the default initialization (on batch size of 64 in float
4.1. Predicting Transport Distances
'If the cost function is not || - =11", the Gaussian initializa- tion not theoretically justified. Empirically, we noted that it behaves similar to the default initialization in these cases.
We compare the convergence of the Sinkhorn algorithm in terms of relative error on the transport distance (H,
'W note that it should he possible finetune UNOT on spe- cifie datasets as well; however, we have not tested this.
for our learned initialization xd = eexp(S v)/e) to
6

---

Universal Neural Optimal Transport
0,90
Ones UNDT
0.85
Ones
0.90
0.80
UNOT
0.85
0.75
0.80
0.70
0.75
0.65
0.65
0.70
0.60
Relative Error
0.60
0.55
0.55
0.50
0.50
0.45
Rela
0.45
0.40
0.40
0.35
0.35
0.30
0.30
0.25
0.25
0.20
0.20
0.10
0.15
0.10
0.15
0.05
0.05
0.00
0.00
MNIST
CIFAR
MNIST CIFAR
LEW-BEAR
LFW
BEAR
MNIST
CIFAR
MNIST LFW-BEAR CIFAR
LFW
BEAR
20x28
28x28
20x20
64x64
64x64
64x64
28x28
28x28
28x28
64x64
64x64
64x64
Figure+ 4. Error on the OT distance after single Sinkhorn iteration with UNOT vs. the default initialization (Ones) for cost || y/l on the square [0, 1J² (left) and arccos( (x. y²) on the sphere {x E IR3 |||||| = 11 (right).
To make this problem tractable, consider the Sinkhorn di- vergence barycenter
6
Ours
Σ X₂ SD,
arg min ('EPs(X)
6
/
8
where the Sinkhorn divergence" between i and v is
True
OT,(u.)) 2 OT,(U,
Now for discrete measures μl, denote by (f.g) the dual potentials for (H, and by p that for OT, (H, p)) Writing μ = E; agdz, for some a E R", the gradient of (9) w.r.t a is given by (cf. (Feydy et al., 2018)):
Figure 5. Sinkhorm divergence barycenters computed with UNOT via eq. (10) (top) vs. ground truth (bottom) of between 5 to 10 MNIST samples of the same digit per barycenter.
VaSD((p.)= = -p.
(10)
on an NVIDIA 4090). We achieve an average speedup of 3.57 28 x 28 datasets and 4.4 on 64 64 datasets." For comparison, the relative speedup achieved in (Amos al., 2023) was 1.96 (for model trained only MNIST).
Hence, we can solve (9) with (projected) gradient descent, where S& predicts f and p in (10). 10 Further details and pseudocode can be found in Appendix A.2. Through- out this section, we set c(x, y) || y1 and always run 200 gradient steps using gradients from (10). Figure shows UNOT barycenters vs. the true barycenters (com- puted with the POT library) of between 5 and 10 MNIST samples of the same digit per barycenter. In Appendix D.7, we also provide quantitative results for barycenters with different initializations. In Figure 6, we show barycenters computed between four shapes. UNOT accurately predict- ing barycenters demonstrates it captures the geometry of the Wasserstein space beyond predicting distances.
4.2. Sinkhorn Divergence Barycenters
The Wasserstein barycenter for of measures "NN P2(X) and X e A is defined as
arg min (X)
Σ x.W3ldd.u.
(8)
"We did not optimize the network S& much for efficiency, and more efficient implementations likely exist. Note that FNOs process complex numbers, but PyTorch is heavily optimized for real number operations. With kernel support for complex numbers, UNOT will likely he much faster. In addition, computation times can vary significantly across hardware, batch sizes, precision, etc.
SIt can he seen as a debiased version of OT, (1,, v), and we use it as approximation of the squared Wasserstein distance.
'If both measures are identical, the dual potentials can be cho- sen to he identical as well.
10 be precise, this solves the barycenter problem on the dis- crete space {#1, ..., J.)
7

---

Universal Neural Optimal Transport
Figure 7. Sinkhorn divergence particle flow between distributions of images, from noise to LFF (64x64). Gradients computed via eq. (11) and (10) with UNOT. Bottom row is target images.
Figure 6. UNOT barycenters computed between four shapes (cor- ners) by linearly interpolating (A1, A2, Aa; Xa) from eq. (9) between the four unit vectors, and solving via eq. (10) with UNOT.
GeONet being trained to predict geodesics MNIST, while UNOT does not train on geodesics, geodesics computed by UNOT are significantly closer to the ground truth than the GeONet geodesic.
4.3. Calculating Geodesics
MNIST, both
Let n E (2) be two measures such that v/ Ta for an optimal transport map T X - X' (which exists for the non-entropic optimal transport problem under certain conditions, see Appendix A. The Wasserstein geodesic between μ and v, also called McCann interpolation, is the constant-speed geodesic between μ and v and given by
4.4. Wasserstein on Wasserstein Gradient Flow
Oftentimes in machine learning, the distributions of interest are not images. but distributions over images, such as in generative modeling. In this experiment, we show that UNOT can successfully transport distributions over images as well. Let μ. & E P2((P2((,, 1J3). c), W2), i.e. the space of distributions over images equipped with the Wasserstein distance (and P[0, 1J2) being equipped with c(æ,, y) x )) Denote by (D, the Sinkhorn divergence between jú and é, where we SD, v) the ground approximation of
Me : [0, 1] + P2()), + [11 t)ld + +T
be interpreted as the shortest path between μ and The Wasserstein barycenter (8) between (n, (1 1)) and (v, f) (i.c. where 1-t and f are the weights X; from equation (8)) turns out to be equal to μe (Agueh & Carlier, 201 I). This gives us two methods to approximate the Wasserstein geodesic between μ and 1:: Either by iteratively computing barycenters as in Section 4.2, or by computing the (entropic) transport plan from equation (4) as an approximation to some technicalities for brevity here,
cost between je, P2( 1J3) W (1, P). Writing á 1. E 8uc. I = 1. E2 dug for Pt, Bj E P2, 1)2). we let UNOT approximate the partiele flow ar = - -Va [SD. 8)), for which we can derive the gradient via (see (Li al., 2024b)):
T (we are leaving which can be found in Appendix A.3). We compare the geodesics computed by UNOT to the ground truth geodesic (obtained from the true OT plan), as well as to GeONet (Gracyk& Chen, 2024), a recently proposed framework that also uses Neural Operators to learn Wasserstein geodesics directly by parametrizing a coupled PDE system encoding the optimality conditions of the dynamic OT problem. Akin to (Amos al., 2023), GeONet is inherently dataset depen- dent. Figure 8 shows the McCann interpolation between two MNIST digits using the ground truth OT plan, the OT plan computed by UNOT, barycenters computed by UNOT, and the GeONet geodesic, where we use the UNOT model trained with c(æ, f/) = || - - again. We see that despite
(S.
SSD. v3) Ikj.
k= 1,,,,,, n,
Opk
Opk
(11)
where 1hg an optimal transport plan between μ4 and v.. These gradients can be approximated by UNOT via equation (10) as before; further details can be found in Appendix A.4. In Figure 7, we plot the particle flow from Gaussian noise P to distribution over 10 images, where we visualize f after every 10 gradient steps (using AdamW (Loshchilov Hutter, 2019) with gradients computed via equation (10)). We can see that the UNOT flow converges quickly.
8

---

Universal Neural Optimal Transport
Thus Kcaam
6
9
9
tion compute distances between images (Zhang al., 2020). In flow matching OT can be used straighten paths (Lipman et al., 2023; Tong et al., 2024; Pooladian al., 2023). Approximating Wasserstein gradient flows with the JKO scheme has been explored in numerous works (Alvarez- Mclis & Fusi, 2021; Alvarez-Melis et al., 2022; Bunne et al., 2022; Choi et al., 2024). The theory of Wasserstein gra- dient flows has also been used study learning dynamies in various settings, such as for overparametrized two-layer networks (Chizat & Bach, 2018) or simplified transformers (Geshkovski et al., 2024).
UNOT McCan (Duns)
6
6
6
6
Baryarraa Grocesc (Gurg)
6
6
6
L
GoCNee
6
Generative Adversarial Networks. GANs (Goodfellow et al., 2014), like other types of generative models, aim at generating samples from a distribution Piams given access to a finite number of samples. In contrast, we do have access to samples from the target distribution. However, our loss function (7) shares similarities with the adversarial GAN loss. Given prior samples P= and data samples x ~ Pdata the GAN objective for a generator G is
t-0
t-0.25
0.5
t-0.75
t-1
Ground Truth
Figure 8. McCann interpolations computed with the true OT plan, UNOT OT plan, UNOT barycenters, and GeONet (top to bottom).
5. Related Work
Neural OT. Typically, neural OT approaches aim at solving individual instances of (tigh-dimensional) OT problems. In (Korotin et al., 2023), a maximin formulation for the dual problem is derived and two networks, parametrizing the transport plan and the dual potential resp., are trained adver- sarially. In (Bunne et al., 2023a), transport maps between continuous input distributions conditioned on context vari- able are learned. Another interesting recent paper (Uscidda & Cuturi, 2023) suggests universal regularizer, called the Monge gap, to learn OT maps and distances. Unlike these works, we focus on generalizing across OT problems.
min max E [lgg((=))+B<..) [log(1 - D(G()))))
G D)
where is the discriminator, which predicts the probability that a sample came from the target distribution rather than the generator. Note that while our generator maximizes the objective, the GAN generator minimizes
6. Discussion
We presented UNOT, neural OT solver capable of solving entropic OT problems universally across datasets, for a given cost function. Leveraging Neural Operators, UNOT can pro- cess distributions of varying resolutions supported on grids. UNOT's training involves generator network Go produc- ing synthetic training samples for the predictive network Sg. where both networks are trained jointly via a self- self-supervised adversarial loss. So predicts the potential of the dual OT problem, and our training objective provably minimizes the loss w.r.t. the ground truth potentials. We show that UNOT is universal in the sense that the generator can create any discrete distributions during training, and empirically verify this through experiments on Euclidean and non-Euclidean image datasets of varying resolutions. UNOT consistently predicts OT distances up to 1-3% relative error, and approx- imates barycenters and geodesics in Wasserstein space by solving for the OT plan. Furthermore, we demonstrate that UNOT can be used as state-of-the-art initialization for the Sinkhorn algorithm, achieving speedups of up 7.4x, Current limitations that UNOT does not extrapolate well to measures with significantly higher resolutions than the training samples, nor generalizes to cost functions other than the training cost. Scaling UNOT to higher resolutions, as well as applying it to other data modalities or non-uniform grids, are interesting directions for future research.
Initializing Sinkhorn. There exists very little literature on initializing the Sinkhorn algorithm. (Thornton & Cuturi, 2022) propose using dual vectors recovered from the unreg- ularized ID optimal transport problem, or from closed-form transport maps in Gaussian (mixture) setting, and were able to significantly speed up convergence. (Amos et al., 2023) propose neural approach, training a single network to predict the optimal dual potential f of the discrete dual problem, and their loss is simply the (negative) dual ob- jective (3). This approach works well when training oow-dimensional datasets such MNIST, and is elegant as it does not require ground truth potentials, i.c. is fully unsu- pervised, but it is not able to generalize to out-of-distribution data, and can only be used for input measures of fixed size.
OT for Machine Learning. Leveraging OT to formulate new machine learning methods has seen a surge in pop- ularity in recent years, and it has been applied to a wide range of problems. Relevant works include the celebrated Wasserstein GAN (Arjovsky et al., 2017), multi-label learn- ing (Frogner et al., 2015), inverse problems in physics (En- gquist & Yang, 2019), point cloud processing (Geuter et al., 2025; Fishman et al., 2025), or few-shot image classifica-

---

Universal Neural Optimal Transport
Acknowledgements
Proceedings of the International Conference on Machine Learning, 2019. doi: 10.48550/ARXIV.1811..00 URL
For this work, VL has been funded by Deutsche Forschungs- gemeinschaft (DFG) Project-ID 318763901 1294. GK acknowledges funding within the BMBF project VI-Screen-PRE (Teilvorhaben: Mathematische Bildverar- beitung und maschinelles Lemen) VC3-23/1NN17309.
https:
00995.
Bonciocat, A.-I. and Sturm, K.-T. Mass transportation and rough curvature bounds for discrete spaces. Journal of
Functional Analysis, 256(9):9444-266, 2009. ISSN 0022-1236. doi: http:://boi.rry/l00011ff........
Impact Statement
029. URL httpp:::/wwww..eeeeeeeeeee// science/artiche /tteacee/arttcceee/a////22222221
This paper presents work whose goal advance the field of Machine Learning. There are many potential societal consequences of our work, optimal transport has vast range of applications, but none of these feel must be specifically highlighted here.
Bonev, B., Kurth, T., Hundt, C., Pathak, J., Baust, M., Kashinath, K., and Anandkumar, A. Spherical Fourier Neural Operators: Learning Stable Dynamics on the Sphere. 2023. URL https: /atps///aaxiv... wow.v.org//bba
2306 03838.
References
Bunne, C., Papaxanthos, L.,, Krause, A., and Cuturi, M. Proximal Optimal Transport Modeling of Popula- tion Dynamics. In Proceedings of The 25th Interna- tional Conference on Artificial Intelligence and Statis- tics, volume 151 of Proceedings of Machine Learn- ing Research, pp. 6511-6528. PMLR, 28-30 Mar 2022. URL https: /pproceeeedigggseeppeepeeee www.mmlr.er.....m. press/ v151 1151/bunne22a. html.
Agueh, M. and Carlier, G. Barycenters in the Wasserstein space. SIAM Journal on Mathematical Analysis, 43(2): 904-924, 2011. doi: 10. 137/1000805411 URL https: //doi. P www.o//ggoooooooo 11 37 100805741.
Alvarez-Melis, and Fusi, N. Dataset Dynamics via Gradi- Flows in Probability Space. In Meila, M. and Zhang, T. (eds.). Proceedings of the 38th International Confer- ence on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 219-230. PMLR, 18- 24 Jul 2021. URL. httppp:://poccce...iiii wtw.:/eddd...mmm.m........mmm .gs.mlr sress/v139/ hwtp/////////e///...a wwwtaarez-elissiele html.
Bunne, C., Krause, A., and Cuturi, M. Supervised Training of Conditional Monge Maps, 2023a. URL https: arxiv. axxiv.org/ai org/abs//////oooooooo//.oor.aebr abs 220222222222 14262.
Bunne, C., Stark, S. G., Gut, G., al. Learning single- cell perturbation responses using neural optimal transport. Nature Methods, 1759-1617 2023b. doi: 10.1038/ $415922-023-019699> URL https: /ttpps:////..ii// 020--4111----22021---0000033310-
Mevass D., Schiff, Y., and Mroueh, Y. Optimizing Functionals on the Space of Probabilities with Input Con- vex Neural Networks. Transactions Machine Learning Research, 2022. URL http://wwpneeeeeeeeeeeee forum? 1d=d2d117dddmr
Carlier, G., Chizat, L., and Laborde, M. Lipschitz Continuity of the Schrödinger Map in Entropic Optimal Transport, 2022.
Amos, B., Luise, G., Cohen, S., and Redko, I. Meta Optimal Transport. In Krause, A., Brunskill, E., Cho, K., Engelhardt, B., Sabato, S., and Scarlett, J. (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Ma- chine Learning Research, pp. 791-813. PMLR, 23-29 Jul 2023. httpttps: http::////ccee//eeeeeecooo ngs mlr press// am0023m html.
Chang, B., Meng, L., Haber, E., Ruthotto, I.,, Begert, D., and Holtham, E. Reversible architectures for arbitrarily deep residual neural networks, 2017.
Chewi, S., Niles-Weed, J., and Rigollet, P. Statistical op- timal transport, 2024. URL https://x///iii.orrr abs 0ttp/2407.11666
Arjovsky, M., Chintala, S., and Bottou, L. Wasserstein Generative Adversarial Networks. In Precup, D. and Teh, Y. W. (eds.), Proceedings of the 34th International Con- ference Machine Learning, volume 70 of Proceedings of Machine learning Research, pp. 214-223. PMLR, 06- mmlr.....e......................
Chizat, L. and Bach, F. On the global convergence of gradi- ent descent for over wer-parameterizer models using optimal transport, 2018. URL https: //tp:////avvviiiiaa www.riii....oooooo.ooooooo abs/ 1805.00 09545.
11 Aug 2017. URL press press/v700 / aarjovsky17a.hhe
Choi, J., Choi, J., and Kang, M. Scalable Wasserstein Gradient Flow for Generative Modeling through Un- balanced Optimal Transport, 2024. URL https: arxiv. wwxiv.org/abs/244000 05443.
Behrmann, J., Grathwohl, W., Chen, R. T. Q., Duvenaud, D., and Jacobsen, J.-H. Invertible Residual Networks.
10

---

Universal Neural Optimal Transport
Courty, N., Flamary, R., Habrard, A., and Rako- tomamonjy, A. Joint distribution optimal trans- portation for domain adaptation. In Advances in Neural Information Processing Systems, vol- ume 30, 2017. URL htttps://proceedings ceurips.cc/pppe////// 0070022306b118665388888 pdf.
Inc., 2015.
URL https://proceedngss
neurips poarrips.cc/paper222555///////// 4bse6222361131225220096333333333 Paper. pdf.
a
Geshkovski, B., Letrouit, C., Polyanskiy, Y., and Rigollet, P. A mathematical perspective on transformers, 2024. URL https: http://awxxvvvoora/////aea
67a Pap-p
Cuturi, M. Sinkhorn Distances: Lightspeed Compu- tation of Optimal Transport. In Burges, C., Bottou, L., Welling, M., Ghahramani, Z.,, and Weinberger, K. (eds.), Advances in Neural Information Pro- cessing Systems, volume 26. Curran Associates, Inc., 2013. URL https: /ptpss////oceeeiggs.... neurips. cc/ hctppapee///1///iiie..// 2r21000f97bbee277e1333cfffffff pdf.
Geuter, J., Bonet, C., Korba, A., and Alvarez-Melis, D. DDEQs: Distributional Deep Equilibrium Models through Wasserstein Gradient Flows. In Proceedings of the 28th International Conference on Artificial In- telligence and Statistics (AISTATS 2025), 2025. URL https: www.aiivvcmmmm h 2503. 200..010300000000 01140.
Goodfellow, I., Pouget-Abadie, J., Mirza, M. Xu, B.. Warde-Farley, D., Ozair, S., Courville, A., and Bengio, Y. Generative Adversarial Nets. In Ghahramani, Z., Welling, M., Cortes, C., Lawrence, N., and Weinberger, K. (eds.), Advances in Neural Information Pro- cessing Systems, volume 27. Curran Associates, Inc., 2014. URL. htttt:::/prcceeee/ neurips. /tt.-p/s.ccpape///2////eeee Caa2e9b12226618fff44444cc.bb...2 pdf.
Paper.
Dadashi, R., Hussenot, L., Geist, M., and Pictquin, O. Primal Wasserstein Imitation Learning, 2020. URL https:/ /tas:///avvv..../ org/ 220. 04678.
Engquist, B. and Froese, B. D. Application of the Wasser- stein metric seismic signals, 2013. URL https: //arxiv. org/ 0ete//111.455 4581.
Engquist, B. and Yang. Y. Seismis imaging and optimal transport. Communications in Informa- tion and Systems, 19(2):55-455, 2019. URL
Goswami, S., Bora, A., Yu, Y., and Karniadakis, G. E.
Physics-informed deep neural operator networks, 2022. URL https:://////. //tp:///aavvvv//aa//eae/eeeaaaaa wwwpiii.orr///aa. hblg///b///2 wttp://..rr//bb/20000077777
http::/wwppeeeee..eeeeeeeeeeeeee https::/wwwtt inw.iriitppsss.. httpp//es//w.ilppsss.s... www.cooc///////oo
Gouk, H., Frank, E., Pfahringer, B., and Cree, M. J. Reg- ularisation of Neural Networks by Enforcing Lipschitz Continuity, 2020.
wols/0019//002//a/ http://///a/0//dde.e.hpp
Fefferman, C., Mitter, S., and Narayanan, H. Test- ing the manifold hypothesis. Journal of the Amer- ican Mathematical Society, 29(4):983-1049 2016. URL hww.h:tpp////cc.tm www. ams jrgg:/ouraaas//eeee 201 2016-29-04/508944-22 001-0000-200-00-------111
Gracyk, A. and Chen, X. GeONet: a neural operator for learning the Wasserstein geodesic, 2024, URL https: / //arxiv. www///g///bbaaao.ae/../eeee 2209. 14440.
Feydy, J., Séjoumé, T., Vialard, F.X.., ichi Amari, S., Trouvé, A.. and Peyré, Interpolating between Optimal Transport and MMD using Sinkhom Divergences, 2018. htttts: http::/:///a.xi...o///bb////////
Hashan, M. Facial expression images, 2022. URL https:/ /ww.. kaggle cww /2366499
He, K., Zhang. X., Ren, S., and Sun, Deep Residual Learn ing for Image Recognition. In 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR). pp. 770-778, 2016. doi: 10. 20.1109CCPP.20116.0
08278.
Fishman, N., Gowri, G., Yin, P., Gootenberg, J., and Abu- dayyeh, O. Generative Distribution Embeddings, 2025.
URL https:///aaaat.
org/ abs 250../2555..0 18150.
Jacobsen, J.-H., Smeulders, A., and Oyallon, i-RevNet: Deep Invertible Networks, 2018.
Franklin, J. and Lorenz, J. On the scaling of multidimen- sional matrices. Linear Algebra and its Applications,
114-115:711-735,, mar-apr 1989, URL https: /doi. org/10.1016/0000
Kingma, D. P. and Ba. J. Adam: A Method for Stochas- tic Optimization, 2017. URL https: http::/aaxvvvvoroee abs / 1412. 598C.
(89) 90490-4.
Frogner, C., Zhang, C., Mobahi, H., Araya, M., and Poggio, T. A. Learning with a Wasserstein Loss. In Advances in Neural Information Pro- cessing Systems, volume 28. Curran Associates,
Kolouri, S., Zou, Y., and Rohde, G. K. Sliced Wasser- stein Kernels for Probability Distributions, 2015. URL https://arx: http:://errii..orr/ba///11111111 0319

---

Universal Neural Optimal Transport
Kolouri, S., Park, S. R., Thorpe, M., Slepcev, D., and Rohde, G. K. Optimal Mass Transport: Signal processing and machine-learning applications. IEEE Signal Processing Magazine, 34(4):43-59 2017. 2695801.
Peyré, G. and Cuturi, M. Computational Optimal Trans- port: With Applications to Data Science. Foundations and Trends@ in Machine Learning, 11(5-6):355-677, 2019. ISSN 1935-8237. doi: 10.1561/2200000007.. URL http: /ttp://xxxooo doi www.aoo........ooooooooooo..oo.. 0 561 /2200000073.
doi: 10.1 10.1109/MPP:207.
Kolouri, S., Nadjahi, K., Simsekli, U., Badeau, R., and Rohde, G. Generalized Sliced Wasserstein Distances. Wallach, H., Larochelle, H., Beygelz- imer, A., d'Alché-Buc, F., Fox, E., and Garnett, R. (eds.). Advances in Neural Information Process- ing Systems, volume 32. Curran Associates, Inc., 2019. URL https: http:::/prcoeeedggg..eeeee oc /e//pper_filleesspprr/////////// [D3544cab22adcc7335a55eeeee1e pdf.
Pooladian, A.-A., Ben-Hamu, H., Domingo-Enrich, C., Amos, B., Lipman, Y., and Chen, R. T. Q. Multisam- ple flow matching: Straightening flows with minibatch couplings, 2023. URL https:/ http///axxvvvv.. /bbs/ 2304.14772.
Santambrogio, F. Optimal transport for applied mathemati- cians. Birkäuser. NY, (3)58,63))94, 2015.
Santambrogio, F. Euclidean, Metric, and Wasserstein gradi- ent flows: an overview, 2016. URL https: org/ / abs 600.. 03890.
Korotin, A., Selikhanovych. D., and Burnaev, E. Neural optimal transport, 2023. htt.httss://arxii..oo
abs/2201, 12220.
Schiebinger, G., Shu, J., Tabaka, M., Cleary, B., Subra- manian, V., Solomon, A., Gould, J., Liu, S., Lin, S., Berube, P., Lee, L.,, Chen, J., Brumbaugh, J., Rigollet, P., Hochedlinger, K., Jaenisch, R., Regev, A., and Lan- der, E. Gene Expression Identifies Developmental Trajectories in Reprogramming. Cell, 176(4):928 000888993,,111 2019.
Kovachki, N., Li, Z., Liu, B.. Azizzadenesheli, K., Bhat- tacharya, K., Stuart, A., and Anandkumar, A. Neural Op- erator: Learning Maps Between Function Spaces, 2024. http:::::// /etii//aaaivvvi /aaaaaaxvvvor////./ abs 2108 08481.
Analysis of Single-Cell
Li, S., Yu, X., Xing, W., Kirby, M., Narayan, A., and Zhe, S. Multi-Resole RelliiResolution Active Learning of Fourier Neural Operators, 2024a. URL https:/: /ww.iv..iiivv.///////e..//aaaa org/ abs/ 2309. 16971.
Schmitz, M. A., Heitz, M., Bonneel, N., Ngolè, F., Coeur- jolly, D., Cuturi, M.. Peyré, G., and Starck. J.-L. Wasser- stein Dictionary Learning: Optimal Transport-Based Un- supervised Nonlinear Dictionary Learning. SIAM Jour- nal on Imaging Sciences, 11(1):643-678, jan 2018. doi: 10.1137/17ml 140431. URL http:::tp////e... /doi org/10. 1 1137 2m 40431.
Li, X., Lu, F., Tao, M. and Ye, F. X. F. Robust first and second-order differentiation for regularized optimal transport, 2024b. URL https::///// / /ttxxxxxvvvv//////////////aaavvt org/ /bs/ 2407.02015,
Li, Z., Kovachki, N., Azizzadenesheli, K., Liu, B., Bhat- tacharya, K., Stuart, A., and Anandkumar, A. Neural operator: Graph kernel network for partial differential equations, 2020. URL /tttts://arxiv..oom abs/
Serrurier, M., Mamalet, F., Fel, T., Béthune, L., and Boissin, T. On the explainable properties of 1-Lipschitz Neural Networks: An Optimal Transport Perspective, 2023.
2003. 03485.
Sinkhorn, R. and Knopp, P. Concerning nonnegative Ma- trices and doubly stochastic Matrices. Pacific Journal of Mathematics, 21)2), 1967.
Li, Z., Kovachki, N., Azizzadenesheli, K., Liu, B., Bhat- tacharya, K., Stuart, A., and Anandkumar, A. Fourier neural operator for parametric partial differential equa- tions, 2021. URL https: / /aaxiviivvv.////////.aaa///avavv org/abs/ 2010.00 08895.
Thornton, J. and Cuturi, M. Rethinking Initialization of the Sinkhorn Algorithm, 2022. htts:tt::/// //arxiv. o org / abs / 220.. 07630.
Lipman, Y., Chen, R. T. Q., Ben-Hamu, H., Nickel, M., and Le, M. Flow matching for generative modeling, 2023. URL httph:tt//// http://xxvvvv./// org/ abcp/200.2220.. 02747.
Tong. A., Fatras, K., Malkin, N., Huguet, G., Zhang, Y.. Brotorr-rroos... J., Wolf, G.. and Bengio, Y. Im- proving and generalizing flow-hased generative models with minibatch optimal transport, 2024. URL https:
Loshchilov, I. and Hutter, F. Decoupled Weight Decay Regularization, 2019. URL httttt:s/aaaxi...ooo 0bs/1711.05100
/ / //axxiv.. org. abs 00482.
Uscidda, T. and Cuturi, M. The Monge Gap: A Regularizer to Learn All Transport Maps, 2023.
Nutz, M. Introduction to Entropic Optimal Transport,
2022. URL http://wwwwwtttta www.. columbia........ooo..mo.....a..a edu/ mmnutz/docs/EOT. wewp:oo___eettttteeeeeee_ee .ddf.
Villani, C. Optimal Transport Old and New. Springer, 2009.
12

---

Universal Neural Optimal Transport
Xu, H., Wang, W., Liu, w. and Carin, L. Distilled Wasserstein Learning for Word Embedding and Topic Modeling, In Bengio, S., Wallach, H., Larochelle, H., Grauman, K., Cesa-Bianchi, N., and Garnett, R. (eds.), Advances in Neural Information Process- ing Systems, volume 31. Curran Associates, Inc., 2018. URL http:://proceee.ngg...r neurips. cc co/paper /c/ppppr / 2 pdf.
Zhang, C., Cai, Y. Lin, G., and Shen, C. DeepEMD:
Few-Shot Image Classification With Differentiable Earth Mover's Distance and Structured Classifiers. In 2020 IEEE/CVF Conference Computer Vision and Pat- Iern Recognition (CVPR), pp. 12200-12210, 2020. doi: 10110999VPP:77600.00000222.
13

---

Universal Neural Optimal Transport
The Appendix is structured as follows: In Appendix A we give thorough background on OT, well additional technical details for experiments computing barycenters and geodesics. Appendix B contains all proofs omitted in the paper. In Appendix C, we provide additional training details, and Appendix D contains additional experiments and materials,
A. Background
A.1. Optimal Transport
In this section, we recall some properties of optimal transport. First, we define the unregularized continuous problems for completeness.
Problem A.1 (Kantorovich Optimal Transport Problem). For μ E P(X), M EP))) and cost e: X + +R {c0} the Kantorovich problem takes the form
/ d(x..)) dr (x.
inf
(12)
xell(p,))
The infimum in (12) called the transport cost, and the minimizer π, if exists, the optimal transport plan.
The contimuous dual problem similar to the regularized dual (3), For more thorough overview of OT, we refer the reader (Villami, 2009; Peyré& & Cuturi, 2019; Chewi et al., 2024).
Problem A.2 (Dual Optimal Transport Problem). For μ1, and before, the dual problem reads
/ f(x)dp(x) J (lu)dd(v).
sup (EL ftg5e
where f g / S is to be understood as f(z) + glu) S c(x,, y) for all x,y.
An important concept in optimal transport are transport maps.
Definition A.3 (Transport Maps). map T X + y is called a transport map between μ and v if v Ty (.. u. there exists an optimal transport plan such that (Id, +) p, T is called an optimal transport map.
Of course, not every transport plan admits transport map; however, every transport map yields optimal transport plan via (Id, u4. For sufficient conditions for the existence of both transport plans and maps, we refer the reader to (Villani, 2009).
In the paper we mentioned that if μ and v are supported on finitely many points, one can rewrite the problems A.1 and A.2 with vectors. We now define the discrete problems carefully.
Problem A.4 (Discrete Optimal Transport Problem). For two discrete measures μ Am-1 and E An-1, and cost matrix C' gmxx, the discrete OT problem is defined
L(,, v) min (C, nell(p.))
Here, I[(,, v) denotes the set of all transport plans between μ and v, i... matrices E mxxx s.t. III, 4 and 11 Im = v. The problem has dual formulation:
Problem A.5 (Discrete Dual Optimal Transport Problem). For two discrete measures μ E A-1 and v E n-1 and a cost matrix C E gmxx, the discrete dual OT problem defined as
D(H,V)= = max ream gean 44050
(f 14) + (g.
Here, f + S C is be understood as f. + 9, S Cig for all E [m]] j E [n]. In the special case where = y and C corresponds metric, i.e. Ciş d(z,, vj). the Wasserstein distance of order p between μ and v for p E [1, oc) defined
W, w)(,vv)=
min sell(u.
ΣdJile
14

---

Universal Neural Optimal Transport
This definition coincides with the definition from the paper for continuous measures, if they supported on finitely many points.
For completeness, we also state the entropically regularized primal and dual problem in the discrete case. The discrete problem is typically formulated with an entropy term instead of the KL divergence as in equation (1), but the two can be shown to be equivalent (Chewi et al., 2024).
Definition A.6 (Entropy). For a matrix = [pslls E Rmxx, we define its entropy (P)) as
H(P) =
Σ Pij logPt;
if all entries are positive, and (P)) == -00 if at least one entry is negative. For entries Púj = 0, we use the convention 0 log 0 = 0, as log r ==*,, 0.
The entropic optimal transport problem is defined as follows.
Problem A.7 (Entropic Discrete Optimal Transport Problem). For e > O, the entropic optimal transport problem is defined as:
OTe( = min (C, IT] (II...
Ielll(p)))
Note that this is identical to the unregularized optimal transport problem, except that the unregularized one does not contain the regularization term -cH (IT.. As the objective in Problem A.7 c-strongly convex, the problem admits a unique solution (Peyré & Cuturi, 2019).
The Gibbs kernel is defined as K exp(-C/e). Then the entropic dual problem reads: Problem A.8 (Entropic Discrete Dual Problem).
max fEER", g. ER (f.. + (g..v) - (chlee Kelle))
Again, without the regularization term -6 (etele, Kea-/e), this equals the regular optimal transport dual; note, however, that the unregularized dual is subject to the constraint f g S C.
In both the continuous, as well as the discrete setting, there is duality, i.e. the optima of the primal and dual problems coincide. In addition, the optimizers are intrinsically linked, akin to Proposition for the discrete entropic problem. We refer the reader to (Villani, 2009; Peyré & Cuturi, 2019) for more details.
A.2. Calculating the Barycenter
Recall the Sinkhorn divergence barycenter for a set of discrete measures (VI...,, C P2(A)),
H arg min xxEP((X)
Σ a (D)','
For a solution (f.g) to the dual problem A.8 between two measures ft = E, ar and v Σ, b56 holds
SD, (n. p) + (v.g q)
see (Feydy et al., 2018). Here, p and are the optimal potentials for (1,, 14) and (v, v) resp. (if both measures in the OT problem are the same, the dual potentials can be chosen to be equal).
From this identity, we immediately get
V (D, f
Note that this is not a gradient with respect to the measure μı; instead, we view μ as a vector, and compute the gradient w.r.t. the entries in that vector. This means we essentially compute the barycenter on the discrete space ..))...;a.))
15

---

Universal Neural Optimal Transport
pseudocode for how approximate the Sinkhorn divergence barycenter with UNOT is given in Algorithm 3. Note that instead of using softmax project back to probability measures, could also just rescale; however, softmax proved better in practice. We also run single Sinkhorn iteration on the output of S& in practice, as it improved visual quality of the barycenters, however, this is not strictly needed.
Algorithm 3 Barycenter Computation
1: in set of measures {ui}e A initial 140 E A"-1, weights À E A
2: μ 140
3: for = 1,,,..... T do
4: p + (p (1)
5:
for v; in (Viti do
6: Si 2 (v(vs,
//switch the order of the arguments to get fi instead of gi
end for
8: _ 4- softmax (u E.Aifi
9: end for
Barycenters can be far when Transport Distances are close. We now give a simple example that illustrates that merely predicting transport distances accurately does not necessarily imply predicting barycenters accurately, at least in the nonregularized setting. Let μ be the measure with mass 1/2 the two points (0,-1) and (0, 1) in IR, Let V1 be the measure with mass 1/2 at each of the points (-1,0) and (1, -)) for small > 0, and 1 measure with mass 1/2 at each of the two points (-1,0) and (1,6). Then as goes to 0, the transport distances between between μ and V1 resp. l and V2 become arbitrarily close. However, the unique Wasserstein-p barycenter (for 1) between and V1 has mass 1/2 at each of the points (-1/2,1/2) and (1/2, -(1 + e)) whereas the barycenter between μ and V/2 has mass 1/2 at each of the points (-1/ 1/2, 1/2) and (1/2, (1 + e)) so no matter how small gets, the barycenters will always be far apart.
A.3. Geodesics
In Section 4.3, we saw that the McCann interpolation between two measures p, (2( is a constant-speed geodesic. In this section, we provide additional background on constant-spest speed geodesics, and establish a connection between constant- speed geodesics in P2( 1)²) and the notion of strong-c quasi-geodesica in the discretized space PP [[m]]², ). This makes approximation of geodesics as in Section 4.3 more rigorous.
First, we recall the definition of constant-speck geodesics.
Definition A.9. A curve : [0, 1] (P2()), W2) is called constant-speed geodesic between (0) and w(1) if it satisfies
(w)w(t), = t- - (())),))), Vt, E [0,
turns out that for convex C Rd, constant-speed geodesics are equivalent to push-for furwads under transport plans, and the starting point (0) is absolutely continuous, this is equal to the McCann interpolation.
Theorem A.10. Let X C. Rd be convex. Then a curve w [0, 1] + (P2(X), W2) is a constant- speed geodesic between (0) and w1) if and only if it is of the form
wtt (pe क
for optimal transport plan between w0) and w()), where the interpolation pt is given by pc(x, y) (1 t)x + ty. in addition, w0) is absolutely continuous, then we can write
[d -dd
for an optimal transport map T from w() to (1).
This theorem holds, in fact, for any Wasserstein-p space for >> 1, see (Santambrogio, 2016).
Now, denote by μ the McCann interpolation between μ and v. As mentioned in Section 4.2, we can express μt as the following barycenter:
P4 arg min (11 (p', p) + IW3 (nl. v))
J (2))
16

---

Universal Neural Optimal Transport
which we approximate by the Sinkhorn Divergence barycenter
Ht arg min (Q- tt / (A)
(13)
which is justified by the fact that Sinkhorn Divergences converge to the OT cost as c + 0, and that they are reliable loss functions, in the sense that weak convergence of a sequence of measures is equivalent to convergence of the Sinkhorn divergence, see (Feydy et al., 2018). As also shown in (Feydy et al., 2018), the gradient of the Sinkhorn Divergence w.r.t. the vector a, when writing discrete measure 1a as μ = E; a262,, is given by
V =f
(14)
However, if we now do a simple gradient descent on (13) using (14), we are not actually computing the barycenter on the space (2(2)) anymore, as we only consider gradients w.r.t. a, which does not allow particles to move, but merely to teleport mass to other particles. particular, if X' a discrete space, there exist no constant-speed geodesics between different points anymore, as can easily be seen from the following example. Let μ40 = 62. and μl 82, be two Dirac measures for some 20, 21 E X. Assume there would exist constant speed geodesic w joining po and μl, Then for t > 0,
W2 (w)(,, = (W2(
However, since the space is discrete, this implies that vg = i.e. the only constant-speed geodesics are constant. We therefore work with the following approximation of geodesics.
Definition A.11 (Quasi-lsometry). Let (Ki, di) and (22,d2) be metric spaces. Xi X2 is called a (A, -)--usii- isometry if there exist A 2 0 and e >> 0 such that for all æ, E Xi
X di (2, - e da((()), s (2, + +6
If in addition there exists C > 0 such that for all z E X2 there exist x Xr such that d2(s(z),)) C, f is called quasi-isometry.
We can then this define quasi-geodesics. (Bonciocat & Sturm, 2009) introduced similar concept called h-rough geodesics, for which they just used the upper bound.
Definition A.12 (Strong-c Quasi-Geedesscs). A strong-c quasi-geodesic metric space (X, d) is map y [0, - X such that for all 8, & [0, 1]. d()o - s| - s d(1) S di)o s) +e.
Now let X = =]0, and denote by |+ C. [0, 1 the discrete space consisting of all x, of the form z; A (L) 0) +j (0, 2)). for j.j 0, .... 1. We can then show that (P2(10, 1J3), W2) is quasi-isometric to (P |m]|
+ W2).
1 A --quast-isommtrii to (P (P[ll)))) W2), i.e. there exist an Proposition A.13. The metric space (P2(10, 1)3), W2) is (1, V2n f: (P2(10, 112), Wa) > (P( [[m]]? ), W₂) such that for all u. P2(10, 11) it holds that
W2(prr) 1
S Wo(f(p)) < W2(m, 1
22n
22n
Proof. We split the space [0, 1]2 into squares via N(xz) = (IF [- 1
2n 12)...
We define f: P((0, 1J2) P [[mll' by
n
) Sz.
f(A)= - Σ
du )a..
2.EX
17

---

Universal Neural Optimal Transport
By triangle inequality, we have
Wa( Wa(m,)) W2(4,,00) Wa(ss)))+ Wa(e, (()) W2(1, Wa(J(ww) + (1((,(()) + (()),
For any measure y E P([0, 1J). denoting by : [0, 1]3 [a the map that sends each point to the corresponding midpoint Ty we get
W20,500) S / 10.11
(T(x) x/d S / 10.13 4n du
2
2
4n
Therefore we have a (1, V2n Jquasi-isoomery between both spaces.
We also need show that there exist C such that for all μ E P( (Iull) there exist v P((0, 1)3) with
Wa(s(v). <C
Choosing C and jr concludes the proof.
v2n
We immediately get the following corollary. Corollary A.14. Constant-speed geodesics P2(10, 1)²) are strong-e quasi-geodesics in P([[n]]]/n))
This justifies doing gradient descent (13) using the discrete space gradient (14) to approximate the geodesic, as we can approximate the constant-speed geodesic with a stronn---quas--eoddssi in the discrete space.
A.4. Wasserstein on Wasserstein Distance
In this section, we provide additional details on how to solve the particle flow
a
Ot f
(15)
from Section 4.4. Recall that fú, E P2((P2((0, 112, c), W2). f = 1 E, 8p., D= 1, 5e, for p4, E P2((0, 1J?). From (Li et al., 2024b), we get that
asD. D)
a (DD vg) 1kas
0pk
where 1eg is an optimal transport plan between H4 and Vj- Now as in the previous section, we can approximate a (SD((HE.5)
Bue fuj Ps. where fky is the dual potential from OT.(%. pj), and ps that of OT, m&)) As before, can approximate these gradients with UNOT, which lets us solve (15) with a simple gradient descent scheme, as shown in Section A.2. As in Section A.2, we add single Sinkhorn iteration on the predictions made by Sg as it improves visual quality, but this is not strictly necessary.
A.5. Fourier Neural Operators
In this section, we describe FNOs in more detail. The main breakthrough for Neural Operators came in the combination with approximating solutions to partial differential equations (PDEs) (Li et al., 2020; 2021; Goswami et al., 2022). Many problems, including PDEs, can be numerically solved by discretizing Neural Operators are a class of neural networks that parametrize functions F' AA + U, where A and U are Banach spaces whose elements are functions
dimnite-dimensionall input and output functions.
Da - R. and u D. + Rd respectively, for bounded domains D. C Rd and D, C Rd., One of the main advantages of Neural Operators is that they can generalize over different grid discretizations, unlike traditional neural networks, which makes them particularly well-suited for solving PDEs, !1 and they are universal approximators for continuous operators acting on Banach spaces (Kovachki et al., 2024). While our space P([0, 1 (?) is not
For example, FNO could be used solve PDEs of the form Au with Dirichlet boundary conditions, for which get a unique solution u for every a. The FNO then maps each a E A to the corresponding solution u E U4.
18

---

Universal Neural Optimal Transport
1,
Layer
Layer 2
Layer L
F [[N(Foo))( b(#)
W
Figure 9. Fourier Neural Operator architecture, adapted from (Kovachki et al., 2024). The input measures (μ, are passed through point-wise lifting operator which is then followed by Fourier operators and point-wise non-linearity operators. After the Fourier layer, we project back to the output potential with point-wise operator Q.
technically a Banach space, the space of finite signed measures with the total variation norm is, and P((0, 1J²) is a subset. We note that approximation theory for Neural Operators usually
neural operator usually has the following form:
F: P:A-1
+ QL o.... B1 o P(a),
which in our setting becomes
x - (10, S (H, - Q B1 o.... B1 P(H, v).
Here, P is lifting map, B, are the kernel layers, and Q is projection back the target space.
Different versions of neural operators have been proposed, which mostly differ in how the kernel layers B, are defined. Our network Sø is parametrized as a Fourier Neural Operator (FNO) (Kovachki et al., 2024), where the kernel layers act on Fourier features of the inputs. We outline details for all the layers in the following.
Lifting (P). The lifting map is a pointwise map {a Da + Rd + {uu : Do + Rd which maps the input a to a function Up by mapping points in R to points in Rd We use 2D convolutional layer for P, and in our setting, D. = [0, 1]2 x 1]2, as we can view elements in P((o, 1)2) as maps (0, 1|2 -> R when dealing with discretizations of measures.
Iterative Fourier Layer (B,). The network has L Fourier layers B,. each of them, we map {uz D: - R. + {vi41 1 Rd+41] by first applying the (discrete) Fourier transform J from which we select fixed number of Fourier features, then a neural network NN on these features, and then the inverse Fourier transform F-1. Note that the Fourier features are complex, hence the network N is also complex (with multiplications in C). Each Fourier layer also contains a bypass layer, which is similar to skip connection, but contains layer W which is typically a 2D convolution; cmp Figure 9. Hence, the output of the Fourier layer is given by a ofF- (NN(())( + b Wo. where a is an activation.
Projection (Q). The projection Q is the analogue to the lifting layer, mapping the hidden representation to the output function {nz : Dr - gd + {u: D. + gd ). In our setting, D, = [0, 1|2.
In contrast to (Kovachki et al., 2024), we found that a Fourier layer containing a two-layer neural network N. instead of just linear layer worked better in practice. Our bypass layer is still linear layer W.
On the unit sphere 82, use Spherical FNOs (SFNOs) (Bonev al., 2023) instead of regular FNOs, which respect the geometry of 52. SFNOs leverage the Fourier transform on the sphere Ss which can be viewed as a change of basis into
19

---

Universal Neural Optimal Transport
an orthogonal basis of (S2), instead of the regular Fourier transform F for flat geometries. Everything else about our architecture remains the same.
Details on hyperparameter choices can be found in Appendix C.
20

---

Universal Neural Optimal Transport
B. Proofs
This section contains proofs, as well as further technical details omitted the paper. For convenience, we restate the statements from the paper.
We start off by rigorously restating Proposition 2. Let X' R be compact set. We start off with natural definition of discretization of a continuous measure, which applies, for example, to discrete images discretizations of an underlying "ground truth" continuous image.
Definition B.1 (Discretization of Measures). Let P(X) be absolutely continuous measure, and let Xn [."....,.; C X. The discretization of on Xn is defined as the measure pn E PP(X)) supported on Xn, where
(n(x)) = / du,
with
2 far E X: 1 x. S | y/ Vy X).
Note that the intersections Q n2g have Lebesgue measure zero, so this is well-defined.
We cannot guarantee that arbitrary sequence of discretizations Ynn converges weakly to as - oc; simply consider the case where all the xe are identical for all and i. Hence, we need to ensure that the discretization is uniform over all of A' in some way.
Definition B.2 (Uniform Discretization). Let X₂ = {-; th be subsets of X' for all n E N. Then we call the sequence (Nn)neN a uniform discretization X if for all x E X.
lim
min || - x1 = 0.
7-000 im
While this may seem like "pointwise discretization" at first, it turns out to be uniform, an Arzelà-Ascoli type argument shows.
Theorem B.3. Let X C Rdd be compact, and let (Na (221 be a sequence finite subsets of X with Nal = n for each The following are equivalent:
1. lim sup min || yll 0. zEd yEa,
2. Vx Ex: lim min || yll 0. +00 yEA,
Proof. (1) - (2). If sux mingex, ||x y/l O, then in particular for each fixed X have
|| wll BEAn
min
sup min || yll zex vEA,
0.
(2) - (1). Define
fa(x) = min
| U|1 I E X.
y,
By hypothesis (2), fa(x) + for every X. Moreover for any xr, 2 E X,
(x) fa(=))
min y.
| wl1 min y,
|| U/ll
x 2||
so {fr} is equicontinuous on the compact set X. Since fa + 0 pointwise, the Arzelà-Ascoli theorem upgrades to uniform convergence of the entire sequence (instead of just a subsequence):
lim sup fa(x) = lim sup min ||x - yll = 0. n ZEX
no yeR,
Hence (1) holds, completing the proof.
21

---

Universal Neural Optimal Transport
Note that condition (1) in Theorem B.3 is equivalent to Definition of "discrete refinement" in (Kovachki et al., 2024).
The following lemma holds.
Lemma B.4 (Weak Convergence of Discretizations of Measures). Let X be absolutely continuous, and (pn) JuN) be sequence of discretizations of μ supported on uniform discretization (Kn)neN of X. Then po converges weakly to p..
Proof. Let f E Cb((2') be test function. We have to show that
J fdji " J fdp.
Since X' is compact and R continuous, by the Heine Cantor theorem f is uniformly continuous. Hence, for such that
every 8 > 0 there exists
x
y/1 <d =
If(z) - f(y)] <
for all x, E X.
Since
sup min ||x xx,l
FE 15:5n
70000
we can choose n' such that for all n 2 n',
sup min || - x < Tex Isisn
In particular, for each E X, there is some x! X, with || x; < 5, giving
f(x)) < whenever || - 2 || < 8.
(16)
Let
|| x. S || [x--="] =Iحll for all ,,..., x E X
as above. These sets form a partition of (up to measure-zero boundaries). Then for all n 2 n', we have (using equation (16)):
1 fdun - - fdu
Σ / (f(x?) f())) du(x))
aap
n2
[f(x7) - f(x)) du(x)
i=1
<p((2))
1=1
6,
and letting - 0 finishes the proof.
In Proposition 2, we used the "canonical extension" for dual potentials. For pair of dual variables (J.g) solving the dual problem (3) between μ and v, their canonical extensions are defined by ff and g satisfying the following conditions:
f(x) eloo / exp
(g(v) - c(z, v)) du(y),
g(x) - / exp
clog
(f(y) c(x, v)) du(y).
We refer to (Santambrogio, 2015; Feydy et al., 2018) for more details.
We can now state and prove formal version of Proposition 2.
22

---

Universal Neural Optimal Transport
Proposition 2. (Formal) Let c(x, X R be Lipschitz continuous in both its arguments, and X R compact. Let (pa)ners nENN be discretization sequences for absolutely continuous Ja, E ((X), supported uniform discretization (NN) of X. Let (frs 9n) be the (unique) extended dual potentials of (pas Vn) such that fn(xo) 0 for some zo X and all n. Let (f.g) g) be the (unique) dual potentials of (p. v) such that fzo) O. Then fn and converge uniformly to and g on all of X.
Proof. By Lemma B.4, we know that Hr, - a2 and Vn v. The statement now follows immediately from Proposition 13 in (Feydy et al., 2018).
Theorem 3. Let and Gg IR + Rd be defined via
Gø(z) = ReLU (NNo(z)) + Xz),
where z P= -(,, and where NÑo Rd + R is Lipschitz continuous with Lip(NNe) = L < X. Then Gg is Lipschitz continuous with Lip(g) < + X, and G(z) == NNo(z) + Az is invertible on Rd. Furthermore, for any E Rg. it holds
(I+ (Go (x0))
1
PGouPa (x) 2
In other words, G0 has positive density any non-negative Re
Proof. Since the Lipschitz constant of the sum of two functions is bounded by the sum of the Lipschitz constants of the two functions, we have
Lip(Go) S LL+1 X.
From Theorem in (Behrmann et al., 2019), it follows that Go is invertible, and Lemma 2 therein implies
Lip(Go ")=
1 X-L
The Lipschitz continuity of Go -1 implies that for any h, z E RRd with f a O, we have
Go(z + th) - Go(z)
lim 10
t
Go -1 (Go(z + th)) - Go "G)((())))
1
lim 100
Lip
-1
|nll
Lip(Go
0,
which shows that VGo is invertible everywhere. Hence, by the inverse function theorem, we get
VGo -1 (x) = Vūg -1 (xx
(Go -1 (x))) = (vGo(Go -1 (x)))-
for any Rd. Furthermore, similar to above, we have
(C)()))
lim r-0
Go(z + te;) Go(z) Lip(Go) lim t->0
te,
<L+>
where e; is the 7t unit vector. Hence, we get from Hadamard's inequality that
[derVGo(=)] < IT (Go())
< IT:(L + X) = (L X)a.
23

---

Universal Neural Optimal Transport
Putting everything together, by change of variables, get for any æ E Rd:
P Pa (z) - P= (x)) deiVC (xr)
(x))
PG (x)) deiV - (z))
2 (L + Aad Px (Go (x))
(L (L+ x N(Go - (.))0,
Now clearly, if x E R then
PGo (x) PGo a (x),
as for any z with G(z) = æ, we also have G(z) = x. Thus, we also have
PGogp= (x) 2 (L+ Xp (I+
N(Go (z:)00 (m00 7),
(0)
which finishes the proof.
Corollary 4. Lel Gg = Go, Go. o... Gog be a composition of functions Go,, each of which is of the form as in Theorem
3. Let ~ Pz = I). Then
PGozpz (x) (L (L+ (1+ (Go - (2))).1)
for any Rd, As in Theorem 3, this also holds for any RS if Go is followed by a ReLU activation.
Proof. Consider the case where Go o 0g =G,,
Then for any x E Rdd, we get from the proof of Theorem 3 above:
PGogPs P2 (x) 2
(L + A) P0g d= ((Ca)
(x))
1
((ch)
> (L + Aaza~V
(x) 10,1
V(G " (z)
The claim now follows by induction over the layers of Go. Note that Go is followed by ReLU activation, this inequality also holds for any æ RRS0 similar Theorem 3.
Next, we prove Proposition The proof is based on the Hilbert projective metric. For two vectors u, 1R. is defined
d (u, v) = max[log( (log(u))) log(n;)) minnlog((s)) log(::));
and can be shown to be distance on the projective cone / ~ where u ~ u' if u = ru' for some > 0 (Peyré & Cuturi, 2019; Franklin & Lorenz, 1989). For f = (u) and log(v), v), we thus define the following loss:
Lu(f.g) max(fi gi) min[f: -9.)-
Lemma B.5. Let f.g E R". Then
LH(f.g) S V2 gl|2-
If. in addition, E. - 0. then
|f S Vn (f.g)
24

---

Universal Neural Optimal Transport
Proof. Let f g. For the first inequality, observe that (f( max, h4 min; h4. Let j" and k* be the indices achieving max; h, and min; h; respectively. Define the vector e such that e-- 1, ex- -1, and e; 0 for all other i. Then:
(f = h - || ||]| = v f1 - 2gl22
Now assume that fi = E: Set M max; h and m min, h., If all h; 0. both statements are trivial. Hence, assume at least one of the h, is not zero. Since Σ; h E; 9: 0, this implies M > 0 and m < 0. For any index i, h, S M, and thus
(h.) - M² - (M m)² = (n
Summing over all indices, we have:
n
|l - gll = |11 = Σ (hi)² (n
Taking the square root yields:
|f - 1-gll2 S vn LH (f.g).
This finishes the proof.
Proposition 5. For two discrete measures (p, v) with n particles, let g be a potential solving the dual problem, 96 So(H,v), and gT. [T, v, (6) the target. Without loss of generality, assume that E. 9: Σ, 9 0. Then
L2(ge, g) S c(K,k, n) L2 (g 9rc)
for some constant of K, n) depending only on the Gibbs kernel K, k and nt
Proof. We first show similar inequality in Proposition 5 for the Hilbert loss. A well-known fact about the Hilbert metric is that positive matrices (in case, the Gibb's 's kernel K) act strict contractions positive vectors with respect to the Hilbert metric (ef. Theorem 4.1 in (Peyré & Cuturi, 2019)). More precisely, we have
(H(KU, S (V,)
for any positive vectors v, v R", where
(K) 1 X(K) = Vn(K) +1 T(K) = max K3t 1. KjtKa
The same inequality also holds for TT place of K. Note that by definition, ry (K) 2 1, hence 0 < X(K) < 1. Now consider starting vector o the Sinkhorn algorithm, and let v' denote the 1th iterate of the vector. Denote by V² the limit lim of of the algorithm. Then (letting' /' denote clement-wise division):
d ( v²) = du (u/KT +1 p/KKTu" u²)
=d (K ult1 (..KKTw u²)
X(K)dn(u't) .u") ,u')
A(K)dH (N/Ke,mp/K'") m/Ko*)
-(())d (Ko', Ko")
S x(K) (d),"),
where we used the Hilbert metric inequality twice, once on and once on KT, Iteratively applying this inequality and translating into log-space notation, this gives
LH g) S X(K)2* LH (g(.9))) g).
For now, assume that E, Io. 0. By triangle inequality,
T. (9..9) ()) < T. (go. Ir.) T. S In + X(K)2 (g g),
25

---

Universal Neural Optimal Transport
which gives us
L (g g)
LH Ln 9..) = c(, (go.
1 \(K)2k
Combining this with Lemma B.5 yields
119. 2ll VnLn(go.9) 9)
S 2/no(K, 119 9r, l2 c(K, 1100 -2ull2 (17)
from which the claim follows by squaring both sides. We are left with proving the general case when E; 0. Write 96 IΦ + 9p. where I6 is equal to 1/2 E. 96; in each entry, s.t. I6 sums to zero. We then get
L2 9) = 1100 - -gll= + 1 1|
(18)
as
(9+ - -.96) =
Similarly, we get
Lg(g (-.)
090 9rx 21 |1
(19)
Combining equations (17), (18) and (19), get
L2 g)
11 g|| + |
S cL2(g0 9n,) + ||
cLollg 9r.) - | 11 + 1 /1?
=cLelgo. 9Th) + <) |!
scLe(gor99 9:.)
where the last inequality follows from the fact that 1
This finishes the proof.
Remark B.6. Looking at the proof of Proposition 5, one might wonder why we didn't 't opt for the Hilbert projective metric the loss directly. We tried using it instead of L2, and it works quite well, but training with L2 seems to have edge, probably because the indifference of the Hilbert projetive metric to constant shifts is not helpful inductive bias for deep learning.
26

---

Universal Neural Optimal Transport
C. Training Details
Generator Architecture. Recall that the generator is of the form
Gg R + P(X) x P(X) P= R RRLUU (NNo(==) + A Al + 6],
where we set À = 1.0, = 1e-6 (note first normalize, then add ö, and then normalize again in practice), and is of size 2 10 x 10. R normalizes and randomly downsizes output distributions to resolutions between 10 x 10 and 64 x 64 (per distribution). This improves generalization of the FNO S& across resolutions, which is true for FNOs in general (Li et al., 2024a). NNo o is a five-layer fully connected MLP, where all hidden layers are of dimension 0.04 - 64", and the output is of dimension 642. All layers except the output layer contain Batch Normalization and ELU activations; the last layer has a sigmoid activation only. We note the architecture might seem strange, as the network is relatively deep, while the hidden layers are relatively narrow. However, this architecture worked best amongst an extensive sweep of architectures.
Applying Theorem 3. In the following, we discuss the relation between our generator Go and Theorem 3 in more detail. Note that Theorem 3 is not directly applicable to our setting for a few reasons: First, we add small constant 1j to the generator's 's output. This constant ensures that all training samples are positive everywhere, and vastly improves learning speed as it ensures that all inputs are active. However, this is not restrictive of the problem, as the Sinkhorn algorithm requires inputs to be positive anyways. Second, in Theorem 3 both in- and outputs to Go have the same dimension. This could be achieved in our setting by choosing the input dimension equal to the output dimension, i.e. Id,n equal to the identity. However, in practice, using lower-dimensional inputs achieves significantly better results. This can be argued for by the manifold hypothesis (Fefferman et al., 2016), i.c. the fact that typically, datasets live on low-dimensional manifolds embedded in high-dimensional spaces. Depending on the application, i.e. the expected target dataset dimension, the dimension of the input can be adjusted accordingly. Finally, note that the theorem assumes that NNg is Lipschitz continuous with Lipschitz constant L < X, where À is the scaling factor of the skip connection. We do not enforce this constraint, as not doing so yields empirically better results. Still, Theorem 3 goes to show that our algorithm's 's performance is not bottlenecked by the generator's 's inability to generalize. We note that a bound on the Lipschitz constant is not necessary for invertibility of ResNets; other approaches have been suggested in the literature, e.g. through the lens of ODEs (Chang et al., 2017) or by partitioning input dimensions (Jacobsen et al., 2018). It is also possible to directly divide by the Lipschitz constant of each layer (Serrurier et al., 2023); these approaches could be studied in future research.
We will now describe how one can bound the Lipschitz constant of the generator. Since = 1.0, we need to make sure that the Lipschitz constant of netg is smaller than in order for Theorem 3 to be applicable. Since the Lipschitz constant of a composition of functions bounded by the product of the Lipschitz constants of each component function, this means we have to bound the product of the Lipschitz constants of components of netg. ELU is Lipschitz continuous with constant 1, whereas sigmoid's Lipschitz constant is 0.25. Furthermore, for batch normalization layer BN, we have
(x) - BN(
x 14b
14b σ6
|| -
σb
σb
where ob and og denote the empirical mean and standard deviation of the batch. Since we draw our data from a standard normal Gaussian, we have E[or]] = 1, i.e. in expectation, the batch normalization layer is Lipschitz with constant 1. Hence, all that remains is to bound the product of Lipschitz constants of the three linear layers by (any number smaller than) 4 (because the constant of sigmoid is 0.25, this will ensure that the network has Lipschitz constant smaller than 1), for which it suffices to bound the operator norms of the weight matrix of each layer. In practice, these can be approximated with the power method as in (Gouk et al., 2020) to find a lower bound on the Lipschitz constant of each linear layer, and these bounds can be used to add a soft constraint to the loss. Empirically, this suffices to bound the Lipschitz constant of the generator. Alternatively, one can use a hard constraint as outlined in (Behrmann et al., 2019). However, empirically, this proved detrimental to training, hence we did not control the Lipschitz constant during our training. Yet, Theorem 3 is still of value, as it goes to show that our algorithm's 's performance is not bottlenecked by the generator's inability to generalize. We leave properly enforcing the Lipschitz constraint for future research.
Architecture of Sg. Our FNO architecture follows the general structure outlined in Section A.5. We set de, = 64 for all é; recall this is the hidden dimension in the Fourier layer. We set the number of Fourier features selected from the Fourier transform to 10x x 10, i.c. 10 along each of the two dimensions of the domain. The (complex) weight matrices of the neural network in Fourier space, i.e. the one acting on the Fourier features, are tensors of shape (de, Adoy; N Nmodes, N Nmoaes,)) =
27

---

Universal Neural Optimal Transport
(64, 256, 10, 10) and (4 du N mode, N modes,)) (256, 64, 10, 10) respectively, i.e. the hidden dimension is four times the hidden dimension of ti. Note that since these complex layers, each layer has two (real) weight tensors of this shape, one for the real and one for the complex part. These layers are the only complex layers in the network Sg. The inputs to the layer are of shape (de.: Noodee, N (odes,,) = (64, 10. 10) (in C) and multiplied along all dimensions by the weights, Lee. for input & E C61,110110 and weight matrix A E C61,256, 10,10 (the first of the two layers):
Do.n.m
Σ Alommmannmm
The activation used within this network, as well as after each Fourier block, GeLU. The lifting layer P. bypass layer W. and projection layer Q are 2D convolutions with kernel size 1.
Hyperparameters. In Table 3 we present all relevant hyperparameters again for convenience.
Table 3. Training hyperparameters.
Hyperparameter
Value 272k
# params Go
# layers G
5 (164, 164, 164, 164) le-6
hidden dims G
8 (eq. (6))
À (eq. (6))
d (dimension of latent z)
2 10 x 10 200
optimizer G
Adam ELU 0.001
activations G
81 (initial learning rate Go learning rate decay Go
weight decay G
0
# params Sp
26M
Number of Fourier layers de, (dim. in Fourier blocks) hidden dim. of Fourier NN # layers in Fourier NN
64
256
N Wodess (# Fourier modes)
10
N Nodess, (# Fourier modes)
10
optimizer S
A
a (activation in S.))
GeLU
Q1 (initial learning rate So)
Ic-4
learning rate decay Sp
0.9999
weight decay Se
lc-4
minimum training sample size maximum training sample size # training samples batch size
10 10
64 64
200M
5000
mini batch size
64
T (number batches)
40k
6 (for Sinkhorn targets)
0.01
k (# Sinkhorn iterations for targets)
Code. Source code for UNOT, including the weights for the model used in the experiments, can be found at https: //github. hrggiaabb..cc/rreepeefoooorrrrte
28

---

Universal Neural Optimal Transport
D. Additional Experiments and Materials
D.I. Test Sets
In Figure 10 we show samples from our test datasets. For some the experiments in the appendix, we included two additional datasets, the "cars" class which is also from the Quick, Draw! dataset, and the Facial Expressions dataset (Hashan, 2022), which consists of 48 48448-dimensional greyscale images. The datasets are very diverse, and range in dimensionality from very low (MNIST) to fairly low (BEARS, CARS), medium high (CIFAR) and very high (EXPRESSIONS, LFW).
Figure 11 shows samples from our spherical datasets (where only part of the sphere is visible here). To create a grid on the sphere, we sample elevation angles 0 uniformly in
2 2. and azimuthal angles y uniformly in [o, 2 Concretely, we set
π 8,
j n
n.
2 = ..... n 1,
2
and form the grid {(0:, {o)) ).."" Each pair (0,, (pj) is mapped to point the sphere by = cos(O;) sin({;));
x = cos(O:) cos(p;))
z == sin(():)
MNIST
6
3
CIFAR
LFW
BEAR
CAR
EXPRESSIONS
Figure 10. Test dataset samples the unit square.
D.2. Comparison with Meta OT
We trained a Meta OT (Amos et al., 2023) network with the official GitHub implementation 12 and compared against UNOT on our test datasets, where we rescaled all datasets to 28 28, as Meta OT does not natively support inputs of varying sizes. In Table D.2, report the relative errors on the OT distance (in %) after single Sinkhorn iteration.
Phttps://ce / thub. com htt.:./cwc.fccoooorrreeeeeeeeeee
29

---

Universal Neural Optimal Transport
MNIST
CIFAR
LFW
BEAR
Figure 11. Test dataset samples on the sphere.
Table Relative Errors the OT distance (in after single Sinkhom iteration with UNOT's initialization, compared to Meta OT (Amos et al., 2023), the Gaussian initialization (Thornton & Cuturi, 2022), and the default initialization. Datasets rescaled to 28 x 28 such that Meta OT network can process them.
MNIST
CIFAR
MNIST-CIFAR
LFW
BEAR
LFW-BEAR
UNOT (ours)
2.7 2.4
1.3 1.1
2.8 2.6
1.5 1.3
2.0 1.6
1.8 1.3
MetaOT
1.8
23.1 15.7
1.44588
24.6
18.
31.0 + 14.8 19.3 6.4
Gauss
18.1 10.0
1.6
32.2
2.5
2.3
Ones
39.5 + 13.4
47.4 + 20.2
74.5 + 6.9
56.9
15.4
54.2 + 13.5
66.4 + 10.8
We see that UNOT outperforms Meta OT on all datasets except MNIST, which is be expected, as Meta OT explicitly trained on MNIST, while UNOT is not trained on any MNIST data. However, surprisingly, we see that UNOT almost matches Meta OT's performance MNIST, suggesting strong coverage of MNIST-like distributions by our generator network during training.
D.3. MLP-UNOT
We mention that in applications of fixed-size distributions, one replace the Neural Operator with MLP and achieve similar results for fraction of the training cost. We note that since the MLP acts fixed discrete space one does not need to have equispaced samples. In experiments, we found the MLP approach to also be very reliable for fixed-size inputs, and to vastly outperform the standard initialization of the Sinkhorn algorithm. Notably, can be trained in just a few minutes to relative errors below 5%.
30

---

Universal Neural Optimal Transport
D.4. Generalization across Resolutions
In this section, we show that UNOT successfully generalizes across resolutions. To this end, we downsample resp. upsample test datasets to resolutions between 10 10 and 64 64. Figure 12 shows the relative errors the transport distance over this range of resolutions after single Sinkhorn iteration, compared against the default and the Gaussian initializations. (In Section D.5, we also provide some results on upsampling the dimension of the data beyond 64 x 64, i.e. beyond the largest resolution that the network saw during training.) We see that UNOT generalizes very well across all resolutions between 10 x 10 and 64 x 64.
MNIST
CIFAR
BEAR
LFW
EXPRESSIONS
CAR
0.8
MNIST
0.6
UNOT
Gauss
Ones
0.4
0.2
0.0
0.8
BEAR CIFAR
0.8
0.4
0.2
0.0
0.8
0.6
Relative Error
0.4
0.2
0.0
0.8
LFW
0.8
0.0
CAR EXPRESSIONS
0.8
0.6
0.4
0.2
0.0
0.8
0.6
0.4
0.2
0.0
20
50
30
20
Resolution
Figure 12. Relative error the transport distance over the image resolution, ranging from 10 x 10 to 64 x 64.
31

---

Universal Neural Optimal Transport
D.5. Variable Epsilon
In this section, we provide experimental results on a variant of UNOT that also receives the parameter as input. Instead of the pair of measures (1,,2) encoded tensor of size (B, n, r2), sample epsilon randomly per sample from distribution with values
imput size of (B,3, n). where the third
channel is equal to everywhere. During training, between 0.01 and 1. Otherwise, training is identical to the training of regular UNOT. In Figure 13, we plot the relative errors over ranging from 0.01 to on the x-axis, and the resolution of the data ranging from 10 x 10 to 70 x 70 on the y-axis (where we downsample resp, upsample the data to these dimensions, cf. Section D.4; note that we still only trained on image resolutions between x and 64 64). This variant of UNOT seems to do surprisingly well across different values of and across wide range of resolutions, with relatively stable performance across different values of However, that when the resolution gets smaller than around 15 15, close 70 70, the error increases.
MNIST
CIFAR
BEAR
LFW
EXPRESSIONS
CAR
70
60
MNIST
50
40
30
20
10
-.14
70
60
CIFAR
50
40
-0.12
30
20
10
70
60
0.10
BEAR
50
40
CAR XXPRSSIONS LFW Resolution
30
20
10
70
60
50
40
0.06
30
20
10
70
60
0.04
50
40
30
20
0.02
10
70
60
40
30
20
20
0.01
0.1
0.01
0.1
1
0.01
0.1
0.01
0.1
0.01
0.1
1 0.01
0.1
Σ Regularization Parameter
Figure 13. Relative the transport distance, the resolution and varying values of e.
32

---

Universal Neural Optimal Transport
D.6. Generated Measures
Figure 14 shows images created by the generator. The generator creates very different images over the course of training. including highly structured distributions, large areas of mass, and distributions with mass concentrated in very small areas.
Figure Pairs training samples before and after 20%, 40%, 60%, 80%, and 100% training, from left to right (lighter=more mass). Top row; actual training images; bottom training samples visualized with smaller skip constant À to accentuate learned features.
In Table D.6, we also report the average OT distance error of samples created by the generator at various stages of training. can see that the generator indeed creates samples that are initially difficult, but that it quickly picks up on them, and by the end of training is capable of predictions for samples from all stages of training.
Table Relative on OT distance for samples created after 10, ... 70% of training. Errors for all samples computed the time of
their creation (i.c., after 10, 20, ...% of training) and at the end of training.
Error after ... % of Training
0%
10%
20 %0%
30% %
40 %
50 %
60%
70 %
At Generation
53.2%
3.1%
2.1%
1.6%
1.8%
1.7%
2.1%
1.9%
At End of Training
2.0%
1.6%
1.4%
1.1%
1.6%
1.5%
2.0%
1.9%
33

---

Universal Neural Optimal Transport
D.7. Additional Experiments
We provide additional results from our experiments in this section. In Table 6, we show the average Wasserstein-2 distance of barycenters computed by gradient descent using equation (10) to the true barycenter, where compute the gradient in equation (10) from the different initializations and single Sinkhorn iteration. Figures 15 and 16 show the relative error on the OT distance over Sinkhorn iterations for c(æ, y) || yll (on the square) and c(x, y) = arccos( (ax, y)) (on the sphere) resp., complementing Figure 3.
In Figure 17, we plot the relative error on the transport distance w.r.t. computation time when initializing the Sinkhorn algorithm with UNOT, and compare against the default initialization. We see that particularly on higher dimensional data, UNOT is significantly faster than Sinkhorn. However, interestingly, on MNIST the default initialization actually seems be faster. We note that these results heavily depend on the hardware used, and that did not optimize our FNO architecture for performance, more efficient architecture would probably lead to even more significant speedups. We have not included the initialization from (Thornton & Cuturi, 2022) in the plots, as it was very slow for us, even slower than the standard initialization, despite our best efforts to implement it as efficiently as possible. However, from (Thornton & Cuturi, 2022; Amos et al., 2023) it seems like the speedup should be somewhere between 1.Ix and 2x, depending on the dataset, which would make significantly slower than UNOT on most of our datasets. We mention again that FNOs process complex numbers, but PyTorch is heavily optimized for real number operations. With kernel support for complex numbers, UNOT will likely be much faster.
Finally, in Figures 18 and 19, we plot the marginal constraint violation (MCV), defined as
+ ||II,, - 1
(20)
2
for a transport plan II, again for single Sinkhorn iteration (Figure 18) and over iterations (Figure 19). The MCV measures how far the transport plan is from the marginals l and v. It is often used as a stopping criterion for the Sinkhorn algorithm, as the ground truth OT distance is unknown in practice. We compute the predicted transport plan for UNOT via equation (4).
Table 6. Average W2 distance from the predicted barycenter to the true barycenter on MNIST after 100 gradient steps.
W2 Distance
UNOT (Ours)
0.021 + 0.011
Gauss
0.033 1 0.018 0.057 = 0.034
Ones
MNIST (28x28)
CIFAR128x28)
MNIST-CIFAR,
0.30
0.6
0.25
0.5
0.25
0.20
0.4
0.20
0.3
0.2
0.10
Relative Error
0.05
0.1
0.05
0.00
0.0
0.00
LFW(64x64)
BEAR (64x64)
WWW-BEAR,6AA
0.6
0.5
0.25
0.25
0.4
0.20
0.20
0.3
6.1
0.05
0.65
6.0
0.00
0.00
0
10
20
30
40
50
0
10
20
30
40
50
10 20
30
40
50
Iteration
Iteration
Iteration
UNOT
Ones
Figure 15. Relative Error on the OT distance on the unit square with c(æ; y) default one, over number of Sinkhorn iterations. Note the y-axis has been rescaled for CIFAR and LFW to fit the curve for the default initialization, and that the Gaussian initialization does exist for the Euclidean cost function.
|| bl for the UNOT initialization compared to the
34

---

Universal Neural Optimal Transport
MNIST (N85T(28x28) (28x28)
(2x2R
MNIST-CIFAR22ak28
0.3
0.1
Relative Error
LFW(64x64)
BEAR(64X64)
LFW (64x6)))(64x64)
0.3
0.1
20 10 50
10 20
40
10 30 40 Iteration
Iteration
Iteration UNOT
Ones
Figure 16. Relative Error on the OT distance on the unit sphere with c(æ, y) arcoos((2:: y)) for the UNOT initialization compared to the default one, over number of Sinkhorn iterations. Note that the Gaussian initialization does exist for the spherical cost function.
MNIST
MN1ST28x28 28x28
CIFAR28x28
WNIST-CIFAR28x28
0.8
0.4
0.4
0.2
Relative Error
0.2
0.0
0.000
0.005
0.010
0.000
0.005
0.010
0.000
0,005
0.010
LFW 64x64
BEAR64x64
LFW-BEARs4x6A
0.4
0.4
0.2
0.2
0.2
0.0
0.0
0.0
0.00
0.01
0.02
0.03
0.00
0.01
0.02
0.03
0.00
0.01
0.02
0.03
Time (s)
Time (s)
Time (s)
UNOT
Ones
Figure 17. Comparison of relative emoor the transport distance over computation time in seconds. Evaluated on an NVIDIA 4090. The z-offset the UNOT curves corresponds the time needed for forward pass through Sg.
35

---

Universal Neural Optimal Transport
Ones
Marginal Constraint Violation
0.6
Gauss
UNOT
0.5
0.4
0.3
0.2
0.1
0.0
MNIST
CIFAR
MNIST CIFAR
LFW-BEAR
LFW
BEAR
28x28
28x28
28x28
64x64
64x64
64x64
0.8
Ones
Marginal Constraint Violation
UNOT
0.7
0.6
0.5
0.4
0.3
0.2
0.1
0.0
MNIST
CIFAR
MNIST CIFAR
LFW-BEAR
LFW
BEAR
28x28
28x28
28x28
64x64
64x64
64x64
Ones
0.7
Marginal Constraint Violation
UNOT
0.6
0.5
0.4
0.3
0.2
0.0
MNIST
CIFAR
MNIST CIFAR
LFW-BEAR
LFW
BEAR
28x28
28x28
28x28
64x64
64x64
64x64
Figure 18. Average marginal constraint violation (sec eq. (20)) after single Sinkhorn iteration, for the unit square domain with c(x, y) = || yl 2 (top) and c(x, y) = || - y/ (middle), and the unit sphere with c(x,, y) = arccos( (xc, y)) (bottom). Note that the Gaussian initialization exists only for the squared Euclidean distance cost.
36

---

Universal Neural Optimal Transport
(2x27
22
MNIST-CIFAR
0.3
Marginal Constraint Violation
0.2
0.1
0.0
(64x64)
(EAR(64x64)
LFW-BEAR (64x64)
0.3
0.2
0.1
0.0
lteration
Iteration Gauss
Iteration
UNOT
Ones
MNIST (28x28)
(262R
MIIST-CIFAR
0.3
Marginal Constraint Violation
0.1
LFW (64x6))
BEAR(64x64)
LFW FFW-BEAR(4xx65 (64x64)
0.3
0.2
0.1
0.0
10
30
50
0
20
30
50
20
40
50
Iteration
Iteration
Iteration
UNOT
Ones
MNIST (28X2(28x28)) (28x28)
CIFAR
MNIST-CIFAR,28.288
Marginal Constraint Violation
0.2
0.0
LFW(64x64)
BEAR
(EW
0.3
0.2
0.1
0.0
20
10
50
10 20
40
10 30
40
Iteration
Iteration
Iteration
UNOT
Ones
Figure Average marginal constraint violation (see eq. (20)) over number Sinkhorn iterations, the unit square domain with <(æ, y) = | y (top) and c(x, y) I|æ yll (middle), and the unit sphere with c(æ, y) = arccos((a, v)) (bottom). Note that the Gaussian initialization exists only for the squared Euclidean distance cost.
37


Reviews:
[{"review_id": "SG6ITQDnKw", "reviewer": "Reviewer_JPz7", "summary": "The paper suggests UNOT (Universal Neural Optimal Transport) a method for (single forward-pass) prediction of OT distances, which are typically computed by iterative methods like the Sinkhorn algorithm. The network architecture is based on Fourier Neural Operators (FNOs), that provide discretization-invariance, which is useful for dealing with discrete distributions of different dimensions. Training is based on a novel GAN-like adversarial scheme, where a generating network is made to provide challenging instances, which the prediction network is optimized over using a self-supervised bootstrapping loss, theoretically proven to minimize the ground truth loss. Extensive experiments show that UNOT outperforms prior methods in several aspects of the problem, including quality of approximation, especially as an initialization to Sinkhorn iterations, but also when considering the implied geodesics and barycenters in the Wasserstein space.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thorough and positive review!\n\nThe plots in Figures 3, 13, and 14 indeed got mixed up, thanks for catching! We will also include more details about FNOs and our architectures in the main text and define $R$ in the generator more clearly.\n\n**Figure 5:** the corners should not be blurry; we will fix the plot.\n\n**Figure 6:** \"Ground truth\" are the input images at times $t=0$ and $t=1$. Each row corresponds to a different transport plan between the ground truth images which transports the ground truth image from time $t=0$ to the predicted image at time $t=1$. This means the interpolation at time $t=0$ is equal to the ground truth by design (at $t=0$, the images should be identical, we will fix the plot), but the interpolation at time $t=1$ can be different from the ground truth if the transport plan is not accurate. This is why we also provide the ground truth for $t=1$. We hope this answers your question.\n\n>**The focus on images [...] how general is the solution? Have you tested it on other domains [...]?**\n\nWe note that the model is “agnostic” to the data modality and should work with data of any modality with the same cost function and grid structure. We ran additional experiments in the Euclidean image setting with different costs ($L_1$ and $L_2$, which are metrics unlike the squared Euclidean cost $L_2^2$ from the paper), and out-of-the-box (without any hyperparameter finetuning) it matches the relative errors on the transport distance for $L_2^2$. We also ran an additional experiment with the **unit sphere in $\\mathbb{R}^3$ as the domain, where we parametrize the network with Spherical FNOs** (https://arxiv.org/abs/2306.03838), and use the spherical distance $c(x,y)=\\text{arccos}(<x, y,>)$ as the cost. Due to time constraints, we have only tested this on our image datasets (where we define an angular grid on the sphere, and then simply map the images onto this grid). Without any hyperparameter tuning, this works very well out-of-the-box and **even outperforms Euclidean UNOT** averaged over the test datasets. With proper hyperparameter tuning, these results can probably be even improved. For the camera-ready version, we will test UNOT on real-world spherical datasets. \n\n>**The choice of adversarial training. Have you considered alternatives to this? [...] Did you encounter stability issues in training [...]?**\n\nInitially, we tried training on synthetically generated data which works, but the adversarial approach works better. Additionally, synthetic data requires significant “fine-tuning” to capture the space of distributions well. Training on a single dataset would destroy generalization capabilities (cf. the comparison between UNOT and MetaOT in response to reviewer **dxme**). This shows that a general-purpose model requires our universal training approach. (However, it is probably possible to fine-tune a pre-trained UNOT model on downstream datasets.)\nWe have not encountered any stability issues during training, and the training loss, as well as the test losses, decrease very stably during training. One stabilizing factor compared to GANs could be that our generator has a skip connection, s.t. its outputs can never completely collapse.\n\n>**I find the explanation as to [the generated pairs’] appearance [...] unclear. Why does it look like it does? [...] Does the predictor eventually solve these instances well? Does it suffer from forgetting?**\n\nSince our generator has a skip connection, part of the appearance of the distributions is indeed due to the architecture. However, the learned network in the generator ensures that training samples are generated in such a way that they are difficult for the predictor. We ran additional experiments which show that **a) generated samples are indeed difficult for the predictor initially, b) over the course of training, the predictor eventually solves them, c) forgetting of training data seems not to be happening.** Specifically, at the start of training, as well as after each 10% of training, we save a training batch and track the relative error on the OT distance on each. The following table shows the initial error on these samples as well as the final error on them at the end of training (with stable improvements on all of them during training).\n\n| Relative OT Distance Error | 0% Training | 10% Training | 20% Training | 30% Training | 40% Training | 50% Training | 60% Training | 70% Training |\n|-|-|-|-|-|-|-|-|-|\n| At Generation| 53.2%| 3.1%|2.1%|1.6%|1.8%|1.7%| 2.1%|1.9%|\n| At End of Training| 2.0%| 1.6%|1.4%|1.1%| 1.6%| 1.5%|2.0%|1.9%|\n\nWe also experimented with keeping a cache of previous training samples and re-feeding them to the model over the course of training to prevent potential forgetting, similar to what is typically done when fine-tuning pre-trained models. However, this did not improve our training, probably because there does not seem to be any forgetting in the first place.\n\nWe hope this answers all your questions!", "claims_and_evidence": "Yes. The work is well motivated, since as claimed, the use of OT is constantly growing in the field of ML, and the iterative natured computations are known to be expensive, especially for large instances of the problem, for which alternative approximations (such as projections to lower-dimensional spaces, or use of closed-form OT formulations) are typically used to save time. Therefore, good approximations, or ones that can be used for starting local optimization, are of high demand.\nThe proposed method indeed shows excellent results, over the different datasets and usage cases, that make significant progress with respect to the goals that were set.", "methods_and_evaluation_criteria": "The solution is well explained and its merits are demonstrated very clearly, with clear improvements across the experimentation.\nThe paper overall is very well written, with a very good separation between the essential parts in the main paper vs. the technicalities that are deferred to the appendix.\nThe solution proposed is elegant, with several new ideas that provide independent contributions. The choice of predicting the dual potential, from which everything else can be recovered, is very natural. The way that the supervision is achieved by running several Sinkhorn iterations over the prediction is a good idea, that is also justified theoretically. \nThere are some design choices (in both method and experimentation setup) that are well explained in the paper, but the choices (among possible alternatives) are not sufficiently discussed and justified. \nOne main point I wonder about is why the method is demonstrated only on the domain of images and as a consequence - How much is it suited (or limited) to this setting (where e.g. ground distances have a very particular structure).\nAnother is the choice to use the adversarial type training over synthetic data. Such an approach has typically stability issues in training, which are a downside of methods like GAN, where specific attention is required to avoid different failures (such as mode collapse). Additionally, it is unclear how the generated samples are related to natural image distributions and therefore how efficient the training and generalization is. Even when inspecting the generated examples - it is hard to understand why they appear as they do and how they develop during training.", "theoretical_claims": "I didn't fully check the proofs. But the appendix does contain the necessary formulated background and proofs of all claims made in the paper. \nOne claim that I have not comfortably understood is regarding the universality of the generator. I understand that any input can be generated, due to the positive density of the generated distribution across the domain. But this density can be extremely non-uniform in practice and furthermore, the adversarial training strategy limits the ability to cover the space even further.", "experimental_designs_or_analyses": ".", "relation_to_broader_scientific_literature": "The contributions are quite general and will have an impact on related research. \nThe proposed method follows a line of works that try to improve Sinkhorn initialization and sets a new standard in this respect. \nOther contributions, such as the theoretical understandings that take care of the discretization issues (with the practical FNO based solution) and the bootstrapping based loss, should be of general interest for designing and training neural networks in this domain.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": ".", "other_comments_or_suggestions": "I would try to clarity (and slightly expand) the description of the architecture (both the generator and FNO), since the current one is somewhat cryptic and I need to refer to the appendix for clarifications. These are main components of the solution and should be more explicit.\nRegarding FNOs - a few sentences on the idea itself (lifting and operations in Fourier space) rather than only referring to the discretization invariance property. Regarding the generator - \" R denotes renormalizing to two probability measures and downsampling them to random dimensions in a set range\" is not very clear, even though it is correct.\n\nFigure 3 and 13/14 - seems like you switched \"Ones\" and \"Gauss\"\n\nFigure 5 - Is it intentional that 3 out of the 4 corner images are blurry. It would be better to interpolated between sharp images.\n\nFigure 6 - What is \"ground truth\"? Only for t=1? And why are don't the interpolations coincide with the inputs at the endpoints (t=0/1)?", "questions_for_authors": "I am very positive about the paper. Given clarifications to the following questions, I might reconsider my score.\n1. The focus on images. I think that it is a very good domain for explaining and demonstrating the method. But how general is the solution? Have you tested it on other domains (especially ones with different ground costs - ones that are not Euclidean, not a metric)?\n2. The choice of adversarial training. Have you considered alternatives to this? Wouldn't training on the (real) images of a single particular dataset (or multiple ones) give better results (perhaps at the cost of generalization)? Did you encounter stability issues in training (e.g. dying/vanishing gradients, mode collapse) and if so what did you do to avoid them? Do the loss dynamics follow those typical in GANs?\n3. Generated training distributions. The example generated pairs are nicely visualized, but I find the explanation as to their appearance and its evolution very unclear. Why does it look like it does? Is it a result of the specific architecture? Or are these actually instances that are difficult for the predictor? Does the predictor eventually solve these instances well? Does is suffer from forgetting? (e.g. solve the first examples at their time, but fail to do so later on)", "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": "7i9Jmqai3h", "reviewer": "Reviewer_dxme", "summary": "The paper presents UNOT- a method to learn universal discrete OT plans/potentials/costs approximator. Interestingly, the method may deal with different discrete resolution scales simultaneously by specific parameterization of the universal learned OT potential (through FNO operator). Several experiments on images (considered as discrete 2d distributions of different resolutions) validate the proposed methodology.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thorough and helpful review! In the following we will try to answer your questions.\n\n>**I do not understand the practical implication of section 4.4. [...] we can do the same by adding noise to the image**\n\nSection 4.4 shows that UNOT can _solve the OT problem between distributions over images_, which e.g. arises in training of generative image models. This means UNOT can a) optimally match the images in the two marginals; b) interpolate the transport along this optimal trajectory, which essentially solves the discretized Wasserstein gradient flow w.r.t. the functional $W_2^2$. Neither marginal has to be noise; we will add another figure where both marginals are image distributions. We hope this clarifies that Section 4.4 is fundamentally different from adding noise to images (adding noise does not solve OT, and in particular, cannot interpolate between two non-noisy marginals).\n\n>**Table 2: time for LFW: 10^-2; time for BEAR: 8\\*10^-3. Then, time for LFW-BEAR should be ~9*10^-3?**\n\nThe hardness of solving OT between two datasets (LFW and BEAR in this case) does not interpolate between the hardness of each of the individual datasets. Imagine one of the two datasets consisted only of Dirac measures, then solving OT between this dataset and any other dataset is also trivial. Hence, the time for LFW-BEAR need not equal 9*10^-3.\n\n>**I think the authors should add MetaOT [Amos et. al., 2022] to the comparison**\n\nWe did not include a comparison to MetaOT as it is inherently restricted to a single dataset of fixed dimension, but for completeness, we trained MetaOT on MNIST with the official implementation (https://github.com/facebookresearch/meta-ot) and compared it against UNOT on our test datasets (rescaled to 28*28 s.t. MetaOT can process them). Here are the relative errors on the OT distance (in %):\n\n| |MNIST| CIFAR|MNIST-CIFAR|LFW|BEAR|LFW-BEAR|\n|-|-|-|-|-|-|-|\n| **UNOT** | 2.7 ± 2.4 | 1.3 ± 1.1 | 2.8 ±2.6 | 1.5 ± 1.3 | 2.0 ± 1.6 | 1.8 ± 1.3 |\n| **MetaOT**| 2.4 ± 1.8   | 23.1 ± 15.7 | 11.4 ± 5.8 | 24.6 ± 15.7 | 11.8 ± 8.3 | 31.0 ± 14.8 |\n| **Gauss** |18.1 ± 10.0|19.7 ± 7.6 | 32.2 ± 8.7 | 21.1 ± 6.5 | 20.4 ± 8.3 | 19.3 ± 6.4 |\n| **Ones** | 39.5 ± 13.4 | 47.4 ± 20.2 | 74.5 ± 6.9 | 56.9 ± 15.4 | 54.2 ± 13.5 | 66.4 ± 10.8 |\n\nSurprisingly, **UNOT is almost on par with MetaOT on MNIST, despite _never having seen a single MNIST sample during training_, whereas MetaOT was _only_ trained on MNIST**, and we see that MetaOT breaks down on all other datasets. More on UNOT vs MetaOT in our response below.\n\n>**How can we use UNOT? Importantly, the UNOT is trained for a fixed cost [...] I think that [...] UNOT has no advantages compared to MetaOT**\n\nWe provide the trained UNOT model in our GitHub (which we cannot link here for anonymity reasons), so it can be used out-of-the box as a general-purpose OT solver without any additional training required. You are right that the model is trained for a fixed cost. However, the squared Euclidean cost function is the most widely used cost function in practice (it also gives rise to $W_2$). We also refer to our response to reviewer **JPz7** for experiments with other costs and on spherical domains (we will provide these models on GitHub once they are ready).\nComparing UNOT to MetaOT, we note that MetaOT can only be trained on a single dataset of fixed dimension. This means the downstream task needs to be known in advance, and crucially, _a training dataset needs to exist for this task_. For a new task, a new model needs to be trained from scratch. UNOT, on the other hand, can be applied to any downstream dataset (for a given cost) once trained, and it even matches MetaOT in terms of performance on the MetaOT training dataset (cmp. our response above).\n\n\n\n>**[...] this speedup in [Amos et. al., 2022] corresponds to MNIST dataset, where your method has a deterioration compared with “ones” initialization. Also [...] your reported time [..] with “ones” [...] is smaller**\n\nRuntimes depend on various factors (e.g. hardware) and are difficult to compare.\nAlso, our FNO uses complex (number) layers, which are known to be sub-optimally implemented in PyTorch. With faster kernels we expect a drastic speedup. Since the submission, **we have also optimized UNOT with JAX and a better architecture, achieving a 2.5x speedup over the previous version without loss in performance.**\n\n>**“... any pair of discrete probability measures $(\\mu,\\nu)$ can be generated by $G_\\theta$.\" - this is not quite right**\n\nThank you for pointing this out, we’ll improve the wording!\n\n>**I found it strange that no information on FNO $S_\\phi$ is given in the main text**\n\nYes, a more detailed description of $S_\\phi$ should be given in the main text, we’ll add this!\n\n>**Why do you need the generator network $G_\\theta$ [...]? Why not just sample $\\mu$ and $\\nu$ at random [...]?**\n\nWe refer you to our response to the same question by Reviewer **JPz7**.\n\nWe hope this answers all your questions!", "claims_and_evidence": "Ok", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "I didn’t check carefully the proofs, but the results seems to be reasonable", "experimental_designs_or_analyses": "For the problem on hand, the experimental section seems to be comprehensive enough. However, I have several comments. \n\n**General**:\n1. I do not understand the practical implication of section 4.4. Figure 7 shows that, given an image, we can transform from noise to this image. Actually, we can do the same by adding noise to the image…\n2. Table 2: time for LFW: $10^{-2}$; time for BEAR: $8.0\\cdot10^{-3}$. Then, time for LFW-BEAR should be ~$9\\cdot10^{-3}$?\n\n**UNOT vs. MetaOT**\n1. I think the authors should add MetaOT [Amos et. al., 2022] to the comparison. Actually, this method could be used in the all provided experimental setups (except, probably, 4.1 with mixed distributions mnist + cifar/lfw+bear: but even in this case we can just learn two MetaOT models with dimensions 28x28 and 64x64)\n2. Quite a strange phrase (line 318): “the relative speedup achieved in [Amos et. al., 2022]” is 1.96. But this speedup in [Amos et. al., 2022] corresponds to MNIST dataset, where your method has a deterioration compared with “ones” initialization. Also, it is interesting that your reported time for MNIST with “ones” initialization is smaller compared to [Amos et. al., 2022].", "relation_to_broader_scientific_literature": "The closest approach for UNOT is MetaOT, It would be good to have a detailed comparison with this approach (see “experimental design and analysis” section). Note on the related works section: a good work on OT in flow matching: [1], an interesting alternative approach for Neural OT: [2].\n\n[1] Optimal Flow Matching, NeurIPS’24\n\n[2] On amortizing convex conjugates for optimal transport, ICLR’23", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "The paper is generally good written. However, to be honest, I do not see serious practical applications of the proposed method. How can we use UNOT? Importantly, the UNOT is trained for a fixed cost matrix. Also, I think that in the vast majority of situations, UNOT has no advantages compared to MetaOT, while being more time-consuming for training.", "other_comments_or_suggestions": "1. lines 43-45, first column: line 44 is skipped. \n2. Line 179, second column: “... any pair of discrete probability measures ($\\mu$, $\\nu$) can be generated by $G_{\\theta}$.’’ - this is not quite right, because $G_{\\theta}$ can generate only discrete distributions with dimensions $n, m$ such that $n + m \\leq d’$.\n3. I found it strange that no information on FNO $S_{\\phi}$ is given in the main text. I think, a quick introduction of how they are constructed (in section 3.2.) would enhance the flow of the manuscript", "questions_for_authors": "Why do you need the generator network $G_{\\theta}$, and, subsequently, adversarial max-min objective? Why not just sample $\\mu$ and $\\nu$ at random, treating this as Meta distribution, similar to MetaOT [Amos et. al., 2022], in eq. (6)?", "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": "1k0J8UCRae", "reviewer": "Reviewer_rgYo", "summary": "This paper introduces the universal neural optimal transport (UNOT) solver, a framework designed to efficiently solve entropic-regularized optimal transport (OT) problems. Unlike existing neural OT methods that can only handle input distributions of a fixed dimension, UNOT leverages Fourier neural operators (FNOs) to predict transport plans of variable resolutions, enabling generalization across datasets and input dimensions.\n\n## Update after rebuttal\n\nI have raised my score after the authors fix the presentation of the theorem, but I still think the significance of Theorem 3 is limited, and the paper needs to be restructured to better highlight its novelty in other aspects.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your helpful review!\n\nWe will include definitions of $\\rho(x)$ and $\\mathcal{N}(\\cdot |0,I)$ in Theorem 3, thank you for catching that. Otherwise, your main concern seems to be regarding the proof of Theorem 3. **The statement of Theorem 3 is indeed correct (if you interpret $G^{-1}(x)$ as a particular point in the preimage, see below).** However, this is only immediately clear for $x\\in\\mathbb{R}^d_{>0}$, as opposed to $x\\in\\mathbb{R}^d_{\\ge 0}$ as stated in the paper; indeed, your counterexample shows that $G_\\theta$ is not invertible if $x=0$ (or at least one entry in $x$ equals $0$). Hence, we could simply rewrite Theorem 3 for $x\\in\\mathbb{R}^d_{>0}$ instead of $x\\in\\mathbb{R}^d_{\\ge 0}$ (on this set, $G$ is indeed invertible), which would be sufficient for our purposes as we use the Sinkhorn algorithm to generate targets from the generated samples, and the inputs to the Sinkhorn algorithm need to be positive everywhere. However, it is also straightforward to adapt the proof of Theorem 3 to work for any non-negative $x\\in\\mathbb{R}^d_{\\ge 0}$:\n\nLet $\\tilde{G}(z)=\\text{NN}(z)+\\lambda z$ (we drop the subscript $\\theta$ for ease of notation here), i.e. $\\tilde{G}$ is equal to $G$ without the ReLU activation. Denote by $\\rho_{\\tilde{G},\\rho_z}$ the density of $\\rho_z$ pushed forward by $\\tilde{G}$ (this interface won't render the pushforward symbol '#' correctly, hence we slightly adapt the notation compared to the paper). Then the exact same proof as for Theorem 3, but applied to $\\tilde{G}$, yields\n$$\n\\rho_{\\tilde{G},\\rho_{z}}(x)\\ge \\frac{1}{(L+\\lambda)^d}\\mathcal{N}\\left({\\tilde{G}}^{-1}(x)|0,I\\right)\n$$\nfor _any_ $x\\in\\mathbb{R}^d$, where we get invertibility of $\\tilde{G}$ by Theorem 1 in [1] (cmp. our proof of Theorem 3 in Appendix B). Now for any non-negative $x\\in\\mathbb{R}^d_{\\ge 0}$, we clearly have\n$$\n\\rho_{G,\\rho_{z}}(x)\\ge \\rho_{\\tilde{G},\\rho_{z}}(x),\n$$\nas for any $z$ with $\\tilde{G}(z)=x$, we also have $G(z)=x$. Combining these two inequalities yields\n$$\n\\rho_{G,\\rho_{z}}(x)\\ge\\frac{1}{(L+\\lambda)^d}\\mathcal{N}\\left({\\tilde{G}}^{-1}(x)|0,I\\right)\n$$\nfor any non-negative $x\\in\\mathbb{R}^d_{\\ge 0}$. **This shows that the statement in Theorem 3 is indeed correct, but there could be many points in the preimage $G^{-1}(x)$, and we need to pick the unique $\\tilde{G}^{-1}(x)$ amongst those for the statement to hold.**\nThank you for pointing out this inaccuracy, and we will adapt the statement and proof accordingly in the updated version of the paper, and also update the wording of Theorem 3 to make it more clear that $\\tilde{G}$ is indeed invertible under the assumptions of Theorem 3. (Corollary 4 will be updated accordingly as well.)\n\nAs this seems to have been your only concern, we would kindly ask you to reconsider your score.\n\n[1] Behrmann, J., Grathwohl ,W., Chen, R.T.Q., Duvenaud, D., and Jacobsen, J.-H. Invertible Residual Networks. Proceedings of the International Conference on Machine Learning, 2019. doi:10.48550/ARXIV.1811.00995. URL https://arxiv.org/abs/1811.00995.", "claims_and_evidence": "I suspect that some of the claims in the paper are not well supported by convincing evidences, especially in the theoretical analysis. Please see the \"Theoretical Claims\" section below for details.", "methods_and_evaluation_criteria": "The numerical evaluation seems reasonable to me.", "theoretical_claims": "I have checked the proofs of the theorems, and I feel some of the claims may be incorrect. In Theorem 3, the inequality involves the term $G_{\\theta}^{-1}(x)$, and I suppose it stands for the inverse of the mapping $G_{\\theta}:\\mathbb{R}^d\\rightarrow\\mathbb{R}^d$. So basically the theorem implicitly claims that $G_{\\theta}$ is invertible, but we can get a simple counterexample. By construction, $G_{\\theta}$ is the output of a ReLU activation, so by carefully choosing $z$, $\\mathrm{NN}\\_\\theta(z)+\\lambda z$ can be made negative, and then the output of $G_{\\theta}(z)$ would be a zero vector. By slightly perturbing the value of $z$ to $z'\\neq z$, $G_{\\theta}(z')$ will still be a zero vector, which shows that $G_{\\theta}(z)$ is not invertible.", "experimental_designs_or_analyses": "The experimental design seems reasonable.", "relation_to_broader_scientific_literature": "The proposed UNOT method may be useful for machine learning models that involve OT computation.", "essential_references_not_discussed": "None.", "other_strengths_and_weaknesses": "None.", "other_comments_or_suggestions": "I think some of the notation can be made clearer and more precise. For example, in Theorem 3, the meaning of the function $\\rho(x)$ is unclear, although it can be inferred that it may stand for the density function. Similarly, it should be made clear that $\\mathcal{N}(\\cdot|0,I)$ stands for the density function of the standard normal distribution.", "questions_for_authors": "I suggest the author(s) checking the theoretical claims and fixing any potential flaws in the proofs.", "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": "SqtNkGQGin", "reviewer": "Reviewer_xg8H", "summary": "The authors introduced UNOT (Universal Neural Optimal Transport), a novel framework designed to efficiently predict entropic optimal transport distances and plans between discrete measures of varying resolutions. Motivated by the universal domain adaptation methodology, they utilized Fourier Neural Operators, which are capable of processing inputs of different sizes, allowing UNOT to generalize across diverse datasets.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thorough review! We will address your concerns in the following.\n\n>**Proposition 5 [...] depends on assumptions that may not hold**\n\nThe only assumption we make in Proposition 5 is that $\\sum_i g_i=\\sum_i g_{\\phi_i}=\\sum_i g_{\\tau_{k_i}}=0$. Since optimal potentials of the dual OT problem can be shifted by constant factors, we simply choose $g$ and $g_{\\tau_k}$ to have zero sum. We could also shift $g_\\theta$ to have zero sum in practice. However, we have updated the proof and show that the same result with the same constant holds _without_ assuming that $\\sum_i g_{\\phi_i}=0$, i.e. only assuming $\\sum_i g_i=\\sum_i g_{\\tau_{k_i}}=0$.\n\n\n>**\"0.01 relative error on the OT distance\" is not the best metric [...] the classic FID [...] is [...] more reasonable**\n\nThe purpose of UNOT is to predict OT distances with high accuracy. Hence, using the OT distance itself for evaluation seems reasonable, and is in line with other works. FID, however, is usually used to evaluate the quality of generated distributions for generative image models. Since UNOT does not generate images, we do not see how FID could be applied in our setting, but if you could explain what you mean, we are happy to include it.\n\n\n>**The selected datasets are toy and low-dimensional**\n\nWhile the datasets are toy datasets, we believe they are not low-dimensional. We evaluate on $64\\times 64$-dimensional images, which corresponds to distributions in $4096$ dimensions (note, in particular, that in this case the cost matrix and the transport plan live in $4096^2\\approx 17M$ dimensions). In comparison, MetaOT only trains and tests on $28\\times 28=784$-dimensional images, more than $5$ times smaller than us. We have also run experiments on $128\\times 128=16384$-dimensional data (here the cost matrix is already ~$268M$-dimensional), and the results look very promising so far, with accuracy almost recovering that of the $64\\times 64$-dimensional setting. We also refer to our response to Reviewer **JPz7** regarding additional experiments with varying costs and on the unit sphere domain, and we will include real-world Euclidean and non-Euclidean test datasets in the evaluation.\n\n>**the authors did not provide a comparison with the universal domain-adaptation methods**\n\nWe assume that you mention this paper because their method also features an adversarial training objective between a predictor and two discriminators (equation (4) therein). However, our objective is more closely related to that of GANs, and the similarities to the paper you link are fairly limited as far as we can tell.\n\n>**The calculation of Wasserstein barycentres and their accuracy was only compared visually**\n\nWe did not include quantitative results as we thought they were difficult to interpret, but we have added them here for completeness and can add them to the camera-ready version too. We report the average $W_2$ distance from the predicted barycenter (after 1 Sinkhorn iteration) to the true barycenter on MNIST.\n\n|||\n|-|-|\n| **UNOT**|0.021 ± 0.011|\n|**Gauss**|0.033 ± 0.018|\n| **Ones**|0.057 ± 0.034|\n\n\n>**The model's performance may degrade with significantly higher resolutions than those encountered during training**\n\nWe evaluate UNOT in much higher dimension than e.g. MetaOT, and we have conducted additional experiments in higher-dimensional settings ($128\\times 128$) with very promising results, suggesting that UNOT can be scaled up. However, you are correct that the performance may degrade with significantly higher-dimensional inputs than those seen during training, but this is not due to an inherent limitation of our approach, but rather an impossible learning task. If a model only sees up to $n\\times n$-dimensional distributions during training, it is impossible to accurately learn $(n+1)\\times (n+1)$ distributions (you can always come up with two different $(n+1)\\times (n+1)$ distributions that look identical when rescaled to $n\\times n$). Empirically, however, UNOT performs well on inputs up to 10% larger than the largest distributions it has seen during training without loss in performance.\n\n>**While UNOT is claimed to generalize across datasets, the paper does not provide sufficient evidence for its effectiveness in highly diverse or out-of-distribution scenarios**\n\nWe respectfully disagree, as we believe our test datasets are highly diverse. Not only are the individual datasets very different (compare, e.g., BEAR with low intrinsic dimension to LFW with high intrinsic dimension, cf. Figure 11), but by adding cross-datasets such as BEAR-LFW, we cover a very diverse subset of the space. Furthermore, our test datasets cover a wide range of dimensions ($28\\times 28$, $48\\times 48$, and $64\\times 64$). UNOT performs very well on _all_ test datasets, suggesting that its generalization capabilities are strong. Also note that _all_ our test datasets are out-of-distribution, as the model does not see any of them during training.\n\nWe hope this answers your questions!", "claims_and_evidence": "The paper claims that the proposed self-supervised bootstrapping loss minimize the ground truth loss. However, the theoretical justification is somewhat vague, and while Proposition 5 suggests a relationship between minimize the bootstrapping loss and the ground truth loss, it depends on assumptions that may not hold in all scenarios.", "methods_and_evaluation_criteria": "The evaluation criteria is questionable. \"0.01 relative error on the OT distance\" is not the best metric because it is very sensitive to the preturbations. On the image domain to calculate the classic FID distance is a more reasonable choice.", "theoretical_claims": "Theoretical claims and supporting evidence are correct and provide the necessary information about the method.", "experimental_designs_or_analyses": "The selected datasets are toy and low-dimensional, and the authors did not provide a comparison with the universal domain-adaptation methods https://openaccess.thecvf.com/content_CVPR_2019/papers/You_Universal_Domain_Adaptation_CVPR_2019_paper.pdf. The calculation of Wasserstein barycentres and their accuracy was only compared visually with the ground truth barycentres.", "relation_to_broader_scientific_literature": "The paper is based on the adaptation of the well known Neural Optimal Transport framework and universal domain adaptation methods.", "essential_references_not_discussed": "The related work section covers the research area.", "other_strengths_and_weaknesses": "The model's performance may degrade with significantly higher resolutions than those encountered during training, indicating a potential scalability issue that needs to be addressed. While UNOT is claimed to generalize across datasets, the paper does not provide sufficient evidence for its effectiveness in highly diverse or out-of-distribution scenarios", "other_comments_or_suggestions": "None", "questions_for_authors": "None", "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:

Meta-review: The paper proposes a one-shot prediction method for optimal transport (OT) distances called UNOT. By applying Fourier Neural Operators, the authors enable the UNOT to work with discrete distributions of varying dimensions. This method gives a state-of-the-art approximation of initial distributions for Sinkhorn iterations, thereby reducing the computational complexity of the Sinkhorn algorithm in high-dimensional cases. Moreover, the authors show that UNOT also captures Wasserstein geometry via barycenters, and approximates geodesics through both barycenters and predicted transport plans.

Regarding the reviewer JPz7, most of the questions and concerns were addressed during the rebuttal phase. JPz7 notes that such tasks as using initializations for solving OT are in demand and contributions are general, mentioning that the idea follows previous works aimed at improving Sinhorn initializations. The authors responded to all of the questions and concerns of the reviewer, including issues with domain preferences in experiments, choice of adversarial training, and the understanding of generated samples. Most of these responses were supported with solid evidence. However, I found that one questions about adversarial training and other domains were addressed without experimental evidence.

The reviewer dxme primarily highlights two issues: the lack of comparison with related work, MetaOT [Amos et. al., 2022], and the limited range of applications. The former concern was addressed by the authors with additional experiments comparing MetaOT [Amos et. al., 2022] and the proposed method. By the table, we can conclude that such a universal approach indeed works comparably, even without being trained on MNIST. Obviously, on out-of-domain datasets, MetaOT [Amos et. al., 2022] gives worse results than the proposed method. The latter concern is more substantial. The reviewer questions whether this method could be applied as the OT solver for mini-batch flow matching setups [Pooladian, A.-A., et al., Tong, A., et al.]. Despite the authors’ explanation of this use case, dxme remains unconvinced. Nevertheless, the authors included additional examples of UNOT application, which, in my opinion, adequately address the reviewer’s concerns. In this case, I believe the issue arises from the unique formulation presented in Section 3, which differs from the formulations commonly used in other articles addressing mini-batch applications of OT. As the author, I do not see a fundamental problem with the discussed application.

The third reviewer, rgYo, had only one question concerning the application of Theorem 3. Despite the response, the reviewer ended up convinced that the significance of Theorem 3 is limited. I agree to some extent that the theorem could benefit from a clearer interpretation to aid understanding. However, this appears to be more of a presentational issue rather than a flaw in the proposed method itself. 

Finally, in my opinion, the authors successfully addressed all the concerns raised by the reviewers. Only minor textual/presentational issues remain, which do not seem to require substantial rewriting. Moreover, it is important to note that the issues that appeared across the reviews are mostly unique, such that both the experimental (JPz7, dxme, xg8H) and theoretical (JPz7, rgYo, xg8H) aspects of the work were comprehensively covered. Based on 4 reviews, no clear or critical problems were found that were repeatedly mentioned in all reviewers.
Decision: Accept (poster)