A Differential Equation Approach for Wasserstein GANs and Beyond

Authors: Zachariah Malik, Yu-Jui Huang
Primary area: generative models
Keywords: Generative modelling, finite elements, gradient flow, persistent training

Abstract:
This paper proposes a new theoretical lens to view Wasserstein generative adversarial networks (WGANs). To minimize the Wasserstein-1 distance between the true data distribution and our estimate of it, we derive a distribution-dependent ordinary differential equation (ODE), which represents the gradient flow of the Wasserstein-1 loss, and show that a forward Euler discretization of the ODE converges. This inspires a new class of generative models that naturally integrates persistent training (which we call W1-FE). When persistent training is turned off, we prove that W1-FE reduces to WGAN. When we intensify persistent training appropriately, W1-FE is shown to outperform WGAN in training experiments from low to high dimensions, in terms of both convergence speed and training results.

Decision: reject
Reviewer scores: 6, 6, 5

Reviews and discussion:
<thread>
creation date:
1730343958189

thread_type:
official_review

<thread_object>
creator: reviewer_y5er

<content>
summary:
This paper provides a new perspective on WGAN from the view of ODE associated with the gradient of Wasserstein derivative, demonstrating that persistent training on a generator can improve WGAN training.

strengths:
1. This paper provides a clear and novel explanation of WGAN's training from the perspective of the gradient flow of Wasserstein distance.

2. Both theoretical explanations and experiments on toy and MNIST datasets demonstrate the effectiveness of persistent training, a common trick for training WGAN.

weaknesses:
1. The main contributions in this paper, i.e. discretization and persistent training, are common tricks for training WGANs[1,2], which are not novel enough in practical implementation. For example, as is shown in the proof of Proposition 4.1, persistent training seems equal to just increasing the generator's iterations in the original WGAN's training. Thus please clarify how discretization and persistent training differ from the above existing methods.

2. Obtaining Kantorovich potential is a challenging and important step in ODE-based WGAN's training, but in this paper, it's still the same as the original WGAN, leaving some problems for persistent training as discussed in Remark 4.2, harming the consistency between theory and practice.

3. Experimental validation is not comprehensive enough, the used datasets are too small scale, like in WGAN and WGAN-GP,  more common and large-scale baseline datasets should also be verified. For example, how does the proposed method perform on CIFAR-10 or CelebA compared to baseline WGAN methods in terms of FID, IS, and convergence speed?

4. Some mathematical explanations and notations should be modified. For example, there is no definition of $m$ in Eq.(2.4). Besides, more introduction of $\nabla \frac{\delta J}{\delta m}$ are suggested to be added, since $\nabla \frac{\delta J}{\delta m}$ is very important in the whole method, how to understand this gradient and how to get Eq(3.2) from Eq(3.1) should be added in the main paper.

[1] Variational Wasserstein gradient flow. ICML 2022.

[2] Scalable Wasserstein Gradient Flow for Generative Modeling through Unbalanced Optimal Transport. ICML 2024.

questions:
1. please clarify how persistent training differs from simply increasing generator iterations.

2. As discussed in the third point of the above weaknesses, more common large-scale datasets should also be verified.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730700017663

thread_type:
official_review

<thread_object>
creator: reviewer_gqdk

<content>
summary:
This work analyzes Wasserstein GAN training through the lens of gradient flow dynamics derived from the optimal transport map between initial and data distributions. Authors demonstrate in Theorem 4.1 that by applying a sequence of updates to the generator distribution according to the gradient of the W1 witness function (Kantorovich potential), the resulting distribution in the limit matches the optimal one. 
This framework motivates the use of *persistent training* of the generator, where at each substep of training the discriminator and generator noise $z$ is frozen while the generator $G_\theta$ is updated for K steps. Lastly, the authors use a few simplified experimental settings  (e.g. 2D mixture of Gaussians) to demonstrate how persistent training can improve the rate of training convergence.

strengths:
* Utilizes tools from optimal transport for analyzing WGANs, providing a generalized WGAN training framework.
* Theory provides good insight into generator training hyperparameters, which is corroborated by experiments.

weaknesses:
* It seems to me that most of the uncertainty about how well WGAN training follows idealized gradient flow dynamics lies with the discriminator. Can you still arrive at a similar conclusion to thoerem 4.1 when the distance between approximated + true potential function is bounded?
* I'd like to see more discussion in the introduction about key differences between contributions of this work and the W2-FE paper.

Minor Notes:

* Eq 2.2 \mu_t, \mu^d not defined initially
* both \varphi and \phi used on line 087
* Eq. 2.3 variable i is not defined

questions:
* What was the number of steps used when training the discriminator in figure 1? Does your theory suggest when it is more or less important to have accurate $\varphi$ approximations during training?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730710597197

thread_type:
official_review

<thread_object>
creator: reviewer_bhdk

<content>
summary:
This paper presents a new variant of Wasserstein generative adversarial networks (WGANs) based on a distribution-dependent ordinary differential equation (ODE). It introduces a method called W1 Forward Euler (W1-FE), which includes persistent training to improve efficiency. When persistent training is not utilized, the approach reverts to standard WGAN algorithms. By appropriately increasing the level of persistent training, the model's performance is enhanced.

strengths:
This paper presents a novel framework for Wasserstein generative adversarial networks (WGANs) based on distribution-dependent ordinary differential equations (ODEs). It utilizes persistent training to enhance the training process of WGANs and provides a solid mathematical foundation for this approach. Additionally, the paper includes numerical results that demonstrate the effectiveness of persistent training at various levels.

weaknesses:
The experiments are insufficient; the paper should include additional datasets or real-world applications. It is better to present empirical results using well-known standard datasets.

The mathematical framework needs further clarification. For instance, the inequality at the top of page 4 should be explained more thoroughly.

The section on "MATHEMATICAL PRELIMINARIES" could benefit from more references. Please include citations for definitions and related concepts.

The contributions seem to lack originality, as mentioned in paragraph 2 of the Introduction. This section suggests that the differential equation approach outlined in the paper is already well-established in the existing GANs literature. Therefore, the question to the author is: Is the approach a variant of Generative Modeling using Wasserstein-1 loss, achieved by minimizing Wasserstein-2 loss (Huang and Malik, 2024)? If the novelty is something else, please clarify this in the Introduction. See references from the papers itself

Y.-J. Huang and Z. Malik. Generative modeling by minimizing the wasserstein-2 loss, 2024. URL
https://arxiv.org/abs/2406.13619.

Y.-J. Huang and Y. Zhang. GANs as gradient flows that converge. Journal of Machine Learning
Research, 24(217):1–40, 2023. URL http://jmlr.org/papers/v24/22-0583.html.

questions:
I have a few questions and suggestions as follows.

1.	For completeness, please include a proof for the statement, “$\mu \rightarrow W_1(\mu,\mu_d)$ is strictly convex on $P_1(X)$.”

2.	In Proposition 3.1, it would be helpful to explain the first inequality (at the top of page 4).

3.	The manuscript mentions the use of persistent training, which is known to have certain limitations. Could you elaborate on how these limitations are addressed or mitigated in your training approach?

4.	The manuscript's writing could benefit from some improvements, particularly in conveying the contributions of the proposed work more clearly and concisely.

5.	Please consider including some recent, relevant citations to enhance the manuscript's context within the field.

6.	The numerical experiments are limited in scope; please include more real-world datasets (such as CIFAR-10 or CIFAR-100) in the empirical results.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732785437759

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Revision Submitted

comment:
We would like to thank all the reviewers for your valuable comments from various angles. We have carefully revised our paper according to every single comment you provided and have just submitted the revived version of our paper. While each one of you can see our detailed response to your comments right below your review, we would like to point out here two major changes made to our paper. 

1. A new training experiment on the CIFAR-10 dataset is added.
2. In the introduction, we now explain in detail our contributions relative to other recent studies that also feature "Wasserstein gradient flows." In short, while almost all other papers rely on the well-developed gradient flow theory under the Wasserstein-2 distance, we focus on Wasserstein-1 gradient flows, which are much less understood and necessary for making a connection to Wasserstein GAN.  


Sincerely,

The authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper is concerned with the Wasserstein generative adversarial network (WGAN). It introduces a notation of gradient flow associated with WGAN leveraging the linear functional derivative. Based on it, the authors propose an algorithm to train WGAN. The major criticism is on the contribution. The proposed method resembles existing methods that utilizes gradient flow associated with Wasserstein-2 metric or f-divergence. The only difference is on the computation of the potential function (discriminator). In addition, the experiments are insufficient to demonstrate the advantages of the proposed method over existing ones. Finally, the theoretical result is weak. It (Theorem 4.1) claims convergence to a continuous trajectory, but the property of this trajectory is unknown. Its convergence to the target distribution is important for it to be useful in applications. The non-uniqueness of the Kantorovich potential is also overlooked in the theoretical development.