PaperlayerView exact revision on arXiv

Source-linked AI summary

Recent Advances in Optimal Transport for Machine Learning

Eduardo Fernandes Montesuma, Fred Ngolè Mboula, Antoine Souloumiac

arXiv:2306.16156v2cs.LGmath.PRstat.ML

TL;DR

Comparing and manipulating probability distributions is a central use of Optimal Transport in machine learning, and the field's rapid growth motivates an updated survey. The paper reviews OT theory, computation, and applications across four learning areas; it concludes that OT has influenced probabilistic ML while high-dimensional estimation and computational cost remain challenges.

  • Problem

    The rapid growth of Optimal Transport for machine learning motivates an updated account of its developments and applications.

    Source: §1 · p. 1

  • Method

    The survey reviews OT theory, computational developments, and applications in supervised, unsupervised, transfer, and reinforcement learning.

    Source: §1 · p. 1

  • Results

    OT is used in probabilistic machine learning to compute losses or manipulate distributions, and machine-learning methods are also influencing OT, including neural approximations.

    Source: §8.3 · p. 16

  • Takeaways & Limitations

    OT offers flexibility and useful theoretical properties, but high-dimensional estimation and computational cost remain relevant challenges for machine learning.

    Source: §8.3 · p. 16

    OT estimation becomes harder in high dimensions, and exact OT computation is computationally heavy.

    Source: §8.1 · p. 15, §8.3 · p. 16

Abstract

from arXiv · show

Recently, Optimal Transport has been proposed as a probabilistic framework in Machine Learning for comparing and manipulating probability distributions. This is rooted in its rich history and theory, and has offered new solutions to different problems in machine learning, such as generative modeling and transfer learning. In this survey we explore contributions of Optimal Transport for Machine Learning over the period 2012 -- 2023, focusing on four sub-fields of Machine Learning: supervised, unsupervised, transfer and reinforcement learning. We further highlight the recent development in computational Optimal Transport and its extensions, such as partial, unbalanced, Gromov and Neural Optimal Transport, and its interplay with Machine Learning practice.

1 INTRODUCTION

Optimal Transport (OT) contributes to machine learning both as a loss for comparing distributions and as a framework for manipulating them. This survey updates the view of OTML and reviews its theory, computational developments, and applications across four machine-learning areas.

1 INTRODUCTION

  • OTML uses OT both to compare distributions with a loss and to manipulate probability distributions.The survey presents OT as a principled tool for understanding the space of probability distributions.Source: §1 · p. 1
  • The Wasserstein distance is one name for OT's distribution metric, which can be preferable to KL divergence in problems such as generative modeling.The passage attributes this preference to topological, statistical, and geometrical properties.Source: §1 · p. 1
  • The survey updates the account of OTML and covers OT theory, computational developments, and applications in supervised, unsupervised, transfer, and reinforcement learning.The authors cite the field's rapid growth as justification for a closer look.Source: §1 · p. 1

2 BACKGROUND

Optimal Transport admits several formulations that provide different ways to represent the movement of probability mass. Its solutions define distributional losses and support metrics whose geometry can matter for machine-learning optimization.

2 BACKGROUND

  • Monge represents transport with a map T, while Kantorovich uses a mass-preserving plan γ; the plan formulation is linear in γ and admits a dual problem.The Monge constraint is described as difficult to analyze; for squared Euclidean cost, Brenier's theorem connects the formulations through T = ∇φ [14].Source: §2 · pp. 1–2
  • Dynamic OT describes transport through a time-dependent distribution ρ(t, x) and a vector field that moves mass between endpoint distributions.The endpoints are ρ(0, ·) = P and ρ(1, ·) = Q.Source: §2 · p. 2, Figure 1 · p. 2
  • OT solutions define losses between distributions, and Wasserstein distances provide a metric when the ground cost comes from a metric [12, Chapter 5].Wasserstein geodesics and barycenters give principled ways to interpolate and average distributions,.Source: §2 · p. 2
  • IPMs compare distributions through differences in expectations, whereas f-divergences measure discrepancies through the ratio between distributions.W1 is an IPM with the 1-Lipschitz function class; MMD is another IPM,.Source: §2 · p. 2
  • IPMs are defined for distributions with disjoint supports and account for sample-space geometry, while the choice of discrepancy changes the geometry and optimization problem in machine learning.The passage contrasts a meaningful IPM metric for disjoint-support GAN distributions with an infinite f-divergence, and contrasts Wasserstein and KL geometry on Gaussian distributions.Source: §2 · p. 2

3 COMPUTATIONAL OPTIMAL TRANSPORT

This section surveys computational strategies and extensions for making Optimal Transport tractable in machine-learning settings, including strong and weak OT formulations. It also discusses stochastic-map implementations using neural networks and mini-batch sampling, along with tradeoffs such as mini-batch plan artifacts and difficulties with non-convex formulations.

3 COMPUTATIONAL OPTIMAL TRANSPORT

  • Entropic OT is GPU-friendly with O(Ln^2) complexity, smoothly approximates Wasserstein distance, and has better sample complexity; its Sinkhorn divergence interpolates between MMD and Wasserstein distance.The stated complexity applies for L ≥ 1 iterations.Source: §3 · p. 3

3.1 Projection-based Optimal Transport

  • Sliced Wasserstein (SW) costs O(Lnd + Ln log n), max-SW is easier to compute with SW-like sample complexity, and subspace-robust Wasserstein optimizes over k-dimensional projections.Source: §3.1 · pp. 3–4

3.2 Structured Optimal Transport

  • Structured OT uses sub-modular costs to obtain structured plans [40], while low-rank transport addresses dimensionality through K-point barycenters and hub-mediated plans.Source: §3.2 · p. 4

3.3 Neural Network-based Solvers

  • Neural OT approximates plans through entropic dual potentials, barycentric maps through neural networks, and convex transport maps through ICNN potentials whose gradients yield OT maps.Source: §3.3 · p. 5

3.4 Mini-batch Optimal Transport

  • Mini-batch OT reduces complexity from O(n^3 log n) to O(Km^3 log m), is scalable and dimension-independent with unbiased gradients, but its plans become less sparse,.The artifact arises because mini-batch OT must transport all mass between batches.Source: §3.4 · p. 5, Figure 5 · p. 5

3.5 Extensions to Optimal Transport

  • Unbalanced OT relaxes mass conservation with f-divergence penalties and is robust to outliers,,; partial OT leaves a fraction 1−s untransported using dummy sinks.Source: §3.5 · p. 6
  • Gromov-Wasserstein compares relational structure across spaces [59], and fused GW adds feature costs; these formulations are non-convex, GW is NP-hard generally, and entropy regularization alleviates complexity,.Source: §3.5 · p. 6
  • Weak OT measures transport effort from a source point to a conditional target distribution; parameterizes that distribution through a map T(x,·) applied to noise S.The original strong OT formulation is recovered when the weak cost is the expected pointwise transport cost under the target distribution.Source: §3.5 · pp. 6–7
  • Stochastic maps T(x1, z) can be parameterized with neural networks, with the optimization carried out by sampling mini-batches from P, Q, and S.The source describes T and φ as neural networks with parameters θ and ξ, respectively.Source: §3.5 · p. 7
  • Weak OT represents each source point with a distribution over target points, unlike strong OT’s point-to-points matching through an OT plan.For a source point x1, the weak-OT figure depicts a conditional distribution γ(·|x1) over points in Q.Source: Figure 8 · p. 7

4 SUPERVISED LEARNING

The survey reviews OT as a classification loss that accounts for label similarities and as a tool for enforcing statistical parity in binary classification.

4.1 OT as a loss

  • Wasserstein loss incorporates label-space similarities, penalizing semantically distant classification errors more than close ones.The label metric defines the transport cost between predictions and ground-truth labels.Source: §4.1 · p. 7

4.2 Fairness

  • For fairness preprocessing, maps both protected groups’ feature distributions toward a shared Wasserstein barycenter while seeking to preserve information.Random repair tunes the trade-off between fairness and classification performance.Source: §4.2 · p. 8
  • In-processing fairness methods regularize representations or output distributions to align groups, with extending the approach to multiple attribute values.[81] uses MMD or Sinkhorn divergence; transports outputs toward a barycenter of group-conditional distributions.Source: §4.2 · p. 8

5 UNSUPERVISED LEARNING

The survey covers OT applications in unsupervised learning, including generative modeling, dictionary learning, and clustering. It describes methods that use Wasserstein-based objectives and extensions, alongside their benefits and constraints.

5.1.1 Generative Adversarial Networks

  • Wasserstein GANs replace the Jensen-Shannon divergence with the Kantorovich-Rubinstein metric, but neural implementations do not estimate the Wasserstein distance,, [90], [91].Enforcing the 1-Lipschitz constraint is difficult; weight clipping, gradient penalties, and spectral normalization are proposed approaches.Source: §5.1.1 · p. 8

5.1.2 Autoencoders

  • Wasserstein Autoencoders generalize Adversarial Autoencoders by using deterministic decoding to simplify the Minimum Kantorovich formulation.The constraint is enforced with a penalty, for which proposed choices include JS divergence, MMD, and Sinkhorn loss.Source: §5.1.2 · pp. 9–10

5.1.3 Continuous Normalizing Flows

  • Regularized Neural ODEs combine kinetic-energy and Jacobian penalties to constrain continuous normalizing flows, with OT properties enforced only on training data.The OT-based penalty favors non-intersecting trajectories and geodesic paths; the Jacobian penalty is proposed to help generalize these properties.Source: §5.1.3 · p. 10

5.2 Dictionary Learning

  • OT-based NMF uses Sinkhorn distances to reduce computational complexity and yield a smooth objective suitable for gradient descent.Wasserstein dictionary learning instead computes barycenters in Wasserstein space, enabling non-linear aggregation of atoms.Source: §5.2 · p. 10
  • Wasserstein Dictionary Learning non-linearly interpolates atom distributions, making learned dictionaries more expressive than linear aggregation [113].The comparison uses Gaussian histograms whose means vary from -6 to 6; the survey says this setting illustrates WDL’s advantage.Source: §5.2 · p. 11, Figure 11 · p. 11

5.3 Clustering

  • Clustering is linked to OT quantization, which frames it as finding a discrete distribution minimizing Wasserstein distance and supports theoretical convergence bounds for K-means.Co-clustering approaches include CCOT and CCOT-GW, the latter using entropic Gromov OT between row, barycenter, and column kernels.Source: §5.3 · p. 11
  • Sliced-Wasserstein Gaussian mixture models are reported to be robust to parameter initialization.The method uses the closed-form sliced-Wasserstein distance between Gaussian distributions.Source: §5.3 · pp. 11–12

6 TRANSFER LEARNING

The section reviews optimal transport for transfer learning, especially unsupervised domain adaptation under source–target distribution shift. It surveys shallow and deep alignment methods, extensions to heterogeneous settings, and OT-based transferability measurement.

6 TRANSFER LEARNING

  • Transfer learning addresses source and target data with different distributions, and UDA uses labeled source data with unlabeled target data.Source: §6 · p. 12

6.1 Shallow Domain Adaptation

  • Barycentric mapping transports source samples toward the target to form training data, but can cross class boundaries and misplace samples.Class-based regularization and group sparsity make transport plans more class-sparse.Source: §6.1 · p. 12, Figure 13 · p. 12
  • Barycentric mappings are defined only on observed source samples; nearest-neighbor extension can quantize, while linear, kernel, and neural parameterizations address new samples.The neural parameterization is described as a large-scale strategy.Source: §6.1 · p. 12
  • For joint-distribution shift in UDA, a proxy target distribution based on a classifier supports minimizing Wasserstein distance over candidate classifiers.Source: §6.1 · p. 12

6.2 Deep Domain Adaptation

  • Deep adaptation aligns learned source and target representations; MMD, DANN, and deep JDOT provide strategies for this alignment.Deep JDOT includes the feature extractor in the ground cost.Source: §6.2 · p. 13
  • WDGRL replaces DANN’s domain loss with Wasserstein distance to maximize domain confusion while minimizing classification error.Unbalanced OT was proposed to allow smaller batches and improve performance when stable JDOT training needs large batches.Source: §6.2 · p. 13

6.3 Extensions to Domain Adaptation

  • Multi-source adaptation uses a Wasserstein barycenter followed by target adaptation, while heterogeneous adaptation uses Gromov-Wasserstein plans to relate incomparable domains.In HDA, sample-level transport supports target-label estimation by label propagation.Source: §6.3 · p. 13
  • OTDD is highly correlated with transferability and can interpolate between datasets with different classes.The distance incorporates a Wasserstein-based metric on the label space.Source: §6.3 · p. 14

7 REINFORCEMENT LEARNING

The section surveys optimal transport across reinforcement-learning settings that model uncertainty over returns, Q-values, or policy parameters. It describes distributional, Bayesian, and policy-optimization approaches.

7 REINFORCEMENT LEARNING

  • Standard reinforcement learning models sequential decisions with a Markov decision process, evaluates policies through returns and value functions, and learns Q-values from experience.Source: §7 · p. 14

7.1 Distributional Reinforcement Learning

  • Distributional RL fits a return distribution by minimizing Wasserstein distance between its parameterized estimate and its Bellman-updated distribution.One parameterization fixes support locations and estimates weights; another fixes weights and estimates support locations, avoiding biased Wasserstein gradients and allowing more flexible support.Source: §7.1 · p. 14

7.2 Bayesian Reinforcement Learning

  • Bayesian RL represents Q-function estimates with posterior distributions and defines state-value posteriors through Wasserstein barycenters.Its Wasserstein temporal-difference update has analytical solutions for specific distribution families such as Gaussians.Source: §7.2 · p. 14

7.3 Policy Optimization

  • Policy optimization can place a distribution over policy parameters, with a prior-regularized objective whose gradient flow admits indirect or direct learning formulations.The indirect formulation is particularly suited to low-dimensional parameter spaces.Source: §7.3 · p. 15

8 FINAL REMARKS

The closing sections identify high-dimensional estimation and computational cost as continuing OT challenges, while surveying applications, future research directions, and OT’s reciprocal influence with machine learning.

8.1 Challenges

  • OT estimation becomes harder in high dimensions, while exact computation is costly; Sinkhorn and sliced methods offer alternatives with trade-offs in iteration or projection counts and optimization smoothness.The stated estimation bound is O(n^-1/d); exact OT has O(n^3 log n) complexity, versus O(n^2) per Sinkhorn iteration and O(n log n) per slice-sorting problem.Source: §8.1 · p. 15

8.2 Future Trends

  • Future directions include learned OT estimators, generalized sliced distances, private and federated Wasserstein computation, and connections among fairness, domain adaptation, and generative modeling.The survey notes that learned methods may perform an ML task without estimating OT, and suggests studying theory for parallels between random repair and OTDA.Source: §8.2 · pp. 15–16
  • Future supervised-learning work could build label embeddings that capture the geometry of class-conditional distributions, while generative models could use partial or unbalanced OT for outliers and FGW for graphs.Source: §8.2 · p. 16
  • Wasserstein losses and barycenters support dictionary learning with probabilistic objects; future work could study atom interpolations and coefficient interpretability, alongside convergence and initialization robustness for sliced-Wasserstein GMMs.Source: §8.2 · p. 16
  • OT-based UDA has shown state-of-the-art performance across several application areas, with large-scale methods and theoretical work establishing OT as a toolbox for transferring knowledge across domains.Source: §8.2 · p. 16
  • Distributional RL studies reward distributions, while WQL uses Wasserstein barycenters to update value-function distributions; policy optimization can also be formulated as gradient flows in distribution space.Source: §8.2 · p. 16

8.3 Conclusion

  • OT is a probabilistic toolkit for machine learning, but high-dimensional estimation and computational cost remain challenges as machine learning methods also influence OT.Source: §8.3 · p. 16

Related papers

Loading 2306.16156v2…