Neural Sinkhorn Gradient Flow

Thank you for your appreciation of our idea, we will answer your questions one by one regarding these weaknesses/problems.

Answer to Weakness 1: "Typos"

We have fixed all the typos, thank you!

Answer to Weakness 2: "There are several papers (also gradient flow-based ones), which evaluate a FID on MNIST. Please provide it as well."

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.

Answer to Weakness 3: "Many of the MNIST digits appear flipped"

We use data augmentation methods during training, including horizontal flipping on image datasets MNIST/CIFAR10. This is a common practice in the training of generative models like diffusion models [D]. In the revised version of our paper, we have provided in the appendix the results of retraining on datasets that were not subjected to data augmentation.

Answer to Weakness 4: "Generated images provided"

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

Answer to Weakness 5: "It is unclear what the trajectories show"

The trajectories show the Neural Sinkhorn Gradient Flow. In NSGF, a $32\times 32$ image can be understood as a particle in the flow. We have provided a clearer description of Figure 4 in the revised version.

Answer to Weakness 6: "The statement of theorem 2"

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

Answer to Weakness 7 and Weakness 8: "The proof of Theorem 2"

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]$.

Answer to Weakness 9: "The code is not runnable"

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

Answer to Weakness 10: "Minibatch"

The regression of vector fields with minibatches serves as an approximation to the potential vector fields of Sinkhorn gradient flows. This method effectively circumvents the problem of repeated pairing, which becomes more prominent as the batch size increases—a challenge commonly encountered in the training of most neural ODE/SDE-based and neural gradient-flow-based generative models, as highlighted in studies like [E], [F], and others. Interestingly, our experiments on toy datasets revealed that, even with small batch sizes, the outcomes of minibatch pairing training were on par with those achieved using larger batches. Furthermore, we can assure that when the samples from the initial and target distributions are identical, we can generate the exact same Sinkhorn gradient flow for training, a consistency that [I] and [J] do not guarantee.

[A] Variational Wasserstein gradient flow, ICML'2022 [B] Deep Generative Learning via Euler Particle Transport, MSML'2022 [C] Deep Generative Learning via Variational Gradient Flow, ICML'2019 [D] Denoising diffusion probabilistic models, NIPS'2020 [E] Multisample Flow Matching: Straightening Flows with Minibatch Couplings, ICML'2023 [F] Improving and Generalizing Flow-Based Generative Models with Minibatch Optimal Transport, ICML Workshop'2023 [G] Sliced-Wasserstein Flows: Nonparametric Generative Modeling via Optimal Transport and Diffusions, ICML'2019 [H] Sliced Iterative Normalizing Flows, ICML'2021 [I] Flow Matching for Generative Modeling, ICLR'2023 [J] Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow, ICLR'2023

While I think this paper has improved I still have some concerns.

1. Not all typos are fixed. In the appendix there are still several "Sinkhron"

2. Regarding 10) please comment on the U-net. As I understand your code (which now includes the model) this has 128 channels, which enables learning the particle dynamics. However, this would become very costly for more particles as the number of parameters in the U-net scales with the channels. So in order to get the mean field limit one would like to take number of particles to $\infty$, however this is probably very computationally expensive.

3. The cifar samples look good, could you also provide l2-nearest neighbors from the training set in order to show how your model generalizes.

4. can you please also upload the checkpoints of the u-nets? The memory requirements are not really feasible for me to train. If i understand correctly, then I think this will take a lot of velocity field/flow saves in the range of terabytes. (Is that correct? Then this should be mentioned in the paper)

Thank you for your thorough review and valuable suggestions, particularly regarding paper revisions. Your insights have greatly enhanced the clarity and detail of our paper. We sincerely appreciate your guidance.

1. We apologize and have promptly rechecked and corrected the spelling errors.

2. Generally speaking, the choice of channel count in a U-net reflects a trade-off between the neural network's expressive capacity and the computational and storage costs. In our case, we use 128 channels, a common setting in generative models that employ U-net as their neural network architecture. It's foreseeable that expressing particle flows under mean field limit conditions requires a very deep neural network. However, we wish to emphasize that our demonstration of particle flows in the mean field limit as spatially continuous flows are to showcase our algorithm generating along paths that are more "information-rich". Furthermore, the theory related to mean field limits has also been developed in Bayesian sampling methods, especially in particle VI. Our insights are derived from advancements in this field.

3. Thank you for your suggestion. We have included a discussion on sample images and their nearest neighbors in the appendix of the revised version.

4. Your understanding is correct. The volume of data stored during the first phase of our algorithm increases with the dimensionality of the dataset. For the MNIST/CIFAR-10 experiments, a considerable amount of storage space is needed to establish the trajectory pool during the first phase of the algorithm. Specifically, for CIFAR-10, setting the batch size to 128 and saving all minibatch Sinkhorn gradient flow trajectories while traversing the entire dataset requires about 115GB of storage space. In cases where storage space is limited, we recommend dynamically adding and removing trajectories in the trajectory pool to accommodate training needs. Identifying a more effective trade-off between training time and storage space utilization is a direction for future improvement. We will add a discussion on this in the updated version of our paper. We note that the supplementary materials are subject to a 100 MB maximum file size limit. Unfortunately, even after compression into a ZIP format, our neural network's size significantly surpasses this limit. With a commitment to reproducibility and aiding the adoption of our approach, we intend to make the code publicly available following the publication of our paper. We try to use "anonymous.4open.science" but meet the same limit: "this exceeds GitHub's file size limit of 100.00 MB"

Answer to Question 1: "Learn the vector field from an unlimited number of minibatches"

Yes, we can follow a similar approach in [B] and [C]. In our experiments, we chose to pre-construct the Sinkhorn WGF training set and perform the vector field matching algorithm. However, in fact, running main_sinkhorn.py repeatedly to build a trajectory pool and then using the stored pre-trained data for vector field matching is equivalent to simulating the gradient flow during training and learning the vector field from an unlimited number of minibatches. In our image experiments, considering the balance between expensive training costs and training quality, we opted to first build a trajectory pool of Sinkhorn WGF and then sample from it to construct the velocity field matching algorithm. This is more thoroughly explained in the revised version of Algorithm 1, hoping to eliminate any confusion regarding the code.

Answer to Question 2: "What is the performance of your method concerning the number of simulation steps"

According to Tables 1 and 2, we can observe that our algorithm achieves better generative effects with fewer simulation steps. Similarly, from Table 1, it can be seen that using more simulation steps can enhance the quality of our generation.

Answer to Question 3: "1. What is the time of the preprocessing step concerning the training time?"

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)$. In practical terms, we observed that the preprocessing step consumed a considerable amount of time. Therefore, in our image experiments, we chose to first build a trajectory pool and subsequently sample from it for constructing the velocity field matching algorithm. This approach reduces the time spent on the preprocessing step, thereby enhancing the overall efficiency of our algorithm's training process.

Answer to Question 4: "Could you compare your method with OT-CFM [B, C] on the synthetic data?"

We have already provided some quantitative comparison data for OT-CFM in Tables 1 and 2. Considering the limited length of the paper, we have provided some quantitative image results (similar to Figure 1) in the revised version.

[A] Learning with minibatch Wasserstein: asymptotic and gradient properties, AISTATS’2020 [B] Multisample Flow Matching: Straightening Flows with Minibatch Couplings, ICML'2023 [C] Improving and Generalizing Flow-Based Generative Models with Minibatch Optimal Transport, ICML Workshop'2023 [D] Learning Generative Models with Sinkhorn Divergences, AISTATS'2018

Thank you for your suggestions, we will answer your questions one by one regarding these weaknesses/problems.

Answer to Weakness 1: "The differences in the empirical implementation between the paper and the code."

We utilized minibatch sampling for both initial and target distribution pairs and calculated Sinkhorn potentials to construct the trajectory of the Sinkhorn gradient flow. Similar minibatch approximations for calculating the optimal transport plan or the corresponding Sinkhorn distance are noted in [A], [B], [C], and [D], which vary from our method of constructing the minibatch Sinkhorn gradient flow. The theoretical impact of minibatch use has been discussed in [A], and [B], [C] demonstrating through experiments that constructing interpolation flows using minibatch OT plans can train flow-based models with approximate OT paths. We plan to conduct further experiments in our future work. We recognize that minibatch was not explicitly mentioned in the main text, and have addressed this in the revised version (see Algorithm 1 for details). Furthermore, we have added a discussion about the minibatch approximation in Appendix A.

Answer to Weakness 2: "Discussion regarding the statistical performance of the method"

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.

Answer to Weakness 3: "Training time between your method and the recent Flow Matching approaches"

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.

Answer to Weakness 4: "typos"

We have fixed all the typos, thank you!

Answer to Weakness 5: "big datasets (Flowers, CelebA) or on other tasks such as single-cell dynamics"

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

Answer to Weakness 6: "The related work on optimal transport is incomplete"

We have added new references on optimal transport, thanks!

Dear authors,

Thank you for your rebuttal which I have read carefully.

[Minibatch OT]: While the authors acknowledge the difference between the original submission and the original provided code, they did not make a careful discussion regarding the minibatch OT. I expect the authors to explain and discuss how it is related and how it differs from existing work [A,B,C]. Authors's approximation is a special case of minibatch OT where they only consider specific terms (ie in the original code, they computed OT between minibatches from one partition of the training dataset). Theoretically, it would correspond to a special Incomplete U-statistics. Therefore, I think that a detailed discussion and comparison to minibatch OT (and its bias) is needed and a remark is not enough.

[Practical implementation] Thank you, Algorithm 1 is much clearer now and reflects what is done in practice. I highly recommend authors to update the methodology section in the main paper to incorpore this practical implementation and to discuss it in a limitation section.

[Training time and motivation] Thank you for having run these experiments. In my opinion, they highlight a (known) limitation of the simulation-based approach like the proposed method. The proposed approach takes more than 10 times than other Flow Matching methods to converge. Indeed, the point of Flow Matching method is to avoid the computation of the interpolation samples that can be computed directly on the fly, resulting in a lower training time. It is the interest of these recent flow matching methods compared to simulation-based method. Furthermore, while authors do (minibatch) OT, it is also the case of (minibatch) OT-CFM [Tong et al., Poolidean et al.] and they do not require any simulation during training. Therefore, it is hard to understand the method interest to this recent line of work and that might require a revised motivation section.

[Inference performance] While I understand authors's points regarding the inference, Flow Matching methods take at most 150 steps to converge towards SOTA FIDs on large datasets [Lipman et al.]. The differences between authors's method and Flow Matching method is 50 steps with a non-competitive FID. Therefore, I disagree with the authors that their method is competitive in that aspect.

[Related work]: I would recommend to include a detailed discussion on minibatch OT and other gradient flow based OT method.

Thank you for your answers regarding my questions. They have been addressed.

I still think that some key elements are missing from the current manuscript (like the theoretical section which should mention the different biases). Therefore, I keep recommending to reject the current manuscript.

Thank you for your appreciation of our work, we will answer your questions one by one regarding these weaknesses/problems.

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.

Answer to Weakness 2: "Typos"

We have fixed all the typos, thank you!

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.

Thank you for your suggestions, we will answer your questions one by one regarding these weaknesses/problems.

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.

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.

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.

[A] Sinkhorn Barycenter via Functional Gradient Descent, NeurIPS’2020 [B] Interpolating between Optimal Transport and MMD using Sinkhorn Divergences, AISTATS’2019 [C] Sinkhorn Barycenters with Free Support via Frank-Wolfe Algorithm, NeurIPS’2019 [D] Diffusion Models Beat GANs on Image Synthesis, NIPS'2021 [E] DPM-Solver: A fast ODE solver for diffusion probabilistic model sampling in around 10 steps, NIPS'2022

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.

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.

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.

Answer to Question 4: "A good work regarding poor expressiveness of ICNNs"

We have added that reference, thanks!

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

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

Answer to Question 6: "What are the differences between 1-RF, 2-RF and 3-RF methods"

Differences between 1-RF, 2-RF, and 3-RF methods depend on the rectified flow technique which recursively applies rectification to obtain a sequence of flows with increasingly straight paths. As the authors show in their experiments, although the rectified flow method produces straighter paths in favor of reducing the number of sampling steps, the accumulation of errors in recursive training may reduce the quality of generation. More details can be found in [M].

[F] Flow Matching for Generative Modeling, ICLR'2023 [H] Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow, ICLR'2023 [I] Large-scale Wasserstein gradient flows, NeurIPS'2021 [J] Optimizing functionals on the space of probabilities with input convex neural networks, TMLR [K] Proximal optimal transport modeling of population dynamics. AISTATS [L] Variational wasserstein gradient flow, ICML'2021 [M] Do Neural Optimal Transport Solvers Work? A Continuous Wasserstein-2 Benchmark, NeurIPS’2021.

Thank you for your appreciation of our work, we will answer your questions one by one regarding these weaknesses/problems.

Answer to Weaknesses 1 & Question 2: “Sample complexity of the method”

Note that the NSGF model is gradient-flow-based, and the training of NSGF can be divided into two stages (see Algorithm 1 for details). In the first stage, NSGF constructs trajectories simulating Sinkhorn WGF using target points and source points and stores trajectories in the trajectory pool. In the second stage, we sample from the trajectory pool for velocity matching algorithms instead of using the source and target points. A similar procedure is widely used in flow-based models such as [A], and [B]. As a result, our method has a similar sample complexity to existing flow-based models. We have added relevant experiments in our 2D experiments, see the table below for details:

Answer to Weaknesses2: "Generative models that output discrete random variable"

Thank you for discussing the application of the NSGF model in generative tasks involving discrete random variables. Although our current focus is on simulating Sinkhorn WGF in continuous state spaces, such as for real-valued images, we recognize the potential for extending the NSGF model to discrete state spaces. This extension would follow methodologies similar to those found in [C] and [D]. Exploring the generative capability of the NSGF model in discrete state spaces is an interesting extension, and we will investigate this in our future work.

Answer to Question 1: "Change the source distribution to be closer to the target distribution"

The NSGF model is designed to simulate Sinkhorn WGF from a source distribution to a target distribution. If the source distribution is closer to the target distribution, Sinkhorn WGF will converge more quickly. This allows us to use fewer steps and lower computational costs during both the training and inference processes of NSGF. [E] considers using gradient-flow-based methods to refine the samples from a generator, which improves the metrics of the generated images. The NSGF model can also be used to achieve similar improvement, and this will be a focus of our future work.

[A] Flow Matching for Generative Modeling, ICLR'2023 [B] Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow, ICLR'2023 [C] Structured Denoising Diffusion Models in Discrete State-Spaces, NIPS'2021 [D] Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions, NIPS'2021 [E] Refining Deep Generative Models via Discriminator Gradient Flow, ICLR'2021