Research map / Probabilistic and scientific modeling / Variational and generative sampling
Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Defines Wasserstein-over-Wasserstein (WoW) gradient flows over labeled datasets viewed as measures over class-conditional distributions, and applies a sliced-Wasserstein MMD functional (MMDSW) to domain adaptation, dataset distillation and transfer learning.
Problem
Designing tractable gradient flows over probability measures on probability measures, so that labeled datasets (mixtures of class-conditional distributions) can be transported for domain adaptation, transfer learning and dataset distillation without Gaussian modeling of classes.
Approach
Lift the Wasserstein gradient-flow machinery to the WoW space (distributions over distributions with W2 as ground cost), define WoW gradients/tangent spaces and a forward Euler (Wasserstein gradient descent) scheme, and minimize an MMD with Sliced-Wasserstein-based kernels between labeled datasets.
Key results
- On dataset distillation, MMDSW outperforms Distribution Matching in the ambient space on MNIST (e.g., 66.5 vs 61.1 at p=1, 93.2 vs 88.2 at p=10).
- In transfer learning from MNIST to Fashion-MNIST with k=1 shot, MMDSW reaches 40.5% accuracy vs 26.0% for training on the target alone.
- MMDSW flow runs in about 14 seconds for k=1 transfer, much faster than OTDD (294.5 s) in the reported runtime table.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Augment target classes with samples generated by flowing MNIST using MMDSW | LeNet5 transfer learning · MNIST to Fashion-MNIST, k-shot | accuracy: 40.5 ± 4.7 vs 26.0 ± 5.3; 61.5 ± 4.6 vs 38.5 ± 6.7; 65.4 ± 1.5 vs 53.9 ± 7.9; 74.7 ± 0.8 vs 71.1 ± 1.5 | Mixed |
| Augment target classes with samples generated by flowing MNIST using MMDSW | LeNet5 transfer learning · MNIST to KMNIST, k-shot | accuracy: 20.9 ± 2.0 vs 18.4 ± 3.1; 37.4 ± 2.2 vs 25.9 ± 4.0; 44.7 ± 1.8 vs 30.9 ± 4.6; 66.8 ± 0.8 vs 60.1 ± 1.1 | Mixed |
| Augment target classes with samples generated by flowing MNIST using MMDSW | LeNet5 transfer learning · MNIST to USPS, k-shot | accuracy: 37.4 ± 6.1 vs 32.4 ± 7.9; 73.0 ± 1.0 vs 51.4 ± 9.8; 77.2 ± 1.2 vs 60.3 ± 10.1; 89.7 ± 0.4 vs 87.5 ± 0.7 | Mixed |
| Use MMDSW rather than training directly on the target data | LeNet5 transfer learning · MNIST-to-target transfer datasets | runtime (seconds) | Helped |
| Use a Sliced-Wasserstein Riesz kernel rather than Distribution Matching in the ambient space | Dataset distillation · MNIST and Fashion-MNIST | classifier accuracy: MNIST, p=10: 88.2 ± 2.8 -> 93.2 ± 0.7; FMNIST, p=10: 74.6 ± 1.0 -> 76.7 ± 1.0 | Helped |
| Add embedding and augmentation to MMDSW dataset distillation (ablation) | MMDSW · MNIST and Fashion-MNIST, p=10 | classifier accuracy: MNIST: 96.4 ± 0.2 -> 96.4 ± 0.3; FMNIST: 78.0 ± 0.9 -> 78.9 ± 1.2 | Hurt |
Surprising findings
- The Riesz SW kernel -SW2 is not positive definite yet works well in practice and needs no bandwidth tuning.
- In dataset distillation, the embedding-free MMDSW does not always beat DM once an embedding is added, with comparable results instead.
Limitations
- Theory assumes compact connected manifolds and absolutely continuous measures, while experiments are on R^d.
- Domain adaptation experiments assume pseudo-labels or class labels for the flowed dataset are available.
- No convergence guarantees for the forward scheme; continuity equation for WoW flows left open.
Experimental setup
Datasets: MNIST Fashion-MNIST KMNIST USPS CIFAR10 SVHN
Benchmarks: MNIST Fashion-MNIST KMNIST USPS CIFAR10 SVHN
Metrics: classifier accuracy runtime in seconds domain adaptation accuracy along flow
Baselines: Distribution Matching (DM) OTDD MMD with product kernel (Alvarez-Melis et al.) Random subset Train on target data
Builds on: Wasserstein gradient flows JKO scheme Riemannian Wasserstein gradient descent distribution matching
Closest papers by meaning
- Dataset Dynamics via Gradient Flows in Probability Space (ICML 2021 · similarity 0.92)
- Wasserstein Flow Matching: Generative Modeling Over Families of Distributions (ICML 2025 · similarity 0.88)
- Generalized Wasserstein Flow Matching: Transport Plans, Everywhere, All at Once (NeurIPS 2026 · similarity 0.88)
- Scalable Wasserstein Gradient Flow for Generative Modeling through Unbalanced Optimal Transport (ICML 2024 · similarity 0.87)
- Wasserstein Transfer Learning (NeurIPS 2025 · similarity 0.86)
- Minimizing $f$-Divergences by Interpolating Velocity Fields (ICML 2024 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.