Flowing Datasets with Wasserstein over Wasserstein Gradient Flows

Thank you for your appraisal and positive comments on our paper. We address your comments below.

Domain Specification of the Functional

For most of ML applications where we aim at minimizing distances, we agree that the codomain is $\mathbb{R}_+\cup \\{+\infty\\}$. Nonetheless, the differential structure holds also for $\mathbb{R}$ as codomain, which may cover e.g. potential energies with $V$ negative.

Higher-Order Terms

The $o(W_{W_2})$ notation can be defined using the Landau notation. The main difficulty when using these terms in the article is whether we can integrate them when using the Taylor expansion. This can be handled using regularity hypothesis as bounded Hessian.

Continuity Equations in WoW Framework

We expect a similar continuity equation $\partial\_t \mathbb{P}\_t = \mathrm{div}(\mathbb{P}\_t \nabla_{W\_{W\_2}}\mathbb{F}(\mathbb{P}\_t))$. In Proposition 3.7, we took a first step by showing that AC curves satisfy a continuity equation. However, to derive the equation of the WoW gradient flow, we would need to show that the minimizing movement scheme converges towards the right equation (see e.g. (Santambrogio, 2017)), and to define an appropriate notion of divergence operator on this space, e.g. the one proposed in (Schiavo, 2020). We leave these questions for future work.

Comparisons with relevant baselines.

Thank you for these suggestions. KALE flows and the flows of the KKL divergence are not designed to optimize over $\mathcal{P}(\mathcal{P}(\mathbb{R}^d))$. In Figure 4, we showed an example of the flow of the MMD with Riesz kernel, which is designed on $\mathcal{P}(\mathbb{R}^d)$ as KALE flows and the KKL divergence.

Related works.

Thank you for pointing to us these related works which we were not aware of.

Limitations and future research directions.

We will add a paragraph in the revised version of the paper indicating future research directions (e.g. using other kernels, f-divergences or continuous labels) and limitations (e.g. theory on compacts manifolds, lack of continuity equation for the flow, dimension of xps).

This work mainly focus on the MMD functional.

We focus on the MMD functional as it can be decomposed as a potential and an interaction energy, which we know how to differentiate in the WoW space (see Section 4.2).

Using this approach for f-divergences could be of great interest to do e.g. sampling. We plan to tackle this in future works, but we note that it is not straightforward, as we would need to identify a base measure w.r.t which the measures need to be AC (see e.g. (Schiavo, 2020)), and then compute the WoW gradient of these functionals.

Adding a GIF for 3 ring case.

Thank you for this suggestion. We will add https://ibb.co/5qgjhgC.

1.Labels not categorical?

Handling continuous labels is an interesting extension of our method. This can be viewed as conditioning with continuous labels, akin to an infinite mixture. More formally, we could define it as $\mathbb{P}=\int \delta_{\mu_y}\mathrm{d}\lambda(y)$ with $\mu_y$ the conditional distribution given the label $y$ and $\lambda$ a distribution over the labels. For a discrete distribution $\lambda$, this would recover the discrete case we use in the paper. In practice, we would need to discretize $\lambda$.

2. Why the authors consider the SW distance rather than wasserstein or sinkhorn?

We consider SW for two reasons:

1. It is much faster to compute than the Wasserstein or Sinkhorn distance (complexity of $O(Ln\log n)$ for SW versus $O(n^3\log n)$ for Wasserstein and $O(n^2\log n/\varepsilon^2)$ for Sinkhorn, see e.g. [1]).
2. SW allows to define valid positive definite kernels [2], which is not the case for the Wasserstein distance [1].

[1] Peyré, G., & Cuturi, M. Computational optimal transport: With applications to data science. Foundations and Trends in Machine Learning, 2019.

[2] Kolouri, S., Zou, Y., & Rohde, G. K. Sliced Wasserstein kernels for probability distributions. IEEE Conference on Computer Vision and Pattern Recognition. 2016.

3. The loss curve in the Rings.ipynb appears to fluctuate.

The loss curve may fluctuate as SW is estimated through a Monte-Carlo approximation. We observed empirically that for a sufficient number of iterations and with well chosen hyperparameters, it converges well and it fluctuates around small values, close to the minimum.

4. Is it possible to extend the proposed approach to constrained support?

The proposed flow might be combined with methods to constrain the support such as Mirror Descent, which has been recently extended to Wasserstein gradient flows in [1], or using barrier methods as in [2].

[1] Bonet, C., Uscidda, T., David, A., Aubin-Frankowski, P. C., & Korba, A. Mirror and preconditioned gradient descent in wasserstein space. NeurIPS 2024.

[2] Li, L., Liu, Q., Korba, A., Yurochkin, M., & Solomon, J. Sampling with mollified interaction energy descent. ICLR 2023.

Thank you for your appraisal and positive comments on our paper. We answer your comments below.

The qualitative results (Figures) are not entirely convincing and do not clearly suggest a continuous flow from one dataset to another.

We observed empirically that the flow goes very fast from the source dataset towards the neighborhood of the target dataset, but takes more time to converge towards clean images of the target dataset. This might explain why the flows do not appear very continuous, e.g. on Figure 2, as the discretization in time is linear. On Figure 9, we reported a finer discretization, which shows slightly better the behaviour.

We will also add in the zip supplementary materials the following gif showing the evolution of the rings (https://ibb.co/5qgjhgC).

How can the method be applied to larger-scale problems?

As the computational complexity of computing the MMD with a Sliced-Wasserstein kernel is in $O(C^2 Ln(\log n + d))$ for datasets with $C$ classes, and $n$ samples in each class, we believe that we can scale this algorithm to larger datasets with more classes and more samples by class.

The bottleneck will probably be for higher dimensional datasets. We tried on a (relatively) higher dimensional dataset CIFAR10 (see Figure 10 in Appendix D.4), and observed that the flow scales well. However, it required a lot more optimization steps to converge (150K steps for CIFAR10 against 18K steps for MNIST) with the same number of Monte-Carlo samples.

To scale to higher dimensional datasets, there could be several solutions. On one hand, one can try to get the number of projections $L$ bigger to get a better approximation of the Sliced-Wasserstein distance, or use better approximations using e.g. control variates as in [1]. We could also use faster optimization algorithms to try to accelerate the speed of convergence or other variants of the Sliced-Wasserstein distance more adapted to images, e.g. using convolution as projections as in [2,3]. We leave these investigations for future works.

We also replicated the domain adaptation experiment (Figure 3) by flowing from the dataset SVHN to CIFAR10, which are both dataset of shape 32x32x3, see https://ibb.co/C3B3Ffty. It demonstrates that it still works in moderate higher dimension. We will put it in Figure 3 in place of the KMNIST to MNIST example.

[1] Leluc, R., Dieuleveut, A., Portier, F., Segers, J., & Zhuman, A. Sliced-Wasserstein estimation with spherical harmonics as control variates. International Conference on Machine Learning (2024).

[2] Nguyen, Khai, and Nhat Ho. "Revisiting sliced Wasserstein on images: From vectorization to convolution." Advances in Neural Information Processing Systems 35 (2022).

[3] Du, C., Li, T., Pang, T., Yan, S., & Lin, M. "Nonparametric generative modeling with conditional sliced-Wasserstein flows." International Conference on Machine Learning (2023).

Typos

Thank you, we corrected the typo.

Thank you for reading the paper and for your feedback. We answer your comments below. Please do not hesitate if you have other questions.

The continuity equation relies on the strong assumption that the velocity field vt is Lipschitz.

In Proposition 3.7, we show that if $(\mathbb{P}\_t)\_t$ is an absolutely continuous curve, then we have $\|v_t\|_{L^2(\mathbb{P}_t)}\le |\mathbb{P}'|(t)$ and the continuity equation. We do not assume that the velocity field is Lipschitz.

The class alignment via 1-NN and majority vote seems heuristic.

Another solution less heuristic to assign the classes is to compute an OT map between $\mathbb{P}=\frac{1}{C}\sum_{c=1}^C \delta_{\mu_c}$ and the target $\mathbb{Q}=\frac{1}{C}\sum_{c=1}^C \delta_{\nu_c}$ with the 2-Wasserstein distance as ground cost. We replicated the results of Figure 3 with this method, and observe that it gives similar results. We will fix this in the revision.

Recent methods like Dataset Condensation with Gradient Matching (DC) are absent from Table 1.

We chose to compare with methods that only optimize a distance to generate new samples, as these are the are the most comparable approaches, and we do not claim to be state of the art for dataset distillation. Thus, we compared in Table 1 with Distribution Matching (DM), introduced in [1]. Notably, [1] already compares DM with DC (see their Table 1), showing that their performances are comparable.

[1] Zhao, B., & Bilen, H. Dataset condensation with distribution matching. WACV 2023.

Ablation study where the target dataset has non-Gaussian class distributions.

As shown in Table 2, our method outperforms in most settings prior works that assume Gaussian class distributions, suggesting that avoiding this assumption leads to better performance on image datasets.

Related work.

Thank you for pointing to us this work, which seems relevant.

Can the method scale to larger datasets beyond MNIST variants?

Please find the answer in the response of Reviewer LUYC.

How sensitive is the method to choices of L, tau, momentum , and kernel bandwidth?

The method is sensitive to the kernel bandwidth when using the Gaussian SW kernel, as we showed in Figure 5 of Appendix D.2. In practice, we use the Riesz SW kernel, which does not require to tune a bandwidth.

As we use a gradient descent method, taking the step size too big will make the scheme diverge.

We use momentum to accelerate the convergence of the scheme, but we notice that it still converges when using no momentum, but slower, see Figure 11 in Appendix D.4.

In low dimension, the method is not very sensitive to the number of projections $L$ (see e.g. the Figure https://ibb.co/SX2gtntd for the ring experiment). In higher dimension such as for MNIST, it is more sensitive, and a relatively big number of projections improves the convergence as it provides a better approximation of the gradient. We will add the following Figure (https://ibb.co/pr2MwMLH) showing results for different values of $L$ for the generation of MNIST samples.

How does the computational cost of your method scale with the number of classes and samples per class?

In our experiments, we focused on the minimization of the MMD with a kernel based on the Sliced-Wasserstein distance, which can be computed in $O(Ln\log n)$ between $\mu_n=\frac{1}{n}\sum_{i=1}^n \delta_{x_i}$ and $\nu_n=\frac{1}{n}\sum_{j=1}^n \delta_{y_j}$. For $C$ classes with each class having $n$ samples, the MMD has therefore a total runtime complexity of $O(C^2 Ln(\log n + d))$. Moreover, note that the pairwise Sliced-Wasserstein distances can be computed in parallel.

In contrast, OTDD requires first to compure $C^2$ pairwise Wasserstein distances which has a complexity of $O(C^2 n^3 \log n)$ if they are computed between the empirical distributions, and of $O(C^2 (nd^2 + d^3))$ if they are computed using the Gaussian approximation. Then, it requires to compute a second OT problem between the $nC$ samples with cost $d\big((x,c), (x,c')\big) = \|x-x'\|^2 + W_2^2(\mu_c,\nu_{c'})$, which has a complexity in $O(n^3C^3\log(nC))$ and can be reduced to $O(n^2C^2\log(nC)/\varepsilon)$ using the entropic regularized problem.

We compared the runtime for the transfer learning experiment in https://ibb.co/4R0GtczT.

The paper assumes compactness and connectivity of the manifold for theoretical results.

We assume compactness and connectivity of the manifold for theoretical purpose, as we notably rely on the seminal results of [1]. Of course, this is not always the case in practical applications, and thus we acknowledge that there is still a gap here to make the theoretical derivations to hold on non-compact spaces. Nonetheless, it is reasonable to assume that the data lie on a compact space.

[1] Schiavo, L. D. A Rademacher-type theorem on L2-Wasserstein spaces over closed Riemannian manifolds. Journal of Functional Analysis, 2020.