Meta Optimal Transport
Authors: Brandon Amos
Samuel Cohen
Giulia Luise
Ievgen Redko
Venue: Submitted to ICLR 2023 (2023)
Keywords: optimal transport
meta learning
amortized optimization

Abstract: We study the use of amortized optimization to predict optimal transport (OT) maps from the input measures, which we call Meta OT. This helps repeatedly solve similar OT problems between different measures by leveraging the knowledge and information present from past problems to rapidly predict and solve new problems. Otherwise, standard methods ignore the knowledge of the past solutions and suboptimally re-solve each problem from scratch. We instantiate Meta OT models in discrete and continuous (Wasserstein-2) settings between images, spherical data, and color palettes and use them to improve the computational time of standard OT solvers by multiple orders of magnitude.

META OPTIMAL TRANSPORT
Anonymous authors
Paper under double-blind review
ABSTRACT
We study the use of amortized optimization to predict optimal transport (OT) maps from the input measures, which we call Meta OT. This helps repeatedly solve sim- ilar OT problems between different measures by leveraging the knowledge and in- formation present from past problems to rapidly predict and solve new problems. Otherwise, standard methods ignore the knowledge of the past solutions and sub- optimally rolvolve each problem from scratch. We instantiate Meta OT models in discrete and continuous (Wasserstein-2) settings between images, spherical data, and color palettes and use them to improve the computational time of standard OT solvers by multiple orders of magnitude.
INTRODUCTION
Optimal transportation (Villani, 2009; Ambrosio, 2003; Santambrogio, 2015; Peyré et al., 2019;
Merigot and Thibert, 2021) is thriving in domains including economies (Galichon, 2016), reinforce- ment learning (Dadashi et al., 2021; Fickinger et al., 2021), style transfer (Kolkin et al., 2019), generative modeling (Arjovsky et al., 2017; Seguy et al., 2018; Huang et al., 2020; Rout et al., 2021), geometry (Solomon et al., 2015; Cohen et al., 2021), domain adaptation (Courty et al., 2017; Redko et al., 2019), signal processing (Kolouri et al., 2017), fairness (Jiang et al., 2020), and cell reprogramming (Schiebinger et al., 2019). A core component in these settings is to couple two mea- sures (a. B) supported on domains (X, y) by solving transport optimization problem such as the primal Kantorovich problem, which is defined by:
T (a, R,c) E arg min / =El((a, xxy
c(x, y)) v),
(1)
where the optimal coupling n² is joint distribution over the product space, U(a, 3) is the set of admissible couplings between a and 8, and c X' x y + IR is the ground cost, that represents a notion of distance between elements in X and elements in V.
Challenges. Unfortunately, solving eq. (1) once is computationally expensive between general mea- sures and computationally cheaper alternatives port (Cuturi, 2013) smooths the transport problem with an entropy penalty, and sliced distances (Kolouri al., 2016; 2018; 2019; Deshpande et al., 2019) solve OT between 1-dimensional projec- tions of the measures, where eq. (1) can be solved easily.
active research topic: Entropic optimal trans-
Furthermore, when an optimal transport method is deployed in practice, eq. (1) is not just solved a single time, but is repeatedly solved for new scenarios between different input measures (a, B). For example, the measures could be representations of images we care about optimally transporting between and in deployment we would receive a stream of new images to couple. Repeatedly solving optimal transport problems also comes up in the context of comparing seismic signals (Engquist and Froese, 2013) and in single-cell perturbations (Bunne et al., 2021; 2022b;a). Standard optimal transport solvers deployed in this setting would re-solve the optimization problems from scratch, but this ignores the shared structure and information present between different coupling problems.
Overview and outline. We study the use of amortized optimization and machine learning methods to rapidly solve multiple optimal transport problems and predict the solution from the input measures (a, .). This setting involves learning a meta model to predict the solution to the optimal transport problem, which we will refer to Meta Optimal Transport. We leam Meta OT models to predict the solutions to optimal transport problems and significantly improve the computational time and number of iterations needed to solve eq. (1) between discrete (sect. 3.1) and continuous (sect. 3.2)

---

measures. The paper is organized as follows: sect. 2 recalls the main concepts needed for the rest of the paper, in particular the formulations of the entropy regularized and unregularized optimal transport problems and the basic notions of amortized optimization; sect. 3 presents the Meta OT models and algorithms; and sect. 4 empirically demonstrates the effectiveness of Meta OT.
Settings that not Meta OT. Meta OT is not useful in OT settings that do not involve repeatedly solving OT problems over a fixed distribution, including 1) standard generative modeling settings, such as Arjovsky
al. (2017) that estimate the OT distance between the data and model distri- butions, and 2) the out-of-sample setting of Seguy et al. (2018); 8); Perrot et al. (2016) that couple measures and then extrapolate the map larger measures containing the original measures.
PRELIMINARIES AND BACKGROUND
2.1 DUAL OPTIMAL TRANSPORT SOLVERS
We review foundations of optimal transportation, following the notation of Peyré et al. (2019) in most places. The discrete setting often favors the entropic regularized version since it can be com- puted efficiently and in a parallelized way using the Sinkhorn algorithm. On the other hand, the continuous setting is often solved from samples using convex potentials. While the primal Kan- torovich formulation in eq. (1) provides an intuitive problem description, optimal transport problems are rarely solved directly in this form due to the high-dimensionality of the couplings and the diffi- culty of satisfying the coupling constraints Ud ((x, B). Instead, most computational OT solvers use the dual of eq. (1), which we build our Meta OT solvers on top of in discrete and continuous settings.
2.1.1 ENTROPIC OT BETWEEN DISCRETE MEASURES WITH THE SINKHORRN ALGORITIIM
Let a Ex= a,O2, and B == Ei=1 b,0g, be discrete measures, where 5₂ is a Dirac at point and a E Am-1 and b An-1 are in the
Algorithm Sinkhorn(a, 8, c,
0)
for iteration i = to N do go c log b log (K" exp((i-1/())
f: f- c clog a c clog (K exp(g:/t))
probability simplex defined by
end for
A E and Σ a; 13. (2)
Compute PN from fN.gN using cq. (6)
return Px x P*
Discrete OT. In the discrete setting, eq. (1) simplifies to the linear program
(,c,, arg argmin PeU((ab)
U(a,b) = E Pl a, P 1, bb
(3)
where P is a coupling matrix, (a B3 is the optimal coupling, and the cost can be discretized as matrix C E Rm with entries Ct. == c(x,, Uj), and (C, P) Σng CipP..j.
Entropic OT. The linear program above can be regularized adding the entropy of the coupling to smooth the objective as in Cominetti and Martín (1994); Cuturi (2013), resulting in:
P* c).) E arg min (C, P) (H()) (a.(),
(4)
where H(P) = -x, (log(())).))) 1) is the discrete entropy of a coupling matrix P.
Entropic OT dual. As presented in Peyré et al. (2019, Prop. 4.4), the dual of eq. (4) is
f" E arg max (f. bg (expf.f/c), K exp(g/e)), == cxpp=cxp(CCag/e) (5)
gER"
where K E Bm is the Gibbs kernel and the dual variabley or potentials f E Rh" and g E R
associated, respectively, with the marginal constraints P1 = a and PT 1, = b. The optimal duals depend on the problem, e.g. fa G,C,,, c)., but omit this dependence for notational simplicity.
Recovering the primal solution from the duals. Given optimal duals f".g" that solve eq. (5) the optimal coupling p* to the primal problem in eq. (4) can be obtained by
(a, == explSi /e) Kij exp(g]/e)
(K is defined in eq. (5))
(6)
2

---

The Sinkhorn algorithm. Algorithm summarizes the log-space version, which takes closed form block coordinate ascent updates on eq. (5) obtained from the first-order optimality conditions (Peyré et al., 2019, Remark 4.21). We will use it to fine-tune predictions made by our Meta OT models.
Computing the error. Standard implementations of the Sinkhorn algorithm, such as Flamary et al. (2021); Cuturi et al. (2022), measure the error of candidate dual solution (f g) by computing the deviation from the marginal constraints, which we will also use in comparing our solution quality:
err(S 93 a, 3, = P - -all + 11I 1, bll
(compute P from eq. (6))
Mapping between the duals. The first-order optimality conditions of eq. (5) also provide an equiv- alence between the optimal dual potentials that we will make use of:
g(f:b.c) elog e log (K explf/e)).
2.1.2 DUAL WASSERSTHIN-2 OT BETWEEN CONTINUOUS (EUCLIDEAN) MEASURES
Let and 3 be continuous measures in Euclidean space X' = V = E (with a absolutely contin- uous with respect to the Lebesgue measure) and the ground cost be the squared Euclidean distance <(x, y) := | Then the minimum of eq. (1) defines the square of the Wasserstein-2 distance:
Algorithm 2 W2GN 8, (P))
for iteration == 1 to N do
Sample from (a, 5) and estimate L(p--1) Update yp; with approximation to vof(fi-11) end for return (w(+) = V. =
|| - vl2d((2))) = min / || - ((=))Bda(x) Jxxy
/
(2
min =El((o.8)
where T transport map pushing to 8. i.e. T a = 3 with the pushforward operator defined by (B)(B) (B) for any measurable set B.
Convex dual potentials. The primal form in eq. (9) is difficult to solve, as in the discrete setting, due to the difficulty of representing the coupling and satisfying the constraints. Makkuva et al. (2020); Taghvaei and Jalali (2019); Korotin et al. (2019; 2021b; 2022) propose to instead solve the dual:
(())d(()) + / (()))B())
( (,;a,,) E arg min ('noovee x
(10)
where vb is convex function referred to a convex potential, and (1) max (z, y) (xx is the Sreendre-Fencher transform or convex conjugate of v (Fenchel, 1949; Rockafellar, 2015). The potential &b often approximated with an input-convex neural network (ICNN) (Amos al., 2017).
Recovering the primal solution from the dual. Given an optimal dual * for eq, (10), Brenier (1991) remarkably shows that an optimal map T for eq. (9) can be obtained with differentiation:
T (x) (x).
(11)
Wasserstein-2 Generative Networks (W2GNs). Korotin et al. (2019) model v/ and thp in eq. (10)
with two separate ICNNs parameterized by sp. The separate model for Vh is useful because the conjugate operation in eq. (10) becomes computationally expensive. They optimize the loss:
((()) == E anwa
[x)(
+
E. ((Vo),))) the (Vei
+y E wwg
V Vih (v) - (12)
Cyclic momatone correlations (dual objective)
cocee-comssteeccyy regularizer
where is a detached copy of the parameters and ry is a hyper- parameter. The first term are the cyclic monotone correlations (Chartrand et al., 2009; Taghvaei and Jalali, 2019), that optimize the dual objective in eq. (10), and the second term provides cycle consistency (Zhu et al., 2017) to estimate the conjugate क. Algorithm shows how is typically optimized using samples from the measures, which
fine-tune Meta OT predictions.
3

---

0
0
B
B
g"
P
T
D
D
General
Discrete (Entropic)
Continuous (Wasserstein-2 Couplings
Input measures and cost
O Dual potentials
Figure 1: Meta OT uses objective-based -bssed amortization for optimal transport. In the general formula- tion, the parameters 0 capture shared structure in the optimal couplings ग between multiple input measures and costs over some distribution DD In practice, we learn this shared structure over the dual potentials which map back the coupling: f* in discrete settings and x in continuous ones.
2.2 AMORTIZED OPTIMIZATION AND LEARNING TO OPTIMIZE
Our paper is an application of amortized optimization methods that predict the solutions of opti- mization problems, surveyed in, e-g-, Chen al. (2021); Amos (2022). We use the basic setup from Amos (2022), which considers unconstrained continuous optimization problems of the form
(1) arg min (=;
(13)
where J is the objective, E E is the domain, and Φ E & is some context or parameterization. In other words, the context conditions the objective but is not optimized over. Given distribution over contexts P()), we learn model &s parameterized by 0 to approximate eq. (13), i.e. 60( Rs (6). will be differentiable for us, so we optimize the parameters using objective-based learning with
min
E d~)(
O()
(14)
which does not require ground ground-truth -truth solutions 2* and can be optimized with gradient- basdd solver. While we focus on optimizing eq. (14) because we do not assume easy access to ground-truth solu- tions 2" (Φ), one alternative regression gggeeesion-based learning if the solutions easily available:
min
E
(2)
2g (0)
(15)
((b)
META OPTIMAL TRANSPORT
Figure illustrates our key contribution of connecting objective-based amortization in eq. (14) to optimal transport. We consider solving multiple OT problems and learning shared structure and correlations between them. We denote joint meta-distribution over the input measures and costs with D(a, 3,c). which we call meta to distinguish it from the measures a, 8.
In general, we could introduce model that directly predicts the primal solution to eq. (1), i.c. (a 33, c) x π (a, 3) for (a, 3, c) ~ DD This is difficult for the same reason why most compu- tational methods do not operate directly the primal space: the optimal coupling often a high- dimensional joint distribution with non-trives trivial marginal constraints. We instead turn to predicting dual variables used by today's ' solvers.
3.1 META OT BETWEEN DISCRETE MEASURES
We build on standard methods for entropic OT reviewed in sect. 2.1.1 between discrete measures a := Em a,or; and 3 := Σn= byd2, with a E Am-1 and E A,-1 coupled using a cost c. the Meta OT setting, the measures and cost are the contexts for amortization and sampled from meta- distribution, i.e. (a, 3, c) ~ D(a, B)c For example, sects. 4.1 and 4.2 considers distrrbbitss over the weights of the atoms, i.e. (a, D, where D is distribution over Am-1 x A.-1-
4

---

Algorithm 3 Training Meta OT
Algorithm 4 Fine-tuning with Sinkhorn
Initialize amortization model with 0g for iteration i do
Predict duals fo 3, return Sinkhorn(ax, 3, c, f,
Sample (a, () ~
Predict duals fo &o on the sample Estimate the loss eq. (17) or eq. (18)
Algorithm 5 Fine-tuning with W2GN Predict dual ICNN parameters (, return (2GN(ax 3, c, T, Ge))
Update 0:+1 with a gradient step
end for
Amortization objective. We will seek to predict the optimal potential. At optimality, the pair of potentials are related to each other via eq. (8), i.e. a (,,,,,c) := e llg e log (K exp(ʃ/e)) where K E RRxx is the Gibbs kernel from eq. (5). Hence, it is sufficient to predict one of the potentials, e.g. f. and recover the other. We thus re-formulate eq. (5) to just optimize over f with
(a, 3,c,6) E arg min a(; fER"
(16)
where -J(f: -;
K exp /c]) the (negated) dual objective. Even though most solvers optimize over ff and g jointly as in eq. (16), amortizing over these would likely need: 1) to have a higher capacity than a model just predicting f. and 2) to learn how and are connected through eq. (8) while in eq. (16) we explicitly provide this knowledge.
Amortization model. We predict the solution eq. (16) with So B, c) parameterized by 0, resulting in computationally efficient approximation fo = f* Here use the notation fo 3, mean that the model fo depends on representations of the input measures and cost. In our settings, we define fg as a fully-connected MI.P mapping from the atoms of the measures to the duals.
Amortization loss. Applying objective-based amortization from eq. (14) to the dual in eq. (16) completes our learning setup. Our model should best-optimize the expectation the dual objective
min
E (o.J.c)~D
a,),,,,,
(17)
which is appealing as it does not require ground- solutions f*. Algorithm 3 shows a basic training loop for eq. (17) using gradient-based optimizer such as Adam (Kingma and Ba, 2014).
Sinkhorn fine-tuning. The dual prediction made by fo with an associated 9 can easily be input as the initialization to a standard Sinkhorn solver shown in algorithm 4. This allows us to deploy the predicted potential with Sinkhorn to obtain the optimal potentials with only a few extra iterations.
On accelerated solvers. Here we have only considered fine-tuning the Meta OT prediction with a log-Sinkhorn solver. Meta OT can also be combined with accelerated variants of entropic OT solvers such Thibault al. (2017); Altschuler et al. (2017); Alaya et al. (2019); Lin et al. (2019) that would otherwise solve every problem from scratch.
3.2 META OT BETWEEN CONTINUOUS MEASURES (WASSERSTEIN-2)
We take an analogous approach predicting the Wasserstein-2 map between continuous measures for Wasserstein-2 as reviewed in sect. 2.1.2. Here the measures a, 3 are supported in continuous space X' y Rd and we focus on computing Wasserstein-2 -2 couplings from instances sampled from a disiiiiiittiittion (a, 8) D(a, .). The cost e is not included in D as it remains fixed to the squared Fuclidean cost everywhere here.
One challenge here is that the optimal dual potential 1/* (.; (3 in eq. (10) is a convex function and not simply a finite-dimeniionl real vector. The dual potentials in this setting are approximated by, e.g., ICNN, We thus propose a Meta ICNN that predicts the parameters yp of an ICNN th. that approximates the optimal dual potentials, which can be seen as a hypernetwork (Stanley et al., 2009; Ha et al., 2016). The dual prediction made by &a can easily be input the initial value standard W2GN solver as shown in algorithm 5. App. B discusses other modeling choices we considered: tried models based MAMI. (Finn et al., 2017) and neural processes (Garnelo al., 201 8b;a).
5

---

Sinkhorn (converged, ground-truth)
Meta OT (initial prediction)
/ ag
e#3eccceeeeee1'
e#EeeCCCeeeeee
Q2
C00
Q2
11
a1
Figure 2: Interpolations between MNIST test digits using couplings obtained from (left) solving the problem with Sinkhorn, and (right) Meta OT model's initial prediction, which is 1100 times computationally cheaper and produces nearly identical coupling.
ResNeta
a
21
MLP8
T(-) - 12 (.)
ResNeta
Parameters
ICNN
Transport map
B
22
Figure 3: A Meta ICNN for image-based input measures. A shared ResNet processes the input measures and 3 into latents that are decoded with MLP into the parameters yp of an ICNN dual potential v.. The derivative of the ICNN provides the transport map T".
Table 1: Sinkhorn runtime (seconds) to reach a marginal error of 10- Meta OT's initial predic- tion takes = 5 10-5 seconds.
Table 2: Color transfer runtimes and values.
Iter
Runtime
Dual Value
Meta OT W2GN
None
3.5 10 =2.7-10 10
0.90
10
Intiializtiion
MNIST
Spherical
Ik
0.93 #2.27 20-22
1.0 -2 3
Zenos (saco))
4.5 4.5-10 as.5 19
0.88 s0.
23.78
1.0 #5.30
Ganssina
10 +1.2 10-3
0.56 10-
Meta (ste))
2.3 10-3 192 10
10 1344
W2GN
22.02-100000 10
0.96 02 0-2.02110" 100
Improvement US Psten)
1.96
11.3
2k
1.81
#3.05 -00
0.99 =111-101 11-10-2
We report the mean and standard deviation across 10 test instances.
Amortization objective. We build on the W2GN formulation (Korotin et al., 2019) and seek pa- rameters so* optimizing the dual ICNN potentials 1i and the with L(y;; or, (3) from eq. (12). We chose W2GN due to the stability, but could also easily use other losses optimizing ICNN potentials.
Amortization model: the Meta ICNN. We predict the solution to eq. (12) with Φo (3 param- eterized by 0, resulting in a computationally efficient approximation to the optimum & R y²" Figure instantiates convolutional Meta ICNN model using ResNet-18 (He et al., 2016) archi- tecture for coupling image-basee based measures. We again emphasize that a, 3 used with the model here representations of measures, which in our cases simply images.
Amortization luss. Applying objective-based amortization from eq. (14) to the W2GN loss in eq. (12) completes our learning setup. Our model should best-optimize the expectation of the loss:
min (a.
(()o(a,88)a,,))
(18)
As in the discrete setting, it does not require ground solutions ;p" and we learn it with Adam.
EXPERIMENTS
We demonstrate how Meta OT models improve the convergence of the state-of-the-art solvers in settings where solving multiple OT problems naturally arises. We implemented our code in JAX (Bradbury et al., 2018) as an extension to the the Optimal Transport Tools (OTT) package (Cuturi al., 2022). App. C covers further experimental and implementation details, and shows that all of our experiments take a few hours to run on our single Quadro GP100 GPU. We will open source the code reproduce all of our experiments.

---

MNIST
Spherical
Brror
Error
0.5
0.1
0.0
0.0
10
15
25
1000
Sinkhorn Iterations
Sinkhorn Iterations
Initialization Zeros
Gaussian (Thornton and Cuturi, 2022)
Meta OT)
Figure 4:
Meta OT successfully predicts warm-start initializations that significantly improve the convergence of Sinkhorn iterations on test data. The error the marginal error defined in eq. (7).
Sinkhorn (converged, ground-truth)
Meta OT (initial prediction)
Figure 5: Test set coupling predictions of the spherical transport problem. Meta OT's initial pre- diction is =37500 times faster than solving Sinkhorn optimality. Supply locations shown as black dots and the blue lines show the spherical transport maps I going to demand locations at the end. The sphere is visualized with the Mercator projection.
4.1 DISCRETE OT BETWEEN MNIST DIGITS
Images provide a natural setting for Meta OT where the distribution over images provide the meta- distribution D over OT problems. Given pair of images ao and a1, each grayscale image is cast as a discrete measure in 2-dimensional space where the intensities define the probabilities of the atoms. The goal is to compute the optimal transport interpolation between the two measures as in, e.g., Peyré et al. (2019, $7). Formally, this means computing the optimal coupling P by solving the entropic optimal transport problem between ag and a and computing the interpolates as a' = (t proj,, +(1 - t) projz) P*, for E [o, 1], where proj (z, y) == & and proj, y) = y. We selected = 10 as app. shows that it gives interpolations that are not too blurry or sharp.
Our Meta OT model fo (sect. 3.1) is an MLP that predicts the transport map between pairs of MNIST digits. We train on every pair from the standard training dataset. Figure 2 shows that even without fine-tuning, Meta OT's predicted Wasserstein interpolations between the measures are close the ground-truth interpolations obtained from rumning the Sinkhorn algorithm to convergence. We then fine-tune Meta OT's prediction with Sinkhorn as in algorithm 4. Figure 4 shows that the near-optimal predictions can be quickly refined in fewer iterations than running Sinkhorn with the default initialization, and table shows the runtime required to reach an error threshold of 10-2 showing that the Meta OT initialization help solve the problems faster by an order of magnitude. We compare learned initialization to the standard zero initialization, well as the Gaussian initial- ization proposed in Thorton and Cuturi (2022), which takes continuous Gaussian approximation of the measures and initializes the potentials be the known coupling between the Gaussians. This Gaussian initialization assumes the squared Euclidean cost, which is not the case in our spherical transport problem, but find it is still helpful over the initialization.

---

B
Tyo
B 3
α
W2GN (converged, ground-truth truth)
Meta OT (Initial prediction)
POLLAN
Figure 6: Color transfers with a Meta ICNN on test pairs of images. The objective to optimally transport the continuous RGB measure of the first image transport map T. Meta OT's prediction is = 1000 times faster than training W2GN from scratch. The image generating is Market in Algiers by August Macke (1914) and B is Argenteuil, The Seine by Claude Monet (1872), obtained from WikiArt.
the second 8, producing invertible
4.2 DISCRETE OT FOR SUPPLY-DEMANT TRANSPORTATION ON SPHERICAL DATA
We next up synthetic transport problem between supply and demand locations where the supply and demands may change locations or quantities frequently, creating another Meta OT setting to be able to rapidly solve the new instances. We specifically consider measures living on the 2-sphere defined by 82 := fir R3 | = 1). i.e. X' ) S2, with the transport cost given by the spherical distance c(æ; y) arccos( (rr, y)). We then randomly sample supply locations uniformly from Earth's s landmass and demand locations from Earth's 's population density to induce class of transport problems on the sphere obtained from the CC-licensed dataset from Doxsey-Whiin Whitfield et al. (2015). Figure 5 shows that the predicted transport maps on test instances close the optimal maps obtained from Sinkhorn to convergence. Similar to the MNIST setting, fig. and table show improved convergence and runtime.
4.3 CONTINUOUS WASSERSTEIN-2 COLOR TRANSFER
The problem of color transfer between two im- ages consists in mapping the color palette of
1.00
Dual Objective
image into the other one. The images are re-
quired to have the same number of channels, for example RGB images. The continuous formula- tion that we use from Korotin et al. (2019). takes i.c. X y [0, 1 with being the squared Fuclidean distance. We collected =200 public domain images from WikiArt and trained Meta ICNN model from sect. 3.2 to predict the color transfer maps between every pair of them. Fig- ure 6 shows the predictions on test pairs and fig. 7 shows the convergence in comparison to the stan- dard W2GN learning. Table reports runtimes and app. E shows additional results.
0.50
0.25
500
1000
1500
2000
W2GN Iterations
Initialization
Standard
Meta OT)
Figure 7: Convergence on color transfer test instances using W2GN. Meta ICNNs predicts warm-start initializations that significantly im- prove the (normalized) dual objective values.
8

---

5 RELATED WORK
Efficiently estimating OT maps. To compute OT maps with fixed cost between pairs of mea- sures efficiently, neural OT models (Korotin et al., 2019; 2021a; Mokrov et al., 2021; Korotin et al., 2021b) leverage ICNNs to estimate maps between continuous high-dimensional measures given samples from these, and Litvinenko et al. (2021); Scetbon et al. (2021); Forrow et al. (2019); Som- merfeld et al. (2019); Scetbon et al. (2022); Muzellee and Cuturi (2019); Bonet et al. (2021) leverage structural assumptions coupling and cost matrices to reduce the computational and memory com- plexity. In the meta-OT setting, we consider learning to rapidly compute OT mappings between new pairs measures. All these works hence potentially benefit from acceleration effect by leveraging amortization similarly.
Embedding measures where OT distances are discriminative. Effort has been invested in learn- ing encoeings//nojeccions of measures through a nested optimization problem, which aims to find discriminative embeddings of the measures to be compared (Genevay al., 2018; Deshpande al., 2019; Nguyen and Ho, 2022). While these works share an encoder and/or projection across task with the aim of leveraging more discriminative alignments (and hence an OT distance with a metric different from the Euclidean metric), our work differs in the sense that we find good initializations to solve the OT problem itself with fixed cost more efficiently across tasks.
Optimal transport and amortization. Few previous works in the OT literature leverage amor- tization. Courty et al. (2018) learn latent space in which the Wasserstein distance between the measure's embeddings is equivalent to the Euclidean distance. Concurrent work (Nguyen and Ho, 2022) amortizes the estimation of the optimal projection in the max-sliced objective, which differs from our work where we instead amortize the estimation of the optimal coupling directly. Also, Lacombe et al. (2021) learns to predict Wasserstein barycenters of pixel images by training a con- volutional networks that, given images generalization of this pixel-based work to general measures - both discrete and continuous. A limi- tation of Lacombe et al. (2021) is that it does not provide alignments, the amortization networks predicts the barycenter directly rather than individual couplings.
input, outputs their barycenters. Our work is hence
6 CONCLUSIONS, FUTURE DIRECTIONS, AND LIMITATIONS
We have presented foundations for modeling and learning to solve OT problems with Meta OT by using amortized optimization to predict optimal transport plans. This works best in applications that require solving multiple OT problems with shared structure. We instantiated it to speed up entropic regularized optimal transport and unregularized optimal transport with squared cost by multiple orders of magnitude. We envision extensions of the work in:
1. Meta OT models. While we mostly consider models based on hypernetworks, other meta- learning paradigms can be connected in. In the discrete setting, we only considered settings where the cost remains fixed, but the Meta OT model can also be conditioned on the cost by considering the entire cost matrix as an input (which may be too large for most models handle), considering lower-dimensional parameterization of the cost that changes between the Meta OT problem instances.
2. OT algorithms. While we instantiated models on top of log-Sinkhorn and W2GN, Meta OT could be built on top of other methods.
3. OT applications that are computationally expensive and repeatedly solved, e.g. in multi- marginal and barycentric settings, or for Commov-aasserstein distances between metric- measure spaces.
Limitations. While we have illustrated successful applications of Meta OT, it is also important to understand the limitations: 1) Meta OT does not make previously intractable problems tractable.
All of the baseline OT solvers we consider solve our problems within milliseconds or seconds. 2) Out-of-distributiin generalization. Meta OT may not generate good predictions are not close to the training OT problems from the distribbtti D over the measures and cost. If the model makes bad prediction, one fallback option is to re-solve the instance from seratch.
instances that

---

International Conference on REPRODUCIBILITY STATEMENT
We have tried to clearly articulate all of the relevant details in the text so that this paper can be completely reproduced from the contents of this paper alone. We will also open-source the full source code behind every experimental result, table, and figure in this paper, and will anonymously send it with this submission.
REFERENCES
Mokhtar Z Alaya, Maxime Berar, Gilles Gasso, and Alain Rakotomamonjy. Screening sinkhorn algorithm for regularized optimal transport. Advances in Neural Information Processing Systems, 32, 2019.
Jason Altschuler, Jonathan Niles- Weed, Philippe Rigollct. Near-linear time approximation algorithms for optimal transport via sinkhorn iteration. Advances in neural information processing systems, 30, 2017.
Luigi Ambrosio. Lecture notes on optimal transport problems. In Mathematical aspects of evolving interfaces, pages 1-52. Springer, 2003.
Brandon Amos.
Tutorial on amortized optimization for learning to optimize over continuous domains. arXiv preprint arXiv:2002. 00665, 2022.
Brandon Amos, Lei Xu, J 7ico Kolter. Input convex neural networks. Machine Learning, pages 146-155. PMLR, 2017.
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. Inter- national conference on machine learning, pages 214-223. PMLR, 2017.
Clément Bonet, Titouan Vayer, Nicolas Courty, François Septier, and Lucas Drumetz. Subspace detours meet gromov-wassertein.. Algorithms, 14(12):366, 2021.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milnee and Qiao Zhang. JAX: composable trans- formations of Nython+NumPy programs. GitHub, 2018. URL http:// //tthuubbbb/////// com/ /googoooooooooooooooooooooooo Le/ jax.
Yann Brenier. Polar factorization and monotone rearrangement of vector-valued functions. Communications on pure and applied mathematics, 44(4):375 1991.
Charlotte Bunne, Stefan G Stark, Gabriele Gut, Jacobo Sarabia del Castillo, Kjong- Van Lehmann, Lucas Pelk- mans, Andreas Krause, and Gunnar Ratsch. Learning single-cell perturbation responses using neural optimal transport. bioRxiv, 2021.
Charlotte Bunne, Andreas Krause, and Marco Cuturi, Supervised training of conditional monge maps. arXiv preprint arXiv:2206. 14262, 2022a.
Charlotte Bunne, Laetitia Papaxanthos, Andreas Krause, and Marco Cuturi. Proximal optimal transport mod- cling of population dynamics. In International Conference on Artificial Intelligence and Statistics, pages 6511-6528. PMLR, 2022b.
Rick Chartrand, Brendt Wohlberg, Kevin Vixie, and Erik Bollt. A gradient descent solution to the monge- kantorovich problem. Applied Mathematical Sciences, 3(22):1071-1000, 2009.
Tianlong Chen, Xiaohan Chen, Wuyang Chen, Howard Heaton, Jialin Liu, Zhangyang Wang, and Wotao Yin. Learming to optimize: primer and benchmark. arXiv preprint arXiv:2 2ot.2/033./28 /2828, 2021.
Samuel Cohen, Brandon Amos, and Yaron Lipman. Riemannian convex potential maps. In International Conference on Machine Learning, pages 2028 2038. PMLR, 2021.
Roberto Cominetti and J San Martín. Asymptotic analysis of the exponential penalty trajectory in linear pro- gramming. Mathematical Programming, 169-)7119-1177 1994.
Nicolas Courty, Rémi Flamary, Amaury Habrard, and Alain Rakotomamonjy. Joint distribution optimal trans- portation for domain adaptation. Advances in Neural Information Processing Systems, 30, 2017.
Nicolas Courty, Rémi Flamary, and Mélanic Ducoffe. Learning wasserstein embeddings. In International
Conference on Learning Representations, 2018.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural informa-
tion processing systems, 26:2292-2300, 2013.
10

---

Marco Cuturi, Laetitia Panng--aaaaantoos Yingtao Tian, Charlotte Bunne, Geoff Davis, and Olivier Teboul. Optimal transport tools (ott): jax toolbox for all things wasserstein. arXiv preprint arXiv:2201122524 2022.
Robert Dadashi, Léonard Hussenot, Matthieu Geist, Olivier Pietquin. Primal wasserstein imitation learning. In ICLR 202t-Ninth International Conference on Learning Representations, 2021.
Ishan Deshpande, Vung-Ting Hu, Ruoyu Sun, Ayis Pyrros, Nasir Siddiqui, Sanmi Koyejo, Zhizhen Zhao, David Forsyth, and Alexander G Schwing. Max-sliced wasserstein distance and its use for gans. In Proceedings of 2019.
the IEEE/CVF Conference Computer Vision and Pattern Recognition, pages
Erin Doxsey-Whiffield Whitfield, Kytt MacManus, Susana B Adamo, Linda Pistolesi, John Squires, Olena Borkovska, and Sandra R Baptista. Taking advantage of the improved availability of census data: first look the gridded population the world, version 4. Papers Applied Geography, 1(3):26-234, 2015.
Bjorn Engquist and Brittany D Froese. Application the wasserstein metrie to seismic signals. arXiv preprint 13Kiv:3311,4811 2013.
Werner Fenchel. On conjugate convex functions. Canadian Journal of Mathematics, 1(1):73-77, 1949.
Arnaud Fickinger, Samuel Cohen, Stuart Russell, and Brandon Amos. Cross-domain imitation learning via optimal transport. In International Conference on Learning Representations, 2021.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In Doina Precup and Yee Whye Teh, editors. Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 126-1135. PMLR, 06-11 Aug 2017. URL https: pttp///oceeenng....o.m mlr pww.ssppeess/ee /t00 .tmll.
Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z Alaya, Aurélie Boishunon, Stanislas Chambon, Lactitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, al. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78):1-8, 2021.
Aden Forrow, Jan-Christian Hütter, Mor Nitzan, Philippe Rigollet, Geoffrey Schiebinger, and Jonathan Weed. Statistical optimal transport via factored couplings. In The 22nd International Conference on Artificial Intelligence and Statistics, pages 2454-2465. PMLR, 2019.
Alfred Galichon. Optimal transport methods in economics. In Optimal Transport Methods in Economics.
Princeton University Press. 2016.
Marta Garnelo, Dan Rosenbaum, Christopher Maddison, Tiago Ramalho, David Saxton, Murray Shanahan, Yee Whye Teh, Danilo Rezende, and SM Eslami. Couditional neural processes. International Con- ference on Machine Learning, pages 1704-1713, PMLR, 2018a.
Marta Garnelo, Jonathan Schwarz, Dan Rosenbaum, Fabio Viola, Danilo Rezende, SM Eslami, and Yee Whye 2018b
Tch. Neural processes. arXiv preprint arXiv:
Aude Genevay, Gabriel Peyre, and Marco Cuturi. Learning generative models with sinkhorn divergences. In Amos Storkey and Fernando Perez-Crus editors, Proceedings of the Twenty-First International Con- ference on Artificial Intelligence and Statistics, volume 84 of Proceedings Machine Learning Re- search, pages 1608 1617. PMLR, 09-11 Apr 2018. URL v8 v84/ 2844/enevay 8a. htm1.
mir press/
David Ha, Andrew Dai, and Quoe Le. Hypernetworks. arXiv preprint arXiv:1609.09106 2016.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, Jian Identity mappings in deep residual networks. In European conference on computer vision, pages 630-645. Springer, 2016.
Chin-We Wei Huang, Ricky TQ Chen, Christos Tsirigotis, and Aaron Courville, Convex potential flows: Universal probability distributions with optimal transport and convex optimization. In International Conference on Learning Representations, 2020.
J. J. Hull. A database for handwritten text recognition research. IEEE Transactions on Puttern Analysis and Machine Intelligence, 16(5):550-$55 1994. doi: 10.1109/34.291440.
John D Hunter. Matplotiib: A 2d graphies environment. Computing in science & engineering, 9(3)990, 2007.
Ray Jiang, Aldo Pacchiano, Tom Stepleton, Heinrich Jiang, and Silvia Chiappa. Wasserstein fair classification. In Uncertainty Artificial Intelligence, pages 862-872. PMLR, 2020.

---

Diederik P Kingma and Jimmy Ba. WrXivv1412.6980 2014.
Adam: A method for stochastic optimization. arXiv preprint
Nicholas Kolkin, Jason Salavon, and Gregory Shakhnarowich. Style transfer by relaxed optimal transport and self-similarity. In Proceedings of the IEEF/CVF Conference Computer Vision and Pattern Recognition, pages 10051-10060, 2019.
Soheil Kolouri, Yang Zou, and Gustavo K Rohde. Sliced wasserstein kernels for probability distributions. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages
2016.
Soheil Kolouri, Se Rim Park, Matthew Thorpe, Dejan Slepcev, and Gustavo K Rohde. Optimal mass transport: Signal processing and machine-lcaming applications. IEEE signal processing magazine, 34(4):43-3) 59, 2017.
Soheil Kolouri, Phillip E Pope, Charles E Martin, and Gustavo K Rohde. Sliced wasserstein auto-encoders. In International Conference on Learning Representations, 2018.
Soheil Kolouri, Kimia Nadjahi, Umut Simsekli, Roland Badeau, and Gustavo K Rohde. Generalized sliced wasserstein distances. arXiv preprint arXiv: 1902.0044, 2019.
Alexander Korotin, Vage Egiazarian, Arip Asadulaev, Alexander Safin, and Evgeny Burnaev. Wasserstein-2 generative networks. arXiv preprint arXiv:1909. 13082, 2019.
Alexander Korotin, Lingxiao Li. Aude Genevay, Justin M Solomon, Alexander Filippov, and Evgeny Burnaev. Do neural optimal transport solvers work? continuous wasserstein-2 benchmark. Advances in Neural Information Processing Systems, 34:14593-14605, 2021a.
Alexander Korotin, Lingxiao Li, Justin Solomon, and Evgeny Burnaev, Continuous wasserstein-2 barycenter estimation without minimax optimization. arXiv preprint arXiv:2102.01752 2021b.
Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. Neural optimal transport. arXiv preprint arXiv:2201. 12220, 2022.
Julien Lacombe. Julie Digne, Nicolas Courty, and Nicolas Bonneel. Learning to generate wasserstein barycen- ters, 2021. URL https:t////. /ttpp//v////peeeeeee.e. net forum? 2id 2: Poaaass 6 vw.
Lingxiao Li, Aude Genevay, Mikhail Yurochkin, and Justin Solomon. Continuous regularized wasserstein barycenters. arXiv preprint arXiv:2008. /2534, 2020.
Tianyi Lin, Nhat Ho, and Michael I Jordan. On the acceleration of the sinkhorn and greenkhorn algorithms for optimal transport. arXiv preprint arXiv:1906.014:: 03377, 2019.
Alexander Litvinenko, Youssef Marzouk, Hermann G Matthies, Marco Scavino, and Alessio Spantini. Com- puting E-divergences and distances of sigh-dimensional probability density functions lowerrons-lowrank tensor ap- proximations. arXiv preprint orXiv:21 1107164, 2021.
Ashok Makkuva, Amirhossein Taghvaci, Sewoong Oh, and Jason Lee. Optimal transport mapping via input convex neural networks. In International Conference on Machine Learning, pages 6672-6681. PMLR, 2020.
Wes McKinney. Python for data analysis: Data wrangling with Pandas, NumPy, and IPython. O'Reilly Media, Inc.", 2012.
Quentin Merigot and Boris Thibert. Optimal transport: ical Analysis. volume 22, pages 133-212. Elsevier, 2021.
discretization and algorithms. In Handbook of Numer-
Petr Mokrov, Alexander Korotin, Lingxiao Li, Aude Genevay, Justin M Solomon, and Evgeny Burnaev. Large- scale wasserstein gradient flows. Advances in Neural Information Processing Systems, 34:15243-1522 2021.
Boris Muzellec and Marco Cuturi. Subspace detours: Building transport plans that are optimal on subspace projections. Advances in Neural Information Processing Syxtems, 32, 2019.
Khai Nguyen and Nhat Ho. Amortized projection optimization for sliced wasserstein generative models. arXiv preprint 2tt.:2203.13117 2022.
Travis E Oliphant. A guide to NumPy, volume 1. Trelgol Publishing USA, 2006.
Travis Oliphant. Python for scientific computing. Computing in Science Engineering, 93)):10-20 2007.
Michaël Perrot, Nicolas Courty, Rémi Flamary, and Amaury Habrard. Mapping estimation for discrete optimal transport. Advances in Neural Information Processing Systems, 29, 2016.
12

---

Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport: With applications to data science. Foun- dations and Trends in Machine Learning, 11(5-6):355 607, 2019.
levgen Redko, Nicolas Courty, Rémi Flamary, and Devis Tuia. Optimal transport for multi-source domain adaptation under target shift. The 22nd International Conference on Artificial Intelligence and Statistics, pages 849-858. PMLR, 2019.
Ralph Tyrell Rockafellar. Convex analysis. In Convex analysis. Princeton university press, 2015.
Litu Rout, Alexander Korotin, and Evgeny Burnee.. Generative modeling with optimal transport maps. In International Conference on Learning Representations, 2021.
Andrei A Rusu, Dushyant Rao, Jakub Sygnowski, Oriol Vinyals, Razvan Pascanu, Simon Osindero, and Raia Hadsell. Meta-learning with latent embedding optimization. In International Conference Learning Rep- resentations, 2018.
Filippo Santambrogio. Optimal transport for applicd mathematicians. Birküuser, NY, 55(58-633))4 2015.
Meyer Scetbon, Marco Cuturi, and Gabriel Peyré. Low-rank sinkhorn factorization. In International Confer- encee on Machine Learning, pages 9344-9354. PMLR, 2021.
Meyer Scetbon, Gabriel Peyré, and Marco Cuturi. Linear-time gromov wasserstein distances using low rank couplings and costs. In International Conference on Machine Learning, pages 19347-19365. PMLR, 2022.
Geoffrey Schiebinger, Jian Shu, Marcin Tabaka, Brian Cleary, Vidya Subramanian, Arych Solomon, Joshua Gould, Siyan Liu, Stacie Lin, Peter Berube, Lia Lec, Jenny Chen, Justin Brumbaugh, Philippe Rigol- let, Konrad Hochedlinger, Rudolf Jaenisch, Aviv Regev, and Eric S. Lander. Optimal-transport analy- sis of single-cell gene expression identifies developmental trajectories in reprogramming. Cell, 176(4): 928-943.e22, 2019. ISSN 0092-8674. doi: https://bbi.ogg//...166........ 201..01.000.. // //www. www.scedieencceeeeeeeeeeee com/ wwwps/cceeeeaaeeeeeeeeeeeeeeeeee htt.ee/eeeaaiiceeeee//////////// 9300099X.
URL https:
Vivien Seguy, Bharath Bhushan Damodaran, Remi Flamary, Nicolas Courty, Antoine Rolet, and Mathieu Blon- del. Large scale optimal transport and mapping cstimation. In International Conference on Learning Repre- sentations, 2018
Justin Solomon, Fernando De Goes, Gabriel Peyré, Marco Cuturi, Adrian Butscher, Andy Nguyen, Tao Du, and Leonidas Guibas. Convolutional wasserstein distances: Efficient optimal transportation geometric domains. ACM Transactions on Graphics (TOG), 34(4):1-11, 2015.
Max Sommerfeld, Jörn Schrieber, Yoav Zemel, and Axel Munk. Optimal transport: Fast probabilistic approxi- mation with exact solvers. J. Mach. Learn. Res., 20: 105-1, 2019.
Kenneth O Stanley, David B D' Ambrosio, and Jason Gauci. A hypercube- encoding for evolving large- scale neural networks. Artificial life. (5(2):185-211 2009.
Amirhossein Taghvaci and Amin Jalali. 2-wasserstein approximation via restricted convex potentials with application to improved training for gans. arXiv preprint arXiv: 190.00990 2019.
Matthew Tancik, Ben Mildenhall, Terrance Wang. Divi Schmidt, Pratul P Srinivasan, Jonathan T Barron, and Ren Ng. Learned initializations for optimizing coordinate-basel neural representations. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pages 2846-2855. 2021.
Alexis Thibault, Lenaic Chizat, Charles Dossal, and Nicolas Papadakis. Overrelaxed sinkhorn -noop algorithm for regularized optimal transport. arXiv preprint arXiv: 2017.
James Thornto and Marco Cuturi. Rethinking initialization of the sinkhorn algorithm. arXiv preprint arXiv:2206.00 07630, 2022.
Stefan Van Der Walt, S Chris Colbert, and Gael Varoquaux. The numpy arry: a structure for fffefficient numerical computation. Computing in Science & Engineering, 13(2):22, 2011.
Guido Van Rossum and Fred L Drake Jr. Python reference manual. Centrum voor Wiskunde en Informatica Amsterdam, 1995.
Cédric Villani. Optimal transport: old and new, volume 338. Springer, 2009,
Omry Yadan. Hydra framework for elegantly configuring complex applications. Github, 2019. URL
https: // thub. com/ /ww.a/eooooreeeerccc//aaaaaa
13

---

Jun-Yan Zhu, Taesung Park, Phillip Isola, and Alexei Efros. Unpaired image-to-image translation using cycle-consistent adversarial networks. In Proceedings of the IEEE international conference on computer vision, pages 2223-2232 2232, 2017.
Luisa Zintgraf, Kyriacos Shiarli, Vitaly Kurin, Katja Hofmann, and Shimon Whiteson. Fast context adaptation via meta-learning. In International Conference Machine Learning, pages 7693-7703 PMLR, 2019.
A
SELECTING E FOR MNIST
0.05
c 0.001 c= < 0.005 00001
Figure 8: We selected e 10 for our MNIST coupling experiments as it results in transport maps that are not too blurry sharp.
B OTHER MODELS FOR CONTINUOUS OT
While developing the hyper-network or Meta ICNN in sect. 3.2 for predicting couplings between continuous measures, we considered alternative modeling formulations briefly documented this section. We finalized only the hyper-network model because is conceptually the most similar to predicting the optimal dual variables the continuous setting and results rapid predictions.
B. OPTIMIZATION-DASED META-LEARNING (MAML-INSPIRED)
The model-agnostic meta-learning setup proposed in MAML (Finn et al., 2017) could also be ap- plied in the Meta OT setting learn adaptable initial parameterization. In the continuous setting, one initial version would take a parameterized dual potential model (1)()) and seek to learn an ini- tial parameterization şo that optimizing loss such the W2GN loss L from eq. (12) results in a minimal L((K) after adapting the model for steps. Formally, this would optimize:
arg min L(px) where Pitt $ V.L((p))
(19)
Tancik et al. (2021) explores similar learned initializations for coordinate-basee neural implicit rep- resentations for 2D images, sean reconstruction, and 3d shape and scene recovery from 2D observations.
14

---

Challenges for Meta OT. The transport maps given by T Vy significantly vary depending on the imput measures a, 3. We found it difficult to learn an initialization that can be rapidly adapted, and optimizing eq. (19) is more computationally expensive than eq. (18) it requires unrolling through many evaluations of the transport loss L. And, we found that only learning to predict the optimal parameters with eq. (18), conditional the input measures, and then fine-tuning with W2GN to be stable.
Advantages for Meta OT. Exploring MAML-inspired methods could further incorporate the knowl- edge that the model's prediction is going to be fine-tuned into the learning process. One promising direction we did not try could be integrate some of the ideas from LEO (Rusu et al., 2018) and CAVIA (Zintgraf et al., 2019), which propose to learn a latent space for the parameters where the initialization also conditional on the input.
B.2 NEURAL PROCESS AND CONDITIONAL MONGE MAPS
The (conditional) neural process models considered in Garnelo et al. (2018b;a) can also be adapted for the Meta OT setting, and is similar to the model proposed in Bunne et al. (2022a). In the continuous setting, this would result in dual potential that also conditioned
representation of the input measures, e.g. &y (x; 2) where 2 == femb(re (ax, 3) is learned embedding the input measures that is learned with the parameters of vb. This could be formulated
arg min (a.B)~>
E
C(v,, 1p (o.3)),
(20)
where & modifies the model used in the loss eq. (12) to also be conditioned on the context extracted from the measures.
Challenges for Meta OT. This raises the issue on bert-foruulating the model to be conditional on the context. One way could be to append to the imput point in the domain. Bunne al. (2022a) proposes to use the Partially Input-Convex Neural Network (PICNN) from (Amos et al., 2017) to make the model convex with respect and not
Advantages for Meta OT. A large advantage is that the representation = of the measures o, B would be significantly lower-dimensional than the parameters yp that our Meta OT models are predicting.
C ADDITIONAL EXPERIMENTAL AND IMPLEMENTATION DETAILS
We have attached the Jax source code necessary to run and reproduce all of the experiments in our paper and will open-source all of it. Here is a basic overview of the files:
meta Meta OT Python library code
conjugate.py
Exact conjugate solver for the continuous setting
data.py
models.py
-utiss.py
config Hydra configuration for the experiments (containing hyper-parameters)
-traindiscrret..py
Train Meta OT models for discrete OT
-train_color_single..pp
Train a single ICNN with W2GN between images (for debugging)
train_color_meta.py
Train a Meta ICNN with W2GN
plot_mmist.py Visualize the MNIST couplings
-plot_world_pair.pe Visualize the spherical couplings
-eval_color.py Evaluate the Meta ICNN in the continuous setting
eva aval_discrete.py Evaluate the Meta ICNN for the discrete tasks
15

---

Connecting to the data is difficulty in running the experiments. The easiest experiment the MNIST one, which will automatically download the dataset:
straee_ Train the model, outputting to <exp_dir> dir> /eval__dscretee <exp_dir> # Rva uatt the learned models /pplot_misst.pp <exp_dir> $ Produce further
oons
C.I
We briefly summarize the hyper-parameters we used for training, which did not extensively tune. In the discrete setting, we use the same hyper-parametens for the MNIST and spherical settings.
Table 4: Continuous OT hyper-parameters.
Name
Value
Table 3: Discrete OT hyper-parameters.
Meta batch size (for a, 9)
8
Inner batch size (to estimate L)
1024
Name
Value
Cycle loss weight (r)
3.
Batch size
128
Adam learning rate
1c-3
Number training iterations
50000
& weight penalty
le-6
MLP Hidden Sizes
[1024, 1024, 1024)
Max grad norm (for clipping)
1.
Adam learning rate
le-3
Number training itcrations
200000
Meta ICNN Encoder
ResNer18
Encoder output size (both measures)
256x2
Meta ICNN Decoder Hidden Sizes
[512]
C.2 SINKHORN CONVERGENCE TIMES, VARYING THRESHOLLD
In the main paper, table reports the runtime of Sinkhorn to reach a convergence threshold of the marginal error being below tolerance of 1023. Tables and report the results from sweeping over other thresholds and show that Meta OT's initialization is consistently able to help.
Table 5: Sinkhorn runtime to reach a thresholded marginal error on MNIST. Threshold=10-2
Initialization
Threshold=10
Threshold=10
Threshold=10-1
Zeros
4.5 10-3 +1.5- 10
7.7 10-3 +1.2-10
1.1 10-2 1188
15 10 42.3 10-2
Gaussian
4.1-10 10
207-10-
114-1003
1.1 21.10-2=1.--00
14-100-2 4244100 10-a
Meta OT
2.3 49.2 10-6
3.9- 10 +1.6100 10"3
6.7. 10 11.4-10
1.0 10-10-2 42.4 10
Table Sinkhorn runtime to reach a thresholded marginal error on the spherical transport problem.
Initialization
Threshold=10- 2
Threshold=10 3
Threshold=10
Threshold=10
Zeros
8.8 10 +1.3 10
1.19 10
2.1 #3.6 10
2.8 15.6 10
Gaussian
5.6 10-1 10-
1.1 12-0-0-1
+3.5--1-10"
2.4 +5.4 10-1
Meta OT
7.8 10 13.4 10-2
0.44 #1.5
0.97 =3.2 10
1.7 17.8-
16

---

C.3 EXPERIMENTAL RUNTIMES AND CONVERGENCE
App. C.3 shows the convergence during training of Meta OT models in the discrete and continuous settings over 10 trials our single Quadro GP100 GPU. The MNIST models consistently trained to optimality within minutes (!) while the continuous model takes a few hours to train.
MNIST
MNIST
Marginal Emrr
0.3
Marginal Error
0.3
0.2
0.2
0.1
0.1
0.0
10
20
30
40
50
2
1k Training Iterations
Training Time (minutes)
Spherical
Spherical
Marginal Error
1.0
Marginal Error
1.0
0.5
0.5
0.0
10
20
30
40
50
1k Training Iterations
Training Time (minutes)
Wasserstein-2
Wasserstein-2
-0.90
-0.90
Dual
Dual
-0.5
100
150
200
1k Training Iterations
Training Time (minutes)
Figure 9: Convergence Meta OT models during training, reported over iterations and wall-clock time. We run each experiment for 10 trials with different seeds and report each trial as a line.
17

---

D OUT-OF-DISTRIBUTION GENERALIZATION
App. D tests the ability of Meta OT to predict potentials for out-of-distribution input data. We consider the pairwise training and evaluation on the following datasets: 1) MNIST; 2) USPS (Hull, 1994) (upscaled to have the same size as the MNIST); 3) Google Doodles dataset" with classes Crab, Cat and Faces; 4) sparsified random uniform data in [0,1] where sparsity (zeroing values below 0.95) is used to mimic the sparse signal in plack-and-white images. For each pair, eg, MNIST-USPS, we train on one dataset and use the other to predict the potentials. The comparison is done using the same metric as before, i.e., the deviation from the marginal constraints defined in eq. (7).
Evaluation Domain
MNIST
Doodles
USPS
Uniform
MNIST
0.3
0.1
0.4
Doodles
0.2
Training Domain
0.1
0.0
0.4
USPS
0.2
0.0
Uniform
0.3
0.1
10
15
20
25
6
1.0
15
20
25
0
10
15
20
25.0
45
Sinkhorn Iterations
Sinkhorn Iterations
Sinkhorn Iterations
Sinkhorn Iterations
Initialization
Zeros
Gaussian (Thornton and Cuturi, 2022)
Meta OT)
Figure 10: Cross-domain experiments.
"https://qui wwdrwwwwwwwuuuuu..oooooo wwt.rrawwwwwitttt witoooo com/ data
18

---

E ADDITIONAL COLOR TRANSFER RESULTS
We next show additional color transfer results from the experiments sect. 4.3 on the following public domain images from WikiArt:
Distant View of the Pyramids by Winston Churchill (1921)
Charing Cross Bridge, Overcast Weather by Claude Monet (1900)
Houses of Parliament by Claude Monet (1904)
October Sundown, Newport by Childe Hassam (1901)
Landscape with House at Ceret by Juan Gris (1913)
Irises in Monet's Garden by Claude Monet (1900)
Crystal Gradation by Paul Klee (1921)
Senecio by Paul Klee (1922)
Váza s květinami by Josef Capek (1914)
Sower with Setting Sun by Vincent van Gogh (1888)
Three Trees in Grey Weather by Claude Monet (1891)
Vase with Daisies and Anemones by Vincent van Gogh (1887)
19

---

B
T
T B
Figure 11: Meta ICNN (initial prediction). The sources are given in the beginning of app. E.
20

---

B
T
T B
Figure 12: Meta ICNN + W2GN fine-tuning. The sources are given in the beginning of app. E.
21

---

B
T
T B
Figure 13: W2GN (final). The sources are given in the beginning of app.
22


Reviews:
[{"review_id": "bTJ_HQRHXR", "reviewer": "Reviewer_zFaG", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "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": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes a meta-learning approach to OT. That is, the authors learn a meta-model on a large train set of OT tasks and then try to predict the solution of the new test OT tasks based on the past experience. Two experimental setups are considered: discrete OT problems for general costs and continuous (small-dimensional) OT problems for the quadratic cost. In both cases, the authors state that the meta-ot predictions notably speed up the OT learning algorithms (meta-prediction is used as the initialization in the OT algorithm). The performance is demonstrated on 2D, 3D optimal transport tasks: interpolation between MNIST digits (considered as 2D distributions), supply-demand transportation and color transfer between RGB images (represented as 3D point clouds).\n", "summary_of_the_review": "The paper has limited scientific novelty – it is a application of Meta Learning to OT which is rather straightforward and its significance is not clear. On the one hand, it seems to indeed boost the underlying OT algorithms thanks to the good predicted initialization of dual variables and overall it is interesting to see that OT problems can be (to some extent) meta-learned. On the other hand, the potential applications prospect of Meta-OT are not convincing enough (do we really need to meta-learn OT?) and there are some questions in the experiments. I currently rate this paper as a borderline, and will wait for the authors feedback to make my final decision.\n\n**Post-rebuttal update:** score +1 (see the discussion under the general authors comment)", "strength_and_weaknesses": "**Strength**\n\n(1) Application of meta-learning to the previously unapproached problem.\n\n(2) Meta-OT allows to predict a good initialization for the OT algorithms which naturally boosts their convergence;\n\n**Weaknesses**\n\n(1) Potential applications of the Meta-OT methodology are questionable;\n\n(2) Meta-OT performance presumably heavily depends on the available meta-distributions (train set of pairs of distributions for OT).\n\n(3) The capabilities of the method are only to predict good initialization for sinkhorn/W2GN, but not the good approximation of the OT plan (not sure here);\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "b5TAEnPUm9p", "reviewer": "Reviewer_eM4a", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "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": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper presents Meta OT, a method for initializing optimal transport problems based on prior solutions that applies in contexts in which the problems share structure. Meta OT is described for discrete and continuous problems. Extensive experiments quantify improvements in speed of convergence on problems including MNIST, transportation problems, and color transfer. Results suggest speed ups of several orders of magnitude are possible and, to the eye, they incur little to no loss in accuracy.  ", "summary_of_the_review": "The paper introduces a nice idea for improving the performance of OT solvers. Given that it only applies to some problems, and an a priori test isn't available to know whether it will help, I would have liked more systematic analysis of the consequences of application to less-well suited domains both in terms of speed and accuracy. ", "strength_and_weaknesses": "Strengths: \n- OT and methods for obtaining solutions to OT problems are a timely and important topic. \n- The paper clearly documents the possible improvements in runtime performance. \n- The math is clear and concise. \n- The writing is clear. \n\nWeaknesses:\n- The specific problem, in which one is solving structurally related problems repeatedly, is a bit niche. \n- I would have like to see more systematic analysis of the decay of performance associated with different degrees of out-of-distribution problems, as well as with greater or less coherence within the domain of generalization. \n- Relatedly, I did not see a clear statement of when this can be applied to the advantage to the user. This seems like an important consideration, given that we wouldn't expect the approach to \"work\" generally, necessarily. \n- Similarly, it seems important to document how the performance decays when applied to domains that are less well suited. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "aH0PMbHzg5", "reviewer": "Reviewer_HF6W", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "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": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper presents a meta-method to obtain OT loss between two distributions.", "summary_of_the_review": "Please read the Strength and Weaknesses part.", "strength_and_weaknesses": "In section 2, the authors recall the OT dual formulation in the case of $W_2$ distance and the entropic OT. Based on simpler formulas of computing OT cost for those cases, they proposed  to use neural network to  modeling the cost function. They utilise some structures like Meta ICNN, Resnet for the neural network function. Section 4 includes experiments of discrete OT on MNIST data, transportation on spherical data and colour transfers using Wasserstein-2 distance. The results show that the meta OT needs significant less number of iterations to converge than the Sinkhorn algorithm.  In the last section, the authors claim the contribution and outline limitations of their methods.\n\nOverall, I think the idea is simple and interesting. I am concerned about the novelty of the method as well as its non-applicability in many important situations  as the authors have stated. The authors also do not show the amount of  time  to build up their meta-OT. It is very important issue because if the number of supports of a discrete distribution is extremely large, then  the main questions are how long does it take to train the meta-OT and can we guarantee the accuracy of its solution? \n\nTo improve the paper, the authors might try to include some approximation theory to quantify the solution's  accuracy and show why the Meta-OT initialisation improve the convergence speed of Sinkhorn algorithm. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "ZIO4AR0ZeXm", "reviewer": "Reviewer_taDB", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "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": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes Meta OT, which can repeatedly solve similar OT problems between different measures by leveraging the knowledge and information present from past problems to rapidly predict and solve new problems.\nThe experiments instantiate that Meta OT can improve the calculation time under various settings.\n", "summary_of_the_review": "This paper proposes Meta OT, which is different from W2GN in that Meta OT takes the cost as the input. \nThe experiments show that it improves the computing time, but it does not explain how to use the existing knowledge", "strength_and_weaknesses": "Strengths: Meta OT significantly improves the computational time and number of iterations needed to solve optimal transport between discrete and continuous measures.\n\nWeaknesses: 1.When comparing the time of algorithms, the time of Meta OT should include the training time.\n2. This paper does not explain why meta can help to quickly predict and solve new problems by using knowledge and information from past problems", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject