Original: Meta Optimal Transport
Revised: Meta Optimal Transport

Original: Anonymous Author(s)
Revised: Anonymous Author(s)

Original: AfﬁliationAddress email
Revised: AffiliationAddress email

Original: Abstract
Revised: Abstract

Original: We study the use of amortized optimization to predict optimal transport (OT) mapsfrom the input measures, which we call Meta OT .
Revised: We study the use of amortized optimization to predict optimal transport (OT) mapsfrom the input measures, which we call Meta OT .

Original: 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.
Revised: This helps repeatedly solve similar OT problems between different measures by leveraging the knowledge andinformation present from past problems to rapidly predict and solve new problems.

Original: Otherwise, standard methods ignore the knowledge of the past solutions and suboptimally re-solve each problem from scratch.
Revised: Otherwise, standard methods ignore the knowledge of the past solutionsand suboptimally re-solve each problem from scratch.

Revised: Meta OT models surpassthe standard convergence rates of log-Sinkhorn solvers in the discrete setting andconvex potentials in the continuous setting.

Original: Weinstantiate Meta OT models indiscrete and continuous (Wasserstein-2) settings between images, spherical data,and color palettes and use them to improve the computational time of standard OTsolvers by multiple orders of magnitude 
Revised: We improve the computational timeof standard OT solvers by multiple orders of magnitude in discrete and continuoustransport settings between images, spherical data, and color palettes

Original:  Introduction
Revised:  Introduction

Original: Optimal transportation [Villani, 2009, Ambrosio, 2003, Santambrogio, 2015, Peyré et al., 2019,
Revised: Optimal transportation [Villani, 2009, Ambrosio, 2003, Santambrogio, 2015, Peyré et al., 2019,

Original: Merigot and Thibert, 2021] is thriving in domains including economics [Galichon, 2016], reinforcement 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, Redkoet al., 2019], signal processing [Kolouri et al., 2017], fairness [Jiang et al., 2020], and cell reprogramming [Schiebinger et al., 2019].
Revised: Merigot and Thibert, 2021] is thriving in domains including economics [Galichon, 2016], reinforcement learning [Dadashi et al., 2020, Fickinger et al., 2021], style transfer [Kolkin et al., 2019],generative modeling [Arjovsky et al., 2017, Seguy et al., 2017, Huang et al., 2020, Rout et al., 2021],geometry [Solomon et al., 2015, Cohen et al., 2021], domain adaptation [Courty et al., 2017b, Redkoet al., 2019], signal processing [Kolouri et al., 2017], fairness [Jiang et al., 2020], and cell reprogramming [Schiebinger et al., 2019].

Original: A core component in these settings is to couple two measures( α, β ) supported on domains ( X , Y ) by solving a transport optimization problem such as the primal
Revised: A core component in these settings is to couple two measures( α, β ) supported on domains ( X , Y ) by solving a transport optimization problem such as the primal

Original: Kantorovich problem , which is deﬁned by:
Revised: Kantorovich problem , which is defined by:

Original: [Equation]
Revised: [Equation]

Original:  where the optimal coupling π (cid:63) is a joint distribution over the product space, U ( α, β ) is the set ofadmissible couplings between α and β , and c : X × Y → R is the ground cost , that represents anotion of distance between elements in X and elements in Y .
Revised: (1) where the optimal coupling π ⋆ is a joint distribution over the product space, U ( α, β ) is the set ofadmissable couplings between α and β , and c : X × Y → R is the ground cost , that represents anotion of distance between elements in X and elements in Y .

Original: Challenges.
Revised: Challenges.

Original: Unfortunately, solving eq.
Revised: Unfortunately, solving eq.

Original: (1) once is computationally expensive between general measures and computationally cheaper alternatives are an active research topic: Entropic optimal transport [Cuturi, 2013] smooths the transport problem with an entropy penalty, and sliced distances[Kolouri et al., 2016, 2018, 2019, Deshpande et al., 2019] solve OT between 1-dimensional projections of the measures, where eq.
Revised: (1) once is computationally expensive between general measures and computationally cheaper alternatives are an active research topic: Entropic optimal transport [Cuturi, 2013] smooths the transport problem with an entropy penalty, and sliced distances[Kolouri et al., 2016, 2018, 2019, Deshpande et al., 2019] solve OT between 1-dimensional projections of the measures, where eq.

Original: (1) can be solved easily.
Revised: (1) can be solved easily.

Original: Furthermore, when an optimal transport method is deployed in practice, eq.
Revised: Furthermore, when an optimal transport method is deployed in practice, eq.

Original: (1) is not just solveda single time, but is repeatedly solved for new scenarios between different input measures ( α, β ) .
Revised: (1) is not just solveda single time, but is repeatedly solved for new scenarios between different input measures ( α, β ) .

Original: For example, the measures could be representations of images we care about optimally transportingbetween and in deployment we would receive a stream of new images to couple.
Revised: For example, the measures could be representations of images we care about optimally transportingbetween and in deployment we would receive a stream of new images to couple.

Original: Repeatedly solvingoptimal transport problems also comes up in the context of comparing seismic signals [Engquistand Froese, 2013] and in single-cell perturbations [Bunne et al., 2021, 2022b,a]. Standard optimaltransport 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
Revised: Standard optimaltransport 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. We note this is not the out-of-sample setting of Seguy et al. [2017], Perrot et al.

Revised: [2016] that seeks tocouple measures and then extrapolate the map to larger measures containing the original measures.

Original: Overview and outline.
Revised: Overview and outline.

Original: We study the use of amortized optimization and machine learning methodsto rapidly solve multiple optimal transport problems and predict the solution from the input measures( α, β ) .
Revised: We study the use of amortized optimization and machine learning methodsto rapidly solve multiple optimal transport problems and predict the solution from the input measures( α, β ) .

Original: This setting involves learning a meta model to predict the solution to the optimal transportproblem, which we will refer to as Meta Optimal Transport .
Revised: This setting involves learning a meta model to predict the solution to the optimal transportproblem, which we will refer to as Meta Optimal Transport .

Original: We learn Meta OT models to predictthe solutions to optimal transport problems and signiﬁcantly improve the computational time andnumber of iterations needed to solve eq.
Revised: We learn Meta OT models to predictthe solutions to optimal transport problems and significantly improve the computational time andnumber of iterations needed to solve eq.

Original: (1) between discrete (sect. 3.1) and continuous (sect. 3.2)measures.
Revised: (1) between discrete (sect. 3.1) and continuous (sect. 3.2)measures.

Original: The paper is organized as follows:  sect.
Revised: The paper is organized as follows: in sect.

Original: 2 recalls the main concepts needed for the restof the paper, in particular the formulations of the entropy regularized and unregularized optimaltransport problems and the basic notions of amortized optimization; sect.
Revised: 2 we recall the main concepts needed for therest of the paper, in particular the formulations of the entropy regularized and unregularized optimaltransport problems and the basic notions of amortized optimization.

Original:  3 presents the Meta OTmodels and algorithms; and  sect.
Revised: In sect. 3 we present the Meta Optimal Transport models and algorithms. In sect.

Original: 4  empirically demonstrates the effectiveness of Meta OT. Settings that are not Meta OT.
Revised: 4 we empirically demonstrate the effectivenessof Meta OT in practice.

Original: Meta OT is not useful in OT settings that do not involve repeatedlysolving OT problems over a ﬁxed distribution, including 1) standard generative modeling settings,such as Arjovsky et al. [2017] that estimate the OT distance between the data and model distributions, and 2) the out-of-sample setting of Seguy et al.
Revised: Finally, sect.

Original: [2018], Perrot et al.
Revised: 5 and sect.

Original: [2016] that couplemeasures and then extrapolate the map to larger measures containing the original measures.
Revised: 6 contain the limitation of the proposed approachand related work respectively.

Original:  Preliminaries and background
Revised:  Preliminaries and background2.1 Dual optimal transport solvers

Original: 2.1 Dual optimal transport solvers

Original: We review foundations of optimal transportation, following the notation of Peyré et al.
Revised: We review foundations of optimal transportation, following the notation of Peyré et al.

Original: [2019] inmost places.
Revised: [2019] inmost places.

Original: The discrete setting often favors the entropic regularized version since it can be computed efﬁciently and in a parallelized way using the Sinkhorn algorithm.
Revised: The discrete setting often favors the entropic regularized version since it can be computed efficiently and in a parallelized way using the Sinkhorn algorithm.

Original: On the other hand, thecontinuous setting is often solved from samples using convex potentials.
Revised: On the other hand, thecontinuous setting is often solved from samples using convex potentials.

Original: While the primal Kantorovich formulation in eq.
Revised: While the primal Kantorovich formulation in eq.

Original: (1) provides an intuitive problem description, optimal transport problemsare rarely solved directly in this form due to the high-dimensionality of the couplings π and the difﬁculty of satisfying the coupling constraints U ( α, β ) .
Revised: (1) provides an intuitive problem description, optimal transport problemsare rarely solved directly in this form due to the high-dimensionality of the couplings π and the difficulty of satisfying the coupling constraints U ( α, β ) .

Original: Instead, most computational OT solvers use thedual of eq.
Revised: Instead, most computational OT solvers use thedual of eq.

Original: (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 Sinkhorn algorithm
Revised: (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 Sinkhorn algorithm

Original: Algorithm  for iteration i = 1 to N do g i ← (cid:15) log b − (cid:15) log (cid:0) K(cid:62) exp { f i − 1 /(cid:15) } (cid:1) f
Revised: Algorithm 1 Sinkhorn( α, β, c, ϵ, f 0 = 0 , g 0 = 0 ) for iteration i = 1 to N do f i ← ϵ log a − ϵ log ( K exp { g i − 1 /ϵ } )

Original:  i ← (cid:15) log a − (cid:15) log ( K  exp { g i /(cid:15) } ) end for
Revised: g i ← ϵ log b − ϵ log (cid:0) K ⊤ exp { f i − 1 /ϵ } (cid:1) end for

Original: Compute P N from f N , g N using eq.
Revised: Compute P N from f N , g N using eq.

Original: (6) return P
Revised: (6) return P

Original: N ≈ P (cid:63)
Revised: N ≈ P ⋆

Original: Let α := (cid:80) mi =1 a i δ x i and β :
Revised: Let α := (cid:80) mi =1 a i δ x i and β :

Original: =
Revised: =

Original: (cid:80) ni
Revised: (cid:80) ni

Original: =1 b i
Revised: =1 b i

Original: δ
Revised: δ

Original: y
Revised: y

Original: i be 62 discrete measures, where δ z is a Dirac at pointz and a ∈ ∆ m − 1 and b ∈ ∆ n − 1 are in theprobability simplex deﬁned by∆ k − 1 :
Revised: i be 60 discrete measures, where δ z is a Dirac at pointz and a ∈ ∆ m − 1 and b ∈ ∆ n − 1 are in theprobability simplex defined by∆ k − 1 :

Original: = { x ∈ R
Revised: = { x ∈ R

Original: k : x ≥ 0 and (cid:88)
Revised: k : x ≥ 0 and (cid:88)

Original: i x i = 1 } .
Revised: [Equation]

Original: Discrete OT.
Revised: Discrete OT.

Original: In the discrete setting, eq.
Revised: In the discrete setting, eq.

Original: (1) simpliﬁes to the linear program
Revised: (1) simplifies to the linear program

Original: P(cid:63) ( α, β, c ) ∈ arg min
Revised: P

Original: [Equation]
Revised: [Equation]

Original: where P is a coupling matrix , P (cid:63) ( α, β ) is the optimal coupling, and the cost can be discretized as amatrix C ∈ R m × n with entries C i,j :
Revised: where P is a coupling matrix , P ⋆ ( α, β ) is the optimal coupling, and the cost can be discretized as amatrix C ∈ R m × n with entries C i,j :

Original: = c ( x i , y j ) , and (cid:104) C, P (cid:105) : = (cid:80) i,j C i,j P i,j ,
Revised: = c ( x i , y j ) , and ⟨ C, P ⟩ := (cid:80) i,j C i,j P i,j ,

Original: Entropic OT.
Revised: Entropic OT.

Original: The linear program above can be regularized adding the entropy of the coupling tosmooth the objective as in Cominetti and Martín [1994], Cuturi [2013], resulting in: 
Revised: The linear program above can be regularized adding the entropy of the coupling tosmooth the objective as in Cominetti and Martín [1994], Cuturi [2013], resulting in: [Equation]

Original: [Equation]
Revised: [Equation]

Original: where H ( P ) :
Revised: where H ( P ) :

Original: = − (cid:80) i,j P i,j (log( P i,j ) − 1) is the discrete entropy of a coupling matrix P .
Revised: = − (cid:80) i,j P i,j (log( P i,j ) − 1) is the discrete entropy of a coupling matrix P .

Original: Entropic OT dual.
Revised: Entropic OT dual.

Original: As presented in Peyré et al.
Revised: As presented in Peyré et al.

Original: [2019, Prop. 4.4], the dual of eq.
Revised: [2019, Prop. 4.4], the dual of eq.

Original: (4) is
Revised: (4) isf ⋆ , g ⋆ ∈ arg max

Original: [Equation]
Revised: [Equation]

Original: where K ∈ R m × n is the Gibbs kernel and the dual variables or potentials f ∈ R n and g ∈ R m areassociated, respectively, with the marginal constraints P 1 m = a and P (cid:62) 1 n = b .
Revised: where K ∈ R m × n is the Gibbs kernel and the dual variables or potentials f ∈ R n and g ∈ R m areassociated, respectively, with the marginal constraints P 1 m = a and P ⊤ 1 n = b .

Original: The optimal dualsdepend on the problem, e.g. f (cid:63) ( α, β, c, (cid:15) ) , but we omit this dependence for notational simplicity.
Revised: The optimal dualsdepend on the problem, e.g. f ⋆ ( α, β, c, ϵ ) , but we omit this dependence for notational simplicity.

Original: Recovering the primal solution from the duals.
Revised: Recovering the primal solution from the duals.

Original: Given optimal duals f (cid:63) , g (cid:63) that solve eq.
Revised: Given optimal duals f ⋆ , g ⋆ that solve eq.

Original: (5) theoptimal coupling P (cid:63) to the primal problem in eq.
Revised: (5) theoptimal coupling P ⋆ to the primal problem in eq.

Original: (4) can be obtained by
Revised: (4) can be obtained by

Original: [Equation]
Revised: [Equation]

Revised: K is defined in eq

Revised: [Equation]

Revised: [Equation]

Revised: [Equation]

Original: The Sinkhorn algorithm .
Revised: The Sinkhorn algorithm .

Original: Algorithm 1 summarizes the log-space version, which takes closed-formblock coordinate ascent updates on eq.
Revised: Algorithm 1 summarizes the log-space version, which takes closed-formblock coordinate ascent updates on eq.

Original: (5) obtained from the ﬁrst-order optimality conditions [Peyréet al., 2019, Remark 4.21].
Revised: (5) obtained from the first-order optimality conditions [Peyréet al., 2019, Remark 4.21].

Original: We will use it to ﬁne-tune predictions made by our Meta OT models.
Revised: We will use it to fine-tune predictions made by our Meta OT models.

Original: Computing the error.
Revised: Computing the error.

Original: Standard implementations of the Sinkhorn algorithm, such as Flamary et al.[2021], Cuturi et al.
Revised: Standard implementations of the Sinkhorn algorithm, such as Flamary et al.[2021], Cuturi et al.

Original: [2022], measure the error of a candidate dual solution ( f, g ) by computing thedeviation from the marginal constraints, which we will also use in comparing our solution quality: 
Revised: [2022], measure the error of a candidate dual solution ( f, g ) by computing thedeviation from the marginal constraints, which we will also use in comparing our solution quality: [Equation]

Original: [Equation]
Revised: [Equation]

Original: Mapping between the duals.
Revised: Mapping between the duals.

Original: The ﬁrst-order optimality conditions of eq.
Revised: The first-order optimality conditions of eq.

Original: (5) also provide an equivalence between the optimal dual potentials that we will make use of:[Equation]
Revised: (5) also provide an equivalence between the optimal dual potentials that we will make use of:

Original: [Equation]
Revised: [Equation]

Original: 2.1.2 Wasserstein-2 OT between continuous (Euclidean) measures with dual potentialsfor iteration
Revised: 2.1.2 Wasserstein-2 OT between continuous (Euclidean) measures with dual potentials

Original:  ≈ T (cid:63) ( · )
Revised: Algorithm for iteration i = 1 to N do Sample from ( α, β ) and estimate L ( φ i − 1 ) Update φ i with approximation to ∇ φ L ( φ i − 1 ) end for return T N ( · ) : = ∇ x ψ φ N ( · ) ≈ T ⋆ ( · )

Original: Let α and β be continuous measures in Euclideanspace X
Revised: Let α and β be continuous measures in Euclideanspace X

Original: = Y =
Revised: = Y =

Original: R d (with α absolutely continuous with respect to the Lebesgue measure) andthe ground cost be the squared Euclidean distancec ( x, y ) :
Revised: R d (with α absolutely continuous with respect to the Lebesgue measure) andthe ground cost be the squared Euclidean distancec ( x, y ) :

Original: = (cid:107) x − y (cid:107) 22 .
Revised: = ∥ x − y ∥ 22 .

Original: Then the minimum of eq.
Revised: Then the minimum of eq.

Original: (1)deﬁnes the square of the Wasserstein-2 distance:
Revised: (1)defines the square of the Wasserstein-2 distance:

Original: [Equation]
Revised: [Equation]

Original: where T is a transport map pushing α to β , i.e. T # α = β with the pushforward operator deﬁnedby T # α ( B ) :
Revised: where T is a transport map pushing α to β , i.e. T # α = β with the pushforward operator definedby T # α ( B ) :

Original: =
Revised: =

Original: α ( T − 1 ( B )) for any measurable set B .
Revised: α ( T − 1 ( B )) for any measurable set B .

Original: Convex dual potentials.
Revised: Convex dual potentials.

Original: The primal form in eq.
Revised: The primal form in eq.

Original: (9) is difﬁcult to solve, as in the discrete setting, dueto the difﬁculty of representing the coupling and satisfying the constraints.
Revised: (9) is difficult to solve, as in the discrete setting, dueto the difficulty of representing the coupling and satisfying the constraints.

Original: Makkuva et al.
Revised: Makkuva et al.

Original: Taghvaei and Jalali [2019], Korotin et al.
Revised: Taghvaei and Jalali [2019], Korotin et al.

Original: [2019, 2021b, 2022] propose to instead solve the dual:
Revised: [2019, 2021b, 2022] propose to instead solve the dual:

Original: [Equation]
Revised: [Equation]

Original: where ψ is a convex function referred to as a convex potential , and ψ ( y ) :
Revised: where ψ is a convex function referred to as a convex potential , and ψ ( y ) :

Original: = max x ∈X (cid:104) x, y (cid:105)− ψ ( x ) isthe Legendre-Fenchel transform or convex conjugate of ψ [Fenchel, 1949, Rockafellar, 2015].
Revised: = max x ∈X ⟨ x, y ⟩− ψ ( x ) isthe Legendre-Fenchel transform or convex conjugate of ψ [Fenchel, 1949, Rockafellar, 2015].

Original: Thepotential ψ is often approximated with an input-convex neural network (ICNN)
Revised: Thepotential ψ is often approximated with an input-convex neural network (ICNN)

Original: [Amos et al., 2017].
Revised: [Amos et al., 2017].

Original: Recovering the primal solution from the dual.
Revised: Recovering the primal solution from the dual.

Original: Given an optimal dual ψ (cid:63) for eq.
Revised: Given an optimal dual ψ ⋆ for eq.

Original: (10), Brenier[1991] remarkably shows that an optimal map T (cid:63) for eq.
Revised: (10), Brenier[1991] remarkably shows that an optimal map T ⋆ for eq.

Original: (9) can be obtained with differentiation:
Revised: (9) can be obtained with differentiation:

Original: [Equation]
Revised: [Equation]

Original: Wasserstein-2 Generative Networks (W2GNs).
Revised: Wasserstein-2 Generative Networks (W2GNs).

Original: Korotin et al.
Revised: Korotin et al.

Original: [2019] model ψ ϕ and ψ ϕ in eq.
Revised: [2019] model ψ φ and ψ φ in eq.

Original: (10)with two separate ICNNs parameterized by ϕ .
Revised: (10)with two separate ICNNs parameterized by φ .

Original: The separate model for ψ ϕ is useful because theconjugate operation in eq.
Revised: The separate model for ψ φ is useful because theconjugate operation in eq.

Original: (10) becomes computationally expensive.
Revised: (10) becomes computationally expensive.

Original: They optimize the loss:
Revised: They optimize the loss:

Original: [Equation]
Revised: [Equation]

Original: where ϕ is a detached copy of the parameters and γ is a hyper-parameter.
Revised: where φ is a detached copy of the parameters and γ is a hyper-parameter.

Original: The ﬁrst term are thecyclic monotone correlations [Chartrand et al., 2009, Taghvaei and Jalali, 2019],that optimize thedual objective in eq. (10), and the second term provides cycle consistency [Zhu et al., 2017] toestimate the conjugate ψ .
Revised: The first term are thecyclic monotone correlations [Chartrand et al., 2009, Taghvaei and Jalali, 2019], and the secondterm provides cycle consistency [Zhu et al., 2017].

Original: Algorithm 2 shows how L is typically optimized using samples from the  α β c π (cid:63) θ
Revised: Algorithm 2 shows how L is typically optimizedusing samples from the measures, which we use to fine-tune Meta OT predictions. α β c π ⋆ θ

Original: D
Revised: D

Original: General α β c f (cid:63) g (cid:63)
Revised: General α β c f ⋆ g ⋆

Original: P (cid:63) θ
Revised: P ⋆ θ

Original: D
Revised: D

Original: Discrete (Entropic) α β ψ (cid:63)
Revised: Discrete (Entropic) α β ψ ⋆

Original: T (cid:63) θ
Revised: T ⋆ θ

Original: D
Revised: D

Original: Continuous (Wasserstein-2)
Revised: Continuous (Wasserstein-2)

Original: Input measures and cost Dual potentials Couplings
Revised: [Figure]

Original: 2.2 Amortized optimization and learning to optimize
Revised: 2.2 Amortized optimization and learning to optimize

Original: Our paper is an application of amortized optimization methods that predict the solutions of optimization problems, as surveyed in, e.g., Chen et al.
Revised: Our paper is an application of amortized optimization methods that predict the solutions of optimization problems, as surveyed in, e.g., Chen et al.

Original: [2021], Amos [2022].
Revised: [2021], Amos [2022].

Original: We use the basic setupfrom Amos [2022], which considers unconstrained continuous optimization problems of the form
Revised: We use the basic setupfrom Amos [2022], which considers unconstrained continuous optimization problems of the form

Original: [Equation]
Revised: [Equation]

Original: (13) where J is the objective, z ∈ Z is the domain , and φ ∈ Φ is some context or parameterization .
Revised: (13) where J is the objective, z ∈ Z is the domain , and ϕ ∈ Φ is some context or parameterization .

Original: Inother words, the context conditions the objective but is not optimized over.
Revised: Inother words, the context conditions the objective but is not optimized over.

Original: Given a distribution overcontexts P ( φ ) , we learn a model ˆ z θ parameterized by θ to approximate eq.
Revised: Given a distribution overcontexts P ( ϕ ) , we learn a model ˆ z θ parameterized by θ to approximate eq.

Original: (13), i.e. ˆ z θ ( φ ) ≈
Revised: (13), i.e. ˆ z θ ( ϕ ) ≈

Original: z (cid:63) ( φ ) .
Revised: z ⋆ ( ϕ ) .

Original: J will be differentiable for us, so we optimize the parameters using objective-based learning with
Revised: J will be differentiable for us, so we optimize the parameters using objective-based learning with

Original: [Equation]
Revised: [Equation]

Original: (14) which does not require ground-truth solutions z (cid:63) and can be optimized with a gradient-based solver.While we focus on optimizing eq.
Revised: (14) which does not require ground-truth solutions z ⋆ and can be optimized with a gradient-based solver.

Original: (14) because we do not assume easy access to ground-truth solutions z (cid:63) ( φ ) , one alternative is regression-based learning if the solutions are easily available:

Original: [Equation]

Original:  Meta Optimal Transport
Revised:  Meta Optimal Transport

Original: Figure 1 illustrates our key contribution of connecting objective-based amortization in eq.
Revised: Figure 1 illustrates our key contribution of connecting objective-based amortization in eq.

Original: (14) tooptimal transport.
Revised: (14) tooptimal transport.

Original: We consider solving multiple OT problems and learning shared structure andcorrelations between them.
Revised: We consider solving multiple OT problems and learning shared structure andcorrelations between them.

Original: We denote a joint meta-distribution over the input measures and costswith D ( α, β, c ) , which we call meta to distinguish it from the measures α, β .
Revised: We denote a joint meta-distribution over the input measures and costswith D ( α, β, c ) , which we call meta to distinguish it from the measures α, β .

Original: In general, we could introduce a model that directly predicts the primal solution to eq.
Revised: In general, we could introduce a model that directly predicts the primal solution to eq.

Original: (1), i.e.π θ ( α, β, c ) ≈ π (cid:63) ( α, β, c ) for ( α, β, c ) ∼ D .
Revised: (1), i.e.π θ ( α, β, c ) ≈ π ⋆ ( α, β, c ) for ( α, β, c ) ∼ D .

Original: This is difﬁcult for the same reason why most computational methods do not operate directly in the primal space: the optimal coupling is often a highdimensional joint distribution with non-trivial marginal constraints.
Revised: This is difficult for the same reason why most computational methods do not operate directly in the primal space: the optimal coupling is often a highdimensional joint distribution with non-trivial marginal constraints.

Original: We instead turn to predictingthe dual variables used by today’s solvers.
Revised: We instead turn to predictingthe dual variables used by today’s solvers.

Original: 3.1 Meta OT between discrete measures
Revised: 3.1 Meta OT between discrete measures

Original: We build on standard methods for entropic OT reviewed in sect.
Revised: We build on standard methods for entropic OT reviewed in sect.

Original: 2.1.1 between discrete measuresα := (cid:80) mi =1 a i δ x i and β :=
Revised: 2.1.1 between discrete measuresα := (cid:80) mi =1 a i δ x i and β :=

Original: (cid:80) ni
Revised: (cid:80) ni

Original: =1 b i δ x i with a ∈ ∆ m − 1 and b ∈ ∆ n − 1 coupled using a cost c .
Revised: =1 b i δ x i with a ∈ ∆ m − 1 and b ∈ ∆ n − 1 coupled using a cost c .

Original: In the
Revised: In the

Original: Meta OT setting, the measures and cost are the contexts for amortization and sampled from a metadistribution , i.e. ( α, β, c ) ∼ D ( α, β, c ) .
Revised: Meta OT setting, the measures and cost are the contexts for amortization and sampled from a metadistribution , i.e. ( α, β, c ) ∼ D ( α, β, c ) .

Original: For example, sects.
Revised: For example, sects.

Original: 4.1 and 4.2 considers meta-distributionsover the weights of the atoms, i.e. ( a, b ) ∼ D , where D is a distribution over ∆ m − 1 × ∆ n − 1 . 
Revised: 4.1 and 4.2 considers meta-distributionsover the weights of the atoms, i.e. ( a, b ) ∼ D , where D is a distribution over ∆ m − 1 × ∆ n − 1 . Amortization objective.

Revised: We will seek to predict the optimal potential.

Revised: At optimality, the pair ofpotentials are related to each other via eq.

Revised: (8), i.e. g ( f ; α, β, c ) :

Original:  Initialize amortization model with θ 0 for iteration i do
Revised: = ϵ log b − ϵ log (cid:0) K ⊤ exp { f/ϵ } (cid:1)where K ∈ R m × n is the Gibbs kernel from eq. Hence, it is sufficient to predict one of the Initialize amortization model with θ 0 for iteration i do

Original: Sample ( α, β, c ) ∼ D
Revised: Sample ( α, β, c ) ∼ D

Original: Predict duals ˆ f θ or ˆ ϕ θ on the sample Estimate the loss in eq.
Revised: Predict duals ˆ f θ or ˆ φ θ on the sample Estimate the loss in eq.

Original: (17) or eq.
Revised: (16) or eq.

Original: Update θ i +1 with a gradient step end for
Revised: Update θ i +1 with a gradient step end for

Original: Predict duals ˆ f θ ( α, β, c )return Sinkhorn( α, β, c, (cid:15), ˆ f θ )
Revised: Predict duals ˆ f θ ( α, β, c )

Original: Predict dual ICNN parameters ˆ ϕ θ ( α, β, c ) return W2GN( α, β, c, T, ˆ ϕ θ ) Amortization objective. We will seek to predict the optimal potential. At optimality, the pair ofpotentials are related to each other via eq.
Revised: Compute ˆ g from ˆ f θ using eq.

Original: (8), i.e. g ( f ; α, β, c ):
Revised: (8) return Sinkhorn( α, β, c, ϵ, ˆ f θ , ˆ g )

Original: = (cid:15) log b − (cid:15) log (cid:0) K (cid:62) exp { f/(cid:15) } (cid:1)where K ∈ R m × n is the Gibbs kernel from eq. Hence, it is sufﬁcient to predict one of thepotentials, e.g. f , and recover the other.
Revised: Predict dual ICNN parameters ˆ φ θ ( α, β, c ) return W2GN( α, β, c, T, ˆ φ θ ) potentials, e.g. f , and recover the other.

Original: We thus re-formulate eq.
Revised: We thus re-formulate eq.

Original: (5) to just optimize over f with
Revised: (5) to just optimize over f with

Original: [Equation]
Revised: [Equation]

Original: (16) where − J ( f ; α, β, c ) : = (cid:104) f, a (cid:105) + (cid:104) g, b (cid:105)− (cid:15) (cid:104) exp { f/(cid:15) } , K exp { g/(cid:15) }(cid:105) is the(negated) dual objective.
Revised: where − J ( f ; α, β, c ) := ⟨ f, a ⟩ + ⟨ g, b ⟩ is the dual objective over f .

Original: Even though most solvers optimize over f and g jointly as in eq.
Revised: Even though most solversoptimizers over f and g jointly as in eq.

Original: (16), amortizing over these wouldlikely need: 1) to have a higher capacity than a model just predicting f , and 2) to learn how f and gare connected through eq.
Revised: (15), amortizing over these would likely need: 1) to have ahigher capacity than a model just predicting f , and 2) to learn how f and g are connected througheq.

Original: (8) while in eq.
Revised: (8) while in eq.

Original: (16) we explicitly provide this knowledge.
Revised: (15) we explicitly provide this knowledge.

Original: Amortization model.
Revised: Amortization model.

Original: We predict the solution to eq.
Revised: We predict the solution to eq.

Original: (16) with ˆ f θ ( α, β, c ) parameterized by θ ,resulting in a computationally efﬁcient approximation ˆ f θ ≈ f (cid:63) .
Revised: (15) with ˆ f θ ( α, β, c ) parameterized by θ ,resulting in a computationally efficient approximation ˆ f θ ≈ f ⋆ .

Original: Here we use the notation ˆ f θ ( α, β, c )to mean that the model ˆ f θ depends on representations of the input measures and cost.
Revised: Here we use the notation ˆ f θ ( α, β, c )to mean that the model ˆ f θ depends on representations of the input measures and cost.

Original: In our settings,we deﬁne ˆ f θ as a fully-connected MLP mapping from the atoms of the measures to the duals.
Revised: In our settings,we define ˆ f θ as a fully-connected MLP mapping from the atoms of the measures to the duals.

Original: Amortization loss.
Revised: Amortization loss.

Original: Applying objective-based amortization from eq.
Revised: Applying objective-based amortization from eq.

Original: (14) to the dual in eq.
Revised: (14) to the dual in eq.

Original: (16)completes our learning setup.
Revised: (15)completes our learning setup.

Original: Our model should best-optimize the expectation of the dual objective
Revised: Our model should best-optimize the expectation of the dual objective

Original: [Equation]
Revised: [Equation]

Revised: [Equation]

Original: which is appealing as it does not require ground-truth solutions f (cid:63) .
Revised: which is appealing as it does not require ground-truth solutions f ⋆ .

Original: Algorithm 3 shows a basictraining loop for eq.
Revised: Algorithm 3 shows a basictraining loop for eq.

Original: (17) using a gradient-based optimizer such as Adam [Kingma and Ba, 2014].
Revised: (16) using a gradient-based optimizer such as Adam [Kingma and Ba, 2014].

Original: Sinkhorn ﬁne-tuning.
Revised: Sinkhorn fine-tuning.

Original: The dual prediction made by ˆ f θ with an associated ˆ g can easily be input asthe initialization to a standard Sinkhorn solver as shown in algorithm 4.
Revised: The dual prediction made by ˆ f θ with an associated ˆ g can easily be input asthe initialization to a standard Sinkhorn solver as shown in algorithm 4.

Original: This allows us to deploy thepredicted potential with Sinkhorn to obtain the optimal potentials with only a few extra iterations.
Revised: This allows us to deploy thepredicted potential with Sinkhorn to obtain the optimal potentials with only a few extra iterations.

Original: On accelerated solvers.
Revised: On accelerated solvers.

Original: Here we have only considered ﬁne-tuning the Meta OT prediction witha log-Sinkhorn solver.
Revised: Here we have only considered fine-tuning the Meta OT prediction witha log-Sinkhorn solver.

Original: Meta OT can also be combined with accelerated variants of entropic OTsolvers such as Thibault et al.
Revised: Meta OT can also be combined with accelerated variants of entropic OTsolvers such as Thibault et al.

Original: [2017], Altschuler et al.
Revised: [2017], Altschuler et al.

Original: [2017], Alaya et al.
Revised: [2017], Alaya et al.

Original: [2019], Lin et al.
Revised: [2019], Lin et al.

Original: [2019]that would otherwise solve every problem from scratch.
Revised: [2019]that would otherwise solve every problem from scratch.

Revised: Convergence rates.

Revised: The knowledge of the hyper-distribution D of problems being solved enables

Revised: Meta OT methods to surpass the convergence rates and computational time of standard algorithmsby restricting the set of problems rather than considering the average- or worst-case performance.

Revised: The model ˆ f θ distills information between the problem instances into the parameters θ .

Original: 3.2 Meta OT between continuous measures (Wasserstein-2)
Revised: 3.2 Meta OT between continuous measures (Wasserstein-2)

Original: We take an analogous approach to predicting the Wasserstein-2 map between continuous measuresfor Wasserstein-2 as reviewed in sect.
Revised: We take an analogous approach to predicting the Wasserstein-2 map between continuous measuresfor Wasserstein-2 as reviewed in sect.

Original: Here the measures α, β are supported in continuousspace X
Revised: Here the measures α, β are supported in continuousspace X

Original: = Y =
Revised: = Y =

Original: R d
Revised: R d

Original: and we focus on computing Wasserstein-2 couplings from instances sampledfrom a meta-distribution ( α, β ) ∼ D ( α, β ) .
Revised: and we focus on computing Wasserstein-2 couplings from instances sampledfrom a meta-distribution ( α, β ) ∼ D ( α, β ) .

Original: The cost c is not included in D as it remains ﬁxed to thesquared Euclidean cost everywhere here.
Revised: The cost c is not included in D as it remains fixed to thesquared Euclidean cost everywhere here.

Original: One challenge here is that the optimal dual potential ψ (cid:63) ( · ; α, β ) in eq.
Revised: One challenge here is that the optimal dual potential ψ ⋆ ( · ; α, β ) in eq.

Original: (10) is a convex function andnot simply a ﬁnite-dimensional real vector.
Revised: (10) is a convex function andnot simply a finite-dimensional real vector.

Original: The dual potentials in this setting are approximated by,e.g., an ICNN.
Revised: The dual potentials in this setting are approximated by,e.g., an ICNN.

Original: We thus propose a Meta ICNN that predicts the parameters ϕ of an ICNN ψ ϕ thatapproximates the optimal dual potentials, which can be seen as a hypernetwork [Stanley et al., 2009,
Revised: We thus propose a Meta ICNN that predicts the parameters φ of an ICNN ψ φ thatapproximates the optimal dual potentials, which can be seen as a hypernetwork [Stanley et al., 2009,

Original: Ha et al., 2016].
Revised: Ha et al., 2016].

Original: The dual prediction made by ˆ ϕ θ can easily be input as the initial value to a standardW2GN solver as shown in algorithm 5. App.
Revised: The dual prediction made by ˆ φ θ can easily be input as the initial value to a standard

Original: B discusses other modeling choices we considered:we tried models based on MAML

Original: [Finn et al., 2017] and neural processes [Garnelo et al., 2018b,a].

Original: Sinkhorn (converged, ground-truth)
Revised: Sinkhorn (converged, ground-truth)

Original: [Equation]
Revised: [Equation]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Equation] [Figure] [Equation] [Figure]

Original:  [Table]
Revised: [Figure] [Figure] [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: Amortization objective.
Revised: We report the mean deviation W2GN solver as shown in algorithm 5. App. B discusses other modeling choices we considered:we tried models based on MAML

Revised: [Finn et al., 2017] and neural processes [Garnelo et al., 2018b,a].

Revised: Amortization objective.

Original: We build on the W2GN formulation [Korotin et al., 2019] and seek parameters ϕ (cid:63) optimizing the dual ICNN potentials ψ ϕ and ψ ϕ with L ( ϕ ; α, β ) from eq.
Revised: We build on the W2GN formulation [Korotin et al., 2019] and seek parameters φ ⋆ optimizing the dual ICNN potentials ψ φ and ψ φ with L ( φ ; α, β ) from eq.

Original: Wechose W2GN due to the stability, but could also easily use other losses optimizing ICNN potentials.
Revised: Wechose W2GN due to the stability, but could also easily use other losses optimizing ICNN potentials.

Original: Amortization model: the Meta ICNN.
Revised: Amortization model: the Meta ICNN.

Original: We predict the solution to eq.
Revised: We predict the solution to eq.

Original: (12) with ˆ ϕ θ ( α, β ) parameterized by θ , resulting in a computationally efﬁcient approximation to the optimumˆ ϕ θ ≈ ϕ (cid:63) .
Revised: (12) with ˆ φ θ ( α, β ) parameterized by θ , resulting in a computationally efficient approximation to the optimum

Original:  We α measures, which in  are simply 
Revised: ˆ φ θ ≈ φ ⋆ . Figure coupling image-based measures We again emphasize that α, β used with the model hereare representations of measures, which in our cases are simply images.

Original: Amortization loss.
Revised: Amortization loss.

Original: Applying objective-based amortization from eq.
Revised: Applying objective-based amortization from eq.

Original: (14) to the W2GN loss ineq.
Revised: (14) to the W2GN loss ineq.

Original: (12) completes our learning setup.
Revised: (12) completes our learning setup.

Original: Our model should best-optimize the expectation of the loss:
Revised: Our model should best-optimize the expectation of the loss:

Original: [Equation]
Revised: [Equation]

Original: As in the discrete setting, it does not require ground-truth solutions ϕ (cid:63) and we learn it with Adam.
Revised: As in the discrete setting, it does not require ground-truth solutions φ ⋆ and we learn it with Adam.

Original:  Experiments
Revised:  Experiments

Original: We demonstrate how Meta OT models improve the convergence of the state-of-the-art solvers insettings where solving multiple OT problems naturally arises.
Revised: We demonstrate how Meta OT models improve the convergence of the state-of-the-art solvers insettings where solving multiple OT problems naturally arises.

Original: We implemented our code in JAX[Bradbury et al., 2018] as an extension to the the Optimal Transport Tools (OTT) package [Cuturiet al., 2022].
Revised: We implemented our code in JAX[Bradbury et al., 2018] as an extension to the the Optimal Transport Tools (OTT) package [Cuturiet al., 2022].

Original: App.
Revised: We have attached the complete source code to reproduce every experiment in our paper [Figure]

Original: C covers further experimental and implementation details, and shows that all ofour experiments take a few hours to run on our single Quadro GP100{+GP
+} GPU.
Revised: [Figure] and will open source it. All experiments take roughly ≈ 2 hours to run on our single Quadro {+GP
+} GPU.

Original: [Figure]
Revised: App.

Original: Gaussian [Thornton and Cuturi, 2022]
Revised: C covers further experimental and implementation details.

Original: Meta OT) Discrete  between 
Revised: 4.1 Discrete OT between MNIST digits

Original: Images provide a natural setting for Meta OT where the distribution over images provide the metadistribution D over OT problems.
Revised: Images provide a natural setting for Meta OT where the distribution over images provide the metadistribution D over OT problems.

Original: Given a pair of images α 0 and α 1 , each grayscale image iscast as a discrete measure in 2-dimensional space where the intensities deﬁne the probabilities ofthe atoms.
Revised: Given a pair of images α 0 and α 1 , each grayscale image iscast as a discrete measure in 2-dimensional space where the intensities define the probabilities ofthe atoms.

Original: The goal is to compute the optimal transport interpolation between the two measuresas in, e.g., Peyré et al.
Revised: The goal is to compute the optimal transport interpolation between the two measuresas in, e.g., Peyré et al.

Original: Formally, this means computing the optimal coupling P (cid:63) bysolving the entropic optimal transport problem between α 0 and α 1 and computing the interpolatesas α t =
Revised: Formally, this means computing the optimal coupling P ⋆ bysolving the entropic optimal transport problem between α 0 and α 1 and computing the interpolatesas α t =

Original: ( t proj y +(1 − t ) proj x ) # P (cid:63) , for t ∈ [0 , 1] , where proj x ( x, y ) :
Revised: ( t proj y +(1 − t ) proj x ) # P ⋆ , for t ∈ [0 , 1] , where proj x ( x, y ) :

Original: = x and proj y ( x, y ) =
Revised: = x and proj y ( x, y ) =

Original: y .
Revised: y .

Original: We selected (cid:15) = 10 − 2 as app.
Revised: We selected ϵ = 10 − 2 as app.

Original: A shows that it gives interpolations that are not too blurry or sharp.
Revised: A shows that it gives interpolations that are not too blurry or sharp.

Original: Our Meta OT model ˆ f θ (sect. 3.1) is an MLP that predicts the transport map between pairs of MNISTdigits.
Revised: Our Meta OT model ˆ f θ (sect. 3.1) is an MLP that predicts the transport map between pairs of MNIST digits.

Original: We train on every pair from the standard training dataset.
Revised: We train on every pair from the standard training dataset.

Original: Figure 2 shows that even withoutﬁne-tuning, Meta OT’s predicted Wasserstein interpolations between the measures are close to theground-truth interpolations obtained from running the Sinkhorn algorithm to convergence.
Revised: Figure 2 shows that evenwithout fine-tuning, Meta OT’s predicted Wasserstein interpolations between the measures are closeto the ground-truth interpolations obtained from running the Sinkhorn algorithm to convergence.

Original: We thenﬁne-tune Meta OT’s prediction with Sinkhorn as in algorithm 4.
Revised: We then fine-tune Meta OT’s prediction with Sinkhorn as in algorithm 4.

Original: Figure 4 shows that the nearoptimal predictions can be quickly reﬁned in fewer iterations than running Sinkhorn with the defaultinitialization, and table 1 shows the runtime required to reach the default threshold, which uses thedefault marginal error threshold of 10 − 3 .
Revised: Figure 4 shows that thenear-optimal predictions can be quickly refined in fewer iterations than running Sinkhorn with thedefault initialization, and table 1 shows the runtime required to reach the default threshold.

Original: We compare our learned initialization to the standard zeroinitialization, as well as the Gaussian initialization proposed in Thornton and Cuturi [2022], whichtakes a continuous Gaussian approximation of the measures and initializes the potentials to be theknown coupling between the Gaussians.
Revised: α β

Original: This Gaussian initialization assumes the squared Euclideancost, which is not the case in our spherical transport problem, but we ﬁnd it is still helpful over thezero initialization.
Revised: T # α T − 1# β [Figure]

Original: 4.2 Discrete OT for supply-demand transportation on spherical data
Revised: 4.2 Discrete OT for supply-demand transportation on spherical data

Original: We next set up a synthetic transport problem between supply and demand locations where the supplyand demands may change locations or quantities frequently, creating another Meta OT setting to beable to rapidly solve the new instances.
Revised: We next set up a synthetic transport problem between supply and demand locations where the supplyand demands may change locations or quantities frequently, creating another Meta OT setting to beable to rapidly solve the new instances.

Original: We speciﬁcally consider measures living on the 2-spheredeﬁned by S 2 :=
Revised: We specifically consider measures living on the 2-spheredefined by S 2 :=

Original: { x ∈ R 3 : (cid:107) x (cid:107) = 1 } , i.e. X =
Revised: { x ∈ R 3 : ∥ x ∥ = 1 } , i.e. X =

Original: Y = S 2 , with the transport cost given by thespherical distance c ( x, y ) =
Revised: Y = S 2 , with the transport cost given by thespherical distance c ( x, y ) =

Original: arccos( (cid:104) x, y (cid:105) ) .
Revised: arccos( ⟨ x, y ⟩ ) .

Original: We then randomly sample supply locations uniformlyfrom Earth’s landmass and demand locations from Earth’s population density to induce a class oftransport problems on the sphere obtained from the CC-licensed dataset from Doxsey-Whitﬁeld etal.[2015].
Revised: We then randomly sample supply locations uniformlyfrom Earth’s landmass and demand locations from Earth’s population density to induce a class of transport problems on the sphere obtained from the CC-licensed dataset from Doxsey-Whitfield et

Original: Figure 5 shows that the predicted transport maps on test instances are close to the optimalmaps obtained from Sinkhorn to convergence.
Revised: Figure 5 shows that the predicted transport maps on test instances are close to the optimalmaps obtained from Sinkhorn to convergence.

Original: Similar to the MNIST setting, ﬁg. 4 and table 1 showimproved convergence and runtime.
Revised: Similar to the MNIST setting, fig. 4 and table 1 showsimproved convergence and runtime.

Original: 4.3 Continuous Wasserstein-2 color transfer
Revised: 4.3 Continuous Wasserstein-2 color transfer

Revised: [Figure]

Revised: warm initializations that significantly im prove the (normalized) dual values

Original: The problem of color transfer between two images consists in mapping the color palette of one imageinto the other one.
Revised: The problem of color transfer between two images consists in mapping the color palette of oneimage into the other one.

Original: The images are required to have the same number of channels, for example RGBimages.
Revised: The images are required to have the same number of channels, forexample RGB images.

Original: The continuous formulation that we use from Korotin et al.
Revised: The continuous formulation that we use from Korotin et al.

Original: [2019], takes i.e. X
Revised: [2019], takesi.e. X

Original: = Y =[0 , 1] 3 with c being the squared Euclidean distance.
Revised: = Y = [0 , 1] 3 with c being the squared Euclidean distance.

Original: We collected ≈ 200 public domain images from WikiArt and trained a Meta ICNN model from sect.
Revised: We collected ≈ 200 publicdomain images from WikiArt and trained a Meta ICNN model from sect.

Original: 3.2 to predict the color transfer maps between [Figure] (converged, ground-truth
Revised: 3.2 to predict the colortransfer maps between every pair of them.

Original: [Figure] [Equation] [Figure] generating
Revised: Figure 6 shows the predictions on test pairs and fig.

Original: Market August Argenteuil every pair of them. Figure 6 shows the predictions on test pairs and ﬁg. 7 shows the convergence incomparison to the standard W2GN learning.
Revised: 7shows the convergence in comparison to the standard W2GN learning.

Original: Table 2 reports runtimes and app.
Revised: Table 2 reports runtimesand app.

Original: E shows additionalresults.
Revised: D shows additional results.

Revised:  Limitations of Meta OT

Revised: While we have illustrated successful applications of Meta OT, it is also important to understand thelimitations: 1)

Original: Related work
Revised: Meta OT does not make previously intractable problems tractable. All of thebaseline OT solvers we consider can solve our problems within milliseconds or seconds. Outof-distribution generalization. Meta OT may not generate good predictions on instances that arenot close to the training OT problems from the meta-distribution D over the measures and cost If

Revised: Meta OT generates a bad prediction in practice, it is possible to measure the error and re-solve theinstance from scratch if necessary.

Revised:  Related work

Original: Efﬁciently estimating OT maps.
Revised: Efficiently estimating OT maps.

Original: To compute OT maps with ﬁxed cost between pairs of measuresefﬁciently, neural OT models [Korotin et al., 2019, Li et al., 2020, Korotin et al., 2021a, Mokrovet al., 2021, Korotin et al., 2021b] leverage ICNNs to estimate maps between continuous high- dimensional measures given samples from these, and Litvinenko et al.
Revised: To compute OT maps with fixed cost between pairs of measures efficiently, neural OT models [Korotin et al., 2019, Li et al., 2020, Korotin et al., 2021a, Mokrov et al., 2021, Korotin et al., 2021b] leverage ICNNs to estimate maps between continuoushigh-dimensional measures given samples from these, and Litvinenko et al.

Original: [2021], Scetbon et al. Forrow et al.
Revised: [2021], Scetbon et al.[2021a], Forrow et al.

Original: [2019], Sommerfeld et al.
Revised: [2019], Sommerfeld et al.

Original: [2019], Scetbon et al.
Revised: [2019], Scetbon et al.

Original: [2022], Muzellec and Cuturi [2019], Bonet et al.
Revised: [2021b], Muzellec and Cuturi[2019], Bonet et al.

Original: [2021] leverage structural assumptions on coupling and cost matrices to reduce thecomputational and memory complexity.
Revised: [2021] leverage structual assumptions on coupling and cost matrices to reducethe computational and memory complexity.

Original: In the meta-OT setting, we consider learning to rapidlycompute OT mappings between new pairs measures.
Revised: In the meta-OT setting, we consider learning to rapidlycompute OT mappings between new pairs measures.

Original: All these works can hence potentially beneﬁtfrom an acceleration effect by leveraging amortization similarly.[Figure]
Revised: All these works can hence potentially benefitfrom an acceleration effect by leveraging amortization similarly.

Original: warm-start initializations that signiﬁcantly improve the (normalized) dual objective values.

Original: Embedding measures where OT distances arediscriminative.
Revised: Embedding measures where OT distances are discriminative.

Original: Effort has been invested in learning encodings/projections of measures througha nested optimization problem, which aims toﬁnd discriminative embeddings of the measuresto be compared [Genevay et al., 2018, Deshpandeet al., 2019, Nguyen and Ho, 2022].
Revised: Effort has been invested in learning encodings/projections of measures through a nested optimization problem, which aims to finddiscriminative embeddings of the measures to be compared [Genevay et al., 2018, Deshpande et al. ,2019, Nguyen and Ho, 2022].

Original: While theseworks share an encoder and/or a projection acrosstask with the aim of leveraging more discriminative alignments (and hence an OT distance with ametric different from the Euclidean metric), ourwork differs in the sense that we ﬁnd good initializations to solve the OT problem itself with ﬁxedcost more efﬁciently across tasks.
Revised: While these works share an encoder and/or a projection across taskwith the aim of leveraging more discriminative alignments (and hence an OT distance with a metricdifferent from the Euclidean metric), our work differs in the sense that we find good initializationsto solve the OT problem itself with fixed cost more efficiently across tasks.

Original: Optimal transport and amortization.
Revised: Optimal transport and amortization.

Original: Few previous works in the OT literature leverage amortization.
Revised: Few previous works in the OT literature leverage amortization.

Original: Courty et al.
Revised: Courty et al.

Original: [2018] learn a latent space in which the Wasserstein distance between themeasure’s embeddings is equivalent to the Euclidean distance.
Revised: [2017a] learn a latent space in which the Wasserstein distance between themeasure’s embeddings is equivalent to the Euclidean distance.

Original: Concurrent work [Nguyen and Ho,2022] amortizes the estimation of the optimal projection in the max-sliced objective, which differsfrom our work where we instead amortize the estimation of the optimal coupling directly.
Revised: Concurrent work [Nguyen and Ho,2022] amortizes the estimation of the optimal projection in the max-sliced objective, which differsfrom our work where we instead amortize the estimation of the optimal coupling directly.

Original: Also,
Revised: Also,

Original: Lacombe et al.
Revised: Lacombe et al.

Original: [2021] learns to predict Wasserstein barycenters of pixel images by training a convolutional networks that, given images as input, outputs their barycenters.
Revised: [2021] learns to predict Wasserstein barycenters of pixel images by training a convolutional networks that, given images as input, outputs their barycenters.

Original: Our work is hence ageneralization of this pixel-based work to general measures – both discrete and continuous.
Revised: Our work is hence ageneralization of this pixel-based work to general measures – both discrete and continuous.

Original: A limitation of Lacombe et al.
Revised: A limitation of Lacombe et al.

Original: [2021] is that it does not provide alignments, as the amortization networkspredicts the barycenter directly rather than individual couplings.
Revised: [2021] is that it does not provide alignments, as the amortization networkspredicts the barycenter directly rather than individual couplings.

Original:  Conclusions, future directions, and limitations
Revised:  Conclusions and future directions

Original: We have presented foundations for modeling and learning to solve OT problems with Meta OT byusing amortized optimization to predict optimal transport plans.
Revised: We have presented foundations for modeling and learning to solve OT problems with Meta OT byusing amortized optimization to predict optimal transport plans.

Original: This works best in applications thatrequire solving multiple OT problems with shared structure.
Revised: This works best in applications thatrequire solving multiple OT problems with shared structure.

Original: We instantiated it to speed up entropicregularized optimal transport and unregularized optimal transport with squared cost by multipleorders of magnitude.
Revised: We instantiated it to speed up entropicregularized optimal transport and unregularized optimal transport with squared cost by multipleorders of magnitude.

Original: We envision extensions of the work in:  Meta OT models .
Revised: We envision extensions of the work in: 1) Meta OT models .

Original: While we mostly consider models based on hypernetworks other meta- learning paradigms can be connected in. In the discrete setting, we only considered settings
Revised: While we mostlyconsider models based on hypernetworks, other meta-learning paradigms can be connected in, and 2)

Original: where the cost remains ﬁxed, but the Meta OT model can also be conditioned on the cost

Original: by considering the entire cost matrix as an input (which may be too large for most models

Original: to handle), or considering a lower-dimensional parameterization of the cost that changes

Original: between the Meta OT problem instances.

Original: OT algorithms .
Revised: OT algorithms .

Original: While we instantiated models on top of log-Sinkhorn and W2GN, Meta OT could be built on top of other methods.
Revised: While we instantiated models on top of log-Sinkhorn and W2GN, Meta OT couldbe integrated with most other methods too.

Original:  OT applications that are computationally expensive and repeatedly solved, e g. in multimarginal and barycentric settings, or for Gromov-Wasserstein distances between metricmeasure spaces.Limitations.
Revised: 3) OT applications that are computationally expensiveand repeatedly solved, e.g. in multi-marginal and barycentric settings, or for Gromov-Wassersteindistances between metric-measure spaces.

Original: While we have illustrated successful applications of Meta OT, it is also important tounderstand the limitations: 1) Meta OT does not make previously intractable problems tractable.

Original: All of the baseline OT solvers we consider solve our problems within milliseconds or seconds.

Original: Out-of-distribution generalization.

Original: Meta OT may not generate good predictions on instances thatare not close to the training OT problems from the meta-distribution D over the measures and cost.

Original: If the model makes a bad prediction, one fallback option is to re-solve the instance from scratch.

Original: Checklist
Revised: Checklist

Original: For all authors...
Revised: For all authors

Original: (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s
Revised: (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s

Original: contributions and scope?
Revised: contributions and scope?

Original: [Yes] We hope so
Revised: [Yes] We hope so

Original: (b) Did you describe the limitations of your work?
Revised: (b) Did you describe the limitations of your work?

Original: [Yes] In sect.
Revised: [Yes] In sect.

Original: (c) Did you discuss any potential negative societal impacts of your work?
Revised: (c) Did you discuss any potential negative societal impacts of your work?

Original: [No] We do not
Revised: [No] We do not

Original: immediately foresee any that our work would add that the broader optimal transport
Revised: immediately foresee any that our work would add that the broader optimal transport

Original: ﬁeld doesn’t already have
Revised: field doesn’t already have

Original: (d) Have you read the ethics review guidelines and ensured that your paper conforms to
Revised: (d) Have you read the ethics review guidelines and ensured that your paper conforms to

Original: them?
Revised: them?

Original: [Yes]
Revised: [Yes]

Original: If you are including theoretical results...
Revised: If you are including theoretical results...

Original: (This is not a theory paper)
Revised: (This is not a theory paper)

Original: (a) Did you state the full set of assumptions of all theoretical results?
Revised: (a) Did you state the full set of assumptions of all theoretical results?

Original: [N/A]
Revised: [N/A]

Original: (b) Did you include complete proofs of all theoretical results?
Revised: (b) Did you include complete proofs of all theoretical results?

Original: [N/A]
Revised: [N/A]

Original: If you ran experiments...
Revised: If you ran experiments...

Original: (a) Did you include the code, data, and instructions needed to reproduce the main experi-
Revised: (a) Did you include the code, data, and instructions needed to reproduce the main experi-

Original: mental results (either in the supplemental material or as a URL)?
Revised: mental results (either in the supplemental material or as a URL)?

Original: [Yes]
Revised: [Yes]

Original: (b) Did you specify all the training details (e.g., data splits, hyperparameters,
Revised: (b) Did you specify all the training details (e.g., data splits, hyperparameters,

Original: how they
Revised: how they

Original: were chosen)?
Revised: were chosen)?

Original: [Yes]
Revised: [Yes]

Original: (c) Did you report error bars (e.g., with respect to the random seed after running experi-
Revised: (c) Did you report error bars (e.g., with respect to the random seed after running experi-

Original: ments multiple times)?
Revised: ments multiple times)?

Original: [Yes] We show results from multiple trials in most places
Revised: [Yes] We show results from multiple trials in most places

Original: (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
Revised: (d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?

Original: [Yes]
Revised: [Yes]

Original: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
Revised: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...

Original: (a) If your work uses existing assets, did you cite the creators?
Revised: (a) If your work uses existing assets, did you cite the creators?

Original: [Yes]
Revised: [Yes]

Original: (b) Did you mention the license of the assets?
Revised: (b) Did you mention the license of the assets?

Original: [Yes]
Revised: [Yes]

Original: (c) Did you include new assets either in the supplemental material or as a URL?
Revised: (c) Did you include new assets either in the supplemental material or as a URL?

Original: [No]
Revised: [No]

Original: (d) Did you discuss whether and how consent was obtained from people whose data[Table]
Revised: (d) Did you discuss whether and how consent was obtained from people whose data

Original: you’re using/curating?
Revised: you’re using/curating?

Original: [N/A]
Revised: [N/A]

Original: [Table] (e) Did you discuss whether the data you are using/curating contains personally identiﬁ- [Table]
Revised: (e) Did you discuss whether the data you are using/curating contains personally identifi-

Original: able information or offensive content?
Revised: able information or offensive content?

Original: [N/A]
Revised: [N/A]

Original: [Table]

Original: If you used crowdsourcing or conducted research with human subjects...
Revised: If you used crowdsourcing or conducted research with human subjects...

Original: [Table] (a) Did you include the full text of instructions given to participants and screenshots, if[Table]
Revised: (a) Did you include the full text of instructions given to participants and screenshots, if

Original: applicable?
Revised: applicable?

Original: [N/A]
Revised: [N/A]

Original: [Table] (b) Did you describe any potential participant risks, with links to Institutional Review[Table]
Revised: (b) Did you describe any potential participant risks, with links to Institutional Review

Original: Board (IRB) approvals, if applicable?
Revised: Board (IRB) approvals, if applicable?

Original: [N/A]
Revised: [N/A]

Original: [Table]

Original: (c) Did you include the estimated hourly wage paid to participants and the total amountspent on participant compensation? [N/A] [Table] A
Revised: (c) Did you include the estimated hourly wage paid to participants and the total amount

Original: Selecting (cid:15) for MNIST

Original: [Figure]

Original: B Other models for continuous OT

Original: While developing the hyper-network or Meta ICNN in sect. 3.2 for predicting couplings betweencontinuous measures, we considered alternative modeling formulations brieﬂy documented in thissection.

Original: We ﬁnalized only the hyper-network model because it is conceptually the most similar topredicting the optimal dual variables in the continuous setting and results in rapid predictions.

Original: B.1 Optimization-based meta-learning (MAML-inspired)

Original: The model-agnostic meta-learning setup proposed in MAML [Finn et al., 2017] could also be applied in the Meta OT setting to learn an adaptable initial parameterization.

Original: [-In the continuous setting,one initial version would take a parameterized dual potential model ψ ϕ ( x ) and seek to learn an initial parameterization ϕ 
-]
Revised: [-In the continuous setting,one initial version would take a parameterized dual potential model ψ ϕ ( x ) and seek to learn an initial parameterization ϕ 
-]

Original: so that optimizing a loss such as the W2GN loss L from eq.

Original: (12) results ina minimal L ( ϕ K ) after adapting the model for K steps.

Original: Formally, this would optimize:

Original: [Equation]

Original: Tancik et al.

Original: [2021] explores similar learned initializations for coordinate-based neural implicit representations for 2D images, CT scan reconstruction, and 3d shape and scene recovery from 2Dobservations.

Original: Challenges for Meta OT.

Original: The transport maps given by T = ∇ ψ can signiﬁcantly vary depending onthe input measures α, β .

Original: We found it difﬁcult to learn an initialization that can be rapidly adapted,and optimizing eq.

Original: (19) is more computationally expensive than eq.

Original: (18) as it requires unrollingthrough many evaluations of the transport loss L .

Original: And, we found that only learning to predictthe optimal parameters with eq.

Original: (18), conditional on the input measures, and then ﬁne-tuning with

Original: W2GN to be stable.

Original: Advantages for Meta OT.

Original: Exploring MAML-inspired methods could further incorporate the knowledge that the model’s prediction is going to be ﬁne-tuned into the learning process.

Original: One promising direction we did not try could be to integrate some of the ideas from LEO [Rusu et al., 2018] and

Original: CAVIA [Zintgraf et al., 2019], which propose to learn a latent space for the parameters where theinitialization is also conditional on the input.

Original: B.2 Neural process and conditional Monge maps

Original: The (conditional) neural process models considered in Garnelo et al.

Original: [2018b,a] can also be adaptedfor the Meta OT setting, and is similar to the model proposed in Bunne et al.

Original: [2022a].

Original: In thecontinuous setting, this would result in a dual potential that is also conditioned on a representationof the input measures, e.g. ψ ϕ ( x ; z ) where z := f emb ϕ ( α, β ) is a learned embedding of the inputmeasures that is learned with the parameters of ψ .

Original: This could be formulated as

Original: [Equation]

Original: where L modiﬁes the model used in the loss eq.

Original: (12) to also be conditioned on the context extractedfrom the measures.

Original: Challenges for Meta OT.

Original: This raises the issue on best-formulating the model to be conditional onthe context.

Original: One way could be to append z to the input point x in the domain.

Original: Bunne et al.

Original: [2022a]proposes to use the Partially Input-Convex Neural Network (PICNN) from [Amos et al., 2017] tomake the model convex with respect to x and not z .

Original: Advantages for Meta OT.

Original: A large advantage is that the representation z of the measures α, β wouldbe signiﬁcantly lower-dimensional than the parameters ϕ that our Meta OT models are predicting.

Original: C Additional experimental and implementation details

Original: We have attached the Jax source code necessary to run and reproduce all of the experiments in ourpaper and will open-source all of it.

Original: Here is a basic overview of the ﬁles:meta ot

Original: Meta OT Python library code conjugate.py

Original: Exact conjugate solver

Original: for the continuous setting data.py models.py utils.py config

Original: Hydra conﬁguration for the experiments (containing hyper-parameters) train_discrete.py

Original: Train Meta OT models for discrete

Original: OT train_color_single.py

Original: Train a single ICNN with W2GN between 2 images (for debugging)

Original: train_color_meta.py

Original: Train a Meta ICNN with W2GN plot_mnist.py

Original: Visualize the MNIST couplings plot_world_ Visualize the spherical couplings eval_color.py Evaluate the Meta ICNN in the continuous setting eval_discrete.py Evaluate the Meta ICNN for the discrete tasks
Revised: spent on participant compensation?

Original: Connecting to the data is one difﬁculty in running the experiments.

Original: The easiest experiment to re-runis the MNIST one, which will automatically download the dataset:557 1 ./ train_discrete.py

Original: # Train the model , outputting to <exp_dir >2 ./eval_discrete.py <exp_dir >

Original: # Evaluate the learned models3 ./plot_mnist.py

Original: <exp_dir >

Original: # Produce further visualizations

Original: [Table]

Original: C.1 Hyper-parameters

Original: We brieﬂy summarize the hyper-parameters we used for training, which we did not extensively tune.

Original: In the discrete setting, we use the same hyper-parameters for the MNIST and spherical settings

Original: [Table]

Original: [Table]

Original: [Table]

Original: C.2 Sinkhorn convergence times, varying thresholds

Original: In the main paper

Original: [Table]

Original: table 1 reports the runtime of Sinkhorn to reach a convergence threshold of themarginal error being below a tolerance of 10 − 3 , which is the default value used in many solvers.app.

Original: C.2 report the results from sweeping over other thresholds and show that Meta OT’s initializa-

Original: [Table]
Revised: [N/A]

Original: tion is consistently able to help

Original: [Table]

Original: [Table]

Original: [Table]

Original: C.3 Experimental runtimes and convergence

Original: App.

Original: C.3 shows the convergence during training of Meta OT models in the discrete and continuoussettings over 10 trials on our single Quadro GP100 GPU.

Original: The MNIST models are consistently trainedto optimality within 2 minutes (!) while the continuous model takes a few hours to train.

Original: [Figure]

Original: M a r g i n a l

Original: E rr o r

Original: MNIST

Original: [Figure]

Original: MNIST

Original: [Figure]

Original: Spherical

Original: [Figure]

Original: Wasserstein-

Original: D Out-of-distribution generalization

Original: App.

Original: D tests the ability of Meta OT to predict potentials for out-of-distribution input data.

Original: Weconsider the pairwise training and evaluation on the following datasets: 1) MNIST; 2) USPS

Original: [Hull,1994] (upscaled to have the same size as the MNIST); 3) Google Doodles dataset * with classes Crab,

Original: Cat and Faces; 4) sparsiﬁed random uniform data in [0,1] where sparsity (zeroing values below 0.95)is used to mimic the sparse signal in black-and-white images.

Original: For each pair, eg, MNIST-USPS, wetrain on one dataset and use the other to predict the potentials.

Original: The comparison is done using thesame metric as before, i.e., the deviation from the marginal constraints deﬁned in eq.

Original: [Figure]

Original: Sinkhorn

Original: [Figure]

Original: Sinkhorn Iterations

Original: Sinkhorn Iterations

Original: Sinkhorn Iterations

Original: M N I

Original: S T

Original: D ood l e s

Original: U

Original: SPS

Original: U n i f

Original: o r m

Original: T r a i n i ng

Original: D o m a i n

Original: MNIST

Original: Doodles

Original: USPS

Original: Uniform

Original: Evaluation Domain

Original: Initialization ( Zeros

Original: Gaussian [Thornton and Cuturi, 2022]

Original: Meta OT)

Original: E Additional color transfer results

Original: We next show additional color transfer results from the experiments in sect.

Original: 4.3 on the followingpublic domain images from WikiArt:

Original: • Distant View of the Pyramids by Winston Churchill (1921)

Original: • Charing Cross Bridge, Overcast Weather by Claude Monet (1900)

Original: • Houses of Parliament by Claude Monet (1904)

Original: [Table]

Original: • October Sundown, Newport by Childe Hassam (1901)

Original: [Table]

Original: • Landscape with House at Ceret by Juan Gris (1913)• Irises in Monet’s Garden by Claude Monet (1900)

Original: [Table]

Original: • Crystal Gradation by Paul Klee (1921)

Original: • Senecio by Paul Klee (1922)

Original: [Table]

Original: • Váza s kvˇetinami by Josef Capek (1914)

Original: • Sower with Setting Sun by Vincent van Gogh (1888)

Original: [Table]

Original: • Three Trees in Grey Weather by Claude Monet (1891)

Original: [Table]

Original: • Vase with Daisies and Anemones by Vincent van Gogh (1887)

Original: α β

Original: T # α

Original: T − 1# β α β

Original: T # α

Original: T − 1# β α β

Original: T # α

Original: T − 1# β