Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
This addresses the need for tractable gradient flows on complex data structures in machine learning, though it appears incremental as it builds on existing optimal transport concepts.
The paper tackles the problem of designing gradient flows for probability distributions over probability distributions, which arise when representing labeled datasets as mixtures of class-conditional distributions. They propose the Wasserstein over Wasserstein (WoW) distance and gradient flows, applying them to transfer learning and dataset distillation tasks.
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.