Amortized Optimal Transport from Sliced Potentials

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.

• The paper introduces a novel amortization framework leveraging sliced OT dual potentials to efficiently predict full OT duals.
• It develops regression-based (RA-OT) and objective-based (OA-OT) strategies that achieve lower RMSE and significant speedup on tasks like MNIST interpolation and color transfer.
• Empirical evaluations demonstrate strong accuracy, robustness, and computational gains, enabling scalable and adaptive OT solutions.

Amortized Optimal Transport from Sliced Potentials: A Technical Analysis

Introduction and Problem Context

Optimal Transport (OT) is central in computational mathematics and machine learning for tasks that measure distributional distances and define couplings between measures. However, iterative OT solvers (e.g., Sinkhorn) remain computationally demanding, particularly when one must repeatedly solve transport over many measure pairs, as in generative modeling, domain adaptation, and flow-matching. The amortization of OT—learning a function that predicts the OT plan or cost—has been proposed to alleviate this inefficiency but remains constrained by high input dimensionality and parameter redundancy.

"Amortized Optimal Transport from Sliced Potentials" (2604.15114) introduces a new paradigm for amortizing OT solutions by leveraging sliced OT dual potentials as compact, information-rich representations. The work proposes two amortization strategies—regression-based (RA-OT) and objective-based (OA-OT)—and demonstrates substantial benefits across canonical OT tasks, outperforming prior approaches such as Meta-OT in both accuracy and efficiency, especially in low-data regimes.

Figure 1: Amortized-OT uses sliced potentials to predict original dual potentials.

Methodological Advances: Sliced Potentials for OT Amortization

The key insight is to substitute direct raw input measure encoding with a projection-based featurization: for each OT problem between measures $\mu$ and $\nu$ with cost $c$, multiple one-dimensional projections (parameterized by $\theta$) are applied, and the corresponding one-dimensional sliced OT dual potentials are computed. These sliced potentials are then used as feature functions for an amortized model to predict the full OT dual potential in the original space.

Regression-Based Amortization (RA-OT)

RA-OT is formulated as a functional regression problem, where ground-truth OT duals are regressed onto the collection of sliced potentials. Practically, this reduces to optimizing a least-squares objective, leading to an analytic solution for linear models or efficient gradient-based routines for more expressive architectures. This approach is structurally robust and model-agnostic with respect to the dimensionality and the discretization of the input measures.

Objective-Based Amortization (OA-OT)

OA-OT extends the amortization philosophy of Meta-OT but replaces raw measure input with sliced potential features. The predictive model is optimized to maximize the Kantorovich dual objective as a function of the predicted dual, where the other dual variable is analytically recovered via the entropic OT duality. The amortization gap for OA-OT is minimized when the true duals lie in the linear span of sliced potentials, making increased projection count $L$ a critical knob for model expressivity.

Slicing yields several essential advantages: 1) data-agnostic featurization, so the amortizer’s parameter count is decoupled from the number of support atoms; 2) fast computation due to closed-form 1D OT; 3) improved ability to generalize across input domains with diverse granularities and structures.

Empirical Evaluation: Canonical and Geometric OT Tasks

A suite of experiments substantiates that sliced-potential amortization delivers superior accuracy, robustness, and efficiency on a broad range of OT tasks.

MNIST Wasserstein Interpolation

In 2D image transport on MNIST, both RA-OT and OA-OT achieve lower RMSE relative to the Sinkhorn ground truth compared to Meta-OT, Min-STP, and min-SWGG, especially in data-limited settings ($M=10, 20$). The amortized interpolations match the converged Sinkhorn barycenters both visually and quantitatively.

Figure 2: Wasserstein interpolations between MNIST test digits. Top: Sinkhorn (ground-truth) and Meta-OT. Bottom: RA-OT and OA-OT. Both methods closely match the ground truth sequence.

Spherical Supply-Demand Matching

On spherical OT (matching world population demand to landmass supply), both RA-OT and OA-OT outperform sliced-OT baselines and Meta-OT in producing smooth, globally coherent assignments that respect the non-Euclidean domain.

Figure 3: RA-OT and OA-OT generate globally coherent spherical matching patterns, outperforming baselines which produce local artifacts.

Color Transfer in RGB Space

For color histogram matching between natural images, RA-OT and OA-OT track the Sinkhorn solution closely and yield transferred images with high color fidelity and preserved texture.

Figure 4: Color transfer results on held-out test pairs $(M=50, L=100)$; both RA-OT and OA-OT adhere to target color palettes and preserve structure.

Figure 5: Amortized solvers robustly align high-dimensional RGB color histograms and preserve image detail in color transfer tasks.

OT-Conditional Flow Matching

In generative modeling (OT-CFM), amortized methods achieve 2.5x–4.5x training speedup over regular OT-CFM with nearly equivalent path straightness (NPE) and $W_2^2$ scores, demonstrating practical scalability—significant given the bottleneck due to OT computation in flow matching.

Figure 6: Flow trajectories—RA-OT and OA-OT replicate the uncrossed, straight integration paths of exact OT-CFM, with lower collision energy than independent baseline.

Figure 7: Both amortized methods generate effective, smooth integration flows in the Moons target transport.

Figure 8: Amortized approaches straighten trajectories for S-Curve targets, bypassing turbulence seen in standard trajectories.

Practical and Theoretical Implications

The presented amortization strategies reframe repeated OT computation as a feature regression or transfer problem, providing a scalable mechanism decoupled from input measure discretization. Theoretically, the approach indicates that OT plans can be efficiently and accurately reconstructed from collections of one-dimensional projections—a property that has implications for both statistical sample complexity and rapid adaptation to variable-sized measures.

In practice, these methods enable efficient and robust deployment of OT routines in large-scale or streaming settings (e.g., within stochastic generative models, real-time color transfer, or multi-agent matching), where latency and resource demands are prohibitive for repeated Sinkhorn or linear program-based solves.

Strong numerical results highlight: 1) consistent subdominant RMSE across all regimes, 2) robust visual preservation of topological structure, 3) strict parameter parsimony, and 4) order-of-magnitude speed gains, particularly in amortized generative modeling.

Future Directions

This work opens several future directions:

• Exploration of richer function classes (e.g., deep non-linear models) for mapping sliced potentials to duals, with study of generalization guarantees and amortization gap minimization.
• Application to continuous OT estimation, domain adaptation on complex manifolds, and broader use as an efficient inner loop in large-scale optimization.
• Extension to unbalanced OT, Gromov-Wasserstein, or semi-coupled matching settings by designing task-customized slicing and amortization architectures.
• Systematic investigation of projection selection strategies (randomized/quasi-Monte-Carlo, adaptive, or learned) to optimally balance computational complexity and expressivity.
• Integration with stochastic and Bayesian modeling frameworks, exploiting amortized OT as a meta-learning subroutine.

Conclusion

The paper "Amortized Optimal Transport from Sliced Potentials" (2604.15114) establishes a theoretically sound and practically effective pipeline for OT amortization. By leveraging sliced OT duals as efficient, low-dimensional surrogates, the proposed regression- and objective-based methods enable accurate and scalable plan prediction, decoupled from measure structure. This work meaningfully advances both the computational and statistical toolkit for repeated OT, with wide-ranging implications for generative modeling, domain adaptation, and geometric data analysis.