REVIEW 4 major objections 6 minor 1 cited by
Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
T0 review · 4 major / 6 minor · reviewed 2026-08-07 · deepseek-v4-flash
Pith's one-line read The paper builds gradient flows on the space of distributions over distributions and moves whole labeled datasets, class by class, toward targets — demonstrated on domain adaptation, dataset distillation, and few-shot learning.
desk verdict A genuinely new dataset-flow framework over Wasserstein space with solid experiments, but the theoretical well-posedness results do not cover the MMD objectives used in practice. read the letter →
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
The reading
What carries the argument
The load-bearing object is the Wasserstein-over-Wasserstein distance $WW_2(P,Q)^2 = \inf_{\Gamma \in \Pi(P,Q)} \int W_2^2(\mu,\nu)\,d\Gamma(\mu,\nu)$ — a Wasserstein distance whose ground cost is the squared Wasserstein distance between measures $\mu,\nu$ on $M$ — together with the lifted inverse-exponential formalism $\exp^{-1}_P(Q)$ that records not only which measure moves to which, but along which geodesic. Two facts carry the argument: for $P$ absolutely continuous with respect to a reference measure, the optimal plan is a transport map $(\mu \mapsto (Id - \nabla\varphi_{\mu,T(\mu)})_\#\mu)_\#P$ giving explicit geodesics, and the WoW gradient of potential and interaction energies is computed level by level ($\nabla_{WW_2}\mathcal{V}(P)(\mu) = \nabla_{W_2}F(\mu)$). For the proposed MMD objective with Sliced-Wasserstein kernel $K(\mu,\nu) = -SW_2(\mu,\nu)$ or the Gaussian SW kernel, the drift at level $\mu$ is $\int \nabla_{W_2}K_\nu(\mu)\,d(P-Q)(\nu)$, evaluated in practice on discrete class measures by rescaling the Euclidean autodiff gradient by $Cn$. The forward scheme $P_{k+1} = \exp_{P_k}(-\tau \nabla_{WW_2}F(P_k))$ then reads, at the particle level, $x^c_{i,k+1} = x^c_{i,k} - \tau \nabla_{WW_2}F(P_k)(\mu^c_k)(x^c_{i,k})$, which is the concrete two-level interaction the experiments exploit.
What would settle it
Run the JKO scheme (13) and the forward-Euler scheme (15) on a target made of genuinely continuous (absolutely continuous) class densities, such as the three smooth rings of the paper's Figure 1, initialized from discrete empirical class measures; if the JKO objective (13) fails to attain its minimum, or the Euler trajectory does not converge to the JKO limit as $\tau \to 0$, the claimed well-posedness fails exactly at the absolute-continuity requirement. The paper's own appendix already shows the complementary failure mode: with four source classes and three target classes, the pure WoW gradient flow cannot match the target because no mass can split.
Extended reading notes
Core claim
The paper establishes that the space $P_2(P_2(M))$ of probability measures over probability measures carries a differential structure for the Wasserstein-over-Wasserstein distance $WW_2$: tangent spaces spanned by Wasserstein gradients of cylinder functions, sub- and super-differentials with a first-order Taylor expansion along optimal transport plans, and velocity fields satisfying a continuity equation for absolutely continuous curves. In this setting, WoW gradient flows of functionals that are $\lambda$-convex along generalized geodesics exist and are unique as limits of the JKO scheme, and a forward-Euler discretization yields a two-level particle system in which each sample moves under both intra-class and inter-class interactions. The paper flows labeled datasets, represented as uniform mixtures of class-conditional empirical distributions, toward target datasets, and reports that classes form, migrate, and match target classes one-to-one. It introduces an MMD objective over distributions with Sliced-Wasserstein based kernels, derives its WoW gradient, and uses the framework for domain adaptation, dataset distillation, and k-shot transfer learning, matching or exceeding prior dataset-flow baselines.
Load-bearing premise
The existence theory only applies when the inner measures are absolutely continuous (density-bearing), while every experiment flows discrete empirical class distributions in $\mathbb{R}^d$ — and the forward Euler scheme used in practice carries no convergence guarantee in that discrete setting.
Editorial extensions
If this is right
- A dataset can be transported as a single object with one-to-one class correspondence preserved; a classifier trained on MNIST reaches roughly 100% accuracy on other *NIST datasets flowed toward MNIST.
- The Sliced-Wasserstein Riesz kernel needs no bandwidth tuning, and the resulting MMDSW objective outperforms per-class Distribution Matching for dataset distillation in the ambient space in the reported experiments.
- For k-shot transfer, flowing a large labeled source toward a small target augments each target class with synthetic samples, improving accuracy over training on the target alone for $k \in \{1, 5, 10, 100\}$.
- Each update costs $O(C^2 L n(\log n + d))$ with $L$ projection directions, and the paper reports runtimes orders of magnitude below the OTDD and product-kernel MMD baselines in its implementations.
- For any functional with a closed form on discrete class measures, the WoW gradient is computable by rescaling the particle autodiff gradient by $Cn$, so the framework inherits standard Wasserstein-gradient-descent tooling.
Reading between the lines
- The same two-level differential structure applies to any mixture family over a Wasserstein space — Gaussian mixtures, topic mixtures, hierarchical clusterings — so the WoW gradient of free energies could serve as a unified optimization recipe for variational inference, barycenters, and generative modeling; the appendix already sketches the Gaussian-mixture variational-inference case.
- All experiments use discrete empirical class distributions, while the well-posedness theorems require absolutely continuous inner measures; a convergence proof for the forward Euler scheme on empirical measures, or a quantitative bound in samples-per-class, would close the gap between algorithm and theory.
- The Riesz Sliced-Wasserstein kernel is not positive definite, so the 'MMD' interpretation is only formal; an ablation that keeps the class-interaction structure but swaps in provably positive-definite SW kernels would separate the effect of interactions from the effect of kernel choice.
- Pure WoW gradient descent cannot split mass, so a source with more classes than a target cannot be matched (the appendix shows this failure on rings); extending the flow with birth-death or weight-adjusting steps would handle unbalanced class counts, and the appendix's semi-relaxed transport scheme is a first step in that direction.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The paper introduces a framework for gradient flows on the space of probability distributions over probability distributions (random measures), endowed with the Wasserstein-over-Wasserstein (WoW) distance. It defines a differential structure (sub/super-differentials, tangent space), derives WoW gradients for potential and interaction energies, and states a JKO-based existence/uniqueness result (Proposition 4.1) for functionals that are λ-convex along generalized geodesics. The paper then proposes an MMD objective with sliced-Wasserstein kernels, computes its WoW gradient, and uses a forward Euler (Riemannian gradient descent) scheme to flow labeled datasets represented as mixtures of class-conditional distributions. Experiments cover synthetic rings, domain adaptation (MNIST/FMNIST/KMNIST/USPS, SVHN/CIFAR10), dataset distillation, and transfer learning, with comparisons to OTDD and distribution-matching baselines.
Significance. If the theoretical gaps are repaired, this is a significant contribution: it is the first systematic differential calculus and gradient-flow construction on the double Wasserstein space, and it demonstrates a practical use for dataset-level optimization, going beyond the Gaussian-conditional models of Alvarez-Melis & Fusi (2021) and Hua et al. (2023). The paper ships code, reports runtimes, and includes useful ablations (kernel bandwidth, number of SW projections, momentum), which are strengths. However, the central well-posedness claim is currently not fully supported for the MMD objectives actually used, and the experiments run in a discrete regime outside the paper's stated absolute-continuity assumptions. The empirical results are promising, but the gap between the theory and the implemented scheme needs to be addressed before the claims as written can be accepted.
major comments (4)
- [§4.1, Appendix C.6] The proof of Proposition 4.1 applies (Ambrosio et al., 2008, Theorem 4.0.4) after checking only λ-convexity of the JKO functional. The remaining hypotheses of that theorem — properness, lower semi-continuity, and crucially the compactness/coercivity needed to ensure existence of minimizers in Eq. (13) — are not verified for (P2(P2(Rd)), WW2). This space is not locally compact, and no tightness or second-moment compactness argument is supplied, so the existence of JKO minimizers does not follow from the cited theorem as stated. Moreover, the λ-convexity notion in Eq. (14) is only defined for generalized geodesics anchored at P ∈ P2(P2,ac(Rd)); the JKO iteration minimizes over P2(P2,ac(Rd)) but the minimizer is not shown to remain absolutely continuous, so the anchoring assumption is not closed under iteration. This is a load-bearing gap for the 'well-posed flows' claim in the abstract and conclusion.
- [§4.2, Proposition 4.1] The MMD functional in Eq. (18) with the Gaussian SW kernel K(μ,ν)=exp(−SW2²(μ,ν)/(2h)) or the Riesz SW kernel K(μ,ν)=−SW2(μ,ν) is never shown to satisfy the λ-convexity condition (14). Proposition B.8 requires the potential F(μ)=−∫K(μ,ν)dQ(ν) to be λ-convex along generalized geodesics of P2(Rd), and Proposition B.9 requires K to be jointly convex along generalized geodesics; neither is established for the SW-based kernels, and for the Riesz kernel the required convexity is plainly false. Consequently, Proposition 4.1 does not cover the objective minimized in Section 5, and the forward Euler scheme (17)—which is not a JKO limit and for which no convergence proof is given—cannot inherit well-posedness from it. The paper should either prove the required convexity for these kernels, give an alternative existence/consistency argument for the discrete forward scheme, or explicitly restrict the well-posedness claim to functionals satisfying (14).
- [§5, Eq. (17); §B.5] All experiments use P = (1/C)∑ δ_{μ_{c,n}} with empirical class distributions μ_{c,n}, which are atomic and hence not in P2(P2,ac(Rd)). Since Proposition 4.1 and Proposition B.10 both require absolute continuity of the inner measures (and the latter also requires P, Q, O ∈ P2(P2,ac(Rd))), the theoretical existence and uniqueness results do not apply to the implemented scheme. The paper's conclusion mentions 'quantitative guarantees' as future work, but the main text should state this limitation explicitly and, ideally, provide a consistency or approximation argument for empirical measures. As written, the central claim that the implemented MMDSW flows are instances of well-posed WoW gradient flows is not supported by the theorems in the paper.
- [Appendix D.1] The statement that the Riesz SW kernel K(μ,ν)=−SW2(μ,ν) is 'not positive definite (but conditionally positive definite)' is asserted without proof or reference. Since the MMDSW loss is interpreted as an MMD in Eq. (18), and conditional positive definiteness is exactly what guarantees nonnegativity of the squared MMD for signed measures of zero total mass such as P−Q, a proof or a precise citation should be provided. If the claim is incorrect, the interpretation of the objective and the minimization results would need to be revisited.
minor comments (6)
- [§4.1, Eq. (13)] The notation 'P2(Pac(Rd))' in Eq. (13) is inconsistent with the notation 'P2(P2,ac(Rd))' used elsewhere; please unify.
- [Appendix C.16, C.17] There are typos in the definitions of generalized geodesics: 'π1,2# Γ ∈ Πo(P, Q) a,d π1,3# Γ' should read 'and' in both places.
- [§5, first paragraph] The claim that 'we are the first to represent labeled datasets this way' is too strong; the conditional-distribution embeddings of Alvarez-Melis & Fusi (2020) and Hua et al. (2023) are closely related. Please qualify the novelty or discuss the difference in representation more precisely.
- [§4.2, paragraph on Riesz kernel] The statement that the Riesz SW kernel 'does not require tuning a bandwidth h' is true but potentially misleading: the step size τ and momentum m still require tuning, and the experiments do tune them.
- [Figure 3] The accuracy curves would benefit from shaded standard-deviation bands; the text reports averages over 3 flows, but the figure appears to show only mean values.
- [§6, Conclusion] The conclusion defers 'quantitative guarantees on the convergence of these schemes'; in the revised paper, please state clearly at the outset which results are formal theorems and which parts of the pipeline are heuristic or empirical.
Circularity Check
No significant circularity: the theoretical construction is self-contained, and the numerical results are benchmarked against external methods and datasets.
full rationale
The paper's derivation chain is: (1) define the Wasserstein-over-Wasserstein space and a differential structure; (2) derive WoW gradients for potential and interaction energies and for the MMD with sliced-Wasserstein kernels; (3) prove existence of JKO gradient flows by verifying the hypotheses of an external theorem (Ambrosio, Gigli, Savaré, Theorem 4.0.4); (4) propose a forward-Euler discretization; and (5) validate on transfer learning and dataset distillation against external baselines such as OTDD and the MMD of Hua et al. Each of these steps is supported either by self-contained proofs in the appendix or by independent, externally published results (McCann, Gigli, Erbar, Emami-Pass, Bonnotte, Hertrich et al.). The self-citations that appear (Bonet et al. 2024, 2025) are used only as references for standard Riemannian Wasserstein gradient descent and sliced-Wasserstein flows; they are parameter-free external tools, not hidden assumptions equivalent to the present conclusions. The MMD objective and its SW gradient are standard constructions, and no fitted parameter is relabeled as a prediction. The acknowledged gaps—lack of a converse continuity-equation result, restriction to compact manifolds, no quantitative convergence guarantees for the forward scheme, and the fact that Proposition 4.1's λ-convexity hypotheses are not verified for the SW kernels used in experiments—are correctness/completeness concerns, not circularity. They indicate that the implemented algorithm is not fully covered by the existence theorem, but this is an unproven bridge rather than an equation reducing to its own input.
Assumptions & free parameters
free parameters (5)
- Gaussian SW kernel bandwidth h =
0.05 for synthetic rings, other values in ablations
- Step size tau =
0.05, 0.1, or 1.0 depending on experiment
- Momentum m =
0.9
- Number of SW projections L =
500 for main image experiments, 100 for 2D rings
- Samples per class n or p =
n=200 or 500, p=1, 10, 50
assumptions (6)
- standard math McCann's theorem and Brenier's theorem guarantee existence and uniqueness of OT maps for absolutely continuous measures on compact Riemannian manifolds.
- standard math The space (P2(M), W2) has a Riemannian structure with tangent spaces, exponential maps, and Wasserstein gradients as developed by Otto, Erbar, and Gigli.
- domain assumption The reference measure P0 satisfies Assumption B.2: no atoms, an integration by parts formula, and quasi-invariance under smooth vector field flows.
- domain assumption The manifold satisfies the smooth transport property (Assumption C.9), a strong Ma-Trudinger-Wang type condition, to ensure WoW gradients can be chosen in the tangent space.
- domain assumption Inner measures are absolutely continuous, i.e. P lies in P2(P2,ac(Rd)), for the JKO scheme and generalized geodesic convexity results.
- domain assumption The functional F is proper, coercive, lower semicontinuous, and lambda-convex along generalized geodesics for the JKO existence theorem.
Cite this review
Pith. "Pith review of Flowing Datasets with Wasserstein over Wasserstein Gradient Flows." pith.science (2026). https://pith.science/paper/YOZ6DQNY
@misc{pith2026250607534,
author = {Pith},
title = {Pith review of: Flowing Datasets with Wasserstein over Wasserstein Gradient Flows},
year = {2026},
howpublished = {\url{https://pith.science/paper/YOZ6DQNY}},
note = {Machine review of arXiv:2506.07534}
}
read the original 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.
Figures
Figures from the paper (15 more)
Forward citations
Cited by 1 Pith paper
-
Totally convex functions, $L^2$-Optimal transport for laws of random measures, and solution to the Monge problem
Using totally convex functionals and their Lagrangian liftings, the authors prove existence and uniqueness of Monge solutions for optimal transport between super-regular laws of random measures.
Reference graph
Works this paper leans on
-
[6]
Wasserstein Flow Matching: Generative modeling over families of distributions
(Cited on p. 1) Gigli, N. On the geometry of the space of probability mea- sures in Rn endowed with the quadratic optimal transport distance. PhD thesis, Scuola Normale Superiore, 2004. (Cited on p. 3) Gigli, N. On the inverse implication of Brenier-McCann theorems and the structure of(P2(M ), W2). Methods and Applications of Analysis, 18(2):127–158, 2011...
work page Pith review arXiv 2004
-
[14]
(i.e. we sample one augmentation among color jittering, cropping, cutout, scaling and a rotation for MNIST, and also add flipping for Fashion MNIST). We optimize (241) by stochastic gradient descent over the particles, sampling one random network and one random augmentation at each step. We trained it for 20K iterations, a learning rate ofτ = 1 and a mome...
work page 2021
-
[338]
Springer, 2009. (Cited on p. 16, 24, 34) Vilnis, L. and McCallum, A. Word Representations via Gaussian Embedding. In Bengio, Y . and LeCun, Y . (eds.), 3rd International Conference on Learning Representa- tions, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015. (Cited on p. 1) von Renesse, M.-K. and Sturm, K.-T. Entropic Me...
arXiv 2009
-
[1725]
PMLR, 2019. (Cited on p. 1) Dusson, G., Ehrlacher, V ., and Nouaime, N. A Wasserstein- type metric for generic mixture models, including location-scatter and group invariant measures. arXiv preprint arXiv:2301.07963, 2023. (Cited on p. 1) Dvurechensky, P., Gasnikov, A., and Kroshnin, A. Com- putational Optimal Transport: Complexity by Acceler- ated Gradie...
work page Pith review arXiv 2019
-
[2006]
(Cited on p. 36) Lim, J. N. and Johansen, A. Particle Semi-Implicit Varia- tional Inference. Advances in Neural Information Pro- cessing Systems, 37:123954–123990, 2024. (Cited on p. 57) Liu, X., Bai, Y ., Lu, Y ., Soltoggio, A., and Kolouri, S. Wasserstein Task Embedding for Measuring Task Simi- larities. Neural Networks, 181:106796, 2025. (Cited on p. 5...
arXiv 2024
- [2010]
-
[2015]
1) Hertrich, J., Gr¨af, M., Beinert, R., and Steidl, G
(Cited on p. 1) Hertrich, J., Gr¨af, M., Beinert, R., and Steidl, G. Wasser- stein Steepest Descent Flows of Discrepancies with Riesz Kernels. Journal of Mathematical Analysis and Applica- tions, 531(1):127829, 2024a. (Cited on p. 2, 3) Hertrich, J., Wald, C., Altekr ¨uger, F., and Hagemann, P. Generative Sliced MMD Flows with Riesz Kernels. In The Twelft...
work page 2017
-
[2017]
1) Rabin, J., Peyr´e, G., Delon, J., and Bernot, M
(Cited on p. 1) Rabin, J., Peyr´e, G., Delon, J., and Bernot, M. Wasserstein Barycenter and its Application to Texture Mixing. In Scale Space and Variational Methods in Computer Vision: Third International Conference, SSVM 2011, Ein-Gedi, Israel, May 29–June 2, 2011, Revised Selected Papers 3, pp. 435–446. Springer, 2012. (Cited on p. 6, 44) Rønning, O., ...
arXiv 2011
Show all 14 references
-
[2018]
7) Courty, N., Flamary, R., Tuia, D., and Rakotomamonjy, A
(Cited on p. 7) Courty, N., Flamary, R., Tuia, D., and Rakotomamonjy, A. Optimal Transport for Domain Adaptation. IEEE transactions on pattern analysis and machine intelligence, 39(9):1853–1865, 2016. (Cited on p. 8) Csisz´ar, I. Information-Type Measures of Difference of Prob...
2016 arXiv
-
[2019]
52) Lu, Y ., Slepˇcev, D., and Wang, L
(Cited on p. 52) Lu, Y ., Slepˇcev, D., and Wang, L. Birth–death dynamics for sampling: global convergence, approximations and their asymptotics. Nonlinearity, 36(11):5731, 2023. (Cited on p. 52) Manupriya, P., Jagarlapudi, S., and Jawanpuria, P. MMD- Regularized Unbalanced Op...
2023 arXiv
-
[2021]
differential
(Cited on p. 8, 50) 14 Flowing Datasets with Wasserstein over Wasserstein Gradient Flows Zhao, B. and Bilen, H. Dataset Condensation with Dis- tribution Matching. In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision, pp. 6514–6523, 2023. (Cited o...
2025
-
[2023]
53) Bellazzi, R., Codegoni, A., Gualandi, S., Nicora, G., and Vercesi, E
(Cited on p. 53) Bellazzi, R., Codegoni, A., Gualandi, S., Nicora, G., and Vercesi, E. The Gene Mover’s Distance: Single- cell similarity via Optimal Transport. arXiv preprint arXiv:2102.01218, 2021. (Cited on p. 1) Bojchevski, A. and G ¨unnemann, S. Deep Gaussian Em- bedding ...
2021 arXiv
-
[2024]
1) Chen, Y ., Georgiou, T
(Cited on p. 1) Chen, Y ., Georgiou, T. T., and Tannenbaum, A. Optimal Transport for Gaussian Mixture Models. IEEE Access, 7: 6269–6278, 2018. (Cited on p. 1, 55) Chen, Z., Mustafi, A., Glaser, P., Korba, A., Gretton, A., and Sriperumbudur, B. K. (De)-regularized Maxi- mum Mea...
2018
-
[2025]
2, 3, 19, 54) Bonnabel, S
(Cited on p. 2, 3, 19, 54) Bonnabel, S. Stochastic Gradient Descent on Riemannian Manifolds. IEEE Transactions on Automatic Control, 58 (9):2217–2229, 2013. (Cited on p. 3, 19) Bonneel, N., Rabin, J., Peyr´e, G., and Pfister, H. Sliced and Radon Wasserstein Barycenters of Meas...
2013 arXiv
Reviewed August 7, 2026 · model on record in the stance chip above.
Discussion (0). Sign in to comment.