Regression-based Amortization (RA-OT)

• RA-OT is defined as a regression-driven method that leverages least-squares estimation on sliced Kantorovich potentials to predict optimal transport dual potentials.
• It builds a linear functional model where one-dimensional sliced potentials serve as predictors and the full Kantorovich potential is the regression response, ensuring parsimony and transferability.
• Empirical evaluations on tasks such as MNIST transport and color transfer demonstrate RA-OT’s efficiency, near-optimal accuracy, and reduced training times compared to fully-parametric methods.

Regression-based amortization (RA-OT) is a regression-driven amortized optimal transport method introduced in "Amortized Optimal Transport from Sliced Potentials" (Truong et al., 16 Apr 2026). It predicts optimal transport solutions across multiple measure pairs by learning a functional relationship between Kantorovich potentials from the original OT problem and closed-form one-dimensional Kantorovich potentials obtained from sliced OT. In this formulation, the original potential is the regression response, the sliced potentials are predictors, and the approximate OT plan is recovered from the estimated dual potentials. The method is presented alongside objective-based amortization (OA-OT), but RA-OT is distinguished by least-squares estimation of a linear functional model. The proposed amortization scheme is described as parsimonious, independent of specific structures of the measures such as the number of atoms in the discrete case, and effective on MNIST digit transport, color transfer, spherical supply-demand transportation, and mini-batch OT conditional flow matching (Truong et al., 16 Apr 2026).

1. Position within amortized optimal transport

RA-OT is defined on repeated OT problems drawn from an unknown meta-distribution $\mathcal{D}$ over triples $(\mu,\nu,c)$. Its purpose is not to solve each OT instance from scratch, but to reuse information learned from prior solved problems so that new instances can be approximated rapidly. In the terminology of the source paper, RA-OT and OA-OT are both amortized OT methods: both predict OT structure from previous instances, but they differ in how the functional model is estimated. RA-OT uses least-squares regression, whereas OA-OT estimates the same type of functional model by optimizing the Kantorovich dual objective (Truong et al., 16 Apr 2026).

The defining structural idea is to leverage sliced OT. For a fixed family of one-dimensional projections, sliced OT provides closed-form one-dimensional Kantorovich potentials. RA-OT uses these sliced potentials as reusable features for predicting the full OT potential. This places the method between direct repeated optimization and fully parametric hypernetwork-style predictors such as Meta-OT. A common misconception is to interpret RA-OT as direct plan regression. In the stated formulation, the learned object is the Kantorovich potential $f^*$; plan recovery is a subsequent step.

The paper characterizes this design as more parsimonious than approaches whose parameterization depends on the detailed structure of the measures. A plausible implication is that the method aims to amortize the most transferable part of the dual OT problem rather than learning an instance-specific plan map end-to-end.

2. Functional regression formulation

For each OT problem $(\mu,\nu,c)\sim\mathcal{D}$, the original Kantorovich dual yields a potential

$f^*[\mu,\nu,c]: x\mapsto f^*(x),$

which RA-OT treats as the response in a regression model. Independently, for a fixed family of one-dimensional projections $P_{\theta_1}^c,\dots,P_{\theta_L}^c$, sliced OT yields closed-form one-dimensional Kantorovich potentials

$$f_{\theta_\ell}^*[\mu,\nu,c]: z\mapsto f_{\theta_\ell}^*(z), \qquad \ell=1,\dots,L,$$

where $z=P_{\theta_\ell}^c(x)$. These $L$ functions serve as predictors.

RA-OT posits the linear functional regression model

$$\hat f_{\boldsymbol\omega}[\mu,\nu,c](x) = \sum_{\ell=1}^L \omega_\ell \, f_{\theta_\ell}^*[\mu,\nu,c]\bigl(P_{\theta_\ell}^c(x)\bigr),$$

with coefficient vector $(\mu,\nu,c)$0. The model is therefore linear in the coefficients $(\mu,\nu,c)$1, while the predictors themselves are OT-derived functions computed from the sliced problems.

On the support of $(\mu,\nu,c)$2, with atoms $(\mu,\nu,c)$3, the method constructs a slice-feature matrix

$(\mu,\nu,c)$4

and the corresponding true potential vector

$(\mu,\nu,c)$5

This formulation makes explicit that RA-OT transfers information through a finite-dimensional coefficient vector $(\mu,\nu,c)$6, while the predictors retain geometric information inherited from sliced OT. The paper’s claim that the model is independent of specific structures of the measures, such as the number of atoms in the discrete case, follows from this emphasis on projection-based features rather than direct instance-sized parameterizations.

3. Least-squares estimation and training workflow

RA-OT measures the squared $(\mu,\nu,c)$7-error between the true potential and its prediction and minimizes its expectation under $(\mu,\nu,c)$8. The population risk is

$(\mu,\nu,c)$9

where $f^*$0 is an optional ridge penalty.

Given a finite training set $f^*$1, the expectation is replaced by the empirical average

$f^*$2

The minimizer satisfies the normal equations

$f^*$3

For moderate $f^*$4, the coefficients may be obtained by direct inversion; otherwise, any standard linear solver may be used.

The reported training workflow has four stages. First, $f^*$5 solved OT problems $f^*$6 are assembled, and for each one the entropic (or exact) Kantorovich potential $f^*$7 is computed. Second, $f^*$8 projection directions $f^*$9 are chosen, with examples including random and quasi-Monte Carlo directions, and the one-dimensional sliced potentials $(\mu,\nu,c)\sim\mathcal{D}$0 are evaluated to build the design matrices $(\mu,\nu,c)\sim\mathcal{D}$1. Third, the ridge regression $(\mu,\nu,c)\sim\mathcal{D}$2 is solved by the normal-equation method or a gradient-based solver, with $(\mu,\nu,c)\sim\mathcal{D}$3 selected by cross-validation on a held-out subset. Fourth, the number of projections $(\mu,\nu,c)\sim\mathcal{D}$4 and regularization $(\mu,\nu,c)\sim\mathcal{D}$5 may be tuned by grid search or Bayesian optimization using a validation plan-prediction error metric (Truong et al., 16 Apr 2026).

This training pipeline makes RA-OT an explicit offline-to-online method: expensive OT solves are incurred during data collection, and subsequent inference uses the learned linear coefficients together with sliced features for new instances.

4. Recovery of approximate transport plans

After estimating

$(\mu,\nu,c)\sim\mathcal{D}$6

RA-OT recovers the dual partner $(\mu,\nu,c)\sim\mathcal{D}$7 in the entropic setting by the Sinkhorn update

$(\mu,\nu,c)\sim\mathcal{D}$8

where $(\mu,\nu,c)\sim\mathcal{D}$9 are the source and target weights and $f^*[\mu,\nu,c]: x\mapsto f^*(x),$0. The approximate OT plan is then recovered from the estimated dual potentials. In the unregularized Monge case, when $f^*[\mu,\nu,c]: x\mapsto f^*(x),$1, a map is recovered via

$f^*[\mu,\nu,c]: x\mapsto f^*(x),$2

This recovery step is central to the method’s interpretation. RA-OT does not learn the transport matrix directly; it predicts a dual object and reconstructs the plan from dual feasibility structure. A plausible implication is that the method benefits from retaining the dual OT architecture, rather than bypassing it with unconstrained plan prediction.

The source description also emphasizes that the same recovery principle applies after amortization: sliced OT supplies predictor functions, regression supplies $f^*[\mu,\nu,c]: x\mapsto f^*(x),$3, and standard OT dual machinery converts $f^*[\mu,\nu,c]: x\mapsto f^*(x),$4 into an approximate primal solution.

5. Theoretical properties and interpretive points

The paper states an exact characterization of amortization error for the regression model: the expected regression error

$f^*[\mu,\nu,c]: x\mapsto f^*(x),$5

vanishes if and only if the true potential $f^*[\mu,\nu,c]: x\mapsto f^*(x),$6 lies in the linear span of the sliced potentials $f^*[\mu,\nu,c]: x\mapsto f^*(x),$7 (Truong et al., 16 Apr 2026). This is the key expressivity statement for RA-OT. It identifies the method’s approximation class explicitly: the quality of amortization is governed by how well the full potential can be represented as a linear combination of sliced potentials.

The paper further states that the dual-objective gap

$f^*[\mu,\nu,c]: x\mapsto f^*(x),$8

is small when $f^*[\mu,\nu,c]: x\mapsto f^*(x),$9. This connects regression fidelity in potential space to OT optimality in the dual objective. It does not assert exact equality between these quantities, but it identifies close potential prediction as the relevant mechanism controlling the amortization gap.

The role of the number of projections $P_{\theta_1}^c,\dots,P_{\theta_L}^c$0 is described as an expressivity-complexity trade-off. Increasing $P_{\theta_1}^c,\dots,P_{\theta_L}^c$1 enlarges the span of predictors and can reduce approximation error. Empirically, the paper reports rapid error decay up to a moderate $P_{\theta_1}^c,\dots,P_{\theta_L}^c$2, beyond which gains saturate. This rules out a simplistic interpretation that arbitrarily many slices are always proportionally beneficial.

For statistical estimation, standard least-squares theory is invoked: under mild design conditions,

$P_{\theta_1}^c,\dots,P_{\theta_L}^c$3

with plan error decaying as $P_{\theta_1}^c,\dots,P_{\theta_L}^c$4. Within the stated assumptions, RA-OT therefore inherits the usual parametric convergence behavior of regularized linear regression, but applied to OT-derived functional features.

Two common misconceptions are directly addressed by these results. First, exact amortization is not guaranteed merely because sliced OT features are available; it requires that $P_{\theta_1}^c,\dots,P_{\theta_L}^c$5 lie in their linear span. Second, adding projections improves the approximation class, but the reported empirical gains saturate beyond moderate $P_{\theta_1}^c,\dots,P_{\theta_L}^c$6, so larger $P_{\theta_1}^c,\dots,P_{\theta_L}^c$7 is not synonymous with unbounded practical improvement.

6. Empirical evaluation and comparative behavior

The empirical study evaluates RA-OT on three standard discrete OT tasks relative to converged Sinkhorn ground truth, using plan-RMSE as the key metric together with training and inference times, and on mini-batch OT for conditional flow matching using Sinkhorn $P_{\theta_1}^c,\dots,P_{\theta_L}^c$8, path straightness NPE, and training speed (Truong et al., 16 Apr 2026).

Task	Setting	Reported outcome
MNIST grayscale transport	$P_{\theta_1}^c,\dots,P_{\theta_L}^c$9, $$f_{\theta_\ell}^*[\mu,\nu,c]: z\mapsto f_{\theta_\ell}^*(z), \qquad \ell=1,\dots,L,$$0, varying training size $$f_{\theta_\ell}^*[\mu,\nu,c]: z\mapsto f_{\theta_\ell}^*(z), \qquad \ell=1,\dots,L,$$1	With $$f_{\theta_\ell}^*[\mu,\nu,c]: z\mapsto f_{\theta_\ell}^*(z), \qquad \ell=1,\dots,L,$$2, RA-OT RMSE $$f_{\theta_\ell}^*[\mu,\nu,c]: z\mapsto f_{\theta_\ell}^*(z), \qquad \ell=1,\dots,L,$$3 vs. Meta-OT $$f_{\theta_\ell}^*[\mu,\nu,c]: z\mapsto f_{\theta_\ell}^*(z), \qquad \ell=1,\dots,L,$$4; inference $$f_{\theta_\ell}^*[\mu,\nu,c]: z\mapsto f_{\theta_\ell}^*(z), \qquad \ell=1,\dots,L,$$5 ms vs. $$f_{\theta_\ell}^*[\mu,\nu,c]: z\mapsto f_{\theta_\ell}^*(z), \qquad \ell=1,\dots,L,$$6 ms
Spherical supply-demand	$$f_{\theta_\ell}^*[\mu,\nu,c]: z\mapsto f_{\theta_\ell}^*(z), \qquad \ell=1,\dots,L,$$7, data on $$f_{\theta_\ell}^*[\mu,\nu,c]: z\mapsto f_{\theta_\ell}^*(z), \qquad \ell=1,\dots,L,$$8 with stereographic projections	RA-OT/OA-OT RMSE $$f_{\theta_\ell}^*[\mu,\nu,c]: z\mapsto f_{\theta_\ell}^*(z), \qquad \ell=1,\dots,L,$$9; single-projection methods fail with RMSE $z=P_{\theta_\ell}^c(x)$0
Color transfer	$z=P_{\theta_\ell}^c(x)$1 RGB clusters, $z=P_{\theta_\ell}^c(x)$2	RA-OT/OA-OT RMSE $z=P_{\theta_\ell}^c(x)$3, Meta-OT $z=P_{\theta_\ell}^c(x)$4
Mini-batch OT for conditional flow matching	2D toy	RA-OT/OA-OT are $z=P_{\theta_\ell}^c(x)$5 faster than OT-CFM, with NPE $z=P_{\theta_\ell}^c(x)$6 and modest $z=P_{\theta_\ell}^c(x)$7 degradation in complex targets

On MNIST grayscale transport, the baselines are Meta-OT, Min-STP, and min-SWGG. With only $z=P_{\theta_\ell}^c(x)$8 training pairs, RA-OT achieves RMSE $z=P_{\theta_\ell}^c(x)$9 versus Meta-OT’s $L$0. Inference is $L$1 ms, whereas Meta-OT requires $L$2 ms; the paper notes that Meta-OT is faster but much less accurate. For $L$3, RA-OT reaches RMSE $L$4, OA-OT $L$5, and Meta-OT $L$6. Training times are reported as RA-OT $L$7 s, OA-OT $L$8 s, and Meta-OT $L$9 s.

On spherical supply-demand transportation, where the data lie on $$\hat f_{\boldsymbol\omega}[\mu,\nu,c](x) = \sum_{\ell=1}^L \omega_\ell \, f_{\theta_\ell}^*[\mu,\nu,c]\bigl(P_{\theta_\ell}^c(x)\bigr),$$0 and stereographic projections are used, RA-OT and OA-OT remain within RMSE $$\hat f_{\boldsymbol\omega}[\mu,\nu,c](x) = \sum_{\ell=1}^L \omega_\ell \, f_{\theta_\ell}^*[\mu,\nu,c]\bigl(P_{\theta_\ell}^c(x)\bigr),$$1, while Meta-OT is at $$\hat f_{\boldsymbol\omega}[\mu,\nu,c](x) = \sum_{\ell=1}^L \omega_\ell \, f_{\theta_\ell}^*[\mu,\nu,c]\bigl(P_{\theta_\ell}^c(x)\bigr),$$2 and is described as slightly worse at small $$\hat f_{\boldsymbol\omega}[\mu,\nu,c](x) = \sum_{\ell=1}^L \omega_\ell \, f_{\theta_\ell}^*[\mu,\nu,c]\bigl(P_{\theta_\ell}^c(x)\bigr),$$3. Single-projection methods fail, with RMSE $$\hat f_{\boldsymbol\omega}[\mu,\nu,c](x) = \sum_{\ell=1}^L \omega_\ell \, f_{\theta_\ell}^*[\mu,\nu,c]\bigl(P_{\theta_\ell}^c(x)\bigr),$$4. Training times are RA-OT $$\hat f_{\boldsymbol\omega}[\mu,\nu,c](x) = \sum_{\ell=1}^L \omega_\ell \, f_{\theta_\ell}^*[\mu,\nu,c]\bigl(P_{\theta_\ell}^c(x)\bigr),$$5 s, OA-OT $$\hat f_{\boldsymbol\omega}[\mu,\nu,c](x) = \sum_{\ell=1}^L \omega_\ell \, f_{\theta_\ell}^*[\mu,\nu,c]\bigl(P_{\theta_\ell}^c(x)\bigr),$$6 s, and Meta-OT $$\hat f_{\boldsymbol\omega}[\mu,\nu,c](x) = \sum_{\ell=1}^L \omega_\ell \, f_{\theta_\ell}^*[\mu,\nu,c]\bigl(P_{\theta_\ell}^c(x)\bigr),$$7 s.

On color transfer with $$\hat f_{\boldsymbol\omega}[\mu,\nu,c](x) = \sum_{\ell=1}^L \omega_\ell \, f_{\theta_\ell}^*[\mu,\nu,c]\bigl(P_{\theta_\ell}^c(x)\bigr),$$8 RGB clusters and $$\hat f_{\boldsymbol\omega}[\mu,\nu,c](x) = \sum_{\ell=1}^L \omega_\ell \, f_{\theta_\ell}^*[\mu,\nu,c]\bigl(P_{\theta_\ell}^c(x)\bigr),$$9, RA-OT and OA-OT achieve RMSE $(\mu,\nu,c)$00, while Meta-OT is at $(\mu,\nu,c)$01. Training times are RA-OT $(\mu,\nu,c)$02 s, OA-OT $(\mu,\nu,c)$03 s, and Meta-OT $(\mu,\nu,c)$04 s.

In mini-batch OT for conditional flow matching on a 2D toy problem, OT-CFM with exact OT attains the best $(\mu,\nu,c)$05 and NPE $(\mu,\nu,c)$06, but is slow at $(\mu,\nu,c)$07 s. RA-OT and OA-OT are $(\mu,\nu,c)$08 faster, at $(\mu,\nu,c)$09 s, with NPE $(\mu,\nu,c)$10 very close to OT-CFM and only modest $(\mu,\nu,c)$11 degradation in complex targets.

Taken together, these results support the paper’s summary characterization of RA-OT as a highly-parsimonious, sample-efficient, and fast amortized predictor for Kantorovich potentials that yields near-optimal transport plans at a fraction of the training cost of fully-parametric hypernetworks such as Meta-OT. A cautious reading is that the empirical evidence is strongest for repeated-problem settings where sliced potentials provide sufficiently expressive predictors and where offline access to solved OT instances is available.