Amortized Optimal Transport from Sliced Potentials

2026-04-16

AI Summary

This work addresses the high computational cost of repeatedly solving optimal transport (OT) problems across multiple measure pairs by introducing two amortized optimization strategies—Regression-based Amortized OT (RA-OT) and Objective-based Amortized OT (OA-OT)—that efficiently predict Kantorovich potentials via function regression or dual objective optimization to rapidly reconstruct OT plans. It is the first to integrate sliced OT with amortized learning, yielding a general-purpose OT prediction model that requires fewer parameters, achieves higher accuracy, and does not rely on the specific structure of input measures. Experimental results demonstrate that the proposed approach significantly accelerates computation while maintaining high fidelity across diverse tasks, including MNIST transport, color transfer, spherical supply-demand transport, and minibatch OT flow matching.

Technology Category

• Machine Learning: Optimization
• Search and Optimization: Learning to Search
• Computer Vision: Learning & Optimization for CV

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications

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.

Problem

• Optimal Transport
• Amortized Optimization
• Kantorovich Potentials
• Sliced Optimal Transport
• Measure Pairs

Innovation

• amortized optimal transport
• sliced optimal transport
• Kantorovich potentials
• functional regression
• dual optimization