Paper: Neural Sinkhorn Gradient Flow (ICLR-2024)

Reviewer comment (Writing):
The article is full with typos. Just to name a few: "piror", "Sinkhron", "Experimrnts", "speedest descent", question mark in the appendix and so on. Please fix those.

Author response:
We have fixed all the typos, thank you!

Reviewer comment (Evaluation):
The authors write "We do not compare with extant neural WGF methods on MNIST because most of the neural WGF methods only show generative power and trajectories on this dataset and lack the criteria to make comparisons." There are several papers (also gradient flow based ones), which evaluate a FID on MNIST. Please provide it as well.

Author response:
In fact, we are very willing to make quantitative comparisons with different methods, but due to limited time and computational resources, we did not replicate related algorithms on MNIST and test FID. The literature we reviewed that conducted MNIST experiments with neural WGF-based methods did not provide relevant data, including references [A] and [B]. However, we found that references [C], [G], and [H] evaluate an FID on MNIST. We have provided both our FID values and the FID values given in reference [C] for your reference. We are also very willing to add more comparisons if you could point us to those papers.

Reviewer comment (Presentation):
Although the CIFAR10 value seems good, there are unfortunately no generated images provided. It is standard practice to sample many images in the appendix.

Author response:
Due to space constraints, we have added sampling images in the Appendix.

Reviewer comment (Theory):
The statement of theorem 2 is incorrect. I guess the authors do not want to sample the Euler scheme (eq 14) but the continuous gradient flow, otherwise the statement would need to depend on the step size $\eta$.

Author response:
We have revised Equation 14 to an ODE (Ordinary Differential Equation) form in the amended version.

Reviewer comment (Theory):
In the proof of Theorem 2: Please provide a proof (or reference) why the mean field limit exists. Or do you mean the gradient flow starting at $\mu_0$ with target $\mu$ (first two sentences).

Author response:
Apologies for the previous lack of clarity. We have re-examined and revised the proof of Theorem 2, and these modifications are included in the revised version of our paper (see Appendix for details). To prove that the empirical distribution $\tilde{\mu}^M_t$ evolving with equation 14 weakly converges to $\mu_t$ which is a solution of the equation 15 when $M\to \infty$ under the assumption $\\tilde{\mu}\_0^M \\rightharpoonup \\mu\_0$, we show that for arbitrary bounded and continuous $h$ and any $t$, 
$\\lim_{M\\to \\infty}\\mathbb{E}\_{\tilde{\mu}^M\_t}[h]\to\mathbb{E}\_{\mu\_t}[h]$. 
We show this property by Induction over the Continuum and analysis to infinitesimal of $\mathbb{E}_{\tilde{\mu}^M_t}[h]$ and $\\mathbb{E}\_{\\mu\_t}[h]$.

Reviewer comment (Theory):
Later in that proof: why does there exists a weakly convergent subsequence of $\mu_t^M$? Further, I cant find the definition of $U_{\mu}$.

Author response:
Apologies for the previous lack of clarity. We have re-examined and revised the proof of Theorem 2, and these modifications are included in the revised version of our paper (see Appendix for details). To prove that the empirical distribution $\tilde{\mu}^M_t$ evolving with equation 14 weakly converges to $\mu_t$ which is a solution of the equation 15 when $M\to \infty$ under the assumption $\\tilde{\mu}\_0^M \\rightharpoonup \\mu\_0$, we show that for arbitrary bounded and continuous $h$ and any $t$, 
$\\lim_{M\\to \\infty}\\mathbb{E}\_{\tilde{\mu}^M\_t}[h]\to\mathbb{E}\_{\mu\_t}[h]$. 
We show this property by Induction over the Continuum and analysis to infinitesimal of $\mathbb{E}_{\tilde{\mu}^M_t}[h]$ and $\\mathbb{E}\_{\\mu\_t}[h]$.

Reviewer comment (Reproducibility):
The code is not runnable, as the model (or any checkpoints) are not provided.

Author response:
We have added a readme and the missing neural network structure components to ensure the completeness of the code.

Reviewer comment (Evaluation):
With the provided code, there are several insights that should be discussed in the paper. In the provided cifar experiments, the number of Gaussian samples used is 50000 samples. This number is extremely low to approximate the semi-discrete OT. Therefore, a discussion regarding the statistical performance of the method is needed in my opinion.

Author response:
In our experiments, we indeed used significantly more than 50,000 Gaussian samples. As detailed in the revised version of our Algorithm 1, we clarified that our method comprises two distinct phases: building the trajectory pool and velocity field matching. The number of Gaussian samples utilized during training is determined in the first phase, while the quantity of trajectory samples used for training the neural velocity fields depends on the second phase. In the first stage, each iteration generates new Sinkhorn WGF paths using freshly resampled Gaussian samples (refer to `main_sinkhorn.py`), which are subsequently employed for vector field matching (see `main_minibatchOT.py`). It's important to note that the parameter of 50,000 Gaussian samples in the code originates from `evaluator.py`, which facilitates the inference process of the NSGF model. For our CIFAR10 experiment, we set this parameter to 50,000 to generate an equivalent number of images for standard FID value testing. We apologize for any confusion this may have caused and have provided detailed descriptions for each file in the supplementary materials of the revised version.

Reviewer comment (Experiments):
As your method requires the simulation of the probability path, I wonder about the training time between your method and the recent Flow Matching approaches which are simulation free.

Author response:
Our algorithm takes longer in training time compared to simulation-free Flow Matching approaches [A]，[B], and [C] (we tested the time required to reach the same $W_2$ in 2D experiments, and the table is attached below), but since our model follows the steepest descent flow, we use fewer steps in the inference process to reduce generation time. We consider this a trade-off, which is discussed in our 2D simulations. 
In our understanding, 'simulation-free' in Flow Matching approaches means that any time snapshot in the Gaussian probability path can be obtained through random interpolation between noise distribution and the target distribution. In the training part of our method, our snapshots come from the discretization of Sinkhorn WGF, obtained through the iterative solving method (see Algorithm 1). In our algorithm, the main complexity of obtaining a snapshot at a specific time $t$ comes from solving the $\mathcal{W}_{\varepsilon}$-potential, with a complexity of $\mathcal{O}(tn^2)$, where $n$ represents the dimension, and $t$ represents the time. Notably, in OT-CFM [B], the authors perform OT matching on the minibatch sample pairs before the stochastic interpolation algorithm, with a complexity of $\mathcal{O}(n^3log(n))$.
It's important to emphasize that although our training process requires more time, this does not impact our inference process. Our experiments demonstrate that we can use fewer steps, which reduces the time needed for generation.

|methods|FM|OTCFM|NSGF|
| ---- | ---- | ---- | -----|
|time|119.11|145.94|1727.59|
Table 1: Time taken by three comparison methods to achieve $W_2<0.2$ in the 8gaussians-moons task, unit: seconds.

Reviewer comment (Writing):
There are many typos in the paper (including in titles: ie ExperimRnts, Notaions) that lead to poor clarity.

Author response:
We have fixed all the typos, thank you!

Reviewer comment (Experiments):
The experiments include two toy datasets (synthetic 2D and MNIST). I would like to know how the method performs on other big datasets (Flowers, CelebA) or on other tasks such as single-cell dynamics [4].

Author response:
Due to time and computational power limitations, we will construct experiments with big datasets and single-cell dynamics in our future work.

Reviewer comment (Writing):
There are minor typos throughout, such as "euclidean" instead of "Euclidean".

Author response:
**Answer to Weakness 2:** "Typos"

We have fixed all the typos, thank you!

Reviewer comment (Writing):
There are minor typos throughout, such as "$lim$" instead of "$\lim$" atop page 15 in the appendix.

Author response:
**Answer to Weakness 2:** "Typos"

We have fixed all the typos, thank you!

Reviewer comment (Writing):
There are minor typos throughout, such as the positive scalar $\delta$ not being defined in the proof of Theorem $1$.

Author response:
**Answer to Weakness 2:** "Typos"

We have fixed all the typos, thank you!

Reviewer comment (Writing):
There are minor typos throughout, such as in the statement of Lemma 3: "teh" should read "the".

Author response:
**Answer to Weakness 2:** "Typos"

We have fixed all the typos, thank you!

Reviewer comment (Theory):
Some references are obscure. For example, for the fact that $\mu + t\delta \mu$ converges weakly to $\mu$, it may be worth simply noting that this is due to the linearity of integration (with respect to the measure term).

Author response:
**Answer to Weakness 3:** "Some references are obscure"

Apologies for the unclear referencing in proof of Theorem 1. We have revised this reference in the revised version of our paper.

Reviewer comment (Theory):
Can it be shown that approximate vector field matching yields approximate solutions for all time $t$?

Author response:
**Answer to Weakness 1 & Question 1:**  "Approximate vector field matching yields approximate solutions for all time $t$"

It is challenging to rigorously analyze the error bounds of neural networks fitting vector fields at all times $t$, as it is related to properties like the smoothness of the flow and the expressive power of neural networks. Qualitatively speaking, as the number of samples used for training and the number of steps increase, the fitting error of the parameterized vector field for all times $t$ will decrease. We will reserve further developments for our future work.

Reviewer comment (Novelty):
Some theoretical results from the paper are known. For example, the statement of Theorem 1 could be found in [B] (eq. 26) or [C] (eq. 8).

Author response:
**Answer to Weakness 1:**  “Some theoretical results from the paper are known”

For the sake of completeness of our paper, we have included the theory of Sinkhorn WGF derivation in our methods section. Obviously, regarding the discussion of $\mathcal{W}_{\varepsilon}$-potential, we are not the first to do so, and we have added more relevant references to enhance the clarity of our paper. It is important to emphasize that our work presents an analysis of our discrete and continuous flows in mean-field limits, which is novel.

Reviewer comment (Reproducibility):
The quality of the code provided is not good. There is no README/or other instruction to run the code. There are imports of non-existing classes. So, there is no possibility of checking (at least, qualitatively) the provided experimental results.

Author response:
**Answer to Weakness 2:** "The quality of the code provided is not good."

Apologies that only the algorithm part was provided in the code. We have added a readme and the missing neural network structure components to ensure the completeness of the code in the latest version.

Reviewer comment (Novelty):
The main weakness of the proposed paper is the limited methodological contribution. The authors simulate the particles of data following Sinkhorn divergence — as already mentioned, this is not a super fresh idea. To make a generative model from these simulated trajectories, the authors simply solve the regression task to learn the local pushforward maps, and that is it. Combined with the fact that the practical performance of the proposed approach is far from being SOTA in generative modelling, the overall contribution of the paper seems limited.

Author response:
**Answer to Weakness 3:** "Limited methodological contribution"

We observed that existing neural SDE/ODE-based models necessitate multiple iterations through a neural network to produce high-quality samples. We noticed that the generative flows used in these models, such as reverse Ornstein-Uhlenbeck processes or linear flows constructed via random matching, tend to incur information loss. Therefore, exploring Wasserstein Gradient Flows (WGFs) that maintain more information by considering the steepest descent direction is a worthwhile direction. Simultaneously, we recognize that current neural gradient-flow-based models also face issues with information loss. Given this context, we considered leveraging the steepest descent property of Sinkhorn WGF to develop the NSGF model. Remarkably, by using a straightforward velocity field matching approach, we achieved generative quality surpassing existing neural gradient-flow-based models. Our experiments also demonstrated other interesting properties, such as improved generative paths and more meaningful interpolations between distributions, as seen in Figures 1 and 2. Notably, we are the first to construct a generative model based on Sinkhorn WGF, achieving results superior to other neural WGF-based works in fewer steps. We believe that further leveraging techniques from existing diffusion models (such as special neural network structures [D], or rapid sampling techniques [E]) could further enhance the capability of the NSGF model. This is a direction for our future work.

Reviewer comment (Reproducibility):
The authors propose to compute certain $\mathcal{W}_{\varepsilon}$ potentials (on discrete support of available samples) and then somehow take the gradients of these potentials w.r.t. the corresponding samples (eq. (13)). From the paper it is not clear how to compute the gradients, because the obtained potentials look like vectors of sample size shape, which are obtained through the iterations of the Sinkhorn algorithm. As understood, in practice, the authors utilize SampleLoss from the geomloss package ([B]). This implies that [B] should be properly cited when deriving the algorithm (section 4.2). It is recommended to explicitly use SampleLoss in the algorithm's listing, as it will contribute to clarity.

Author response:
**Answer to Question 1:**  "explicitly use SampleLoss in the algorithm's listing"

We have added the detailed computation procedure of $\mathcal{W}_{\varepsilon}$-potentials in section 4.2 of the revised version. 

In the original paper, we only explained the calculation of $\mathcal{W}_{\varepsilon}$-potentials and the use of the geomloss package in the Appendix. We apologize for this affecting the completeness of the paper.

Reviewer comment (Theory):
The vector field of the Sinkhorn gradient flow is estimated by empirical samples. It is not clear how well this sample estimate approximates the true vector field. This point should be clarified. Note that Theorem 2 works only for the mean-field limit.

Author response:
**Answer to Question 2:** "How well does this sample estimate approximate the true vector field"

In our paper, we concentrate on the mean-field limit situation and plan to address the analysis of the exact approximation error in our future work. It is important to note that conducting error analysis and theoretical proofs about the use of empirical sampling to approximate vector fields in flow-based models is generally challenging. For example, relevant discussions are not provided in references [F] and [H]. Even the analysis of sample complexity for Sinkhorn distance and Sinkhorn potential has gradually developed over time, following the introduction of these concepts.

Reviewer comment (Theory):
In the Introduction section, the authors consider a taxonomy of divergences used for gradient flow modelling, namely, "divergences [...] with the same support" and "divergences [...] with possible different support". As understood, the first class is about $f$-divergences and the second class is about the other types (like Sinkhorn, MMD, etc.). There is a question regarding the provided examples of works which deal with the former or the latter type of divergences: the works [D], [E], [F], [G] deal with KL-divergence (or f-divergence) minimization, so there is uncertainty as to why they were classified as the second class.

Author response:
**Answer to Question 3:** "The taxonomy in Introduction section"

Based on our understanding, papers [I] and [J] develop algorithms designed for a broad range of functionals, not limited to KL divergence. For instance, [I] uses the Free Energy Functional in the experiment section. [K] utilizes the Free Energy Functional and neural network parametrization techniques to construct functionals for WGF and JKO algorithms. The Free Energy Functional relies solely on a single distribution and does not require the “same support” condition, leading us to categorize it under the second class, “With Possible Different Supports.” We acknowledge that our description of these algorithms lacked clarity, and we have made revisions to our revised version. Regarding [L], our reference in the second paragraph was initially to illustrate the “poor expressiveness” of ICNNs. We have modified this reference in the revised version to address any ambiguity it may have caused.

Reviewer comment (Writing):
What is the “ground” set ($\S$ 3.1, first line).

Author response:
**Answer to Question 5:** "What is the “ground” set"

The "ground set" is simply the underlying set or space under consideration.