Source-linked AI summary
Large-Scale Optimal Transport and Mapping Estimation
Vivien Seguy, Bharath Bhushan Damodaran, Rémi Flamary, Nicolas Courty, Antoine Rolet, Mathieu Blondel
TL;DR
The paper addresses how to learn an optimal map between distributions when direct OT computation is difficult at large scale or for continuous measures. It combines stochastic dual regularized OT with neural estimation of a Monge map, and proves consistency while demonstrating applications in domain adaptation and generative modeling.
Problem
Computing optimal maps is difficult in high-dimensional or continuous settings, while standard discrete OT solvers have super-cubic complexity.
Method
The approach first computes a regularized OT plan with stochastic dual optimization, then learns a neural Monge map by approximating the plan’s barycentric projection.
Results
The method scales better than a related semi-dual approach on very large sample sets, achieves the best reported adaptation results, and its regularized plans converge to continuous OT solutions.
Takeaways & Limitations
The framework provides tractable large-scale or continuous OT computation and supports applications in unsupervised domain adaptation and generative modeling.
Takeaways & Limitations
The paper leaves convergence rates for empirical regularized optimal plans as an open problem.
Abstract
from arXiv · showhide
This paper presents a novel two-step approach for the fundamental problem of learning an optimal map from one distribution to another. First, we learn an optimal transport (OT) plan, which can be thought as a one-to-many map between the two distributions. To that end, we propose a stochastic dual approach of regularized OT, and show empirically that it scales better than a recent related approach when the amount of samples is very large. Second, we estimate a \textit{Monge map} as a deep neural network learned by approximating the barycentric projection of the previously-obtained OT plan. This parameterization allows generalization of the mapping outside the support of the input measure. We prove two theoretical stability results of regularized OT which show that our estimations converge to the OT plan and Monge map between the underlying continuous measures. We showcase our proposed approach on two applications: domain adaptation and generative modeling.
1 INTRODUCTION
The paper studies learning maps that match one distribution to another, with optimal transport providing a principled cost-minimizing criterion. It proposes a scalable two-step approach that computes a regularized OT plan and learns a neural Monge map from its barycentric projection.
- Mapping one distribution to another: Distribution-matching maps support domain adaptation and generative modeling by transforming samples from a source distribution into the target distribution.In generative modeling, sampling consists of drawing from a continuous latent distribution and applying the learned map.
- Mapping with optimality: Optimal transport selects among distribution-matching maps by minimizing the total cost of moving mass between the source and target.This extends the basic distribution-matching problem with an explicit optimality criterion.
- Large-scale OT: Direct map computation becomes difficult in high dimensions, while discrete OT plans remain feasible but require super-cubic or O(n3 log n) computation.Space discretization makes some low-dimensional methods intractable at higher dimensions, and network-simplex solutions scale poorly with support size.
- Contributions: The proposed method first computes a regularized OT plan with stochastic dual optimization, then learns a neural Monge map by approximating its barycentric projection.The neural parameterization supports learning beyond the input measure’s support.
- Contributions: The paper proves stability results for regularized plans and their barycentric projections, showing convergence toward the underlying continuous OT plan and Monge map.The approach is evaluated in domain adaptation and generative modeling.
2 BACKGROUND ON OPTIMAL TRANSPORT
The background contrasts Monge maps, which send each source point to one destination, with Kantorovich plans, which distribute mass across destination points. Regularization makes the relaxed OT problem easier to optimize and differentiable with respect to input weights.
- The Monge Problem: The Monge problem seeks a cost-minimizing map that transports source measure µ to target measure ν.Its feasibility depends on the measures and cost; for discrete sources, a valid map may not exist.
- The Monge Problem: For continuous µ on R^d with squared Euclidean cost, Brenner’s result guarantees existence of an optimal Monge map, while high-dimensional computation remains difficult.The Gaussian-to-Gaussian case has a closed-form affine solution, but general cases remain open computationally.
- Kantorovich Relaxation: Kantorovich relaxation replaces maps with couplings whose marginals equal µ and ν, making the OT problem always feasible.A coupling can send mass from one source point to several target locations.
- Kantorovich Relaxation: The relaxed coupling formulation is a linear program, but specialized discrete solvers have super-cubic complexity in the support sizes.This computational cost limits OT use in large-scale settings.
- Regularized OT: Entropy or squared-L2 regularization adds a strictly convex penalty to the primal problem, smoothing the dual constraints and enabling unconstrained optimization.The paper uses the resulting dual formulation for stochastic gradient algorithms.
3 LARGE-SCALE OPTIMAL TRANSPORT
The paper develops stochastic dual optimization for regularized OT across discrete and continuous settings, then analyzes its computational trade-offs and consistency. Its convergence results connect empirical regularized plans to OT plans between underlying continuous measures.
- Dual stochastic approach: Stochastic gradient ascent maximizes the unconstrained dual by sampling batches from the independent coupling µ × ν.First-order optimality conditions recover the regularized primal plan from the dual solution.
- Regularized OT dual: Regularizing the primal OT problem relaxes the hard coupling constraint into a smooth concave penalty in the dual objective.The framework includes entropy and L2 regularization and can extend to decomposable strictly convex regularizers.
- Dual stochastic approach: Neural networks parameterize continuous dual variables, while discrete measures use finite-dimensional vectors over their support points.This allows the same algorithmic framework to cover discrete and continuous measures.
- Convergence rates and computational cost comparison: Mini-batch stochastic optimization costs O(p2) per iteration, compared with O(n) for the semi-dual method and O(k2) for kernel-expansion approaches.The trade-off is lower per-iteration dependence on the full support size, with non-convexity in neural parameterizations.
- Convergence rates and computational cost comparison: In the semi-discrete setting, the proposed algorithm is known only to converge to a stationary point because neural-network parameterization makes the problem non-convex.The semi-dual alternative retains an O(1/k) convergence rate in that setting.
- Convergence of regularized OT plans: When empirical measures converge weakly to continuous measures and regularization decreases sufficiently fast, regularized OT plans converge weakly to an unregularized OT solution.The result holds up to extraction of a subsequence and is presented as analogous to universal consistency.
4 OPTIMAL MAPPING ESTIMATIONS
The paper recovers an optimal map by taking the barycentric projection of a regularized OT plan and fitting that projection with a deep neural network. Under stated convergence assumptions, the estimated projection approaches the Monge map and its pushforward approaches the target measure.
- Barycentric projection: The barycentric projection averages the target locations assigned by an OT plan, producing a map from each source point to an averaged target location.For the squared Euclidean cost, this projection is an optimal map to its own pushforward measure.
- Neural map estimation: A deep neural network parameterizes the barycentric projection so the learned map is defined beyond the finite support of a discrete input measure.Its parameters are learned by stochastic gradient descent on a weighted mapping objective.
- Theoretical guarantees: The barycentric projection generally matches the target only approximately when the OT plan is not induced by a map, although the approximation improves with large sample size and small regularization.This limitation motivates the paper’s stability analysis.
- Theoretical guarantees: As n → ∞ and ε → 0 under the theorem’s assumptions, the barycentric projection converges weakly to the Monge map between the underlying continuous measures.The result assumes a continuous source measure, weakly convergent empirical measures, uniqueness of each discrete OT solution, and sufficiently fast decay of ε.
- Theoretical guarantees: The pushforward of the source measure by the estimated projection converges weakly to the target measure ν.Equivalently, mapped samples from the empirical source distribution converge in distribution to the target distribution.
5 NUMERICAL EXPERIMENTS
The experiments assess scalability, domain adaptation, and generative modeling. The proposed stochastic dual method converges faster and enables large-scale adaptation, while the learned Monge map improves performance over barycentric projections.
- 5.1 DUAL VS SEMI-DUAL SPEED COMPARISONS: 250000 samples in R31 were used for spectral transfer, and the proposed dual algorithm converged much faster than the stochastic semi-dual approach across ε = {0.025, 0.1, 1}.The comparison used two 500 × 500 CAVE images with 31 spectral bands and squared Euclidean cost.
- 5.2 LARGE SCALE DOMAIN ADAPTATION: The domain-adaptation evaluation used MNIST, USPS, and SVHN, including 60000 MNIST, 9298 USPS, and 73212 SVHN samples.Adaptation directions included MNIST → USPS, USPS → MNIST, and SVHN → MNIST.
- 5.2 LARGE SCALE DOMAIN ADAPTATION: The method was compared with 1-NN source-only classification and prior OTDA methods using discrete or entropy-regularized barycentric projections.The proposed approach learned a neural-network Monge map with entropy or L2 regularization and used squared Euclidean ground cost.
- 5.2 LARGE SCALE DOMAIN ADAPTATION: The proposed approach outperformed previous OTDA algorithms in every domain-adaptation setting, including 20% and 10% accuracy-point gains over OT and regularized OT on MNIST→USPS.On SVHN→MNIST, only the proposed method handled the whole datasets and exceeded source-only results.
- 5.2 LARGE SCALE DOMAIN ADAPTATION: Learning a parametric Monge map improved performance relative to the barycentric projection, acting as a form of regularization.
- 5.3 GENERATIVE OPTIMAL TRANSPORT (GOT): The generative experiment used a Gaussian measure and the preprocessed MNIST dataset, with a fully connected barycentric-projection network producing samples shown in Figure 4.MNIST images were treated as 784-dimensional vectors without convolutional architectures.
6 CONCLUSION
The paper introduces algorithms for large-scale regularized OT and Monge-map learning, demonstrating applications in domain adaptation and generative modeling. It establishes theoretical consistency while leaving convergence-rate analysis for empirical regularized transport plans as future work.
- The paper proposes algorithms for large-scale regularized optimal transport and for learning a Monge map between probability distributions.
- The approach is applied to unsupervised domain adaptation and generative modeling, where the Monge map acts directly as a generator.
- The consistency results support the theoretical grounding of the approach, while convergence rates for empirical regularized optimal plans remain an open direction.
A PROOFS
The proofs establish weak convergence of regularized empirical OT plans and their barycentric projections under suitable decay of the regularization parameter. They use OT stability, entropy-regularization convergence, and discrete assignment structure to connect empirical solutions with continuous OT plans and Monge maps.
- Regularized OT-plan convergence: OT stability yields a subsequence of empirical maximum-entropy plans converging weakly to a continuous OT solution.This stability applies regardless of whether the empirical plans are maximum-entropy solutions.
- Regularized OT-plan convergence: Choosing ε_n sufficiently small makes the regularization error vanish, proving weak convergence of πεn_n to π.The argument uses convergence results for entropy-regularized OT and bounded continuous test functions.
- Monge-map convergence: Absolute continuity of µ and the cost assumptions guarantee existence of a Monge map between µ and ν.This provides the continuous target map for the empirical barycentric-projection analysis.
- Monge-map convergence: For equally sized uniform discrete measures, the empirical OT plan is an optimal assignment T_n by the Birkhoff theorem.This identifies the discrete transport plan with a map, enabling comparison with the barycentric projection.
- Monge-map convergence: A Lipschitz test function bounds the barycentric-projection error by the distance between the empirical regularized plan and the discrete assignment.The bound uses the Lipschitz constant of g, the discrete barycentric-projection formula, and Cauchy–Schwarz.
- Monge-map convergence: An appropriate decay rate for ε_n makes the error vanish and proves weak convergence of (id, ¯πεn_n)#µ_n to (id, f)#µ.The conclusion holds for Lipschitz functions g ∈ C_l(R^d × R^d).