Amortized Optimal Transport from Sliced Potentials
| Authors | Minh-Phuc Truong & Khai Nguyen |
| Year | 2026 |
| Field | Statistics / ML |
| arXiv | 2604.15114 |
| Download | |
| Categories | stat.ML, cs.AI, cs.LG |
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.
Engineering Breakdown
Plain English
This paper addresses the problem of computing optimal transport (OT) plans repeatedly across many different pairs of probability distributions—a computationally expensive task in practice. The authors propose an amortized learning approach that uses simpler sliced OT computations to train neural networks that can quickly predict high-quality OT solutions for new distribution pairs. They introduce two strategies: RA-OT (regression-based) which treats sliced potentials as features to predict original OT potentials, and OA-OT (objective-based) which directly optimizes the Kantorovich dual objective to learn these potentials. Both methods allow practitioners to solve OT problems fast at test time by feeding through learned models rather than solving expensive optimization problems from scratch.
Core Technical Contribution
The key innovation is using sliced OT—a computationally cheaper approximation—to provide training signal for amortized models that predict full OT solutions. Specifically, the authors formulate this as a functional regression problem where Kantorovich potentials (the dual solutions to the OT problem) computed from sliced OT serve as predictors for the original OT potentials. This is novel because prior amortized OT methods typically learn end-to-end mappings from measures to transport maps, whereas this approach explicitly leverages the mathematical structure of OT duality. The two amortization strategies (RA-OT via least-squares fitting and OA-OT via dual objective optimization) offer different trade-offs between computational simplicity and theoretical alignment with the OT problem.
How It Works
The pipeline works as follows: given pairs of measures, compute their sliced OT potentials (which are fast to obtain because slicing reduces dimensionality). For RA-OT, train a functional regression model where the input features are the sliced potentials and the output targets are the true OT potentials from the original high-dimensional problem, using standard least-squares fitting. For OA-OT, instead of fitting to targets, optimize the learned potential parameters directly by evaluating the Kantorovich dual objective on the OT problem and backpropagating through it. Once trained, at test time the learned model takes sliced potentials as input and outputs predicted OT potentials, from which the final OT plan is recovered using standard OT recovery methods. The model architecture itself is a functional mapping (essentially a regression function) that transforms low-dimensional sliced solution information into high-dimensional dual potentials.
Production Impact
For engineers building production systems that solve OT problems repeatedly—such as domain adaptation pipelines, generative modeling with Wasserstein losses, or optimal transport-based data matching—this approach offers substantial speedup at inference time. Instead of solving a new optimization problem each time (which scales poorly and may require hundreds of iterations), you run a single forward pass through a learned model, reducing latency from seconds/minutes to milliseconds. The trade-off is upfront training cost: you need a dataset of measure pairs and their ground-truth OT solutions to train the amortized model, which requires solving the full OT problem on each training sample. This makes sense when you have a distribution of OT problems at test time that resembles the training distribution; if test problems are highly out-of-distribution, the amortized model may be inaccurate. Integration is relatively straightforward—drop in the trained model as a preprocessing step before your OT-dependent computations, though you'll need to validate accuracy on your specific domain.
Limitations and When Not to Use This
The method assumes that sliced OT potentials contain sufficient information to predict full OT potentials, which may not hold for highly structured or multimodal distribution pairs where 1D projections lose critical information. The approach requires access to ground-truth OT solutions during training, which itself is expensive to compute at scale—this limits the method to relatively small training datasets. The paper does not provide theoretical guarantees on how well the amortized predictions approximate true OT solutions, nor does it characterize when RA-OT versus OA-OT should be preferred empirically. Generalization to out-of-distribution measure pairs (different domains, different data modalities than training) remains unclear, and the computational savings only materialize if amortization overhead (model inference) is genuinely cheaper than the OT solver it replaces, which depends on problem dimension and solver efficiency.
Research Context
This work builds directly on the growing literature of amortized optimization in machine learning, where neural networks learn to predict solutions to repeatedly-solved optimization problems. It extends prior amortized OT methods by explicitly using sliced OT—a well-established technique for fast OT approximation—as the bridge between expensive full OT and learned mappings. The paper likely benchmarks against existing amortized OT approaches and classical OT solvers on standard datasets (likely images, point clouds, or synthetic measures). This research opens directions toward combining other OT approximations (entropic regularization, neural OT solvers) with amortization, and toward learning functional mappings in higher-dimensional OT applications where both solving and amortizing are challenging.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::