---
title: Regression-based Amortization (RA-OT)
url: https://www.emergentmind.com/topics/regression-based-amortization-ra-ot
type: topic
---

# Regression-based Amortization (RA-OT)

Regression-based amortization (RA-OT) is a regression-driven amortized optimal transport method introduced in "Amortized Optimal Transport from Sliced Potentials" [2604.15114]. 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 [2604.15114].

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

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 \(\boldsymbol\omega\in\mathbb{R}^L\). The model is therefore linear in the coefficients \(\boldsymbol\omega\), while the predictors themselves are OT-derived functions computed from the sliced problems.

On the support of \(\mu\), with atoms \(x_1,\dots,x_n\), the method constructs a slice-feature matrix
\[
X_\Theta[\mu,\nu,c]
=
\begin{bmatrix}
f_{\theta_1}^*(P_{\theta_1}^c(x_1)) & \cdots & f_{\theta_L}^*(P_{\theta_L}^c(x_1))\\
\vdots & & \vdots\\
f_{\theta_1}^*(P_{\theta_1}^c(x_n)) & \cdots & f_{\theta_L}^*(P_{\theta_L}^c(x_n))
\end{bmatrix}
\in \mathbb{R}^{n\times L},
\]
and the corresponding true potential vector
\[
Y[\mu,\nu,c]
=
\bigl(f^*(x_1),\dots,f^*(x_n)\bigr)^\top \in \mathbb{R}^n.
\]

This formulation makes explicit that RA-OT transfers information through a finite-dimensional coefficient vector \(\boldsymbol\omega\), 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 \(L^2(\mu)\)-error between the true potential and its prediction and minimizes its expectation under \(\mathcal{D}\). The population risk is
\[
\mathcal{L}(\boldsymbol\omega)
=
\mathbb{E}_{(\mu,\nu,c)\sim\mathcal{D}}
\bigl[
\|X_\Theta[\mu,\nu,c]\boldsymbol\omega - Y[\mu,\nu,c]\|_2^2
\bigr]
+\lambda \|\boldsymbol\omega\|_2^2,
\]
where \(\lambda\ge 0\) is an optional ridge penalty.

Given a finite training set \(\{(\mu_i,\nu_i,c_i)\}_{i=1}^N\), the expectation is replaced by the empirical average
\[
\hat{\mathcal{L}}(\boldsymbol\omega)
=
\frac{1}{N}\sum_{i=1}^N
\|X_\Theta[\mu_i,\nu_i,c_i]\boldsymbol\omega - Y[\mu_i,\nu_i,c_i]\|_2^2
+\lambda\|\boldsymbol\omega\|_2^2.
\]
The minimizer satisfies the normal equations
\[
\Bigl(\frac{1}{N}\sum_i X_i^\top X_i + \lambda I\Bigr)\boldsymbol\omega
=
\frac{1}{N}\sum_i X_i^\top Y_i,
\qquad
X_i\equiv X_\Theta[\mu_i,\nu_i,c_i],\;
Y_i\equiv Y[\mu_i,\nu_i,c_i].
\]
For moderate \(L\), the coefficients may be obtained by direct inversion; otherwise, any standard linear solver may be used.

The reported training workflow has four stages. First, \(N\) solved OT problems \((\mu_i,\nu_i,c_i)\) are assembled, and for each one the entropic (or exact) Kantorovich potential \(f^*[\mu_i,\nu_i,c_i]\) is computed. Second, \(L\) projection directions \(\{\theta_\ell\}_{\ell=1}^L\) are chosen, with examples including random and quasi-Monte Carlo directions, and the one-dimensional sliced potentials \(f_{\theta_\ell}^*[\mu_i,\nu_i,c_i]\) are evaluated to build the design matrices \(X_i\). Third, the ridge regression \(\min_\omega \hat{\mathcal{L}}(\omega)\) is solved by the normal-equation method or a gradient-based solver, with \(\lambda\) selected by cross-validation on a held-out subset. Fourth, the number of projections \(L\) and regularization \(\lambda\) may be tuned by grid search or Bayesian optimization using a validation plan-prediction error metric [2604.15114].

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
\[
\hat f
=
\sum_{\ell} \omega_\ell\, f_{\theta_\ell}^* \circ P_{\theta_\ell}^c,
\]
RA-OT recovers the dual partner \(\hat g\) in the entropic setting by the Sinkhorn update
\[
\hat g
=
\epsilon \ln \beta
-
\epsilon \ln\Bigl(\exp(-C^\top/\epsilon)\,\exp(\hat f/\epsilon)\Bigr),
\]
where \(\alpha,\beta\) are the source and target weights and \(C_{ij}=c(x_i,y_j)\). The approximate OT plan is then recovered from the estimated dual potentials. In the unregularized Monge case, when \(c(x,y)=\|x-y\|^2\), a map is recovered via
\[
T(x)=\nabla f(x).
\]

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 \(\hat f\), and standard OT dual machinery converts \(\hat f\) 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
\[
\mathbb{E}\bigl[\|\hat f-f^*\|_{L^2(\mu)}^2\bigr]
\]
vanishes if and only if the true potential \(f^*\) lies in the linear span of the sliced potentials \(\{f_{\theta_\ell}^*\}\) [2604.15114]. 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
\[
\mathbb{E}\bigl[J(\hat f)-J(f^*)\bigr]
\]
is small when \(\hat f\approx f^*\). 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 \(L\) is described as an expressivity-complexity trade-off. Increasing \(L\) enlarges the span of predictors and can reduce approximation error. Empirically, the paper reports rapid error decay up to a moderate \(L\approx 20\text{–}100\), 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,
\[
\|\hat\omega-\omega^*\| = O_p(1/\sqrt{N}),
\]
with plan error decaying as \(O_p(1/\sqrt{N})\). 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 \(f^*\) lie in their linear span. Second, adding projections improves the approximation class, but the reported empirical gains saturate beyond moderate \(L\), so larger \(L\) 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 \(W_2^2\), path straightness NPE, and training speed [2604.15114].

| Task | Setting | Reported outcome |
|---|---|---|
| MNIST grayscale transport | \(n=784\), \(\epsilon=0.1\), varying training size \(M\) | With \(M=10\), RA-OT RMSE \(\approx 8.2\times 10^{-6}\) vs. Meta-OT \(16.2\times 10^{-6}\); inference \(\approx 40\) ms vs. \(2\text{–}6\) ms |
| Spherical supply-demand | \(n=100, m=10{,}000, \epsilon=0.5\), data on \(\mathbb{S}^2\) with stereographic projections | RA-OT/OA-OT RMSE \(\approx 4\text{–}8\times 10^{-7}\); single-projection methods fail with RMSE \(\approx 1.1\times 10^{-5}\) |
| Color transfer | \(K=500\) RGB clusters, \(\epsilon=0.005\) | RA-OT/OA-OT RMSE \(\approx 9\text{–}11\times 10^{-6}\), Meta-OT \(\approx 33\text{–}37\times 10^{-6}\) |
| Mini-batch OT for conditional flow matching | 2D toy | RA-OT/OA-OT are \(2.5\text{–}4.5\times\) faster than OT-CFM, with NPE \(\approx 0.014\text{–}0.06\) and modest \(W_2^2\) degradation in complex targets |

On MNIST grayscale transport, the baselines are Meta-OT, Min-STP, and min-SWGG. With only \(M=10\) training pairs, RA-OT achieves RMSE \(\approx 8.2\times 10^{-6}\) versus Meta-OT’s \(16.2\times 10^{-6}\). Inference is \(\approx 40\) ms, whereas Meta-OT requires \(\approx 2\text{–}6\) ms; the paper notes that Meta-OT is faster but much less accurate. For \(M=50\), RA-OT reaches RMSE \(\approx 7.8\times 10^{-6}\), OA-OT \(\approx 6.0\times 10^{-6}\), and Meta-OT \(\approx 15.5\times 10^{-6}\). Training times are reported as RA-OT \(\approx 1\text{–}3\) s, OA-OT \(\approx 16\) s, and Meta-OT \(\approx 37\) s.

On spherical supply-demand transportation, where the data lie on \(\mathbb{S}^2\) and stereographic projections are used, RA-OT and OA-OT remain within RMSE \(\approx 4\text{–}8\times 10^{-7}\), while Meta-OT is at \(\approx 4\text{–}7\times 10^{-7}\) and is described as slightly worse at small \(M\). Single-projection methods fail, with RMSE \(\approx 1.1\times 10^{-5}\). Training times are RA-OT \(\approx 1\text{–}3\) s, OA-OT \(\approx 19\) s, and Meta-OT \(\approx 52\) s.

On color transfer with \(K=500\) RGB clusters and \(\epsilon=0.005\), RA-OT and OA-OT achieve RMSE \(\approx 9\text{–}11\times 10^{-6}\), while Meta-OT is at \(\approx 33\text{–}37\times 10^{-6}\). Training times are RA-OT \(\approx 2\text{–}7\) s, OA-OT \(\approx 18\) s, and Meta-OT \(\approx 33\) s.

In mini-batch OT for conditional flow matching on a 2D toy problem, OT-CFM with exact OT attains the best \(W_2^2\approx 0.20\text{–}0.27\) and NPE \(\approx 0.01\text{–}0.02\), but is slow at \(\approx 1{,}000\text{–}1{,}500\) s. RA-OT and OA-OT are \(2.5\text{–}4.5\times\) faster, at \(\approx 300\text{–}410\) s, with NPE \(\approx 0.014\text{–}0.06\) very close to OT-CFM and only modest \(W_2^2\) 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.

Source: https://www.emergentmind.com/topics/regression-based-amortization-ra-ot