Flowing Datasets with Wasserstein over Wasserstein Gradient Flows

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows

Speakers: Clément Bonet, Christophe Vauthier, Anna Korba

Conference: ICML 2025

YouTube: https://slideslive.com/39043996

Overview

This talk introduces a novel and computationally efficient framework for comparing and transforming labeled datasets, a critical challenge in modern machine learning. Presented by Clément Bonet, Christophe Vauthier, and Anna Korba at ICML 2025, the work proposes representing labeled datasets as probability distributions over probability distributions, a hierarchical structure that naturally accounts for both individual sample characteristics and inter-class relationships. To enable meaningful comparisons and transformations within this complex space, the authors introduce the Wasserstein over Wasserstein (WoW) distance, an optimal transport metric that endows the space with a formal Riemannian structure. This structure, in turn, allows for the definition of WoW gradient flows, enabling principled optimization for tasks such as dataset distillation and data augmentation.

The significance of this research lies in its ability to facilitate coherent manipulation of dataset structure across various machine learning applications, including domain adaptation, transfer learning, and synthetic data generation. By leveraging the WoW distance and Maximum Mean Discrepancy (MMD) with Sliced Wasserstein (SW) kernels, the proposed method offers a significant improvement in computational efficiency compared to prior optimal transport-based approaches like OTDD, scaling as O(c^2 n log n) instead of O(n^2 c^2), where 'n' is the number of samples and 'c' is the number of classes. This efficiency, combined with its capacity to preserve both intra-class coherence and inter-class dynamics during transformations, positions this work as a substantial advancement in the field of dataset manipulation and generative modeling.

Background

▶ Watch: Problem: Comparing and flowing labeled datasets coherently (0:00)

Working with labeled datasets, composed of samples (X_i, Y_i) where X_i is the data point and Y_i is its assigned class, is fundamental to many machine learning tasks. The goal often extends beyond simple classification to generating new samples that adhere to the dataset's structure, or even flowing one dataset into another while maintaining the integrity of its classes. For instance, one might want to transform images of 'fours' from MNIST into images resembling 'shoes' from Fashion MNIST, ensuring that the transformation is coherent within each class and does not mix with other classes along the way. Such capabilities are crucial for applications like domain adaptation, where models trained on a source domain need to perform well on a related but different target domain, or dataset distillation, where a large dataset is condensed into a smaller, representative one.

A core challenge in these scenarios is how to meaningfully compare labeled datasets, particularly when the labels themselves lack a clear, a priori correspondence. For example, there's no inherent reason why class 'zero' in MNIST should semantically align with class 'zero' in Fashion MNIST. Addressing this, David Alvarez-Melis and Nicolò Fuzi proposed a method to embed labels into the space of probability distributions. They represented each label Y_i by the empirical distribution of samples belonging to that class. A labeled dataset was then viewed as an empirical distribution over the product space Rd x PRd (where Rd is the sample space and PRd is the space of probability distributions over Rd). To compare datasets, they introduced the Optimal Transport Dataset Distance (OTDD), which uses a ground cost defined as the sum of the squared Euclidean distance between samples and the squared Wasserstein distance between their corresponding class distributions. While conceptually sound, computing OTDD, even with approximations like Gaussian measures and entropic regularization, scales with a complexity of O(n^2 c^2), which can be prohibitively expensive for large datasets, limiting its practical applicability. This computational bottleneck and the specific structure of the product space Rd x PRd motivated the search for more efficient and flexible alternatives.

Key Findings

▶ Watch: Our model: Datasets as distributions over distributions (2:30)

The presented work introduces a powerful and efficient framework for flowing labeled datasets, built upon several key findings:

1. Hierarchical Dataset Representation: The authors propose a novel and intuitive representation of labeled datasets as elements of the space of probability distributions over probability distributions (P(PRd)). Each class is represented by the empirical distribution of its samples, and the entire dataset is then an empirical distribution over these class-specific distributions. This hierarchical view naturally captures both intra-class and inter-class structure.

1. Wasserstein over Wasserstein (WoW) Distance: To equip this P(PRd) space with a structure suitable for optimization, the paper defines the WoW distance. This is an optimal transport distance where the ground cost between two "points" (which are themselves probability distributions representing classes) is the squared Wasserstein distance. The WoW distance endows the space of distributions over distributions with a formal Riemannian structure, enabling the application of geometric concepts like geodesics, tangent spaces, and exponential maps.

1. WoW Gradient Flows for Optimization: Crucially, the Riemannian structure allows for the definition of a WoW gradient for functionals defined on P(PRd). This gradient provides a direction of steepest descent, facilitating Riemannian gradient descent to minimize discrepancies between datasets. This enables a principled approach to flowing datasets by iteratively moving them along the gradient direction.

1. Efficient Discrepancy Minimization with MMD and Sliced Wasserstein: For practical optimization, the authors focus on minimizing the Maximum Mean Discrepancy (MMD), a popular distance in machine learning. They demonstrate that the WoW gradient of MMD can be computed, either in closed form or via auto-differentiation, under certain smoothness assumptions on the kernel. By employing kernels based on the Sliced Wasserstein (SW) distance, which is computationally efficient (O(n log n) for two distributions) and Hilbertian, the overall computational complexity for evaluating MMD is reduced to O(c^2 n log n). This is a significant improvement over the O(n^2 c^2) complexity of OTDD.

1. Multi-level Coherent Flow: Experiments demonstrate that the proposed method facilitates a multi-level coherent flow. Samples within each class move together, maintaining their internal structure (e.g., forming rings), while the classes themselves interact and converge towards target classes, preserving inter-class relationships.

1. Effective Applications in ML Tasks: The framework is successfully applied to practical machine learning tasks such as dataset distillation and data augmentation. In distillation, it achieves comparable performance to existing methods while distilling all classes simultaneously. For data augmentation, it outperforms OTDD in improving classifier performance on small target datasets, all with significantly reduced computational cost.

Technical Deep Dive

▶ Watch: Mathematical framework: Wasserstein over Wasserstein gradient flows (4:00)

The core of this work lies in its novel mathematical framework for representing and manipulating labeled datasets. Unlike traditional approaches that treat labeled data as tuples (X_i, Y_i) or prior methods like OTDD that embed labels into a product space Rd x PRd, Bonet et al. introduce a more hierarchical representation.

Dataset Representation:

A labeled dataset, consisting of c classes, is modeled as an empirical probability distribution over probability distributions. Specifically, for each class Y_i, its samples are aggregated into an empirical distribution nu_i. This nu_i itself is a probability distribution over the sample space Rd (e.g., the space of images). The entire dataset is then represented as a uniformly weighted sum of Dirac measures centered at each of these class distributions: P = (1/c) * sum_{i=1 to c} Dirac(nu_i). This means the "points" in our dataset representation are not individual samples, but entire class distributions.

Wasserstein over Wasserstein (WoW) Distance:

To define a meaningful distance and structure on this space P(PRd), the authors introduce the WoW distance. This is an optimal transport distance where the "ground cost" between any two "points" in P(PRd) (which are themselves probability distributions, i.e., nu_i and nu_j) is defined by the squared 2-Wasserstein distance, W_2^2. That is, if we want to transport one dataset P to another Q, we are looking for an optimal plan that transports class distributions nu_i in P to class distributions mu_j in Q, minimizing the total cost, where the cost of transporting nu_i to mu_j is W_2^2(nu_i, mu_j).

The space of probability distributions over probability distributions, when endowed with the WoW distance, acquires a formal Riemannian structure. This is a crucial theoretical contribution, as it allows the definition of standard differential geometry concepts:

• Geodesics: Shortest paths between two distributions P and Q in this space.
• Tangent Space: For every distribution P, a tangent space can be defined, which is a vector space approximating the manifold locally.
• Exponential Map (exp_P): A mapping that projects elements from the tangent space at P back onto the manifold, allowing for "steps" in the manifold direction.

WoW Gradient and Riemannian Gradient Descent:

The Riemannian structure is particularly valuable for optimization. For a functional F defined on P(PRd) (e.g., a discrepancy measure between a current dataset and a target dataset), the WoW gradient (nabla_W2 F P) can be defined. This gradient is an element of the tangent space at P and points in the direction of the steepest descent for F at P.

This enables a Riemannian gradient descent algorithm:

Starting from an initial distribution P_k, a step with size tau is taken by computing P_{k+1} = exp_{P_k}(-tau * nabla_W2 F P_k).

In practice, since P_k is a discrete distribution of c class empirical distributions, each with n samples, this optimization scheme translates to a gradient descent over each individual sample (e.g., image). However, it's not a simple independent update; the gradient for a given sample depends on other samples within the same class and on the overall distribution of classes. This intricate dependency ensures that samples within a class move coherently, and classes themselves interact to converge towards target distributions.

Maximum Mean Discrepancy (MMD) with Sliced Wasserstein Kernels:

For the functional F to be minimized, the authors primarily use the Maximum Mean Discrepancy (MMD). MMD is a popular and versatile distance metric between probability distributions, requiring only a suitable positive definite kernel on the underlying space. In this context, the "underlying space" is PRd (the space of class distributions).

The paper shows that the WoW gradient of MMD can be computed, either in closed form or using auto-differentiation, provided the chosen kernel is sufficiently smooth. For the kernel itself, they leverage the Sliced Wasserstein (SW) distance. SW distance is a computationally efficient approximation of the Wasserstein distance, obtained by projecting high-dimensional distributions onto one-dimensional lines, computing 1D Wasserstein distances (which are trivial), and averaging the results. SW distance also possesses desirable properties such as being Hilbertian.

Specifically, the authors explore two types of SW-based kernels:

1. Gaussian Sliced Wasserstein Kernel: k(nu_i, nu_j) = exp(-gamma * SW^2(nu_i, nu_j))
2. Riesz Sliced Wasserstein Kernel: k(nu_i, nu_j) = -SW(nu_i, nu_j) (a conditionally positive definite kernel).

The computational complexity for evaluating the MMD between two datasets (each with c classes) involves computing O(c^2) Sliced Wasserstein distances between individual class distributions. Since each SW distance computation between two empirical distributions of n samples takes O(n log n) time, the total complexity for one MMD evaluation (and thus one gradient step) is O(c^2 n log n). This represents a substantial computational advantage over the O(n^2 c^2) complexity of OTDD.

Experimental Setup & Results

▶ Watch: Synthetic experiment: Three Rings dataset visualization (8:00)

The effectiveness of the proposed WoW gradient flow framework was demonstrated across a range of synthetic and real-world machine learning tasks.

1. Synthetic Task: Minimizing MMD on the Three Rings Dataset

• Setup: The goal was to minimize the squared MMD, targeting a synthetic "Three Rings" dataset where each ring represented a distinct class. The flow was initiated from a simple distribution, and the evolution was observed.
• Kernels: Both the Gaussian Sliced Wasserstein kernel and the Riesz Sliced Wasserstein kernel were evaluated.
• Results: In both kernel cases, the classes (initially unstructured particles) progressively formed distinct rings, and these rings then coherently moved towards the target ring positions. This experiment clearly illustrated the two-level interaction mechanism:
• Intra-class interaction: Particles within each class self-organized to form and maintain the ring structure throughout the flow.
• Inter-class interaction: The distinct rings (classes) interacted with each other, ensuring they converged to their respective target rings without mixing, demonstrating the preservation of inter-class relationships.

2. Image Dataset Flow with Class Structure Preservation

• Setup: To verify the ability to maintain class structure on complex, high-dimensional data, experiments were conducted using image datasets like MNIST, Fashion MNIST, SVHN, and CIFAR-10.
• A classifier was pre-trained on either MNIST or CIFAR-10.
• A flow was initiated, minimizing MMD with the Riesz Sliced Wasserstein kernel, starting from a source dataset (Fashion MNIST or SVHN) and targeting the classifier's training dataset (MNIST or CIFAR-10).
• At each iteration of the flow, the accuracy of the pre-trained classifier was measured on the currently flowed data.
• Results: As the data flowed for sufficient iterations, the classifier accuracy on the flowed data consistently reached 100%. This demonstrated that the flow successfully transformed the source dataset into the target domain while preserving the discriminability of its classes. Visual inspection of the flowed images confirmed that they adopted the aesthetic characteristics of the target dataset (e.g., t-shirts from Fashion MNIST were transformed into images resembling MNIST '1's), and images from the same original class converged to the same target class, showcasing the coherence of the transformation.

3. Application: Dataset Distillation

• Setup: Dataset distillation aims to generate a smaller, synthetic dataset from a large one, such that a neural network trained on the synthetic subset achieves comparable performance to one trained on the full original dataset.
• The proposed method minimized MMD with the Riesz Sliced Wasserstein kernel to distill all classes simultaneously.
• Comparison: The results were compared against a "distribution matching" method, which minimizes MMD between each class independently, often incorporating random augmentation and embeddings.
• Results: The proposed WoW gradient flow method achieved "pretty comparable" results to the distribution matching baseline. This is noteworthy as the WoW method performs distillation across all classes simultaneously, which is generally considered a more complex and integrated approach than class-by-class distillation.

4. Application: Data Augmentation

• Setup: This task addresses scenarios where a small target dataset leads to poor neural network performance. The idea is to augment this small dataset by flowing a large, known source dataset towards it and then combining the original and flowed samples for training.
• The proposed method again minimized MMD with the Riesz Sliced Wasserstein kernel to perform this flow.
• Comparison: Results were compared against the OTDD method and other related techniques.
• Results: The WoW gradient flow approach yielded "better performances" compared to OTDD, while simultaneously maintaining a significantly lower computational complexity (O(c^2 n log n) vs. O(n^2 c^2)). This highlights the practical advantage of the new framework in improving model performance on data-scarce tasks without incurring prohibitive computational costs.

Practical Implications

▶ Watch: Scaling to image datasets: MNIST and CIFAR-10 (8:40)

The introduction of Wasserstein over Wasserstein (WoW) gradient flows for labeled datasets holds significant practical implications across several domains of machine learning, particularly for practitioners, infrastructure teams, and model builders dealing with complex data transformations.

1. Enhanced Domain Adaptation and Transfer Learning: The framework provides a principled and efficient mechanism for transforming data from a source domain to a target domain while explicitly preserving the semantic structure of individual classes and their inter-relationships. This is crucial for domain adaptation scenarios where models need to generalize across different data distributions. By flowing a source dataset towards a target, practitioners can generate synthetic data that bridges the domain gap, potentially reducing the need for extensive retraining or fine-tuning.

2. More Efficient Dataset Distillation: The ability to distill all classes simultaneously into a smaller, representative dataset offers a compelling alternative to existing methods. For large-scale datasets where training neural networks is computationally expensive, this method can generate compact surrogates that retain the essential information, leading to faster model prototyping, reduced storage requirements, and more efficient experimentation cycles. The demonstrated comparable performance to "distribution matching" methods, combined with a more integrated approach, suggests a powerful tool for dataset summarization.

3. Advanced Data Augmentation Strategies: When faced with small target datasets, traditional data augmentation techniques often involve simple transformations or generative models that might not fully capture the semantic coherence with the target data. This work enables the generation of highly specific, semantically relevant synthetic data by flowing a large source dataset towards the small target. This can significantly boost model performance in low-data regimes, offering a structured way to enrich training sets.

4. Scalability and Computational Efficiency: One of the most critical practical advantages is the improved computational complexity. By leveraging Maximum Mean Discrepancy (MMD) with Sliced Wasserstein (SW) kernels, the proposed method achieves a complexity of O(c^2 n log n), which is a substantial reduction compared to the O(n^2 c^2) of OTDD. This makes the approach far more scalable and applicable to larger datasets and more numerous classes, which is a common requirement in real-world ML deployments. Infrastructure teams can benefit from this efficiency, requiring fewer computational resources for dataset manipulation tasks.

Tradeoffs and Limitations:

• Formal vs. "True" Riemannian Manifold: As highlighted in the Q&A, the space of probability distributions over probability distributions, when endowed with the WoW distance, possesses a formal Riemannian structure, allowing for the definition of geometric concepts like exponential maps and gradients. However, it might not be a "true" Riemannian manifold in the classical sense, which could imply certain theoretical limitations or nuances in its geometric properties compared to simpler manifolds.
• Kernel Selection: The performance of MMD is dependent on the choice of a suitable positive definite kernel. Selecting the optimal kernel for a given task and data distribution can require expertise and experimentation.
• Class Correspondence: While the flow maintains class structure, determining the explicit one-to-one correspondence between classes of the source and target datasets after the flow still requires an additional optimal transport problem to be solved at the end, as noted by the speakers.
• Smoothness Assumptions: The computation of the WoW gradient relies on certain smoothness assumptions regarding the chosen kernel, which might restrict the choice of kernels in some applications.

Despite these considerations, the framework offers a robust and theoretically grounded approach to handling complex labeled dataset transformations, paving the way for more sophisticated and efficient data-centric AI methodologies.

Key Takeaways

• Labeled datasets can be effectively represented as probability distributions over probability distributions, offering a hierarchical view that captures both intra-class and inter-class relationships.
• The Wasserstein over Wasserstein (WoW) distance endows this complex space with a formal Riemannian structure, enabling the definition of gradients and the application of Riemannian gradient descent.
• Minimizing Maximum Mean Discrepancy (MMD) with Sliced Wasserstein (SW) kernels allows for efficient computation of WoW gradient flows, achieving a significantly improved complexity of O(c^2 n log n) compared to prior methods like OTDD.
• The framework facilitates multi-level coherent transformations, ensuring samples within classes maintain structure while classes interact and converge towards target distributions.
• It demonstrates strong performance in practical machine learning tasks, including dataset distillation and data augmentation, offering a scalable and principled approach for structured dataset manipulation.
• This work provides a powerful tool for tasks like domain adaptation and transfer learning, enabling the generation of semantically consistent synthetic data.

About the Speaker(s)

The talk "Flowing Datasets with Wasserstein over Wasserstein Gradient Flows" was presented by Clément Bonet, Christophe Vauthier, and Anna Korba. Based on the content of the transcript, their expertise lies in machine learning, optimal transport theory, and the application of geometric methods to data analysis. Further specific biographical details regarding their affiliations or academic backgrounds were not provided within the scope of the talk transcript.