Amortized Optimal Transport from Sliced Potentials

TL;DR

This paper introduces amortized optimization methods for efficient repeated optimal transport plan prediction across multiple measure pairs, leveraging sliced Kantorovich potentials for improved accuracy and computational speed.

Contribution

It presents two novel amortization strategies, regression-based and objective-based, that utilize sliced OT to efficiently predict OT plans with high accuracy across diverse tasks.

Findings

Both methods enable rapid approximation of OT plans in multiple scenarios.

The approaches are effective on tasks like MNIST digit transport and color transfer.

They achieve high accuracy while being independent of measure structure complexity.

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.