---
title: Objective-based Amortization (OA-OT)
url: https://www.emergentmind.com/topics/objective-based-amortization-oa-ot
type: topic
---

# Objective-based Amortization (OA-OT)

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” [2604.15114], while the same principle is instantiated in “Meta Optimal Transport” [2206.05262]. 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 [2503.12633]. 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 [2604.15114], 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 [2206.05262].

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) [2206.05262]. 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 [2604.15114].

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 [2604.15114]. 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 [2206.05262].

## 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
$$
\sup_{\varphi,\psi}\ \mathbb E_{x\sim\mu}[\varphi(x)] + \mathbb E_{y\sim\nu}[\psi(y)]
$$
subject to
$$
\varphi(x)+\psi(y)\le c(x,y)\quad \forall x,y.
$$
In the entropic setting with $\varepsilon>0$, the optimal plan has Gibbs form and one potential determines the other via the dual mapping [2604.15114].

In the discrete entropic case, for $\mu=\sum_{i=1}^n \alpha_i\delta_{x_i}$ and $\nu=\sum_{j=1}^m \beta_j\delta_{y_j}$ with cost matrix $C_{ij}=c(x_i,y_j)$, the plan is
$$
P_{ij}^\star = \exp\!\Big(\frac{\mathbf f_i^\star+\mathbf g_j^\star-C_{ij}}{\varepsilon}\Big),
$$
and the one-potential update is
$$
\mathbf g = \varepsilon \log \boldsymbol\beta - \varepsilon \log\!\big(\exp(-C^\top/\varepsilon)\exp(\mathbf f/\varepsilon)\big).
$$
Meta OT writes the discrete entropic dual through a single predicted potential and defines the amortized loss
$$
\min_\theta\ \mathbb E_{(a,b,C)\sim D} J(\hat f_\theta(a,b,C);a,b,K(C)),
$$
where $K_{ij}=\exp\{-C_{ij}/\varepsilon\}$ and $g(f)$ is recovered from the optimality relation [2206.05262].

The explicit OA-OT formulation in [2604.15114] parameterizes a Kantorovich potential by combining sliced OT-derived predictors:
$$
\hat f_{\boldsymbol\omega}[\mu,\nu,c](x)=\sum_{l=1}^L \omega_l\, f_{\theta_l}^\star[\mu,\nu,c]\!\big(P_{\theta_l}^c(x)\big),
$$
where $\{\theta_l\}_{l=1}^L$ are projection parameters and $f_{\theta_l}^\star$ are slicewise $1$D Kantorovich potentials. The parameters are then learned by minimizing the empirical entropic dual objective over training pairs:
$$
\min_{\boldsymbol\omega}\ \frac{1}{N}\sum_{i=1}^N J\big(\hat f_{\boldsymbol\omega}[\mu_i,\nu_i,c_i];\mu_i,\nu_i,c_i\big).
$$
The second potential is obtained through the entropic dual mapping, and the transport plan is recovered from the Gibbs kernel [2604.15114].

A central structural property of this formulation is parsimony. Because the potential is parameterized by $L$ sliced predictors, the parameter vector has size $O(L)$ and is independent of the number of atoms $n,m$ in the discrete measures. This contrasts with Meta-OT models whose neural parameterization often scales with problem structure such as fixed atom counts [2604.15114].

## 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 [2206.05262]. Its discrete formulation predicts a dual potential $\hat f_\theta$ from the input measures and cost structure, computes $\hat g_\theta$ via the entropic dual relation, and reconstructs the plan
$$
\hat P_\theta=\mathrm{diag}(e^{\hat f/\varepsilon})\,K\,\mathrm{diag}(e^{\hat g/\varepsilon}).
$$
Training minimizes the OT objective directly rather than the distance to a reference solution:
$$
J(f)= - \langle f,a\rangle - \langle g(f),b\rangle + \varepsilon \langle e^{f/\varepsilon},K e^{g(f)/\varepsilon}\rangle.
$$
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 [2206.05262].

In the continuous $W_2$ 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 $T_\theta(x)=\nabla\psi_{\phi_\theta}(x)$, where $\psi_{\phi_\theta}$ is represented by an ICNN. Training uses the W2GN loss
$$
L(\phi) = \mathbb E_\alpha[\psi_\phi(x)] + \mathbb E_\beta[\langle \nabla \bar\psi_\phi(y),y\rangle - \psi_\phi(\nabla \bar\psi_\phi(y))] + \gamma \mathbb E_\beta[\|\nabla\psi_\phi(\nabla\bar\psi_\phi(y))-y\|_2^2],
$$
and the amortized objective is
$$
\min_\theta \mathbb E_{(\alpha,\beta)\sim D} L(\phî_\theta(\alpha,\beta);\alpha,\beta).
$$
This makes objective-based amortization applicable to both discrete entropic OT and continuous Monge-map prediction [2206.05262].

Empirically, Meta OT demonstrates large warm-start gains. On MNIST, warm-starting Sinkhorn reduces time to reach $\mathrm{err}\le 10^{-2}$ from $4.5\times 10^{-3}\,\mathrm s$ with zero initialization to $2.3\times 10^{-3}\,\mathrm s$ with Meta OT, while prediction alone takes approximately $5\times 10^{-5}\,\mathrm s$ [2206.05262]. On spherical supply-demand transport, warm-starting reduces time to $\mathrm{err}\le 10^{-2}$ from $0.88\,\mathrm s$ to $0.078\,\mathrm s$, reported as an $11.3\times$ speedup, and the initial prediction is approximately $37{,}500\times$ faster than converged Sinkhorn maps [2206.05262]. In continuous color transfer, Meta OT prediction is approximately $3.5\times 10^{-3}\,\mathrm s$ versus approximately $0.9$–$1.8\,\mathrm s$ for $1000$–$2000$ W2GN steps from scratch, and warm-started fine-tuning reaches normalized dual value approximately $1.0$ faster than training from scratch [2206.05262].

## 4. OA-OT from sliced potentials

“Amortized Optimal Transport from Sliced Potentials” gives the most explicit modern formulation of OA-OT [2604.15114]. Its starting point is sliced OT: for projection directions $\theta\in\mathbb S^{d-1}$, the measures are projected via $P_\theta^c(x)=\langle \theta,x\rangle$, yielding one-dimensional pushforwards
$$
\mu_\theta=(P_\theta^c)_\sharp\mu,\qquad \nu_\theta=(P_\theta^c)_\sharp\nu.
$$
For costs of the form $c(x,y)=h(x-y)$ with strictly convex $h$, the one-dimensional OT problem admits quantile coupling, and in the discrete $1$D case the sliced potentials are computed in $O((n+m)\log(n+m))$ time per slice [2604.15114]. 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 $L$ 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 $\omega\in\mathbb R^L$, the model size is independent of the numbers of atoms $n,m$, which allows generalization across measure pairs whose atom counts vary [2604.15114].

Training proceeds pairwise. For each training pair $(\mu_i,\nu_i,c_i)$, one computes sliced pushforwards, solves $1$D OT on each slice to obtain $f^\star_{i,\theta_l}$, composes them with the projections, forms the parameterized potential $\hat f_{\omega}$, computes the entropic dual objective $J$, and updates $\omega$ by gradient-based optimization. The paper reports Adam with learning rate $10^{-3}$ and $5{,}000$ iterations as the standard setting [2604.15114]. Inference on a new pair requires computing the sliced predictors, forming $\hat f_\omega$, recovering $\hat g$ 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 [2604.15114].

The empirical profile is strong across several geometries. With $M=50$ training pairs and $L=100$ slices, OA-OT on MNIST achieves transport-plan RMSE approximately $6.02\times10^{-6}$, compared with approximately $7.77\times10^{-6}$ for RA-OT and approximately $15.54\times10^{-6}$ for Meta-OT; training time is approximately $15$–$16\,\mathrm s$ for OA-OT versus approximately $37\,\mathrm s$ for Meta-OT, while inference is approximately $39\,\mathrm{ms}$ per pair [2604.15114]. On spherical transport, OA-OT reports RMSE approximately $3.93\times10^{-7}$–$4.80\times10^{-7}$ depending on $M$, compared with approximately $7.5$–$7.9\times10^{-7}$ for RA-OT and approximately $4.4$–$6.8\times10^{-7}$ for Meta-OT; single-projection baselines are reported at high RMSE approximately $1.12\times10^{-5}$ [2604.15114]. On color transfer, OA-OT reaches RMSE approximately $9.00\times10^{-6}$, compared with approximately $9.99\times10^{-6}$ for RA-OT and approximately $33.16\times10^{-6}$ for Meta-OT, with training approximately $18\,\mathrm s$ and inference approximately $17$–$19\,\mathrm{ms}$ [2604.15114].

The method also extends to mini-batch OT conditional flow matching. On $2$D toy data, OA-OT and RA-OT are reported to accelerate training by $2.5\times$–$4.5\times$ over OT-CFM while maintaining near-straight trajectories, with NPE approximately $0.014$–$0.063$ versus approximately $0.011$–$0.024$ for OT-CFM [2604.15114]. On high-dimensional CIFAR-10 fine-tuning, OA-OT achieves NFE/sample approximately $146.00$, while RA-OT achieves the best FID approximately $3.543$; both amortized variants significantly reduce fine-tuning time compared to exact OT-CFM, reported as approximately $744\,\mathrm s$ [2604.15114].

## 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 [2503.12633]. The filtering problem is defined by the discrete-time state-space model
$$
X_t \sim a_t(\cdot|X_{t-1}),\qquad Y_t\sim h_t(\cdot|X_t),
$$
with posterior recursion
$$
\pi_t = B_{t,Y_t}(A_t\pi_{t-1}).
$$
The optimal transport filter replaces the conditioning operator by a conditional Brenier map $T_t$ satisfying
$$
T_t(\cdot,y)_\#(A_t\pi_{t-1}) = B_{t,y}(A_t\pi_{t-1}),
$$
so that
$$
\pi_t = T_t(\cdot,Y_t)_\# A_t\pi_{t-1}.
$$
The computational burden arises because the conditional map is learned by solving a stochastic saddle-point problem online at every time step [2503.12633].

A-OTF amortizes this cost through a library of pre-trained OT maps and a mixture-of-experts aggregation. Offline, one generates $M$ pre-trained OTF maps $\tilde T_m$, computes the pairwise distance matrix
$$
D_{u,v}:=d(\tilde S_u,\tilde S_v),
$$
and clusters the training instances with K-medoids using one of three distances: $d_{W_2}$, $d_{MMD}$, or $d_T$. The medoids define representative expert maps $\{T_k^\ast\}_{k=1}^K$ [2503.12633].

Online, the current predictive context is compared against the medoids using a distance $\rho$, and the gating weights are
$$
w_t^k := \frac{\exp(-\lambda \rho(S_k^\ast,S_t))}{\sum_{k'} \exp(-\lambda \rho(S_{k'}^\ast,S_t))},
$$
with $\lambda\ge 0$ and $\lambda=\infty$ yielding nearest-neighbor selection. The amortized map is then
$$
\hat T_t(x,y)=\sum_{k=1}^K w_t^k T_k^\ast(x,y),
$$
and the particle update is
$$
X_{t|t}^i=\hat T_t(X_{t|t-1}^i,Y_t).
$$
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 [2503.12633].

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 [2503.12633]. 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,
$$
\min_{w\in \Delta^K} D(T_w(\cdot,Y_t)_\#\mu_t,\nu_t),
$$
with $\mu_t=A_t\pi_{t-1}$ and $\nu_t=B_{t,Y_t}(A_t\pi_{t-1})$. This would make the online gating an explicit posterior-matching OA-OT objective rather than a similarity-based surrogate [2503.12633].

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 $N=250$ particles for all methods and SIR with $N=10^5$ as ground truth, A-OTF with $d_{W_2}$, $\rho_{W_2}$, $K=20$, $\lambda=1$ performs comparably to OTF while faster [2503.12633]. In Lorenz-63 with scalar observation, observing $X(3)$ yields bimodality in $X(1),X(2)$ due to symmetry, and A-OTF with the same configuration and $N=250$ captures multimodality in unobserved coordinates similarly to OTF; weighted selection with $\rho_{W_2}$ is reported as robust, while $d_T$ and $d_{MMD}$ exhibit overfitting for $K>10$ [2503.12633].

## 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 [2206.05262; 2604.15114; 2503.12633]. 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 [2604.15114]. The same point underlies Meta OT, which learns by minimizing the OT objective rather than an MSE to precomputed plans or potentials [2206.05262].

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 [2604.15114]. 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 [2206.05262; 2503.12633]. 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 [2604.15114; 2206.05262].

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 [2604.15114], 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 $K$ [2604.15114; 2503.12633].

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 [2604.15114]. In a broader sense, it includes objective-driven amortized predictors and warm-start mechanisms for repeated OT solves [2206.05262]. 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 [2503.12633].

Source: https://www.emergentmind.com/topics/objective-based-amortization-oa-ot