PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Kunwoong Kim, Insung Kong, Yongdai Kim

arXiv:2609.06598v1cs.LGcs.AIstat.ML

TL;DR

Existing OT-map estimators either offer sharp minimax guarantees with demanding nonparametric implementation or provide practical parametric training without equally established guarantees. BROT computes an unregularized OT plan and fits a DNN to its barycentric targets by least-squares regression. Its estimator is nearly minimax optimal for Lipschitz OT maps and performs strongly on map estimation, distribution matching, transport cost, and downstream tasks.

  • Problem

    Existing OT-map estimators separate sharp statistical guarantees from practically feasible parametric training.

  • Method

    BROT computes an unregularized empirical OT plan, forms barycentric targets, and fits a DNN to them by least-squares regression.

  • Results

    BROT’s DNN estimator matches the minimax lower bound up to logarithmic factors and achieves strong empirical map, transport, and downstream performance.

  • Takeaways & Limitations

    BROT provides a parametric OT-map estimator that combines standard regression training with statistical optimality for Lipschitz ground-truth maps.

  • Takeaways & Limitations

    The theoretical guarantee is limited to Lipschitz ground-truth OT maps, leaving smoother classes such as Hölder classes for future analysis.

Abstract

from arXiv · show

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.

1 Introduction

BROT addresses the gap between statistically optimal but demanding OT-map estimators and practical parametric methods with limited guarantees. It combines an unregularized OT plan with DNN barycentric regression and achieves minimax-rate estimation alongside strong empirical transport and downstream performance.

  • Motivation: Existing minimax-optimal OT-map estimators are highly nonparametric, demanding to implement and costly at inference.They may require density smoothing, basis construction, nearest-neighbor search, or regularization tuning.
  • Proposed approach: BROT first computes an unregularized OT plan, then trains a DNN by least-squares regression on induced barycentric targets.The framework is designed to retain statistical optimality while using standard regression training.
  • Theory: Under standard regularity conditions, the DNN estimator attains the minimax convergence rate when the ground-truth OT map is Lipschitz.The estimator is presented as simultaneously minimax optimal and practically feasible.
  • Empirical results: BROT achieves lower estimation error, the lowest Wasserstein distance, comparable transport cost, and best or competitive downstream performance.The evaluations cover synthetic and high-dimensional image datasets, single-cell perturbation prediction, and unsupervised domain adaptation.

2 Preliminaries

The OT map deterministically transports a source distribution to a target while minimizing expected cost, whereas the Kantorovich relaxation optimizes over joint transport plans. For empirical discrete distributions, the relaxed problem becomes a linear program, but converting its plan into an out-of-sample map requires an extension.

  • OT map: An OT map transports the source distribution to the target while minimizing expected transport cost.It provides both a distributional discrepancy and an explicit transportation.
  • Kantorovich relaxation: The Kantorovich relaxation replaces a deterministic map with a joint distribution whose marginals are the source and target distributions.This relaxation can remain feasible when no deterministic map exists.
  • Discrete formulation: For discrete empirical measures, the Kantorovich problem reduces to a linear program defining an optimal transport plan.The associated optimal transport cost defines the squared 2-Wasserstein distance.
  • Map estimation: An OT map is generally not well-defined for discrete measures, so estimating one requires extending the empirical OT plan to unseen data.This out-of-sample extension is the setting addressed by BROT.

3 BROT: barycentric regression for OT

BROT reduces OT-map estimation to regression by converting an unregularized empirical OT plan into barycentric targets and fitting a Lipschitz-constrained DNN to them. This preserves a standard least-squares training objective while providing a parametric out-of-sample map.

  • Algorithm: BROT computes an OT plan between observed source and target data, then trains a DNN on the plan’s barycentric targets.The barycentric targets are weighted averages of target data induced by the OT plan.
  • Motivation: Compared with nearest-neighbor extension, a DNN map avoids storing the entire training dataset and performing a search for every test point.Prior DNN fitting approaches used regularized OT plans without a known minimax guarantee for the resulting estimator.
  • Regression formulation: Barycentric regression converts OT-map estimation into a standard regression problem with source data as inputs and barycentric targets as outputs.The resulting least-squares estimator is designed to attain statistical minimax optimality.
  • DNN design: BROT uses a Lipschitz-constrained DNN class, motivated by the theory of Lipschitz OT maps and practical DNN applications.The constraint is mainly theoretical; experiments enforce it with a Jacobian penalty.

4 Theoretical study: minimax optimality

The theoretical analysis studies BROT under bounded source and target densities and regularity of the OT potential. Its specially designed DNN estimator matches the minimax rate for Lipschitz ground-truth OT maps up to logarithmic factors.

  • Assumptions: The analysis assumes absolutely continuous source and target distributions with densities bounded above and below, plus a regular convex OT potential.These are the stated regularity conditions for the minimax result.
  • Minimax benchmark: The minimax lower bound for estimating a Lipschitz OT map is ˜n^-2/d ∨ ˜n^-1, where ˜n := n ∧ m.The risk is measured in the squared L2(P) norm.
  • Proof design: The proof uses a sample-size-dependent DNN architecture and a Lipschitz constraint to control approximation error.Experiments instead use a fixed MLP with Jacobian-penalty regularization.
  • Main theorem: BROT’s DNN risk upper bound matches the minimax lower bound up to logarithmic factors.The result establishes that the DNN estimator is nearly minimax optimal.

5 Experiments

The experiments evaluate BROT’s statistical convergence, map quality, computational practicality, and downstream utility against representative OT-map estimators. Across these settings, BROT achieves strong estimation and distribution-matching performance while remaining competitive in transport cost and runtime.

  • Statistical convergence: BROT’s estimation error decreases monotonically with training size and remains lower than all baselines across every tested size.OTP is omitted from the convergence plot because its error is approximately 2.6 for all training sizes.
  • Statistical convergence: Approximately 15% relative RMSE supports the predicted e_n^-1(log e_n)^2 convergence-rate trend for BROT.The comparison uses a linear regression between estimation error and the theoretical minimax rate.
  • Map quality: BROT attains the lowest Wasserstein distance on all three datasets while maintaining comparable transport cost.The evaluation covers two synthetic datasets and the AFHQV2 wild-to-cat image dataset; lower Wasserstein distance and transport cost are preferred.
  • Map quality: BROT trains faster than min-max DNN baselines, has lower inference cost than 1NN, and can train faster than regularized LSOT.It also achieves smaller Wasserstein distance than LSOT in the reported comparison.
  • Single-cell perturbation prediction: Across 35 drug perturbations, BROT achieves the best MMD, r-std, and TC, and the second-best L2 among methods with MMD ≤0.03.Compared with 1NN, BROT reaches a smaller TC at the same MMD level.
  • Unsupervised domain adaptation: BROT achieves the best per-class mean accuracy in all four unsupervised domain-adaptation cases and remains competitive with end-to-end OT methods.The comparison includes fixed representations and multiple foundation visual encoders and datasets.

6 Conclusion and discussion

BROT combines unregularized OT plans with least-squares DNN regression, achieving minimax-rate estimation under Lipschitz-map conditions and improved empirical performance. Its theory remains limited to Lipschitz OT maps, leaving broader smoothness classes for future work.

  • The DNN estimator attains the minimax convergence rate when the ground-truth OT map is Lipschitz.
  • The theoretical guarantee is limited to Lipschitz OT maps; extending it to Hölder classes could provide broader statistical guarantees.

A.4 Proofs

The proof develops approximation and gradient-representation tools for constrained DNNs, then uses them to control the estimator under the paper’s regularity assumptions.

  • Regularity conditions: The proof assumes compact convex support, bounded positive densities, and a twice-differentiable convex potential whose Hessian bounds make the OT map Lipschitz.
  • DNN approximation: A ReQU DNN can approximate sufficiently smooth functions and their derivatives with error decreasing as K^-(β−ℓ).
  • Gradient representation: The gradient of a ReQU DNN is exactly represented by a mixed ReLU/ReQU DNN with depth at most 3L + 3 and width at most 6W.
  • Lipschitz control: The constructed approximating network has Lipschitz constant bounded by λ + cK^-(β−2).
  • Approximation rate: Choosing network size according to the smoothness parameter yields approximation error of order K^-(β−1).

A.4.2 Proof of Theorem 4.1

The proof decomposes estimation error over Voronoi cells, controls regression and approximation terms using Lipschitz properties, and concludes the minimax rate up to logarithmic factors.

  • Regression objective: The OT plan induces a weighted least-squares objective for fitting the transport map.
  • Error decomposition: The error is decomposed over Voronoi cells into network variation, empirical-map error, and ground-truth-map variation.
  • Error control: Lipschitz bounds on the estimated and true maps control the within-cell variation terms.
  • Network selection: A DNN approximator is selected with (K̃_n)^d ≍ ñ^1/(β−1), placing it in the admissible Lipschitz class.
  • Conclusion: The resulting estimator achieves the minimax convergence rate up to logarithmic factors.

B.1 Implementation details

The experiments compare BROT with regularized and adversarial-style OT estimators using documented hyperparameter searches and shared training protocols. BROT computes an unregularized plan and trains a DNN by least-squares regression.

  • Overall protocol: All methods use documented hyperparameter searches, up to 10,000 outer gradient steps when applicable, and normalized inputs where appropriate.
  • BROT: BROT solves the unregularized empirical Kantorovich problem and selects training hyperparameters using validation Wasserstein distance or task-appropriate MMD.
  • BROT: BROT trains a DNN with selectable activations, architectures, optimizers, learning rates, batch sizes, epochs, and gradient-penalty weights.
  • BROT algorithm: The BROT algorithm first builds the empirical cost matrix and transport plan, then fits the DNN by least-squares regression on mini-batches.
  • Baselines: 1NN extends barycentric targets to unseen source points through nearest-neighbor interpolation.
  • Baselines: LSOT uses regularized barycentric targets and searches downward for the smallest regularization strength whose optimization remains numerically stable.

B.2 Statistical convergence to the ground-truth OT map (Section 5.1)

The synthetic construction uses a truncated Gaussian source and a target generated by a known Lipschitz OT map. The experiments compare map quality, runtime, and downstream performance across methods.

  • Synthetic setup: The synthetic source is a truncated standard Gaussian, while the target is generated as Q = (T0)#P using specified coefficients and R = 3.5.The construction uses d = 2, (a1, a2) = (1.5, 0.8), and (b1, b2) = (0.3, 0.1).
  • Synthetic setup: Figure 3 visualizes the source and target distributions, with point colors indicating the radius of each data point.
  • Synthetic setup: The constructed map T0 is the unique OT map from P to Q and is Lipschitz.For the stated parameter setting, the Lipschitz constant is approximately 12.52.
  • Evaluation: BROT is evaluated against existing methods using map quality, computation time, and downstream prediction performance.The reported materials include distribution matching and transport cost, runtime, and single-cell perturbation prediction.

B.5 Application 2: unsupervised domain adaptation (Section 5.4)

The domain-adaptation experiments compare BROT with standard and end-to-end OT-based methods using target-domain per-class accuracy. BROT performs competitively with strong frozen encoders, while fixed ResNet-50 representations reveal a smaller advantage over some jointly trained methods.

  • Evaluation design: The classifier head is linear for VISDA-17 and an MLP for USPS→MNIST, based on the performance of the Identity baseline.Linear classifiers on VISDA-17 source representations already attain 60–68% and remain within 5%p of MLP heads.
  • Method comparison: BROT estimates the OT map from source-target Euclidean distances without jointly training the representation encoder or classifier head.The end-to-end baselines jointly train the encoder, classifier head, and transport plan.
  • End-to-end comparison: 79.34% with CLIP-ViT/B-32 and 81.59% with DINOv2-ViT/B-14 exceed DeepJDOT 68.0, JUMBOT 72.5, and ELOT 76.32.These results use per-class mean accuracy and compare BROT with end-to-end OT-based domain-adaptation methods.
  • Fixed-encoder comparison: 70.55% with fixed pretrained ResNet-50 representations exceeds DeepJDOT but falls below JUMBOT and ELOT.The passage attributes the gap to BROT not jointly training the encoder.

B.6 Sensitivity analysis: DNN architecture and Jacobian penalty

The sensitivity analysis examines whether DNN architecture and Jacobian-penalty choices materially affect BROT’s estimation error. Errors are broadly similar across tested settings, with the largest penalty producing moderately higher error at large sample sizes.

  • Interpretation: The sample-size-dependent architecture in Theorem 4.1 is primarily a theoretical device for controlling approximation error, not a practical architectural prescription.The experiments instead use a fixed MLP architecture with a Jacobian penalty.
  • Sensitivity results: Estimation errors are broadly similar across three DNN architectures and three Jacobian penalty weights.The experiments use five random seeds and sample sizes from 500 to 10,000.
  • Sensitivity results: The largest Jacobian penalty weight, 10^-2, produces moderately larger error at large n.The tested penalty weights are 10^-4, 10^-3, and 10^-2.

C Background on statistically optimal estimators

Statistically optimal OT-map estimators include wavelet, nearest-neighbor, and general-function-class approaches, but their computational requirements can limit practical use. Neural-network plug-in methods also face challenging optimization and unresolved approximation guarantees.

  • Wavelet-based estimator: Wavelet estimators achieve minimax rates but require basis selection, coefficient estimation, and high-dimensional constrained optimization.Their complexity grows quickly with dimension, limiting applicability.
  • Practical constraints: Theoretical estimator constructions can therefore attain minimax convergence while imposing substantial computational or inference burdens.These burdens include constrained optimization, nearest-neighbor search, and convex-conjugate evaluation.
  • Nearest neighbor estimator: Nearest-neighbor estimators achieve the optimal rate for Lipschitz OT maps but require inference-time searches whose memory and latency scale with n.This scaling can be prohibitive for large-scale applications.
  • Estimators in general classes of functions: General function-class estimators derive oracle inequalities separating approximation and complexity terms, and suitable ReQU DNN classes attain minimax rates.The complexity term is governed by the metric entropy of the candidate class.
  • Neural-network plug-in estimators: Neural-network plug-in estimators may require evaluating convex conjugates through high-dimensional inner optimization problems that are nonconvex in general.Input convex neural networks make the inner problem convex, but their statistical optimality remains unclear.
Loading 2609.06598v1…