Deep Barycentric Regression for Optimal Transport Map Estimation and its Statistical Optimality

Kunwoong Kim; Insung Kong; Yongdai Kim

cs.LGSep 6, 2026

Deep Barycentric Regression for Optimal Transport Map Estimation and its Statistical Optimality

Authors: Kunwoong Kim, Insung Kong, Yongdai Kim

Organizations: KAIST AI · UNIST · Seoul National University

Abstract

The optimal transport (OT) map provides a geometric transformation for aligning probability distributions and has become a useful tool in machine learning. However, existing estimators of the OT map still exhibit a gap between sharp statistical guarantees and practical parametric estimation based on stable training objectives. Theoretical estimators achieve minimax optimal convergence rates, but they are typically nonparametric and can incur demanding implementation design or inference costs. Practical estimators are parametric and scalable, but their statistical guarantees remain underexplored, and their min-max, adversarial-like training objectives can be sensitive to optimization algorithms. We propose BROT (Barycentric Regression for OT), a simple two-step method that first computes the unregularized OT plan and then fits a deep neural network (DNN) to the induced barycentric targets by least-squares regression. Under standard regularity conditions, we prove that the DNN estimator of BROT attains the minimax convergence rate, when the ground-truth OT map is Lipschitz. Numerical studies on synthetic datasets and an image dataset show that BROT provides accurate map estimates, strong target distribution matching, and competitive transport costs, compared to existing estimation methods. Experiments on two downstream tasks, single-cell perturbation prediction and unsupervised domain adaptation, further suggest that the accurate estimation of BROT can translate into stronger task performance.

Figures & tables

Figure 1: Outline of our proposed framework BROT. We first compute the unregularized OT plan, and train a DNN to fit source data to barycentric targets obtained from the OT plan.

Figure 2: Statistical convergence to the ground-truth OT map. (Left) training data size

\widetilde{n}

vs. estimation error. Points are the means across

10

random seeds and the shaded bands show the

10

-th to

90

-th percentile. (Right) theoretical rate

\widetilde{n}^{-1}(\log\widetilde{n})^{2}

vs. estimation error for BROT, with a linear regression line and its relative RMSE (rel. RMSE).

	Gaussian $\to$ Rectangle		Gaussian $\to$ Ellipse		AFHQv2 wild $\to$ cat
Method	Wass	TC	Wass	TC	Wass	TC
1NN Manole et al. (2024)	0.126	1095.67	0.207	1122.30	54.819	3296.54
LSOT Seguy et al. (2018)	0.111	1095.69	0.172	1124.35	53.263	2627.29
ICNN Makkuva et al. (2020)	0.211	1097.51	0.171	1122.54	$\ast$	$\ast$
OTP Choi et al. (2025b)	0.258	1099.17	0.596	1129.88	50.651	1408.23
DIOTM Choi et al. (2025a)	0.123	1098.21	0.228	1123.99	51.000	1453.99

Table 1: Comparison of map quality: distribution matching and transport cost. Best values of Wass are bolded and ‘

\ast

’ marks the failed cases due to algorithm instability. Results are averaged over five random trials, and standard deviations are reported in Table 4 of Section B.3 .

Method	MMD $\downarrow$	$L_{2}$ difference $\downarrow$	$r$ -std $\uparrow$	TC $\downarrow$
Identity	0.028	1.213	0.917	0.000
1NN ( Manole et al., 2024 )	0.002	0.189	0.990	6.745
LSOT ( Seguy et al., 2018 )	0.012	0.287	0.982	4.381
ICNN ( Makkuva et al., 2020 )	0.006	0.258	0.987	4.329
OTP ( Choi et al., 2025b )	0.035	1.435	0.917	0.514
DIOTM ( Choi et al., 2025a )	0.002	0.255	0.990	3.993

Table 2: Comparison of single-cell perturbation prediction performance. The results are averaged across the

35

drug perturbations from 4i dataset Gut et al. (2018) . Among the methods that satisfy

\texttt{MMD}\leq 0.03

(the MMD of Identity), the best value is bolded and the second-best is underlined for each metric. Standard deviations are reported in Table 6 of Section B.4 .

Method	USPS $\to$ MNIST		VisDA-17 (synthetic $\to$ real)
Method	CLIP-ViT/B-32	DINOv2-ViT/B-14	CLIP-ViT/B-32	DINOv2-ViT/B-14
Identity	52.05	58.07	64.04	60.66
1NN Manole et al. (2024)	80.34	75.25	74.17	74.39
LSOT Seguy et al. (2018)	30.12	52.20	76.87	81.57
ICNN Makkuva et al. (2020)	75.76	54.70	74.12	77.77
OTP Choi et al. (2025b)	49.98	57.36	76.15	64.51

Table 3: Comparison of unsupervised domain adaptation performance. The performance is evaluated by the per-class mean accuracy (%) on the target-domain test data. Best results are bolded and the second places are underlined , and the standard deviations are in Table 7 of Section B.5 .

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 3: Source $P$ (left) and target $Q=(\mathbf{T}_{0})_{\#}P$ (right) used in Section 5.1 . The color of each point indicates the radius of data $\|\bm{x}\|.$

Figure 4: Source (purple) and target (red) data for the two synthetic cases used in Section 5.2 .

Source $\to$ Target	Method	Wass $\downarrow$	TC $\downarrow$
Gaussian $\to$ Rectangle	1NN	$0.126\pm 0.008$	$1095.67\pm 0.61$
	LSOT	$0.111\pm 0.003$	$1095.69\pm 0.14$
	ICNN	$0.211\pm 0.005$	$1097.51\pm 0.09$
	OTP	$0.258\pm 0.008$	$1099.17\pm 2.39$
	DIOTM	$0.123\pm 0.007$	$1098.21\pm 2.35$
	BROT	$0.110\pm 0.003$	$1095.65\pm 0.05$

Appendix

Table 4: Comparison of map quality: distribution matching and transport cost. Results are reported as mean

\pm

standard deviation across five random trials. Best values of Wass are bolded .

Method	Synthetic cases		AFHQv2 wild $\to$ cat
	Train (%)	Inference (%)	Train (%)	Inference (%)
1NN [ 36 ]	20%	339%	17%	486%
LSOT [ 48 ]	267%	95%	154%	108%
ICNN [ 35 ]	2129%	87%	4985%	229%
OTP [ 8 ]	1078%	105%	785%	286%
DIOTM [ 7 ]	1445%	81%	2958%	205%

Appendix

Table 5: Comparison of computation time. Relative training and inference computation times (BROT

=100\%

). We report mean relative training and inference computation times (BROT

=100\%

) over five random seeds, and synthetic cases are the average over the two datasets.

Method	MMD $\downarrow$	$L_{2}$ difference $\downarrow$	$r$ -std $\uparrow$	TC $\downarrow$
Identity	0.028 $\pm$ 0.029	1.213 $\pm$ 0.899	0.917 $\pm$ 0.081	0.000 $\pm$ 0.000
1NN [ 36 ]	0.002 $\pm$ 0.000	0.189 $\pm$ 0.062	0.990 $\pm$ 0.009	6.745 $\pm$ 6.108
LSOT [ 48 ]	0.012 $\pm$ 0.002	0.287 $\pm$ 0.121	0.982 $\pm$ 0.013	4.381 $\pm$ 5.437
ICNN [ 35 ]	0.006 $\pm$ 0.001	0.258 $\pm$ 0.096	0.987 $\pm$ 0.008	4.329 $\pm$ 5.751
OTP [ 8 ]	0.035 $\pm$ 0.029	1.435 $\pm$ 0.852	0.917 $\pm$ 0.081	0.514 $\pm$ 0.141
DIOTM [ 7 ]	0.002 $\pm$ 0.001	0.255 $\pm$ 0.101	0.990 $\pm$ 0.010	3.993 $\pm$ 5.875

Appendix

Table 6: Comparison of single-cell perturbation prediction performance. The results are reported as mean

\pm

standard deviation across the

35

drug perturbations. Among the methods that satisfy

\texttt{MMD}\leq 0.03

, the best value is bolded and the second-best is underlined for each metric.

	USPS $\to$ MNIST		VisDA-17 (synthetic $\to$ real)
Method	CLIP-ViT/B-32	DINOv2-ViT/B-14	CLIP-ViT/B-32	DINOv2-ViT/B-14
1NN [ 36 ]	80.34 $\pm$ 0.34	75.25 $\pm$ 0.77	74.17 $\pm$ 0.08	74.39 $\pm$ 0.01
LSOT [ 48 ]	30.12 $\pm$ 2.84	52.20 $\pm$ 0.88	76.87 $\pm$ 0.63	81.57 $\pm$ 0.18
ICNN [ 35 ]	75.76 $\pm$ 4.63	54.70 $\pm$ 21.08	74.12 $\pm$ 1.07	77.77 $\pm$ 0.41
OTP [ 8 ]	49.98 $\pm$ 4.79	57.36 $\pm$ 5.18	76.15 $\pm$ 0.42	64.51 $\pm$ 0.37
DIOTM [ 7 ]	82.71 $\pm$ 1.00	77.46 $\pm$ 0.51	76.45 $\pm$ 0.55	27.05 $\pm$ 35.31

Appendix

Table 7: Comparison of unsupervised domain adaptation performance. Per-class mean accuracy (%) on the target-domain test data with standard deviations across three random seeds, reported as mean

\pm

standard deviation.

Method	Train		Backbone	Accuracy
Method	Encoder	Head	Backbone	Accuracy
DeepJDOT [ 13 ]	✓	✓	ResNet-50	68.00
JUMBOT [ 17 ]	✓	✓	ResNet-50	72.50
ELOT [ 53 ]	✓	✓	ResNet-50	76.32
	✗	✓	CLIP-ViT/B-32	79.34
	✗	✓	DINOv2-ViT/B-14	81.59

Appendix

Table 8: Comparison of domain adaptation performance for BROT and end-to-end OT-based domain adaptation methods. The performance is evaluated by the per-class mean accuracy (%) on the target-domain test data, following [ 40 ] . The results of BROT are averaged over three random seeds, while baseline results are copied from the original papers.

Figure 5: Sensitivity of DNN architecture and Jacobian penalty to the statistical convergence to the ground-truth OT map. (Left) Three MLP architectures (Jacobian penalty weight fixed to

10^{-2}

). (Right) Three Jacobian penalty weights (architecture fixed).

Optimal Transport

Optimal transport compares and transforms probability distributions by finding low-cost ways to move mass between them.

See all 284 papers

The Fundamental Limits of Valid Transport Map Estimation
Jun 29, 2026Sivaraman BalakrishnanMinimax Estimation Generative Modeling
Department of Statistics and Data Science† · Machine Learning Department‡
Simultaneous Neural Optimal Transport
Sep 29, 2026Milena Gazdieva, Kirill Sokolov, Jiawei Chen +2All-in-One Image Restoration Optimal Transport
Applied AI Institute AXXX Moscow, Russia · Lomonosov Moscow State University Moscow, Russia · MIRIAI Moscow, Russia
Sliced-Regularized Optimal Transport
Apr 27, 2026Khai NguyenOptimal Transport
The University of Texas at Austin

Jun 29, 2026J/K move · Enter open · S save

The Fundamental Limits of Valid Transport Map Estimation

Sivaraman Balakrishnan

Department of Statistics and Data Science† · Machine Learning Department‡

Minimax Estimation Generative Modeling Optimal Transport

Open paper AskPDF arXiv

Many modern generative modeling methods, including diffusion models, normalizing flows, and flow matching, estimate transport maps or plans between distributions without explicitly targeting an optimal transport (OT) map. In applications like generative modeling, the transport cost itself is irrelevant, and this makes it natural to target maps which are more tractable from either a statistical or computational standpoint. In this short note, we formalize the task of estimating any valid transport map in a rigorous minimax framework. One consequence of this framing is that it yields sample complexity lower bounds for any method whose learned object is evaluated as a transport map or plan, including flow matching and diffusion-based generative models, in settings where direct analysis would be challenging due to the analytic complexity of the methods and their target maps. We observe that, under standard, though strong, stability assumptions from the OT literature, estimating any valid transport map is statistically as hard as estimating the OT map. We complement these results with some examples showing that when these stability assumptions fail, alternative transport maps can be learned substantially more accurately than the OT map. Our minimax framing provides a rigorous foundation for understanding the statistical limits of modern transport-based generative methods and clarifies when targeting sub-optimal maps can provide real statistical advantages.

Abstract

Optimal Transport

Explore similar work

The Fundamental Limits of Valid Transport Map Estimation

Simultaneous Neural Optimal Transport

Sliced-Regularized Optimal Transport