Paper deep dive
Amortized Optimal Transport from Sliced Potentials
Minh-Phuc Truong, Khai Nguyen
Intelligence
Status: succeeded | Model: google/gemini-3.1-flash-lite-preview | Prompt: intel-v1 | Confidence: 95%
Last extracted: 4/18/2026, 1:49:31 AM
Summary
The paper introduces two novel amortized optimization strategies, Regression-based Amortization (RA-OT) and Objective-based Amortization (OA-OT), for predicting optimal transport (OT) plans. By leveraging Kantorovich potentials derived from sliced OT as inputs, the proposed models are more parsimonious and independent of the specific structure of the measures (e.g., number of atoms) compared to existing methods like Meta-OT. The approach demonstrates high accuracy and efficiency in tasks such as MNIST digit transport, color transfer, and spherical supply-demand transportation.
Entities (5)
Relation Signals (3)
OA-OT → optimizes → Kantorovich dual objective
confidence 95% · In OA-OT, we estimate the parameters of the functional model by optimizing the Kantorovich dual objective.
RA-OT → uses → Kantorovich Potentials
confidence 95% · In RA-OT, we formulate a functional regression model that treats Kantorovich potentials from the original OT problem as responses
RA-OT → improvesupon → Meta-OT
confidence 90% · The proposed models are more parsimonious... unlike in Meta-OT.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:We propose a novel amortized optimization method for predicting optimal transport (OT) plans across multiple pairs of measures by leveraging Kantorovich potentials derived from sliced OT. We introduce two amortization strategies: regression-based amortization (RA-OT) and objective-based amortization (OA-OT). In RA-OT, we formulate a functional regression model that treats Kantorovich potentials from the original OT problem as responses and those obtained from sliced OT as predictors, and estimate these models via least-squares methods. In OA-OT, we estimate the parameters of the functional model by optimizing the Kantorovich dual objective. In both approaches, the predicted OT plan is subsequently recovered from the estimated potentials. As amortized OT methods, both RA-OT and OA-OT enable efficient solutions to repeated OT problems across different measure pairs by reusing information learned from prior instances to rapidly approximate new solutions. Moreover, by exploiting the structure provided by sliced OT, the proposed models are more parsimonious, independent of specific structures of the measures, such as the number of atoms in the discrete case, while achieving high accuracy. We demonstrate the effectiveness of our approaches on tasks including MNIST digit transport, color transfer, supply-demand transportation on spherical data, and mini-batch OT conditional flow matching.
Tags
Links
- Source: https://arxiv.org/abs/2604.15114v1
- Canonical: https://arxiv.org/abs/2604.15114v1
Trouble viewing inline? Open PDF directly →
Full Text
62,469 characters extracted from source content.
Expand or collapse full text
Amortized Optimal Transport from Sliced Potentials Minh-Phuc Truong † Khai Nguyen ⋄⋆ ⋄ The University of Texas at Austin † Hanoi University of Science and Technology April 17, 2026 Abstract We propose a novel amortized optimization method for predicting optimal transport (OT) plans across multiple pairs of measures by leveraging Kantorovich potentials derived from sliced OT. We introduce two amortization strategies: regression-based amortization (RA-OT) and objective-based amortization (OA-OT). In RA-OT, we formulate a functional regression model that treats Kantorovich potentials from the original OT problem as responses and those obtained from sliced OT as predictors, and estimate these models via least-squares methods. In OA-OT, we estimate the parameters of the functional model by optimizing the Kantorovich dual objective. In both approaches, the predicted OT plan is subsequently recovered from the estimated potentials. As amortized OT methods, both RA-OT and OA-OT enable efficient solutions to repeated OT problems across different measure pairs by reusing information learned from prior instances to rapidly approximate new solutions. Moreover, by exploiting the structure provided by sliced OT, the proposed models are more parsimonious, independent of specific structures of the measures, such as the number of atoms in the discrete case, while achieving high accuracy. We demonstrate the effectiveness of our approaches on tasks including MNIST digit transport, color transfer, supply–demand transportation on spherical data, and mini-batch OT conditional flow matching. 1 Introduction Optimal transport (OT) [73,74] has been widely recognized as a fundamental and powerful tool in statistics, machine learning, and data science. It plays a key role across a broad range of applications. For instance, OT has improved the quality of generative models, including autoencoders [70,52], generative adversarial networks [6,33], and diffusion and flow-based models [39,54,71], as well as drifting models [35]. It is also widely used in domain adaptation, where it helps establish correspondences between source and target domains [19,20,23,76]. Beyond these areas, OT has found numerous applications in computational biology [15,62], chemistry [75], image processing [29], fairness [36], signal processing [38], computer graphics [67,66,12], statistical inference [8,9], and dependency measurement [17, 16], among others. Computational optimal transport (OT) [53] has emerged as a central yet challenging area within OT. Except for special cases such as OT between Gaussian measures, which admits a closed form solution [24], computing OT is generally expensive. In the discrete setting, where both measures are supported on finite sets, OT can be formulated as a linear programming problem with super cubic complexity in the number of support points [51]. Various approaches have been proposed to mitigate this computational burden, including stochastic approximation with mini-batches [68,28,27,49] and low rank approximations [61,60]. For general measures, OT is typically solved through its dual formulation, which involves a functional optimization problem over Kantorovich potentials [74]. ⋆ Corresponding author: khainb@utexas.edu 1 arXiv:2604.15114v1 [stat.ML] 16 Apr 2026 This dual problem is constrained and often requires parameterizing the potentials for practical computation [6, 43], which can be challenging. One widely used approach is entropic regularization [21], which smooths the transport problem by adding an entropy penalty. In the discrete setting, this regularization ensures a unique transport plan that can be efficiently computed via matrix scaling algorithms such as Sinkhorn–Knopp [65], or, equivalently, through iterative Bregman projections [7]. Beyond reducing computational complexity to quadratic time, entropic regularization also improves the statistical rates for estimating both the transport cost [31] and the transport plan [44,57]. In the general cases, it removes constraints in the dual formulation, enabling unconstrained optimization [32]. For these reasons, entropic regularized OT is widely used in practice as a surrogate for classical OT. In many scenarios involving numerous pairs of measures, such as mini-batch generative models [33,35], dataset comparison [3], 3D point-cloud autoencoders [1], point-cloud nearest neighbor classification and regression [58], seismic signal analysis [26], and single-cell perturbation modeling [14], optimal transport (OT) must be solved repeatedly. To accelerate computation in such settings, amortized optimization [64,4] has been employed to predict OT solutions directly from input measures. Two primary targets arise for amortization: the transportation cost and the transportation plan. Predicting the transportation cost is relatively straightforward, as it reduces to estimating a nonnegative scalar. This direction has been widely explored, for example through learned embeddings or embedding networks [45,18,37,34], as well as regression onto sliced optimal transport metrics [50]. In contrast, amortizing the transportation plan is significantly more challenging, as it involves predicting a joint measure over the input distributions and implicitly determines the cost. To the best of our knowledge, the only existing approach is Meta-OT [5], which learns to predict entropically regularized Kantorovich potentials from the input measures as a proxy for the plan. In the discrete setting, Meta-OT constructs an amortized model, typically a neural network, that maps input measures, represented by their atomic supports and corresponding weights, to a vector whose dimension matches the number of atoms, thereby approximating one of the Kantorovich potentials. In the continuous setting, Meta-OT instead employs a hypernetwork as the amortized model to predict the parameters of a parametric Kantorovich potential. In both cases, Meta-OT adopts an objective-based amortization strategy, where the amortized models are learned by directly optimizing the Kantorovich dual objective. The main limitation of Meta-OT is that it directly takes the raw representation of input measures, hence, its amortized models contains many parameters. In this work, we propose a new approach to amortizing the transportation plan. The key idea is to use Kantorovich potentials derived from sliced OT [56,46] as inputs to the amortized model, rather than the raw input measures. This representation substantially reduces the dimensionality of the input, and consequently lowers the number of parameters required for the amortized model. Contribution. In summary, our main contributions are three-fold: 1. We introduce a new method to tackle amortized OT problem using sliced OT. In particular, we suggest to predict Kantorovich potentials of the original OT problem from Kantorovich potentials of the corresponding sliced OT problems. Due to the closed-form of sliced OT, the resulting method is computationally fast. Moreover, leveraging expressive sliced OT’s potentials, we are able to use parsimonious amortized models with a small number of parameters and are independent of the structure of measures such as the number of atoms of measures in discrete cases. Finally, our proposal can also be seen as the first attempt to reconstruct OT plan from sliced OT as existing works only 2 form a transporation plan from sliced OT without any notion of approximating the original OT one. 2. We consider two amortization strategies: regression-based amortization (RA-OT) and objective- based amortization (OA-OT). In RA-OT, we formulate a functional regression model, serving as the amortized model, where Kantorovich potentials from the original OT problem are treated as responses and those derived from sliced OT as predictors. The model is then estimated using least-squares methods. In OA-OT, we extend the Meta-OT framework by designing an amortized model that predicts one Kantorovich potential from the sliced OT potentials, and train it by optimizing the Kantorovich dual objective. 3. While the proposed method is well defined for both discrete and continuous measures, we focus on the discrete setting in this work. This choice is motivated by the fact that discrete OT is more commonly used in practice, whereas computing continuous OT has no exact solutions. We conduct experiments on MNIST digit transport, color transfer, and spherical supply–demand transportation to show that the proposed amortized OT approach leads to more favorable prediction of OT plan than the predicted plan from Meta-OT and sliced OT’s plans. 2 Background on Optimal Transport and Entropic Regularization We start this section by reviewing the definition of OT, Kantorovich duality, and Kantorovich potentials. Given two measures (μ,ν) supported on domains (X,Y) and a ground metricc: X ×Y → R + , the primal OT [74] problem is defined as follows: π ⋆ ∈ arg min π∈Π(μ,ν) Z X×Y c(x,y)dπ(x,y),(1) where Π(μ,ν) is the set of admissible transportation plans (joint measures) betweenμandν, andπ ⋆ is the optimal transportation plan. equation 1 admits a dual problem which is known as Kantorovich duality: (f ⋆ ,g ⋆ )∈ arg max (f,g)∈U c (X,Y) Z X f (x)dμ(x) + Z Y g(y)dν(y),(2) whereU c =(f,g)∈ C(X)×C(Y)| ∀(x,y)∈ X ×Y,f(x) +g(y)≤ c(x,y)is the constraint set, C(X) andC(Y) are sets of continuous functions onXandYrespectively, andf ⋆ ,g ⋆ are known as Kantorovich potentials. The connection betweenπ ⋆ andf ⋆ ,g ⋆ is from complementary slackness condition: f ⋆ (x) + g ⋆ (y) = c(x,y) for π ⋆ - almost everywhere (x,y). Entropic regularization. Entropic OT [21] smoothens the OT problem and it can be conveniently written in a single convex optimization problem as follows: π ⋆ ε = arg min π∈Π(μ,ν) Z X×Y c(x,y)dπ(x,y) + εKL(π | μ⊗ ν),(3) whereKL(π | μ⊗ ν) = R X×Y log dπ dμ⊗ν (x,y) − 1 dμ⊗ ν(x,y) andμ⊗ νis the product measure of μ and ν. The dual problem becomes: (f ⋆ ,g ⋆ )∈arg max f∈C(X ),g∈C(Y) Z X f (x)dμ(x) + Z Y g(y)dν(y) − ε log Z X×Y exp f (x) + g(y)− c(x,y) ε dμ(x)ν(y) ,(4) 3 Forε >0, the optimal transportation plan is recovered from any (f,g) by solving: dπ ⋆ ε (x,y) = exp f ⋆ (x)+g ⋆ (y)−c(x,y) ε dμ(x)dν(y),and we can also infer one potential from the other i.e.,g(y) = −ε log R X exp f (x)−c(x,y) ε dμ(x) and f (x) =−ε log R Y exp g(y)−c(x,y) ε dν(y) . Discrete Case. Whenμ= P n i=1 α i δ x i andν= P m j=1 β j δ y j with P n i=1 α i = P m j=1 β j andα i > 0,β j > 0∀i,j, the entropic OT problem becomes: P ⋆ ε = arg min P∈Γ(α,β) ⟨C,P⟩− εH(P ),(5) where Γ(α,β) =P ∈ R n×m + | P1 =α,P ⊤ 1 =βis the set of discrete plans,H(P) = − P n i=1 P m j=1 P ij (logP ij − 1) is the entropy, and C ij = c(x i ,y j ). The entropic OT dual becomes: (f ⋆ ,g ⋆ )∈ arg max f∈R n ,g∈R m ⟨f,α⟩ +⟨g,β⟩− ε n X i=1 m X j=1 exp f i + g j − C ij ε .(6) We can recover a discrete plan from any pairs (f,g) as follows: P ⋆ ij = exp f ⋆ i + g ⋆ j − C ij ε .(7) The mapping between the dual can be rewritten as: g = ε logα− ε log exp(−C ⊤ /ε) exp(f/ε) ,(8) f = ε logβ− ε log (exp(−C/ε) exp(g/ε)).(9) The Sinkhorn algorithm [21] is an efficient iterative method for solving the entropic optimal transport problem. It alternates between updating the dual potentials f and g according to equation 8 and equation 9 until the updates converge. A key property of entropic OT is that knowing only one of the optimal Kantorovich potentials, f ⋆ or g ⋆ , is sufficient: the other potential can be directly computed from it, allowing the optimal plan to be fully determined. The time complexity of Sinkhorn algorithm is O(nm) [2] compared to O((n + m) 3 log(n + m)) [53] of the conventional OT. 3 Amortized Optimization of Optimal Transport In this section, we first review the amortized optimal transport (OT) problem in Section 3.1, including the Meta-OT framework, as well as the definitions of amortized models and amortized losses. We then introduce two new amortization strategies based on sliced OT potentials: regression-based amortized OT (RA-OT) and objective-based amortized OT (OA-OT), and discuss their properties in Section 3.2. 3.1 Amortized Optimal Transport Problem In the amortized OT problem, we observe multiple pairs of probability measures, potentially associated with different ground cost functions. We assume that these triplets are drawn from an unknown meta-distributionD, yielding observations (μ 1 ,ν 1 ,c 1 ),...,(μ N ,ν N ,c N ) (N >1). Our 4 objective is not only to compute entropic optimal transport for the givenNpairs, but also to leverage these observations to accelerate computation for future pairs (μ ′ 1 ,ν ′ 1 ,c ′ 1 ),...,(μ ′ M ,ν ′ M ,c ′ M ) (M >1). Meta-OT [5] provides a solution by adopting amortized optimization. We now discuss how Meta-OT address the problem by reviewing the definition of the amortized model and amortized loss. Amortized model. For each optimal transport (OT) problem, denoted (μ i ,ν i ,c i ), the goal is to obtain one of the corresponding Kantorovich potentials, e.g.,f ⋆ [μ i ,ν i ,c i ]. An amortized model is a mapping that takes a pair of measures and the associated cost as input and predicts a Kantorovich potential. Formally, a parametric mapping ˆ f φ is considered in Meta-OT [5], whereφbelongs to a parameter space Φ. The model is trained such that ˆ f φ [μ,ν,c]≈ f ⋆ [μ,ν,c] for D-almost every (μ,ν,c). Whenμandνare discrete andcis fixed across all realizations fromD, the inputs to ˆ f φ can be written as (α,X) and (β,Y), representing the weights and atoms ofμandν, respectively. In Meta-OT, the authors consider cases where the atoms are fixed, so only the weightsαandβare used as inputs. In this setting, ˆ f φ can be implemented as a neural network (e.g., an MLP) that takes two vectors of sizesnandmas input and outputs a vector of sizen. We recall that Meta-OT requires the number of atoms,nandm, to be the same across all realizations ofD, since designing a neural network that accommodates variable input sizes is nontrivial. Amortized loss. Meta-OT adapts an objective-based optimization approach to estimate the amortized model. In particular, it solves the following optimization problem: min φ∈Φ E (μ,ν,c)∼D [J ( ˆ f φ [μ,ν,c];μ,ν,c)],(10) whereJis the Kantorovich dual objective in equation 6 with the plugged in form of the other potential using equation 8. By replacing the expectation with the finite sum across samples (μ 1 ,ν 1 ,c 1 ),...,(μ N ,ν N ,c N ), we obtain the empirical objective, which can be optimized using gradient descent algorithms or stochastic gradient algorithms for a large-scale dataset. 3.2 Amortized Optimal Transport from Sliced Potentials We start by reviewing some essential definitions, such as one-dimensional OT plan, sliced potentials, and projection functions in sliced OT. One-dimensional Optimal Transport. In one-dimension, whenc=h(x− y) withhis a strictly convex function, OT between μ and ν admits a closed-form [59]: π ⋆ = (F −1 μ ,F −1 ν )♯U ([0, 1]),(11) whereU([0,1]) is the uniform distribution on the interval [0,1],F −1 μ andF −1 ν are quantile functions ofμandνrespectively. We then can access to the Kantorovich potential using complementary slackness condition: f ⋆ (F −1 μ (u)) + g ⋆ (F −1 ν (u)) = h(F −1 μ (u)− F −1 ν (u)),(12) for U ([0, 1])-almost every u∈ [0, 1], this fact allows the OT plan and Kantorovich potentials in the discrete case to be computed extremely efficiently, with a time complexity ofO((n+m)log(n+m)), 5 Meta-OTAmortized-OT from Sliced Potentials Input Measures and CostProjection ParametersProjected Measures Sliced PotentialsDual PotentialsTransportation plan Figure 1: Amortized-OT uses sliced potentials to predict original dual potentials. wherenandmdenote the number of atoms in the two measures. A detailed algorithm for computing the OT plan and Kantorovich potentials in one-dimensional discrete settings can be found in Algorithm 1 of [63]. We also note that, alternatively, Kantorovich potentials in discrete cases can be obtained by taking the gradients of the optimal transport cost with respect to the weights of the two measures [22]. Sliced Optimal Transport. Sliced OT [56,46] leverages the closed-form solution of optimal transport in one dimension through projections. Given the geometry of the space of atomsX ∪Yof the measures and the ground metricc, a projection functionP c θ :X∪Y →Zmaps the atoms to a one- dimensional spaceZ, e.g., the real lineR, whereθ ∈Θ denotes the projection parameters. Designing effectiveP c θ remains an active area of research. For instance, in Euclidean spaces (X ∪Y ⊂ R d ) with the Euclidean ground metricc(x,y) =∥x− y∥ 2 , one can use linear projectionsP c θ (x) =⟨θ,x⟩ withθ ∈ S d−1 (the unit (d−1)-sphere). In spherical settings (X ∪Y ⊂ S d−1 ), spherical [10,55] or stereographic [72] projections are applicable. Extensions to other domains, such as functions [30], manifolds, and products of manifolds [11,48], have also been explored. Using the projection function, sliced OT projectsμandνintoP c θ ♯μ andP x θ ♯ν , then obtains the 1D OT plan between them. Recent works have explored lifting the 1D OT plan back to the original space [42,41,69]. However, these methods focus on fast approximations and are not intended to recover the exact original OT. Furthermore, they do not operate in an amortized OT setting, as they are applied at the level of individual measure pairs. Amortized model with sliced potentials. Given a set of projection parametersθ 1 ,...,θ L (L ≥1 is referred to as the number of projections), which can be created using Monte Carlo sampling [13] or quasi Monte Carlo sampling [47], we are able to obtain a set of sliced potentials f ⋆ θ 1 [μ,ν,c],...,f ⋆ θ L [μ,ν,c] using closed-form of one-dimensional OT problems betweenP c θ l ♯μand P c θ l ♯ν for l = 1,...,L. We then define an amortized model as follows: ˆ f φ [μ,ν,c] = γ φ (f ⋆ θ 1 [μ,ν,c],...,f ⋆ θ L [μ,ν,c]),(13) whereγ φ is an operator that maps from the space ofLfunctions in one dimension to the space of functions in the dimension of measuresμandν. Sincef ⋆ θ 1 [μ,ν,c],...,f ⋆ θ L [μ,ν,c] are "handcraft fea- 6 tures" with a lot of information about transportation betweenμandνthrough one-dimensional slices, we propose to consider a simple linear model as follows: ˆ f ω [μ,ν,c](x) = P L l=1 ω l f ⋆ θ l [μ,ν,c](P c θ l (x)), whereω ∈ R L is the linear coefficients. We recall that the above construction of the amortized model is well-defined for both discrete and continuous cases. Importantly, the parameterωorφof the proposed amortized model is independent of the input measures (e.g., their number of atoms), unlike in Meta-OT. Therefore, we are able to handle settings where the number of atoms varies. Regression-based amortization. We now discuss the first amortization strategy: regression-based amortized OT (RA-OT). In this strategy, we assume that we can afford the computation of (entropic) OT during training. In particular, we solve the following optimization problem: min ω∈R L E (μ,ν,c)∼D [∥ ˆ f ω [μ,ν,c]− f ⋆ [μ,ν,c]∥ 2 2 ],(14) where∥ · ∥ 2 2 denotes the functionalL 2 norm. Whenμis discrete withnatomsx 1 ,...,x n , the above optimization problem admits a closed-form. In particular, we denoteX θ l [μ,ν,c] = (f ⋆ θ l [μ,ν,c](P c θ l (x 1 )),...,f ⋆ θ l [μ,ν,c](P c θ l (x n )) ⊤ ∈ R n ,X Θ [μ,ν,c] = (X θ 1 [μ,ν,c],...,X θ L [μ,ν,c])∈ R n×L , andY[μ,ν,c] = (f ⋆ [μ,ν,c](x 1 ),...,f ⋆ [μ,ν,c](x n )) ⊤ ∈ R n , we can rewrite the optimization problem as: min ω∈R L E (μ,ν,c)∼D [∥X Θ [μ,ν,c]ω− Y [μ,ν,c]∥ 2 2 ],(15) which has the optimal solution is the root of the following equation: E[X Θ [μ,ν,c] ⊤ X Θ [μ,ν,c]]ω = E[X Θ [μ,ν,c] ⊤ Y [μ,ν,c]],(16) which has the closed-form solution: ω ⋆ = E[X Θ [μ,ν,c] ⊤ X Θ [μ,ν,c]] −1 E[X Θ [μ,ν,c] ⊤ Y [μ,ν,c]],(17) where the expectation is underD. When we observe a finite set of (μ 1 ,ν 1 ,c 1 ),...,(μ N ,ν N ,c N ), we can replace the expectation by the finite sum 1 N P N i=1 as a Monte Carlo estimate for tractable computation. In practice, we can use other linear solvers of equation 16 rather than 17 for computational stability. The time complexity for training RA-OT involves the complexity for computing sliced potentials O(ML(n+m)(log(n+m) +d)) (P c θ has the time complexity ofO(Ld(n+m))), computing ground truth potentialO(Nnm) for entropic OT orO((n+m) 3 log(n+m)) for OT, and computing the optimal coefficientO(NnL 2 +L 3 ). For prediction on a new set (μ ′ 1 ,ν ′ 1 ,c ′ 1 ),...,(μ ′ M ,ν ′ M ,c ′ M ), the time complexity of RA-OT involves only the complexity for computing sliced potentialsO(ML(n+ m)(log(n+m) +d)) and applying linear predictionO(MnL). In the case of conventional OT, we can set up another regression problem for the other potentialg ⋆ , while predictingf ⋆ is sufficient for entropic OT as discussed in Section 2. Objective-based amortization. For the objective-based amortized OT (OA-OT) strategy, we follow the same training scheme as Meta-OT with the new proposed amortized model. We consider the following optimization problem: min ω∈R L E (μ,ν,c)∼D [J ( ˆ f ω [μ,ν,c];μ,ν,c)],(18) 7 Sinkhorn (converged, ground-truth) α 0 α 1 α 2 Meta-OT α 0 α 1 α 2 RA-OT α 0 α 1 α 2 OA-OT α 0 α 1 α 2 Figure 2: Wasserstein interpolations between MNIST test digits. (Top) Sinkhorn ground truth run to convergence and baseline Meta-OT. (Bottom) RA-OT and OA-OT (M =50,L=100). Both proposed methods produce interpolation sequences nearly identical to the ground truth. whereJis either the Kantorovich dual objective in equation 4 or in equation 6 in the plugged-in form of the other potential. Similar to Meta-OT, OA-OT is designed to address the entropic OT problem only. We also use a gradient-based algorithm for finding an estimate ofωin OA-OT. For a set ofNsamples fromD, after computing sliced potentials with the time complexity of O(NL(n+m)(log(n+m) +d)), an estimate of gradient forωhasO(Nmn) in time complexity. The time complexity is then scaled byT >0 iterations of gradient updates. Compared to Meta-OT, OA-OT has a much smaller number of parameters; hence, it enjoys faster training. Given an estimate of ω, the time complexity of OA-OT during inference time is the same as RA-OT. Remark 1 (Amortization Error and Amortization Gap). For linear amortized models, the amortiza- tion errorE (μ,ν,c)∼D [∥ ˆ f ω [μ,ν,c]−f ⋆ [μ,ν,c]∥ 2 2 ] and the amortization gapE (μ,ν,c)∼D [(J( ˆ f ω [μ,ν,c];μ,ν,c)]− J (f ⋆ [μ,ν,c];μ,ν,c)) 2 ] are minimized whenf ⋆ [μ,ν,c] lives in the linear spans off ⋆ θ 1 [μ,ν,c],...,f ⋆ θ L [μ,ν,c] forD-almost every (μ,ν,c). Therefore, increasing the number of projectionsLcan potentially enlarge the span generated by the sliced potentials, thereby improving the expressivity of the amortized model. 4 Experiments Following the experimental pipeline in [5], we evaluate the practical effectiveness of our proposed methods, RA-OT and OA-OT, across three representative settings that span a range of discrete optimal transport applications. First, in MNIST grayscale transport, we compute Wasser- stein interpolations between digit images defined on a 2D Euclidean grid. Second, in spherical supply–demand transport, we consider a large-scale, non-Euclidean setting that maps global landmass distributions to population density over a sphere. Finally, in color transfer, we align image color palettes within a 3D RGB space. For these tasks, all methods are quantitatively assessed using the Root Mean Square Error (RMSE) of the predicted transport plan relative to the converged Sinkhorn ground truth, alongside total training and per-pair inference times. We compare our approaches against three recent baselines: Meta-OT [5], which predicts the plan via an MLP applied to raw measure weights, min-SWGG [42], a training-free baseline, and Min-STP [40], which learns an optimized projection direction. To ensure a rigorous evaluation, all trainable methods share an identical pool of 1,000 measure pairs per task with a 70/30 train-test split. Additionally, we evaluate the performance of our methods when applied to mini-batch OT conditional flow matching, focusing on low-dimensional 2D toy datasets to demonstrate accelerated pairing in generative modeling. Since this task deviates from the standard plan-prediction pipeline, we instead measure the training 8 Table 1: RMSE, training and inference time on the MNIST grayscale transport task across varying number of training dataM(ε=0.1,N =300 test pairs,L=100). The results are reported as the mean ±standard deviation. The label “no train” indicates that no prior training is required. Bold values denote the best performance in each block. M MethodRMSE (×10 −6 , ↓) Train (s) Infer (ms,↓) M MethodRMSE (×10 −6 , ↓) Train (s) Infer (ms,↓) 10 Meta-OT [5]16.16± 4.9038.05 2.46± 0.20 50 Meta-OT [5]15.54± 4.7437.11 2.39± 0.29 Min-STP [40]93.51± 10.18199.526.18± 0.37Min-STP [40]93.50± 10.18198.016.25± 0.48 min-SWGG [42]89.85± 8.64no train10.01± 0.73min-SWGG [42]89.85± 8.64no train9.72± 0.41 RA-OT8.23± 3.551.3644.33± 2.91RA-OT7.77± 3.063.0339.36± 3.61 OA-OT6.16± 2.5316.6542.54± 2.95OA-OT6.02± 2.5215.7838.92± 2.23 20 Meta-OT [5]16.68± 4.9837.27 2.44± 0.39 200 Meta-OT [5]14.12± 4.7037.22 2.37± 0.40 Min-STP [40]93.51± 10.18198.566.21± 0.41Min-STP [40]93.51± 10.18197.866.17± 0.49 min-SWGG [42]89.85± 8.64no train9.70± 0.47min-SWGG [42]89.85± 8.64no train9.73± 0.61 RA-OT7.49± 3.191.4440.29± 2.74RA-OT7.74± 3.0711.1240.30± 4.39 OA-OT6.04± 2.5015.6040.65± 2.95OA-OT6.11± 2.5416.3940.48± 4.16 time and the NPE, benchmarking against Independent CFM (I-CFM) and OT-CFM [71]. Across all experiments, the number of sliced projections is fixed atL=100. For the first three standard OT tasks, we report results for varying training sizesM ∈10,20,50,200, focusing onM =50 for qualitative visualizations. In contrast, for the OT-CFM task, we maintain a fixed pre-training size of M =50. Experiment details are provided in Appendix A. Additional ablation results on the number of projectionsLare given in Appendix B, and further visualizations for all tasks are in Appendix D. 4.1 MNIST Pairs of MNIST digit images offer a standard test for amortized OT, sharing a common 28×28 grid geometry but varying in mass distributions. Each image is modeled as a discrete measure where pixel intensities dictate the weights of then=784 atoms. The transport cost is the squared Euclidean distance between coordinates, and the regularization parameter is set toε=0.1. As shown in Table 1, RA-OT and OA-OT remain accurate even with very limited training data (M= 10,20), outperforming Meta-OT in this low-data regime. They also train significantly faster due to their lightweight parameterization. In contrast, single-projection baselines such as Min-STP and min- SWGG struggle to capture the spatial structure of image transport. This improvement is also reflected in Figure 8 (Appendix D), where the interpolations better match the Sinkhorn reference. 4.2 Spherical Transport To evaluate the proposed methods in a non-Euclidean domain, we construct a supply-demand transportation problem on the unit sphereS 2 :=x∈R 3 :∥x∥=1. Supply capacities are assumed to be uniform acrossn=100 locations sampled from Earth’s landmass, while demand weights are drawn from Earth’s population density acrossm=10,000 locations based on the C-licensed dataset [25]. The ground metric is defined by the geodesic costc(x,y)= arccos(⟨x,y⟩) with an entropic regularization parameter ofε=0.5. To properly accommodate the spherical geometry within our framework, we utilize stereographic projections [72] to define the slicing family for both RA-OT and OA-OT. As reported in Table 2, both methods generalize well with few training samples and converge quickly in practice. This task highlights the importance of geometry: while Min-STP and min-SWGG fail to respect the spherical structure, our approach remains close to the Sinkhorn solution. The transport patterns shown in Figure 9 (Appendix D) clearly illustrate this behavior. 9 Sinkhorn (converged, ground-truth)Meta-OT RA-OTOA-OT Figure 3: Prediction on a held-out spherical supply–demand instance (M =50,L=100,ε=0.5). (Top) Sinkhorn ground truth and baseline Meta-OT. (Bottom) our methods: RA-OT and OA-OT Black dots denote active supply locations; blue arcs show great-circle geodesics from each demand location to its assigned supply location under the predicted plan. Table 2: RMSE, training and inference time on the spherical supply-demand transport task across varying number of training dataM(ε=0.5,n=100,m=10,000,N =300,L=100). The results are reported as the mean±standard deviation. The label “no train” indicates that no prior training is required. Bold values denote the best in each block. M MethodRMSE (×10 −7 , ↓) Train (s) Infer (ms,↓) M MethodRMSE (×10 −7 , ↓) Train (s) Infer (ms,↓) 10 Meta-OT [5]6.83± 2.4952.0815.94± 2.14 50 Meta-OT [5]4.42± 1.5552.0716.09± 2.17 Min-STP [40]115.07± 0.5446.06 7.51± 7.23Min-STP [40]115.10± 0.5446.38 7.88± 1.22 min-SWGG [42]112.15± 0.99no train34.62± 2.70min-SWGG [42]112.15± 0.99no train33.25± 1.71 RA-OT7.92± 2.771.0541.16± 3.11RA-OT7.82± 1.892.5341.96± 6.29 OA-OT4.80± 1.9618.6039.92± 3.12OA-OT3.93± 1.9219.3741.03± 2.76 20 Meta-OT [5]6.03± 2.2051.9915.57± 1.22 200 Meta-OT [5]3.62± 1.1452.1215.17± 2.54 Min-STP [40]115.06± 0.5446.39 7.95± 1.41Min-STP [40]115.05± 0.5446.41 7.93± 1.22 min-SWGG [42]112.15± 0.99no train33.85± 2.49min-SWGG [42]112.15± 0.99no train33.59± 1.99 RA-OT7.82± 2.251.3140.70± 3.25RA-OT7.55± 1.888.7439.40± 1.43 OA-OT4.24± 1.9818.7539.65± 3.26OA-OT3.66± 1.5621.3240.03± 3.83 4.3 Color Transfer Color transfer is formulated as an OT problem by representing images as weighted point clouds in the normalized RGB space [0,1] 3 . Following the experimental setup of Meta-OT, we collect approximately 200 public domain images from WikiArt to construct the dataset. To retain a discrete formulation, images are quantized intoK=500 distinct color clusters using mini-batchk-means, where centroids act as atoms and normalized cluster occupancies as weights. The ground metric 10 SourceTargetSinkhorn (GT)RA-OTOA-OT Figure 4: Color transfer results on held-out test pairs (M =50,L=100,ε=0.005,K=500 clusters). Each row corresponds to one source–target pair. Columns show: sourceα, targetβ, Sinkhorn ground truth, RA-OT, and OA-OT. Table 3: RMSE, training and inference time on the color transfer task across varying number of training dataM(ε=0.005,K=500 clusters,N =300 test pairs,L=100). The results are reported as the mean±standard deviation. The label “no train” indicates that no prior training is required. Bold values denote the best performance in each block. M MethodRMSE (×10 −6 , ↓) Train (s) Infer (ms,↓) M MethodRMSE (×10 −6 , ↓) Train (s) Infer (ms,↓) 10 Meta-OT [5]36.83± 10.9733.3216.83± 1.33 50 Meta-OT [5]33.16± 10.4532.9115.16± 0.64 Min-STP [40]80.93± 3.2449.51 4.02± 0.27Min-STP [40]81.04± 3.4250.00 4.10± 0.20 min-SWGG [42]77.69± 3.19no train4.25± 0.21min-SWGG [42]77.69± 3.19no train4.34± 0.17 RA-OT10.73± 6.351.9518.41± 1.21RA-OT9.99± 5.617.4017.40± 0.83 OA-OT9.78± 5.7817.5218.77± 1.31OA-OT9.00± 5.0218.1317.76± 0.79 20 Meta-OT [5]34.73± 10.9434.6315.03± 0.70 200 Meta-OT [5]33.23± 10.8933.1115.44± 1.20 Min-STP [40]80.97± 3.3349.92 4.07± 0.23Min-STP [40]81.10± 3.3849.68 4.08± 0.22 min-SWGG [42]77.69± 3.19no train4.30± 0.17min-SWGG [42]77.69± 3.19no train4.27± 0.19 RA-OT9.74± 5.683.1118.02± 1.34RA-OT10.21± 5.5326.5718.08± 1.40 OA-OT9.31± 5.2417.9218.45± 1.20OA-OT9.00± 5.0520.3718.25± 0.71 is the squared Euclidean distance withε=0.005. Transferred images are generated by updating source pixels via the target centroids weighted by the row-normalized predicted transport plan. As shown in Table 3, RA-OT and OA-OT achieve strong performance across all data regimes, with clear advantages when training data is limited. They also require less training time compared to Meta-OT. The gap with Min-STP and min-SWGG suggests that single projections are insufficient for modeling color distributions. As seen in Figure 4, this leads to more faithful color reconstruction. 4.4 Mini-batch Optimal Transport Conditional Flow Matching OT-CFM [71] leverages mini-batch OT to straighten flow trajectories in flow matching, thereby reducing inference cost. However, computing exact OT remains expensive and constitutes a major training bottleneck. To address this, we introduce a more efficient alternative by integrating RA-OT 11 Table 4: Quantitative results on 2D Conditional Flow Matching tasks. We report both training and pretraining time. Gaussian → 8gaussiansGaussian → moonsGaussian → scurve Method W 2 2 ↓ NPE ↓ Train (s) Pretrain W 2 2 ↓ NPE ↓ Train (s) Pretrain W 2 2 ↓ NPE ↓ Train (s) Pretrain I-CFM0.2800 0.2590 168.70.0s0.3901 0.538045.60.0s0.1113 1.923044.50.0s OT-CFM 0.1956 0.0108 1054.70.0s0.2668 0.0142 1513.10.0s0.0991 0.0243992.00.0s RA-OT0.5072 0.0144411.82.6s0.5569 0.0619331.71.0s0.4675 0.0400299.36.3s OA-OT0.5041 0.0142413.814.0s0.5459 0.0631319.714.6s0.4672 0.0425294.014.4s and OA-OT into the conditional flow matching framework, significantly accelerating training while preserving near-straight trajectories. Following [71], we mapN(0,I) to 2D targets (8gaussians, moons, scurve) using a network with three 64-unit SELU layers and a linear head, trained for 20000 steps (batch size 512,σ= 0.1) with AdamW (learning rate 10 −3 , weight decay 10 −5 ). We compare RA-OT and OA-OT against I-CFM and OT-CFM, where amortized methods useM= 50 pretraining batches andL= 100 projections. All methods employ fixed-step RK4 (101 steps); performance is evaluated using SinkhornW 2 2 att= 1 and path straightness via NPE. As shown in Table 4, our methods achieve 2.5×–4.5×training speedups over OT-CFM while maintaining comparable trajectory quality, with NPE close to OT-CFM and up to 1–2 orders of magnitude lower than I-CFM. Although slight degradation in SinkhornW 2 2 is observed in more complex settings, the overall trade-off highlights the scalability of our approach. The visualizations in Figures 5, 6, and 7 (Appendix D) further corroborate the quality of the learned trajectories. Additional experiments demonstrating the efficacy of our approach on high-dimensional CIFAR-10 data are provided in Appendix C. 5 Conclusion The paper presents two novel amortized OT methods, RA-OT and OA-OT, which leverage sliced OT to efficiently predict OT plans across multiple measure pairs. RA-OT uses regression to map sliced OT potentials to original OT potentials, while OA-OT optimizes the Kantorovich dual within the span of sliced OT potentials. Both approaches enable rapid and accurate solutions for repeated OT problems, are parsimonious, and do not depend on the structure or size of the measures. Experiments on MNIST digit transport, color transfer, spherical supply–demand problems, and conditional flow matching demonstrate the practical effectiveness and generality of the methods. A Experiment Details In this section, we detail the architectures, hyperparameter settings, and optimization constraints for all methods utilized in the main experiments. All configuration values are dynamically loaded from a centralized configuration file to ensure exact reproducibility across multiple experimental runs. To guarantee a fair evaluation, methods requiring gradient-based optimization are allocated an equivalent training budget (5,000 steps). Optimizers (predominantly Adam) and learning rates werefixed per algorithm based on their respective convergence properties. The entropic regularization parameter (ε) defines the inherent Sinkhorn OT problem and is maintained uniformly across all baselines for consistent ground truth definition (ε= 0.1 for Grayscale,ε= 0.5 for spherical transport, and ε = 0.005 for color transfer). 12 Table 5: Hyperparameters for RA-OT (Regression-based Amortized OT). ParameterMNIST Grayscale World Pair Color Transfer Number of Projections (L)100100100 Ridge Coefficient (λ)1× 10 −3 1× 10 −3 1× 10 −3 Table 6: Hyperparameters for OA-OT (Objective-based Amortized OT). ParameterMNIST Grayscale World Pair Color Transfer Number of Projections (L)100100100 Training Iterations (T)5,0005,0005,000 Learning Rate1× 10 −3 1× 10 −3 1× 10 −3 B Ablation Study: Effect of the Number of Projections L To study the effect of the number of projections, we varyL ∈ 3,5,10,20,50,100. As shown in Table 9, both RA-OT and OA-OT benefit from increasingL, with clear improvements when moving away from very small values such asL= 3 or 5. With too few projections, the model cannot capture enough structural information, which leads to weaker performance. This is most noticeable in the Color Transfer task, where the data distribution is more complex and requires a richer representation. AsLincreases, the error decreases quickly at first, and the results also become more stable across runs. However, after aroundL= 20, the improvements start to level off, and adding more projections only brings small gains. This suggests that a moderate number of projections is already enough to capture most of the important structure in the data. Importantly, increasingLdoes not significantly affect efficiency. Both training and inference times remain nearly the same across all values ofL, showing that the method scales well with the number of projections. In practice, this means we can safely use a reasonably largeLto ensure good performance without worrying about additional computational cost. C High-dimensional Flow Matching Fine-tuning on CIFAR-10 In addition to 2D experiments, we investigate the scalability and practical advantages of our amortized methods on high-dimensional image generation. Specifically, we focus on fine-tuning a converged Continuous Normalizing Flow model on the CIFAR-10 dataset using an large batch setting. We employ a widely adopted U-Net architecture containing approximately 35M parameters. Following the standard protocol [71], we utilize a converged Independent CFM (I-CFM) checkpoint natively Table 7: Hyperparameters for the Meta-OT baseline [5]. ParameterMNIST Grayscale World PairColor Transfer Architecture ModelMLPMLPPointCloud Encoder Hidden Layers / Units3 Layers3 Layers, 512 Units256 Enc Dim, 512 Head Training Iterations (T)5,0005,0005,000 Learning Rate1× 10 −3 1× 10 −3 1× 10 −3 13 Table 8: Hyperparameters for the amortized Min-STP baseline [40].αdenotes the inherent temperature parameter used during continuous soft-sorting. ParameterMNIST Grayscale World Pair Color Transfer Encoder/Context Dim323232 Hidden Dim12812864 Training Iterations (T)5,0005,0005,000 Learning Rate1× 10 −3 1× 10 −3 1× 10 −3 Soft-Sort Temperature (α)0.050.050.05 Table 9: Ablation study on the effect of the number of sliced projectionsL∈3,5,10,20,50,100. The performance of RA-OT and OA-OT is evaluated across all three tasks: MNIST grayscale transport, Spherical supply-demand, and Color Transfer. Results are reported in terms of Plan RMSE (×10 −6 ), training time, and inference time. Bold indicates the best metric between the two methods for a given task and L. L TaskMethod RMSE (×10 −6 , ↓) Train (s) Infer (ms)L TaskMethod RMSE (×10 −6 , ↓) Train (s) Infer (ms) 3 MNIST RA-OT9.25± 3.733.2237.47± 3.91 20 MNIST RA-OT7.62± 3.113.1838.11± 2.45 OA-OT8.50± 3.5515.35 37.38± 2.65OA-OT6.05± 2.5315.35 37.83± 2.23 Spherical RA-OT0.71± 0.172.17 30.51± 4.43 Spherical RA-OT0.74± 0.192.3332.06± 4.97 OA-OT0.41± 0.1618.6234.20± 5.70OA-OT0.38± 0.1618.23 30.89± 4.24 Color RA-OT25.60± 10.426.68 16.48± 0.96 Color RA-OT9.39± 5.507.0317.64± 1.28 OA-OT 23.96± 6.7916.8617.29± 1.15OA-OT9.11± 5.1117.31 17.48± 1.37 5 MNIST RA-OT8.80± 3.533.6240.89± 3.02 50 MNIST RA-OT7.66± 3.052.9837.73± 3.34 OA-OT8.01± 3.4116.10 40.76± 3.06OA-OT6.03± 2.5215.19 37.71± 2.40 Spherical RA-OT0.70± 0.182.37 31.90± 4.62 Spherical RA-OT0.78± 0.192.4635.53± 4.06 OA-OT0.41± 0.1618.6033.11± 6.03OA-OT0.39± 0.1918.47 33.04± 3.95 Color RA-OT 13.16± 6.957.2518.69± 1.75 Color RA-OT9.76± 5.626.9117.52± 1.51 OA-OT14.83± 6.2218.93 18.67± 0.89OA-OT9.04± 5.0217.48 17.51± 1.17 10 MNIST RA-OT7.85± 3.223.20 37.56± 2.80 100 MNIST RA-OT7.77± 3.063.0339.36± 3.61 OA-OT6.34± 2.6314.4437.90± 2.74OA-OT6.02± 2.5215.78 38.92± 2.23 Spherical RA-OT0.67± 0.182.19 32.01± 4.35 Spherical RA-OT0.78± 0.192.5341.96± 6.29 OA-OT0.42± 0.1718.5435.13± 5.99OA-OT0.39± 0.1919.37 41.03± 2.76 Color RA-OT9.63± 5.476.4817.14± 1.59 Color RA-OT9.99± 5.617.40 17.40± 0.83 OA-OT9.38± 5.2317.14 17.08± 1.03OA-OT9.00± 5.0218.1317.76± 0.79 pretrained from scratch for 400,000 steps with a learning rate of 2×10 −4 and a batch size of 128. Rather than training all configurations from scratch (which incurs extreme computational cost), we establish a rigorous fine-tuning protocol. All comparative methods natively load the equivalent I-CFM converged checkpoint (400000 steps training) and fine-tune for strictly 10 epochs. This setup establishes an absolutely fair ground where every method originates from an identical prior capacity and uniquely differs strictly by the assigned mini-batch coupling strategies (I-CFM, OT-CFM, OA- OT, and RA-OT). To avoid catastrophic forgetting, we reduce the learning rate to 5×10 −5 , apply a short 100-step warmup, tighten gradient clipping to 0.5, and lower the EMA decay from 0.9999 to 0.999 for faster adaptation. For amortized methods, predictors are constructed fromM= 50 mini-batches withL= 100 sliced projections. All runtime results are reported under a large batch size (B= 2048). About evaluation, we sample 50,000 outputs per model and evaluate the FID and NFE via the adaptivedopri5numerical solver (with tolerance 10 −5 ). As shown in Table 10, directly applying the same optimal transport setup (OT-CFM) to large-batch training only increases the 14 I-CFMOT-CFMRA-OTOA-OT Figure 5: Generated flow trajectories moving from a 2D Gaussian prior to the 8-Gaussians target. (Left to Right) Baseline I-CFM, OT-CFM, and our proposed amortized methods. Black dots signify source samples. The olive tracks map the predicted integration paths. Our solutions successfully reproduce the straight, uncrossed optimal transport paths of OT-CFM, yielding significantly lower collision energy than standard I-CFM. computational burden. In contrast, our precomputation strategies (RA-OT and OA-OT) effectively bypass this bottleneck by generating coupling plans much more efficiently, significantly reducing fine-tuning time. More interestingly, amortizing optimal transport does more than just speed things up, it also improves performance. OA-OT produces straighter trajectories, achieving the best NFE of 146.00, while RA-OT delivers the strongest perceptual quality, reaching an FID of 3.543. Table 10: Adaptive sampling generation results (dopri5) fine-tuning Flow Matching models on CIFAR-10. Despite operating on an extreme scaling batch (B= 2048), our amortized paradigms (OA-OT, RA-OT) drastically accelerate iterative epochs relative to exact OT-CFM while yielding highly superior FID scores and tightly constrained NFE parameters. Method FID ↓ NFE/sample ↓ Training Time (s) ↓ Pretrain Time (s) ↓ ICFM3.638146.61481.90.0 OT-CFM 3.630146.86744.20.0 OA-OT3.575146.00607.012.0 RA-OT 3.543147.10618.116.4 D Additional Visualizations 15 I-CFMOT-CFMRA-OTOA-OT Figure 6: Flow trajectories transforming a 2D Gaussian prior to the Moons target distribution. (Left to Right) Baselines I-CFM and OT-CFM, followed by our proposed RA-OT and OA-OT methods. Both amortized variants effectively map unentangled optimal integration flows (olive paths) across the dual clusters, matching exact OT capabilities. I-CFMOT-CFMRA-OTOA-OT Figure 7: Visualized flow trajectories mapping a 2D Gaussian prior into the S-Curve distribution. (Left to Right) I-CFM, OT-CFM, RA-OT, and OA-OT. Our amortized approaches robustly straighten the generative pathways, practically bypassing the severe integration turbulence traditionally observed in uncoupled I-CFM trajectories. 16 Sinkhorn (converged, ground-truth) α 0 α 1 α 2 Meta-OT α 0 α 1 α 2 Min-STP α 0 α 1 α 2 min-SWGG α 0 α 1 α 2 RA-OT α 0 α 1 α 2 OA-OT α 0 α 1 α 2 Sinkhorn (converged, ground-truth) α 0 α 1 α 2 Meta-OT α 0 α 1 α 2 Min-STP α 0 α 1 α 2 min-SWGG α 0 α 1 α 2 RA-OT α 0 α 1 α 2 OA-OT α 0 α 1 α 2 Sinkhorn (converged, ground-truth) α 0 α 1 α 2 Meta-OT α 0 α 1 α 2 Min-STP α 0 α 1 α 2 min-SWGG α 0 α 1 α 2 RA-OT α 0 α 1 α 2 OA-OT α 0 α 1 α 2 Figure 8: Additional qualitative visualizations of the predicted optimal transport plans on the MNIST dataset (M= 50,L= 100). The figure compares the structural alignment and digit matching quality across different amortized and sliced OT solvers (RA-OT, OA-OT, Meta-OT, Min-STP, and min-SWGG). Our proposed methods consistently preserve the topological structure of the digits. 17 Sinkhorn (converged, ground-truth) Meta-OT Min-STPmin-SWGG RA-OTOA-OT Figure 9: Extended qualitative results on the spherical world supply-demand transport task (M = 50,L= 100). The visualizations illustrate the transport assignments from dense global population distributions (demand) to uniformly scattered supply points across the globe. Observations indicate that RA-OT and OA-OT generate smoother and more globally coherent matching patterns, effectively avoiding the localized structural artifacts often observed in standard linear sliced estimation methods. 18 SourceTargetSinkhorn (GT)Meta-OT Min-STPmin-SWGGRA-OT (ours)OA-OT (ours) SourceTargetSinkhorn (GT)Meta-OT Min-STPmin-SWGGRA-OT (ours)OA-OT (ours) Figure 10: Further qualitative examples of the color transfer task between natural images using different OT solvers (M= 50,L= 100). The results demonstrate the capability of our amortized solvers to robustly align high-dimensional RGB color histograms. The images transported by RA-OT and OA-OT strictly adhere to the target color palettes while faithfully preserving the underlying textural and geometric details of the source images, exhibiting highly competitive visual fidelity against established exact and amortized baselines. 19 SourceTargetSinkhorn (GT)Meta-OT Min-STPmin-SWGGRA-OT (ours)OA-OT (ours) SourceTargetSinkhorn (GT)Meta-OT Min-STPmin-SWGGRA-OT (ours)OA-OT (ours) Figure 11: Additional qualitative examples of the color transfer task (M= 50,L= 100). The figure presents further comparisons of the visual outputs generated by the evaluated methods when mapping the color palette of the target images onto the source images. 20 References [1]P. Achlioptas, O. Diamanti, I. Mitliagkas, and L. Guibas. Learning representations and generative models for 3d point clouds. In International Conference on Machine Learning, pages 40–49. PMLR, 2018. (Cited on page 2.) [2]J. Altschuler, J. Niles-Weed, and P. Rigollet. Near-linear time approximation algorithms for optimal transport via Sinkhorn iteration. In Advances in Neural Information Processing Systems, pages 1964–1974, 2017. (Cited on page 4.) [3] D. Alvarez-Melis and N. Fusi. Geometric dataset distances via optimal transport. Advances in Neural Information Processing Systems, 33:21428–21439, 2020. (Cited on page 2.) [4] B. Amos. Tutorial on amortized optimization. Foundations and Trends in Machine Learning, 16(5):592–732, 2023. (Cited on page 2.) [5] B. Amos, G. Luise, S. Cohen, and I. Redko. Meta optimal transport. In International Conference on Machine Learning, pages 791–813. PMLR, 2023. (Cited on pages 2, 5, 8, 9, 10, 11, and 13.) [6]M. Arjovsky, S. Chintala, and L. Bottou. Wasserstein generative adversarial networks. In International Conference on Machine Learning, pages 214–223, 2017. (Cited on pages 1 and 2.) [7]J.-D. Benamou, G. Carlier, M. Cuturi, L. Nenna, and G. Peyré. Iterative Bregman projections for regularized transportation problems. SIAM Journal on Scientific Computing, 37(2):A1111– A1138, 2015. (Cited on page 2.) [8] E. Bernton, P. E. Jacob, M. Gerber, and C. P. Robert. Approximate Bayesian computation with the Wasserstein distance. Journal of the Royal Statistical Society Series B: Statistical Methodology, 81(2):235–269, 2019. (Cited on page 1.) [9] E. Bernton, P. E. Jacob, M. Gerber, and C. P. Robert. On parameter estimation with the Wasserstein distance. Information and Inference: A Journal of the IMA, 8(4):657–676, 2019. (Cited on page 1.) [10] C. Bonet, P. Berg, N. Courty, F. Septier, L. Drumetz, and M.-T. Pham. Spherical sliced- Wasserstein. International Conference on Learning Representations, 2023. (Cited on page 6.) [11]C. Bonet, L. Drumetz, and N. Courty. Sliced-Wasserstein distances and flows on Cartan- Hadamard manifolds. Journal of Machine Learning Research, 26(32):1–76, 2025. (Cited on page 6.) [12]N. Bonneel and J. Digne. A survey of optimal transport for computer graphics and computer vision. In Computer Graphics Forum, volume 42, pages 439–460. Wiley Online Library, 2023. (Cited on page 1.) [13]N. Bonneel, J. Rabin, G. Peyré, and H. Pfister. Sliced and Radon Wasserstein barycenters of measures. Journal of Mathematical Imaging and Vision, 1(51):22–45, 2015. (Cited on page 6.) [14]C. Bunne, L. Papaxanthos, A. Krause, and M. Cuturi. Proximal optimal transport modeling of population dynamics. In International Conference on Artificial Intelligence and Statistics, pages 6511–6528. PMLR, 2022. (Cited on page 2.) 21 [15]C. Bunne, S. G. Stark, G. Gut, J. S. Del Castillo, M. Levesque, K.-V. Lehmann, L. Pelkmans, A. Krause, and G. Rätsch. Learning single-cell perturbation responses using neural optimal transport. Nature methods, 20(11):1759–1768, 2023. (Cited on page 1.) [16]M. Catalano, H. Lavenant, A. Lijoi, and I. Prünster. A Wasserstein index of dependence for random measures. Journal of the American Statistical Association, 119(547):2396–2406, 2024. (Cited on page 1.) [17]M. Catalano, A. Lijoi, and I. Prünster. Measuring dependence in the Wasserstein distance for Bayesian nonparametric models. The Annals of Statistics, 49(5):2916–2947, 2021. (Cited on page 1.) [18]N. Courty, R. Flamary, and M. Ducoffe. Learning Wasserstein embeddings. In International Conference on Learning Representations, 2018. (Cited on page 2.) [19]N. Courty, R. Flamary, A. Habrard, and A. Rakotomamonjy. Joint distribution optimal transportation for domain adaptation. In Advances in Neural Information Processing Systems, pages 3730–3739, 2017. (Cited on page 1.) [20]N. Courty, R. Flamary, D. Tuia, and A. Rakotomamonjy. Optimal transport for domain adaptation. IEEE transactions on pattern analysis and machine intelligence, 39(9):1853–1865, 2016. (Cited on page 1.) [21]M. Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In C. Burges, L. Bottou, M. Welling, Z. Ghahramani, and K. Weinberger, editors, Advances in Neural Information Processing Systems, volume 26. Curran Associates, Inc., 2013. (Cited on pages 2, 3, and 4.) [22] M. Cuturi and A. Doucet. Fast computation of wasserstein barycenters. In International Conference on Machine Learning, pages 685–693. PMLR, 2014. (Cited on page 6.) [23] B. B. Damodaran, B. Kellenberger, R. Flamary, D. Tuia, and N. Courty. Deepjdot: Deep joint distribution optimal transport for unsupervised domain adaptation. In Proceedings of the European Conference on Computer Vision (ECCV), pages 447–463, 2018. (Cited on page 1.) [24] D. Dowson and B. Landau. The Fréchet distance between multivariate Normal distributions. Journal of Multivariate Analysis, 12(3):450–455, 1982. (Cited on page 1.) [25] E. Doxsey-Whitfield, K. MacManus, S. B. Adamo, L. Pistolesi, J. Squires, O. Borkovska, and S. R. Baptista. Taking advantage of the improved availability of census data: A first look at the gridded population of the world, version 4. Papers in Applied Geography, 1(3):226–234, 2015. (Cited on page 9.) [26]B. Engquist and B. D. Froese. Application of the Wasserstein metric to seismic signals. Communications in Mathematical Sciences, 12(5):979–988, 2014. (Cited on page 2.) [27]K. Fatras, T. Sejourne, R. Flamary, and N. Courty. Unbalanced minibatch optimal transport; applications to domain adaptation. In M. Meila and T. Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 3186–3197. PMLR, 18–24 Jul 2021. (Cited on page 1.) 22 [28]K. Fatras, Y. Zine, R. Flamary, R. Gribonval, and N. Courty. Learning with minibatch Wasser- stein: asymptotic and gradient properties. In AISTATS 2020-23nd International Conference on Artificial Intelligence and Statistics, volume 108, pages 1–20, 2020. (Cited on page 1.) [29]J. Feydy, B. Charlier, F.-X. Vialard, and G. Peyré. Optimal transport for diffeomorphic registration. In Medical Image Computing and Computer Assisted Intervention- MICCAI 2017: 20th International Conference, Quebec City, QC, Canada, September 11-13, 2017, Proceedings, Part I 20, pages 291–299. Springer, 2017. (Cited on page 1.) [30] R. C. Garrett, T. Harris, Z. Wang, and B. Li. Validating climate models with spherical convolutional Wasserstein distance. Advances in Neural Information Processing Systems, 37:59119–59149, 2024. (Cited on page 6.) [31]A. Genevay, L. Chizat, F. Bach, M. Cuturi, and G. Peyré. Sample complexity of sinkhorn divergences. In The 22nd international conference on artificial intelligence and statistics, pages 1574–1583. PMLR, 2019. (Cited on page 2.) [32] A. Genevay, M. Cuturi, G. Peyré, and F. Bach. Stochastic optimization for large-scale optimal transport. Advances in neural information processing systems, 29, 2016. (Cited on page 2.) [33] A. Genevay, G. Peyré, and M. Cuturi. Learning generative models with Sinkhorn divergences. In International Conference on Artificial Intelligence and Statistics, pages 1608–1617. PMLR, 2018. (Cited on pages 1 and 2.) [34]D. Haviv, R. Z. Kunes, T. Dougherty, C. Burdziak, T. Nawy, A. Gilbert, and D. Pe’er. Wasserstein wormhole: Scalable optimal transport distance with Transformer. In Forty-first International Conference on Machine Learning, 2024. (Cited on page 2.) [35]P. He, O. Khangaonkar, H. Pirsiavash, Y. Bai, and S. Kolouri. Sinkhorn-drifting generative models. arXiv preprint arXiv:2603.12366, 2026. (Cited on pages 1 and 2.) [36]R. Jiang, A. Pacchiano, T. Stepleton, H. Jiang, and S. Chiappa. Wasserstein fair classification. In Uncertainty in Artificial Intelligence, pages 862–872. PMLR, 2020. (Cited on page 1.) [37]S. Kolouri, N. Naderializadeh, G. K. Rohde, and H. Hoffmann. Wasserstein embedding for graph learning. In International Conference on Learning Representations, 2021. (Cited on page 2.) [38]S. Kolouri, S. R. Park, M. Thorpe, D. Slepcev, and G. K. Rohde. Optimal mass transport: Signal processing and machine-learning applications. IEEE signal processing magazine, 34(4):43–59, 2017. (Cited on page 1.) [39] Y. Lipman, R. T. Q. Chen, H. Ben-Hamu, M. Nickel, and M. Le. Flow matching for generative modeling. In The Eleventh International Conference on Learning Representations, 2023. (Cited on page 1.) [40]X. Liu, E. Akbari, R. Diaz Martin, N. NaderiAlizadeh, and S. Kolouri. Efficient transferable optimal transport via min-sliced transport plans. arXiv e-prints, pages arXiv–2511, 2025. (Cited on pages 8, 9, 10, 11, and 14.) 23 [41]X. Liu, R. D. Martin, Y. Bai, A. Shahbazi, M. Thorpe, A. Aldroubi, and S. Kolouri. Expected sliced transport plans. In The Thirteenth International Conference on Learning Representations, 2025. (Cited on page 6.) [42] G. Mahey, L. Chapel, G. Gasso, C. Bonet, and N. Courty. Fast optimal transport through sliced Wasserstein generalized geodesics. In Proceedings of the 37th International Conference on Neural Information Processing Systems, pages 35350–35385, 2023. (Cited on pages 6, 8, 9, 10, and 11.) [43] A. Makkuva, A. Taghvaei, S. Oh, and J. Lee. Optimal transport mapping via input convex neural networks. In International Conference on Machine Learning, pages 6672–6681. PMLR, 2020. (Cited on page 2.) [44] T. Manole, S. Balakrishnan, J. Niles-Weed, and L. Wasserman. Plugin estimation of smooth optimal transport maps. The Annals of Statistics, 52(3):966–998, 2024. (Cited on page 2.) [45]C. Moosmüller and A. Cloninger. Linear optimal transport embedding: provable Wasserstein classification for certain rigid transformations and perturbations. Information and Inference: A Journal of the IMA, 12(1):363–389, 2023. (Cited on page 2.) [46]K. Nguyen. An introduction to sliced optimal transport: foundations, advances, extensions, and applications. Foundations and Trends® in Computer Graphics and Vision, 17(3-4):171–391, 2025. (Cited on pages 2 and 6.) [47]K. Nguyen, N. Bariletto, and N. Ho. Quasi-Monte Carlo for 3d sliced Wasserstein. In The Twelfth International Conference on Learning Representations, 2024. (Cited on page 6.) [48]K. Nguyen and P. Mueller. Summarizing nonparametric Bayesian mixture posteriors–sliced optimal transport metrics for Gaussian mixtures. Journal of Computational and Graphical Statistics, (just-accepted):1–22, 2026. (Cited on page 6.) [49] K. Nguyen, D. Nguyen, T. Pham, and N. Ho. Improving mini-batch optimal transport via partial transportation. In Proceedings of the 39th International Conference on Machine Learning, 2022. (Cited on page 1.) [50]K. Nguyen, H. Nguyen, and N. Ho. Fast estimation of Wasserstein distances via regression on sliced Wasserstein distances. In The Fourteenth International Conference on Learning Representations, 2026. (Cited on page 2.) [51] J. B. Orlin. A polynomial time primal network simplex algorithm for minimum cost flows. Mathematical Programming, 78(2):109–129, 1997. (Cited on page 1.) [52] G. Patrini, R. van den Berg, P. Forre, M. Carioni, S. Bhargav, M. Welling, T. Genewein, and F. Nielsen. Sinkhorn autoencoders. In Uncertainty in Artificial Intelligence, pages 733–743. PMLR, 2020. (Cited on page 1.) [53] G. Peyré, M. Cuturi, et al. Computational optimal transport: With applications to data science. Foundations and Trends® in Machine Learning, 11(5-6):355–607, 2019. (Cited on pages 1 and 4.) 24 [54]A.-A. Pooladian, H. Ben-Hamu, C. Domingo-Enrich, B. Amos, Y. Lipman, and R. T. Chen. Multisample flow matching: Straightening flows with minibatch couplings. In International Conference on Machine Learning, pages 28100–28127. PMLR, 2023. (Cited on page 1.) [55]M. Quellmalz, R. Beinert, and G. Steidl. Sliced optimal transport on the sphere. Inverse Problems, 39(10):105005, 2023. (Cited on page 6.) [56]J. Rabin, G. Peyré, J. Delon, and M. Bernot. Wasserstein barycenter and its application to texture mixing. In Scale Space and Variational Methods in Computer Vision: Third International Conference, SSVM 2011, Ein-Gedi, Israel, May 29–June 2, 2011, Revised Selected Papers 3, pages 435–446. Springer, One New York Plaza, Suite 4600, New York, NY 10004-1562, 2012. (Cited on pages 2 and 6.) [57]P. Rigollet and A. J. Stromme. On the sample complexity of entropic optimal transport. The Annals of Statistics, 53(1):61–90, 2025. (Cited on page 2.) [58]Y. Rubner, C. Tomasi, and L. J. Guibas. A metric for distributions with applications to image databases. In Sixth International Conference on Computer Vision (IEEE Cat. No. 98CH36271), pages 59–66. IEEE, 1998. (Cited on page 2.) [59]F. Santambrogio. Optimal transport for applied mathematicians. Birkäuser, NY, 55(58-63):94, 2015. (Cited on page 5.) [60] M. Scetbon and M. Cuturi. Low-rank optimal transport: Approximation, statistics and debiasing. Advances in Neural Information Processing Systems, 35:6802–6814, 2022. (Cited on page 1.) [61]M. Scetbon, M. Cuturi, and G. Peyré. Low-rank Sinkhorn factorization. In International Conference on Machine Learning, pages 9344–9354. PMLR, 2021. (Cited on page 1.) [62]G. Schiebinger, J. Shu, M. Tabaka, B. Cleary, V. Subramanian, A. Solomon, J. Gould, S. Liu, S. Lin, P. Berube, et al. Optimal-transport analysis of single-cell gene expression identifies developmental trajectories in reprogramming. Cell, 176(4):928–943, 2019. (Cited on page 1.) [63]T. Séjourné, F.-X. Vialard, and G. Peyré. Faster unbalanced optimal transport: Translation invariant sinkhorn and 1-d frank-wolfe. In International Conference on Artificial Intelligence and Statistics, pages 4995–5021. PMLR, 2022. (Cited on page 6.) [64]R. Shu. Amortized optimizationhttp://ruishu.io/2017/11/07/amortized-optimization/, 2017. (Cited on page 2.) [65]R. Sinkhorn and P. Knopp. Concerning nonnegative matrices and doubly stochastic matrices. Pacific Journal of Mathematics, 21(2):343–348, 1967. (Cited on page 2.) [66]J. Solomon, F. De Goes, G. Peyré, M. Cuturi, A. Butscher, A. Nguyen, T. Du, and L. Guibas. Convolutional Wasserstein distances: Efficient optimal transportation on geometric domains. ACM Transactions on Graphics (ToG), 34(4):1–11, 2015. (Cited on page 1.) [67] J. Solomon, G. Peyré, V. G. Kim, and S. Sra. Entropic metric alignment for correspondence problems. ACM Transactions on Graphics (TOG), 35(4):72, 2016. (Cited on page 1.) 25 [68]M. Sommerfeld, J. Schrieber, Y. Zemel, and A. Munk. Optimal transport: Fast probabilistic approximation with exact solvers. Journal of Machine Learning Research, 20:105–1, 2019. (Cited on page 1.) [69]E. Tanguy, L. Chapel, and J. Delon. Sliced optimal transport plans. arXiv preprint arXiv:2508.01243, 2025. (Cited on page 6.) [70] I. Tolstikhin, O. Bousquet, S. Gelly, and B. Schoelkopf. Wasserstein auto-encoders. In International Conference on Learning Representations, 2018. (Cited on page 1.) [71]A. Tong, K. FATRAS, N. Malkin, G. Huguet, Y. Zhang, J. Rector-Brooks, G. Wolf, and Y. Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. Transactions on Machine Learning Research, 2024. Expert Certification. (Cited on pages 1, 9, 11, 12, and 13.) [72]H. Tran, Y. Bai, A. Kothapalli, A. Shahbazi, X. Liu, R. D. Martin, and S. Kolouri. Stereographic spherical sliced Wasserstein distances. International Conference on Machine Learning, 2024. (Cited on pages 6 and 9.) [73]C. Villani. Topics in optimal transportation. Number 58. American Mathematical Soc., 2003. (Cited on page 1.) [74]C. Villani. Optimal transport: old and new, volume 338. Springer, One New York Plaza, Suite 4600, New York, NY 10004-1562, 2009. (Cited on pages 1 and 3.) [75]F. Wu, N. Courty, S. Jin, and S. Z. Li. Improving molecular representation learning with metric learning-enhanced optimal transport. Patterns, 4(4), 2023. (Cited on page 1.) [76] J. Zhu, A. Guha, D. Do, M. Xu, X. Nguyen, and D. Zhao. Functional optimal transport: regularized map estimation and domain adaptation for functional data. Journal of Machine Learning Research, 25(276):1–49, 2024. (Cited on page 1.) 26