Objective-based Amortization (OA-OT)

• OA-OT is an amortized optimization paradigm that trains predictors by directly optimizing optimal transport objectives, avoiding regression to precomputed solutions.
• It leverages sliced potentials and parametric ensembles to reuse structure across repeated OT tasks, enabling rapid warm-starts and scalable inference.
• Empirical results show significant speedups and improved RMSE in discrete entropic and continuous OT settings, with extensions to filtering applications.

Searching arXiv for the cited papers to ground the article in the relevant literature. Objective-based Amortization (OA-OT) is an amortized optimization paradigm for optimal transport (OT) in which a parametric predictor is trained by directly optimizing an OT objective over a distribution of transport tasks, rather than by regressing to precomputed solutions. In the literature summarized here, the term is used explicitly in “Amortized Optimal Transport from Sliced Potentials” (Truong et al., 16 Apr 2026), while the same principle is instantiated in “Meta Optimal Transport” (Amos et al., 2022). A related, but terminologically distinct, objective-guided amortization mechanism also appears in the amortized optimal transport filter (A-OTF), where pre-trained conditional OT maps are selected and weighted by discrepancy objectives during nonlinear/non-Gaussian filtering (Al-Jarrah et al., 16 Mar 2025). Across these settings, OA-OT is characterized by reuse of structure across repeated OT problems, direct training against dual or primal-surrogate OT criteria, and rapid approximation of new transport plans or maps.

1. Definition and conceptual scope

OA-OT addresses repeated OT problems drawn from a task distribution. In the multi-pair setting of (Truong et al., 16 Apr 2026), one observes many OT problems sampled from a meta-distribution over triples $(\mu,\nu,c)\sim\mathcal D$ and seeks to predict transport plans for new pairs by learning reusable structure across prior instances. In Meta OT, the corresponding task distribution is described as a meta-distribution over OT problems, with discrete tasks of the form $(a,b,C)$ and continuous tasks over measures $(\alpha,\beta)$ under squared Euclidean cost (Amos et al., 2022).

The defining feature of OA-OT is objective-based training. Rather than solve each new OT instance from scratch, or fit a predictor to ground-truth potentials by least squares, OA-OT optimizes the OT objective itself. In discrete entropic OT, this means minimizing the negative dual objective evaluated at predicted potentials; in continuous $W_2$ settings, Meta OT uses a dual surrogate based on W2 generative networks (W2GN) (Amos et al., 2022). In the sliced-potential formulation, OA-OT estimates the coefficients of a functional potential model by optimizing the entropic Kantorovich dual across multiple training pairs (Truong et al., 16 Apr 2026).

This separates OA-OT from regression-based amortization. “Amortized Optimal Transport from Sliced Potentials” introduces both regression-based amortization (RA-OT) and OA-OT, with the distinction that RA-OT fits original OT potentials by least squares, whereas OA-OT never needs ground-truth potentials and instead optimizes the dual objective directly (Truong et al., 16 Apr 2026). The same distinction is implicit in Meta OT, which does not supervise on exact OT solutions but trains predictors by the OT task loss itself (Amos et al., 2022).

2. Mathematical formulation

The OT background used by OA-OT is standard Kantorovich primal-dual theory, together with entropic regularization. For measures $\mu$ on $\mathcal X$ and $\nu$ on $\mathcal Y$ with cost $c:\mathcal X\times\mathcal Y\to\mathbb R_+$, the primal problem is

$$\pi^\star \in \arg\min_{\pi \in \Pi(\mu,\nu)} \int_{\mathcal X\times \mathcal Y} c(x,y)\, \mathrm d\pi(x,y),$$

and the classical dual is

$(a,b,C)$0

subject to

$(a,b,C)$1

In the entropic setting with $(a,b,C)$2, the optimal plan has Gibbs form and one potential determines the other via the dual mapping (Truong et al., 16 Apr 2026).

In the discrete entropic case, for $(a,b,C)$3 and $(a,b,C)$4 with cost matrix $(a,b,C)$5, the plan is

$(a,b,C)$6

and the one-potential update is

$(a,b,C)$7

Meta OT writes the discrete entropic dual through a single predicted potential and defines the amortized loss

$(a,b,C)$8

where $(a,b,C)$9 and $(\alpha,\beta)$0 is recovered from the optimality relation (Amos et al., 2022).

The explicit OA-OT formulation in (Truong et al., 16 Apr 2026) parameterizes a Kantorovich potential by combining sliced OT-derived predictors:

$(\alpha,\beta)$1

where $(\alpha,\beta)$2 are projection parameters and $(\alpha,\beta)$3 are slicewise $(\alpha,\beta)$4D Kantorovich potentials. The parameters are then learned by minimizing the empirical entropic dual objective over training pairs:

$(\alpha,\beta)$5

The second potential is obtained through the entropic dual mapping, and the transport plan is recovered from the Gibbs kernel (Truong et al., 16 Apr 2026).

A central structural property of this formulation is parsimony. Because the potential is parameterized by $(\alpha,\beta)$6 sliced predictors, the parameter vector has size $(\alpha,\beta)$7 and is independent of the number of atoms $(\alpha,\beta)$8 in the discrete measures. This contrasts with Meta-OT models whose neural parameterization often scales with problem structure such as fixed atom counts (Truong et al., 16 Apr 2026).

3. OA-OT in Meta Optimal Transport

“Meta Optimal Transport” presents a general amortized optimization framework for OT that matches the OA-OT principle even though the paper uses the name Meta OT (Amos et al., 2022). Its discrete formulation predicts a dual potential $(\alpha,\beta)$9 from the input measures and cost structure, computes $W_2$0 via the entropic dual relation, and reconstructs the plan

$W_2$1

Training minimizes the OT objective directly rather than the distance to a reference solution:

$W_2$2

The approach does not differentiate through Sinkhorn iterations; instead, it optimizes the amortized objective with autodiff and uses the prediction directly or as a warm start (Amos et al., 2022).

In the continuous $W_2$3 setting, Meta OT predicts the parameters of convex potentials through a hypernetwork, using input measures such as grayscale images or color palettes. The induced map is $W_2$4, where $W_2$5 is represented by an ICNN. Training uses the W2GN loss

$W_2$6

and the amortized objective is

$W_2$7

This makes objective-based amortization applicable to both discrete entropic OT and continuous Monge-map prediction (Amos et al., 2022).

Empirically, Meta OT demonstrates large warm-start gains. On MNIST, warm-starting Sinkhorn reduces time to reach $W_2$8 from $W_2$9 with zero initialization to $\mu$0 with Meta OT, while prediction alone takes approximately $\mu$1 (Amos et al., 2022). On spherical supply-demand transport, warm-starting reduces time to $\mu$2 from $\mu$3 to $\mu$4, reported as an $\mu$5 speedup, and the initial prediction is approximately $\mu$6 faster than converged Sinkhorn maps (Amos et al., 2022). In continuous color transfer, Meta OT prediction is approximately $\mu$7 versus approximately $\mu$8–$\mu$9 for $\mathcal X$0–$\mathcal X$1 W2GN steps from scratch, and warm-started fine-tuning reaches normalized dual value approximately $\mathcal X$2 faster than training from scratch (Amos et al., 2022).

4. OA-OT from sliced potentials

“Amortized Optimal Transport from Sliced Potentials” gives the most explicit modern formulation of OA-OT (Truong et al., 16 Apr 2026). Its starting point is sliced OT: for projection directions $\mathcal X$3, the measures are projected via $\mathcal X$4, yielding one-dimensional pushforwards

$\mathcal X$5

For costs of the form $\mathcal X$6 with strictly convex $\mathcal X$7, the one-dimensional OT problem admits quantile coupling, and in the discrete $\mathcal X$8D case the sliced potentials are computed in $\mathcal X$9 time per slice (Truong et al., 16 Apr 2026). OA-OT uses these fast slicewise potentials as predictors for the full-dimensional potential.

The model is deliberately parsimonious. A single potential is parameterized as a linear combination of $\nu$0 sliced predictors, while the second potential is obtained by the entropic dual mapping and therefore adds no extra trainable degrees of freedom. Because the parameter vector is $\nu$1, the model size is independent of the numbers of atoms $\nu$2, which allows generalization across measure pairs whose atom counts vary (Truong et al., 16 Apr 2026).

Training proceeds pairwise. For each training pair $\nu$3, one computes sliced pushforwards, solves $\nu$4D OT on each slice to obtain $\nu$5, composes them with the projections, forms the parameterized potential $\nu$6, computes the entropic dual objective $\nu$7, and updates $\nu$8 by gradient-based optimization. The paper reports Adam with learning rate $\nu$9 and $\mathcal Y$0 iterations as the standard setting (Truong et al., 16 Apr 2026). Inference on a new pair requires computing the sliced predictors, forming $\mathcal Y$1, recovering $\mathcal Y$2 via the one-potential mapping, and assembling the plan through the Gibbs kernel. One or two Sinkhorn scaling iterations may optionally be used to improve marginal matching if numerical drift is observed (Truong et al., 16 Apr 2026).

The empirical profile is strong across several geometries. With $\mathcal Y$3 training pairs and $\mathcal Y$4 slices, OA-OT on MNIST achieves transport-plan RMSE approximately $\mathcal Y$5, compared with approximately $\mathcal Y$6 for RA-OT and approximately $\mathcal Y$7 for Meta-OT; training time is approximately $\mathcal Y$8–$\mathcal Y$9 for OA-OT versus approximately $c:\mathcal X\times\mathcal Y\to\mathbb R_+$0 for Meta-OT, while inference is approximately $c:\mathcal X\times\mathcal Y\to\mathbb R_+$1 per pair (Truong et al., 16 Apr 2026). On spherical transport, OA-OT reports RMSE approximately $c:\mathcal X\times\mathcal Y\to\mathbb R_+$2–$c:\mathcal X\times\mathcal Y\to\mathbb R_+$3 depending on $c:\mathcal X\times\mathcal Y\to\mathbb R_+$4, compared with approximately $c:\mathcal X\times\mathcal Y\to\mathbb R_+$5–$c:\mathcal X\times\mathcal Y\to\mathbb R_+$6 for RA-OT and approximately $c:\mathcal X\times\mathcal Y\to\mathbb R_+$7–$c:\mathcal X\times\mathcal Y\to\mathbb R_+$8 for Meta-OT; single-projection baselines are reported at high RMSE approximately $c:\mathcal X\times\mathcal Y\to\mathbb R_+$9 (Truong et al., 16 Apr 2026). On color transfer, OA-OT reaches RMSE approximately $$\pi^\star \in \arg\min_{\pi \in \Pi(\mu,\nu)} \int_{\mathcal X\times \mathcal Y} c(x,y)\, \mathrm d\pi(x,y),$$0, compared with approximately $$\pi^\star \in \arg\min_{\pi \in \Pi(\mu,\nu)} \int_{\mathcal X\times \mathcal Y} c(x,y)\, \mathrm d\pi(x,y),$$1 for RA-OT and approximately $$\pi^\star \in \arg\min_{\pi \in \Pi(\mu,\nu)} \int_{\mathcal X\times \mathcal Y} c(x,y)\, \mathrm d\pi(x,y),$$2 for Meta-OT, with training approximately $$\pi^\star \in \arg\min_{\pi \in \Pi(\mu,\nu)} \int_{\mathcal X\times \mathcal Y} c(x,y)\, \mathrm d\pi(x,y),$$3 and inference approximately $$\pi^\star \in \arg\min_{\pi \in \Pi(\mu,\nu)} \int_{\mathcal X\times \mathcal Y} c(x,y)\, \mathrm d\pi(x,y),$$4–$$\pi^\star \in \arg\min_{\pi \in \Pi(\mu,\nu)} \int_{\mathcal X\times \mathcal Y} c(x,y)\, \mathrm d\pi(x,y),$$5 (Truong et al., 16 Apr 2026).

The method also extends to mini-batch OT conditional flow matching. On $$\pi^\star \in \arg\min_{\pi \in \Pi(\mu,\nu)} \int_{\mathcal X\times \mathcal Y} c(x,y)\, \mathrm d\pi(x,y),$$6D toy data, OA-OT and RA-OT are reported to accelerate training by $$\pi^\star \in \arg\min_{\pi \in \Pi(\mu,\nu)} \int_{\mathcal X\times \mathcal Y} c(x,y)\, \mathrm d\pi(x,y),$$7–$$\pi^\star \in \arg\min_{\pi \in \Pi(\mu,\nu)} \int_{\mathcal X\times \mathcal Y} c(x,y)\, \mathrm d\pi(x,y),$$8 over OT-CFM while maintaining near-straight trajectories, with NPE approximately $$\pi^\star \in \arg\min_{\pi \in \Pi(\mu,\nu)} \int_{\mathcal X\times \mathcal Y} c(x,y)\, \mathrm d\pi(x,y),$$9–$(a,b,C)$00 versus approximately $(a,b,C)$01–$(a,b,C)$02 for OT-CFM (Truong et al., 16 Apr 2026). On high-dimensional CIFAR-10 fine-tuning, OA-OT achieves NFE/sample approximately $(a,b,C)$03, while RA-OT achieves the best FID approximately $(a,b,C)$04; both amortized variants significantly reduce fine-tuning time compared to exact OT-CFM, reported as approximately $(a,b,C)$05 (Truong et al., 16 Apr 2026).

5. Objective-guided amortization in optimal transport filtering

In nonlinear/non-Gaussian Bayesian filtering, objective-guided amortization appears in the amortized optimal transport filter (A-OTF), although the paper does not use the term OA-OT explicitly (Al-Jarrah et al., 16 Mar 2025). The filtering problem is defined by the discrete-time state-space model

$(a,b,C)$06

with posterior recursion

$(a,b,C)$07

The optimal transport filter replaces the conditioning operator by a conditional Brenier map $(a,b,C)$08 satisfying

$(a,b,C)$09

so that

$(a,b,C)$10

The computational burden arises because the conditional map is learned by solving a stochastic saddle-point problem online at every time step (Al-Jarrah et al., 16 Mar 2025).

A-OTF amortizes this cost through a library of pre-trained OT maps and a mixture-of-experts aggregation. Offline, one generates $(a,b,C)$11 pre-trained OTF maps $(a,b,C)$12, computes the pairwise distance matrix

$(a,b,C)$13

and clusters the training instances with K-medoids using one of three distances: $(a,b,C)$14, $(a,b,C)$15, or $(a,b,C)$16. The medoids define representative expert maps $(a,b,C)$17 (Al-Jarrah et al., 16 Mar 2025).

Online, the current predictive context is compared against the medoids using a distance $(a,b,C)$18, and the gating weights are

$(a,b,C)$19

with $(a,b,C)$20 and $(a,b,C)$21 yielding nearest-neighbor selection. The amortized map is then

$(a,b,C)$22

and the particle update is

$(a,b,C)$23

The paper explicitly describes the offline clustering as minimizing within-cluster discrepancies under the chosen metric and the online weighting as a softmax of negative distances akin to Nadaraya–Watson kernel interpolation or scattered data approximation (Al-Jarrah et al., 16 Mar 2025).

This supports an OA-OT interpretation, but only in a qualified sense. The paper states that the term “Objective-based Amortization (OA-OT)” is not used there. What is explicit is that both clustering and online weighting are objective-driven, in the sense that they are determined by discrepancy objectives rather than by direct regression to a target map (Al-Jarrah et al., 16 Mar 2025). A plausible extension, stated in the source material as consistent with the framework but not in the paper, is to choose the mixture weights by directly minimizing a discrepancy between the amortized pushforward and the target posterior,

$(a,b,C)$24

with $(a,b,C)$25 and $(a,b,C)$26. This would make the online gating an explicit posterior-matching OA-OT objective rather than a similarity-based surrogate (Al-Jarrah et al., 16 Mar 2025).

Empirically, A-OTF preserves the ability of OTF to capture non-Gaussian structure while significantly reducing runtime. In the quadratic observation example, OTF and A-OTF capture bimodality while EnKF/KF fail; with $(a,b,C)$27 particles for all methods and SIR with $(a,b,C)$28 as ground truth, A-OTF with $(a,b,C)$29, $(a,b,C)$30, $(a,b,C)$31, $(a,b,C)$32 performs comparably to OTF while faster (Al-Jarrah et al., 16 Mar 2025). In Lorenz-63 with scalar observation, observing $(a,b,C)$33 yields bimodality in $(a,b,C)$34 due to symmetry, and A-OTF with the same configuration and $(a,b,C)$35 captures multimodality in unobserved coordinates similarly to OTF; weighted selection with $(a,b,C)$36 is reported as robust, while $(a,b,C)$37 and $(a,b,C)$38 exhibit overfitting for $(a,b,C)$39 (Al-Jarrah et al., 16 Mar 2025).

6. Strengths, limitations, and recurring misconceptions

A consistent strength of OA-OT is computational reuse across related tasks. In all three lines of work, the expensive part of solving OT is shifted from per-instance optimization to a learned amortized mechanism: a parametric dual predictor in Meta OT, a low-dimensional sliced-potential model in explicit OA-OT, or a pre-trained map library in A-OTF (Amos et al., 2022, Truong et al., 16 Apr 2026, Al-Jarrah et al., 16 Mar 2025). This reuse can provide either direct approximate solutions or warm starts that substantially reduce solver iterations.

A second strength is alignment between training and the eventual task objective. This is explicit in the contrast between RA-OT and OA-OT: RA-OT performs least-squares regression to ground-truth potentials, whereas OA-OT optimizes the entropic Kantorovich dual directly and therefore avoids the bias that can arise from fitting potentials rather than optimizing transport optimality (Truong et al., 16 Apr 2026). The same point underlies Meta OT, which learns by minimizing the OT objective rather than an MSE to precomputed plans or potentials (Amos et al., 2022).

Several limitations recur. First, the explicit sliced-potential OA-OT targets entropic OT; extending it to classical unregularized OT would require explicit inequality-constraint handling and plan recovery on contact sets (Truong et al., 16 Apr 2026). Second, amortized predictions may suffer under strong distribution shift: Meta OT notes degradation under OOD tasks, and A-OTF depends on the coverage and diversity of the offline library, with insufficient coverage degrading accuracy in highly non-stationary regimes (Amos et al., 2022, Al-Jarrah et al., 16 Mar 2025). Third, OA-OT does not automatically guarantee exact feasibility at inference time. In discrete entropic OT, predicted potentials can yield small marginal errors; both the sliced-potential method and Meta OT allow a short Sinkhorn refinement to improve marginal matching (Truong et al., 16 Apr 2026, Amos et al., 2022).

A common misconception is that OA-OT is simply fast regression to previously solved OT problems. The literature here does not support that equivalence. OA-OT, in the explicit sense of (Truong et al., 16 Apr 2026), is defined by optimization of the Kantorovich dual objective. Regression-based amortization is treated as a separate strategy. Another misconception is that amortization removes the need for geometry-aware modeling. In fact, the sliced-potential framework emphasizes projection design, including stereographic or spherical projections for non-Euclidean tasks, and A-OTF shows sensitivity to the choice of discrepancy metric and the number of experts $(a,b,C)$40 (Truong et al., 16 Apr 2026, Al-Jarrah et al., 16 Mar 2025).

Taken together, these works suggest a broad but technically coherent notion of OA-OT. In its strictest form, OA-OT denotes direct optimization of an OT objective for an amortized potential predictor over a task distribution (Truong et al., 16 Apr 2026). In a broader sense, it includes objective-driven amortized predictors and warm-start mechanisms for repeated OT solves (Amos et al., 2022). In filtering, the same logic appears as discrepancy-based selection and weighting of pre-trained conditional OT maps, even though the paper frames the method as A-OTF rather than OA-OT (Al-Jarrah et al., 16 Mar 2025).