[Paper Review] Recent Advances in Optimal Transport for Machine Learning

Eduardo Fernandes Montesuma, Fred Ngolè Mboula+1 more|arXiv (Cornell University)|Jun 28, 2023

Machine Learning and ELM4 citations

TL;DR

This survey reviews recent advances in Optimal Transport (OT) for machine learning from 2012 to 2022, covering its use as a metric, loss function, regularizer, and transformation tool across supervised, unsupervised, transfer, and reinforcement learning. It highlights novel computational methods like projection-robust and mini-batch OT, and presents OT’s role in generative modeling, domain adaptation, and distributional reinforcement learning with theoretical and empirical insights.

ABSTRACT

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.

Motivation & Objective

• To provide a comprehensive, up-to-date review of Optimal Transport's role in machine learning from 2012 to 2022.
• To identify and analyze recent computational advances in Optimal Transport, including projection-robust, structured, neural, and mini-batch OT methods.
• To explore OT's applications in four key ML subfields: supervised, unsupervised, transfer, and reinforcement learning.
• To highlight novel contributions not covered in prior surveys, such as OT in fairness, dictionary learning, and distributional RL.
• To clarify the interplay between OT theory and practical ML implementations, especially in deep learning and probabilistic modeling.

Proposed method

• Uses the Monge-Kantorovich (MK) formulation of Optimal Transport as the primary theoretical foundation, defining the Wasserstein distance as the minimal cost of transporting mass from one distribution to another.
• Applies the MK formulation to define OT as a metric (e.g., Wasserstein distance), loss function (e.g., in WGANs and WDGRL), and regularizer (e.g., in WAE and WQL).
• Introduces the use of the Knothe-Rubinstein (KR) formulation for differentiable OT in generative models, enabling backpropagation through OT layers.
• Employs the Bures-Bernoulli (BB) formulation for modeling continuous-time dynamics in policy optimization, framing it as gradient flows in the space of probability measures.
• Utilizes structured and projection-robust OT to improve scalability and statistical efficiency in high-dimensional and complex data settings.
• Applies OT to data aggregation via Wasserstein barycenters and geodesics, enabling interpolation and alignment of heterogeneous probability distributions in domain adaptation and clustering.

Experimental results

Research questions

• RQ1 How has Optimal Transport advanced machine learning across supervised, unsupervised, transfer, and reinforcement learning from 2012 to 2022?
• RQ2 What are the key computational innovations in Optimal Transport that have enabled its integration into deep learning and large-scale ML?
• RQ3 In what ways does OT serve as a principled alternative to traditional losses like KL divergence in generative modeling and domain adaptation?
• RQ4 How does OT enable fairer and more robust learning in structured data and multi-source domain adaptation?
• RQ5 What is the role of OT in distributional reinforcement learning, and how does it improve uncertainty propagation and policy optimization?

Key findings

• Optimal Transport provides a statistically and topologically well-behaved metric (Wasserstein distance) that outperforms KL divergence in generative modeling due to its continuity and stability under weak convergence.
• The use of projection-robust and mini-batch OT methods significantly improves computational efficiency and scalability for large-scale and high-dimensional data.
• Wasserstein GANs (WGANs) and WDGRL achieve state-of-the-art performance in generative modeling and domain adaptation by using OT as a training loss, enabling more stable training dynamics.
• Wasserstein barycenters and geodesics enable principled interpolation and aggregation of probability distributions, with applications in domain adaptation and clustering.
• In reinforcement learning, the WQL algorithm uses Wasserstein barycenters to update Q-distributions, achieving strong performance with theoretical grounding in distributional dynamics.
• The survey identifies a growing trend of using neural networks to parametrize OT maps (e.g., via ICNNs), enabling end-to-end differentiable OT layers in modern deep learning architectures.