# Flowing Datasets with Wasserstein over Wasserstein Gradient Flows

arXiv: https://arxiv.org/abs/2506.07534 · HTML: https://mlsanity.com/paper/2506.07534/

## Key facts

- Cited by 11 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 9 Jun 2025.

## Built on by (most cited first, 11 of 11)

- [A Novel Sliced Fused Gromov-Wasserstein Distance](https://mlsanity.com/paper/2508.02364/) (2025)
- [Slicing Wasserstein Over Wasserstein Via Functional Optimal Transport](https://mlsanity.com/paper/2509.22138/) (2025)
- [Generalized Wasserstein Flow Matching: Transport Plans, Everywhere, All at Once](https://mlsanity.com/paper/2605.08424/) (2026)
- [Dense associative memory for Gaussian distributions](https://mlsanity.com/paper/2509.23162/) (2025)
- [Busemann Functions in the Wasserstein Space: Existence, Closed-Forms, and Applications to Slicing](https://mlsanity.com/paper/2510.04579/) (2025)
- [Mean-Field Model for Two-Layer Neural Networks Trained with Consensus-Based Optimization](https://mlsanity.com/paper/2511.21466/) (2025)
- [HOT-POT: Optimal Transport for Sparse Stereo Matching](https://mlsanity.com/paper/2601.12423/) (2026)
- [Rethinking the Flow-Based Gradual Domain Adaptation: A Semi-Dual Optimal Transport Perspective](https://mlsanity.com/paper/2602.01179/) (2026)
- [Highly Data Parallelizable Estimation of the Sliced-Wasserstein Distance Using Cumulative Distribution Functions](https://mlsanity.com/paper/2606.30310/) (2026)
- [General Quantification of Covariate and Concept Shifts](https://mlsanity.com/paper/2609.11918/) (2026)
- [When Riemann flows with Wasserstein: Generative Modeling of Probability Distributions on Manifolds](https://mlsanity.com/paper/2609.25659/) (2026)

## Abstract (arXiv)

Many applications in machine learning involve data represented as probability distributions. The emergence of such data requires radically novel techniques to design tractable gradient flows on probability distributions over this type of (infinite-dimensional) objects. For instance, being able to flow labeled datasets is a core task for applications ranging from domain adaptation to transfer learning or dataset distillation. In this setting, we propose to represent each class by the associated conditional distribution of features, and to model the dataset as a mixture distribution supported on these classes (which are themselves probability distributions), meaning that labeled datasets can be seen as probability distributions over probability distributions. We endow this space with a metric structure from optimal transport, namely the Wasserstein over Wasserstein (WoW) distance, derive a differential structure on this space, and define WoW gradient flows. The latter enables to design dynamics over this space that decrease a given objective functional. We apply our framework to transfer learning and dataset distillation tasks, leveraging our gradient flow construction as well as novel tractable functionals that take the form of Maximum Mean Discrepancies with Sliced-Wasserstein based kernels between probability distributions.