Flowing Datasets with Wasserstein over Wasserstein Gradient Flows

TL;DR

This paper introduces a novel framework for modeling and optimizing datasets as probability distributions over distributions using Wasserstein metrics, enabling advanced gradient flows for transfer learning and dataset distillation.

Contribution

It proposes the Wasserstein over Wasserstein (WoW) distance and gradient flows on this space, facilitating new techniques for dataset manipulation and learning tasks.

Findings

Effective transfer learning via WoW gradient flows

Improved dataset distillation using Wasserstein-based functionals

Novel tractable metrics for probability distributions over distributions

Abstract

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.