# Dataset Dynamics via Gradient Flows in Probability Space

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

## Key facts

- Cited by 18 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 24 Oct 2020.

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

- [Recent Advances in Optimal Transport for Machine Learning](https://mlsanity.com/paper/2306.16156/) (2023)
- [Neural Optimal Transport with General Cost Functionals](https://mlsanity.com/paper/2205.15403/) (2022)
- [FlashSinkhorn: IO-Aware Entropic Optimal Transport on GPU](https://mlsanity.com/paper/2602.03067/) (2026)
- [Dynamic Flows on Curved Space Generated by Labeled Data](https://mlsanity.com/paper/2302.00061/) (2023)
- [Lightspeed Geometric Dataset Distance via Sliced Optimal Transport](https://mlsanity.com/paper/2501.18901/) (2025)
- [Universal Neural Optimal Transport](https://mlsanity.com/paper/2212.00133/) (2022)
- [Generating Synthetic Datasets by Interpolating along Generalized Geodesics](https://mlsanity.com/paper/2306.06866/) (2023)
- [Module-wise Training of Neural Networks via the Minimizing Movement Scheme](https://mlsanity.com/paper/2309.17357/) (2023)
- [Neural McKean-Vlasov Processes: Distributional Dependence in Diffusion Processes](https://mlsanity.com/paper/2404.09402/) (2024)
- [A Universal Metric of Dataset Similarity for Cross-silo Federated Learning](https://mlsanity.com/paper/2404.18773/) (2024)
- [Interpolation for Robust Learning: Data Augmentation on Wasserstein Geodesics](https://mlsanity.com/paper/2302.02092/) (2023)
- [Forward-Euler time-discretization for Wasserstein gradient flows can be wrong](https://mlsanity.com/paper/2406.08209/) (2024)
- [Block-wise Training of Residual Networks via the Minimizing Movement Scheme](https://mlsanity.com/paper/2210.00949/) (2022)
- [Busemann Functions in the Wasserstein Space: Existence, Closed-Forms, and Applications to Slicing](https://mlsanity.com/paper/2510.04579/) (2025)
- [Wasserstein Gradient Flows for Scalable and Regularized Barycenter Computation](https://mlsanity.com/paper/2510.04602/) (2025)
- [From Outcomes to Actions: Leveraging Hindsight for Long-Horizon Language Agent Training](https://mlsanity.com/paper/2607.16257/) (2026)
- [Implicit Neural Optimal Transport via Fixed-Point Optimization](https://mlsanity.com/paper/2605.10792/) (2026)
- [Highly Data Parallelizable Estimation of the Sliced-Wasserstein Distance Using Cumulative Distribution Functions](https://mlsanity.com/paper/2606.30310/) (2026)

## Abstract (arXiv)

Various machine learning tasks, from generative modeling to domain adaptation, revolve around the concept of dataset transformation and manipulation. While various methods exist for transforming unlabeled datasets, principled methods to do so for labeled (e.g., classification) datasets are missing. In this work, we propose a novel framework for dataset transformation, which we cast as optimization over data-generating joint probability distributions. We approach this class of problems through Wasserstein gradient flows in probability space, and derive practical and efficient particle-based methods for a flexible but well-behaved class of objective functions. Through various experiments, we show that this framework can be used to impose constraints on classification datasets, adapt them for transfer learning, or to re-purpose fixed or black-box models to classify -- with high accuracy -- previously unseen datasets.