Parallel Wasserstein Generative Adversarial Nets with Multiple Discriminators

1 Introduction

Generative Adversarial Networks (GANs) [Goodfellow et al., 2014] have recently achieved significant progress for numerous applications such as computer vision [Wu et al., 2018], information retrieval [Wang et al., 2017; Su et al., 2017] and recommender system [Zhao et al., 2016; 2018; Zhang et al., 2019]. GANs consist of two networks: a generator network that creates plausible data given some noise as the latent seed variable, and a discriminator network that is trained to distinguish between the generator's fake output and real data sample. The critical issue in the GAN training process is to measure the difference between the real data distribution and the generated fake data distribution. For better measurement between the two distributions, Wasserstein GANs [Arjovsky et al., 2017] are proposed with better training stability over prior Jensen-Shannon (JS) divergence and Kullback-Leibler (KL) divergence GANs [Goodfellow et al., 2014].

∗Corresponding Author

Wasserstein GANs have attracted much interest in the community recently. The original Wasserstein GAN and its variants employ Kantorovich-Rubinstein duality, which involves a saddle-point objective with a strong 1-Lipschitz constraint for the discriminator function. On the other hand, some researchers approximate the Wasserstein distance in the primal form rather than solving a dual problem to improve stability. [Deshpande et al., 2018] introduces a random projection to estimate the Wasserstein distance from samples directly. [Salimans et al., 2018] employs the Sinkhorn algorithm to calculate the primal form of Wasserstein distance with entropic regularization. Although the performance of the state-of-theart Wasserstein GAN is better than the original version, the computation complexity of the approximation to dual or primal Wasserstein measures keeps increasing.

In general, the Wasserstein GAN framework consists of two time-consuming components: the forward and backward updates of the discriminator and the generator, and the approximation of Wasserstein distance. Since the discriminator employs the approximated Wasserstein distance to estimate the distance between the real distribution and the generated distribution, we should guarantee the correctness of the Wasserstein distance per generator update. Otherwise, the discriminator may mislead the generator during the training process. Overall, the computation cost is much heavier when the inaccurate approximation involves the frequent update of the discriminator. In practice, we usually spend several days training a generator model with edging NVIDIA GPUs.

In this paper, we explore the parallel computation to speed up the Wasserstein GAN training. A straightforward solution is based on data-parallelism. A set of worker unit containing model replicas is fed with different partitions of the input dataset, and the model replicas are synchronized via a parameter server. Unfortunately, the above approach is not ideal because of the skew parameter updates in GAN frameworkmost of GAN algorithms update the discriminator several times for every generator update. The frequent update of discriminator easily becomes a bottleneck due to the limited communication bandwidth in all distributed systems.

In order to eliminate the communication bottleneck for the synchronization of the discriminator, we intend to remove the synchronization requirements by allowing the discriminators different among worker units. Besides this, we also follow the data-parallelism to split the whole dataset into different worker unit. In this case, we can consider that each worker unit contains different real distribution supported by portions of the whole dataset. Then we formulate the objective function as minimizing the averaged Wasserstein distance between the generated distribution and the multiple real distributions. Inspired by [Genevay et al., 2016; Staib et al., 2017], we optimize the objective function by the stochastic algorithm for discrete optimal transport problem, which can be efficiently paralleled.

In summary, our contributions to parallel Wasserstein GANs are highlighted as follows:

• To our best knowledge, we are the first to propose a parallel architecture for Wasserstein GANs framework to speed up the training of GANs from the computational perspective.
• Different from common Wasserstein GANs, we develop an efficient stochastic algorithm to approximate the Wasserstein distance with higher accuracy for the newly proposed parallel framework.
• In experiments, we show that the proposed parallel architecture enjoys the superior convergence speed and comparable stability.

2 Preliminary and RelatedWork

The GAN defines two competing networks: The generator network G produces a data x from a source of noise z ∼ Pz. The discriminator network D is trained to distinguish between the generated fake sample and a real sample. Formally, the GAN defines the following minimax objective:

where Pr and Pg denote the real distribution and the generated distribution respectively. Practically, the expectations are empirically evaluated using samples. Unfortunately, welllearned discriminators may suppress the training of generators. We need to well tune the rounds of discriminator updates after every generator update to ease this training instability.

To enhance the training stability, [Arjovsky et al., 2017] introduces Wasserstein-1 distance to GAN framework. The Wasserstein distance is derived from the optimal transport problem, which is defined as follows:

where Π(Pr, Pg) represents all joint distributions, and c (xr, xg) is the cost to move a unit of mass from xr to xg. When c (xr, yg) = dX (xr, xg)p in Eq. (2), W 1/p c is called Wasserstein-p distance between probability measures. The Wasserstein distance is made known as a more geometricaware cost function for learning the distributions supported by the low-dimensional manifold, which is a widely adopted assumption in the feed-forward neural network.

Wasserstein GAN (WGAN) employs KantorovichRubinstein duality of Wasserstein-1 distance, which relies on 1-Lipschitz continuity of the discriminator:

where Pg ∼ Gϑ (Pz). To approximate the supremum in Eq. (3), [Arjovsky et al., 2017] proposes to clip the parameters θ to enforce the discriminator Dθ to be 1-Lipschitz.

Currently, the Wasserstein GAN approach is considered as the state-of-the-art method due to the theoretical contributions and competitive performance. However, to approximate the 1-Lipschitz constraint is very challenging in the Kantorovich-Rubinstein dual form of the Wasserstein-1 metric. [Gulrajani et al., 2017] introduces a soft penalty for the violation of 1-Lipschitzness (WGAN-GP). The gradient is evaluated as a linear interpolation between the training data and generated samples as a proxy to the optimal coupling. The gradient penalty only takes effect on the observed data {x}, the other support domain is not covered. [Wei et al., 2018] follows WGAN-GP to enforces the Lipschitz continuity over all the data manifold and its surrounding regions.

An alternative way to solve the problem in Eq. (2) is to minimize the primal of optimal transport. The primal formulation is numerically stable because it does not involve differentiating the dual solution. [Bousquet et al., 2017] proposes to minimize a regularized primal form of optimal transport problem. [Petzka et al., 2018] explores the theoretical properties of such regularization under Wasserstein GANs framework. [Deshpande et al., 2018] employs random projection to approximate the Wasserstein distance directly. The discriminator is not mandatory in their approach. [Genevay et al., 2018b] proposes a divergence measurement based on the Sinkhorn algorithm, which is originally designed for discrete optimal transport with entropic regularization.

To tackle with parallel computation, [Chavdarova and Fleuret, 2018] proposes an ensemble method of pairs of the generator and discriminator, which can be trained in parallel naturally. However, the communication demand is severe, because the "messenger" discriminators and generators with a massive number of parameters need to be synchronized among workers.

3 Proposed Approach

In all Wasserstein GANs framework, the Wasserstein distance plays a pivotal role to estimate the distance between the real distribution Pr and the generated distribution Pg. In realworld applications, we estimate the distributions by the empirical probabilistic measure with sampled data. Formally, we formulate the empirical measure of the real distribution µr and the artificially generated distribution µg by Dirac mass as follows:

where δxri and δxgj are Dirac function at location xri and xgj , pri and pgj are probability masses associated to the i-th real and j-th generated samples respectively, and nr and ng are the cardinality of the set of real data and generated fake data in a mini-batch respectively.

When the number of data in each mini-batch increases, the empirical distribution µr approximates the real distribution Pr with higher accuracy [Genevay et al., 2018a]. However, the size of the mini-batch is constrained due to the limited memory in each worker (e.g., GPU). To solve this problem, we turn to a parallel training mechanism in a distributed environment.

3.1 Framework

We propose a communication-efficient framework for parallel Wasserstein GAN. The details are shown in Figure 1.

Figure 1 (a) shows the traditional parallel approach [Dean and et al., 2012; Li et al., 2016]. In this framework, we split the training dataset across several workers first. Then, for each iteration, each worker with a model replica conducts forward-backward computation to calculate the gradients in parallel. Finally, the model parameters are synchronized with a global parameter server before the next iteration starts. This framework is widely used in distributed deep learning systems like Tensorflow [Abadi et al., 2016].

However, the traditional parallel approach in Figure 1 (a) contains two drawbacks. The parameter server is not proper to train the discriminator for GAN framework. Because the frequent synchronization of the parameters in the discriminator and the generator leads to massive communication demand, which degrades the training speed in any parallel solutions. Another drawback of the parameter server comes from the gradient aggregation for a large-scale batch. Traditionally, the gradient used for back-propagation is the arithmetic mean of gradients from different portions of mini-batch. The averaged gradient can not preserve the geometry information, which is vital to depict the semantic information.

To overcome the drawbacks in traditional parallel approaches, we propose a new parallel framework for Wasserstein GAN training, shown in Figure 1 (b). In our approach, we assign K worker unit with a discriminator Dθk and a master unit with a generator Gϑ. θk and ϑ represent the parameters of the neural networks in the discriminator k and the generator respectively. Note that θk ̸= θk′ in our proposed architecture generally. Our proposed framework relaxes the communication constraints by eliminating the synchronization of the discriminator among workers. In addition, this relaxed requirement could improve the diversity of the discriminator model and intuitively enhance the generator by increasing the difficulty of cheating the discriminator.

In our proposed architecture, the training real dataset R is split into K workers. Then, for worker k, the empirical measure µrk is supported by the subset of real data Rk.

Similar to the original Wasserstein GAN, we try to minimize the sum of the Wasserstein distance between the generated distribution µg and all real distribution µrk with empirical measures as follows:

3.2 Parallel Wasserstein GANAlgorithm

We develop a novel parallel Wasserstein GAN algorithm to solve the objective function in the our proposed architecture, namely Eq. (6). The loss function defined in Eq. (6) is similar to the Wasserstein barycenter problem introduced in [Agueh and Carlier, 2011]. Intuitively, the optimal generated fake distribution should be approximated to the barycenter of the real distributions in all workers. Specifically, we employ the empirical distribution of the barycenter of all real distributions as the supervised information to train the generator network. In order to solve the problem in Eq. (6) with more accuracy, we consider to calculate the Wasserstein distance directly, which is different from the recently popular solutions with entropic regularization [Genevay et al., 2016; Cuturi, 2013].

In the proposed framework, we follow [Genevay et al., 2016; Claici et al., 2018] to consider the semi-discrete Kantorovich's dual formulations of the primal optimal transport problem in Eq. (2):

where ϕc (xr) = c (xr, xgj ) -ϕj is the c-transform of the Kantorovich potential ϕj [Villani, 2008]. Therefore, L is reformulated as:

Algorithm 1: Generator as Master Unit

Input: K: number of worker unit, ng: batch size in generator, nr: batch size in worker, T: number of iteration for the estimation of the probability Dirac mass.

Output: Gϑ: the generator model.

while ϑ has not converged do

Sample Gaussian noise z1, . . . , zng from N(0, 1)

Broadcast {xgj }ng j=1 to workers

for t=1,. . . ,T do

Broadcast {ϕj , pgj }ng j=1 to workers

Receive {ϕkj , N k j }ng j=1 fromworkers

pgj ← 1K

∑Kk=1 N k

j ∀j ∈ [1, ng]

end

Conduct back-propagation with the loss:

end

From Eq. (8), we can know that F(ϕ) is a concave function. Then we can derive Wc from the optimization using the stochastic gradient ascent method. Since F(ϕ) is derivative with respective to each individual ϕj:

where Vorϕj is Voronoi cells of the point xgj [Santambrogio, 2015] defined in the following:

Therefore, Wc is obtained when ∂F ∂ϕj = 0 and the optimal p∗j is calculated as:

Algorithm 2: Discriminator as Worker Unit

Input: K: number of worker unit, nr: batch size in worker unit, η : learning rate for Kantorovich potential, T: number of iterations for the Dirac mass estimation.

for all worker k ∈ [1, K] do in parallel Sample {xi}nr i=1 a batch from the real dataset Rk

Receive fake images {xgj }ng j=1 from the master unit

Evaluate cost cij = ∥∥Dθk (xri ) -Dθk

(xgj

)∥∥ ∀i ∈

[1, nr], j ∈ [1, ng]

for t = 1, . . . , T do

Receive {ϕj , pgj }ng j=1 frommaster

for j = 1, . . . , ng do

Compute Vorϕj from Eq. (10)

end

end

Conduct back-propagation with the loss:

end

where I is denoted as the indicator function of a given set. p∗j is the valuable supervised information of our GAN framework for the generated fake point xgj since pgj represents the weight of the point xgj in the empirical distribution.

In our framework, the computation of p∗j in Eq. (11) is easily decomposed to all workers in parallel, because the Voronoi cells Vorkϕj in each worker k are independent with each other. When the empirical estimation in each worker is finished, p∗j is reduced to the master worker with xgj . Algorithm 1 and 2 detail the proposed parallel algorithm with mini-batch.

In many machine learning algorithms armed with Wasserstein distance, the ground cost function on the given metric is relatively simple, such as the ℓ2 norm in Euclidean space. However, ℓ2 norm is not proper for high-dimensional image data in GAN framework, because it lacks complex geometric information. Therefore, we parameterize the cost function c (xr, xg) from the discriminator as follows:

where Dθ is the discriminator network with parameters θ. When p∗ is obtained, we update the discriminator network Dk θ in all workers and the generator network Gϑ by backpropagation.

It is important to note that our approach does not contain any cumbersome communication policy that may increase the negative effect of mismatching communication bandwidth in the distribution environment. For each worker unit, they only receive generated data and send a simple vector of size ng to optimize the proposed objective function collaboratively. Moreover, our approach is compatible with different configurations for parallel computation: (a) All GPUs are installed via PCI Express in a single workstation. In this setting, each GPU is considered as an isolated computation worker in the parallel concept. (b) There are multiple machines, each of which may contain more than one GPU, deemed as distributed computation. In the following, we take experiments in case (b) to evaluate our approach.

4 Experiments

In the experimental part, we evaluate the proposed parallel approach on two widely used image datasets CIFAR-10 and LSUN [Yu et al., 2015]. CIFAR-10 contains 60,000 small color images with size 32 × 32. LSUN is a large-scale scene understanding dataset. Following the setting in many Wasserstein GANs papers, we mainly use the bedroom category from LSUN dataset. The bedroom category contains 3 million color images with a size of 128 × 128.

All experiments in this section are conducted in a cluster with four machines with 2 NVIDIA GTX 1080 GPUs each. Therefore, we could scale our approach up to 8 GPUs setting, which contains one master unit for the generator and seven workers for the discriminators.

We implement our algorithms with PyTorch [Paszke et al., 2017], and utilize the 101-layer ResNet block in generator and discriminator. We assign the number of iteration for per mini-batch update T = 5 for all experiments. Intuitively, the probability of Dirac mass for each generator update should be more accurate when T increases. However, larger T means that the communication rounds between the master and workers will increase as well.

4.1 Convergence Speed for Scalability

In a parallel environment, we assume that all workers have the same batch size nr. Therefore, in the proposed parallel approach, the total number of real images sampled from the dataset for a single generator update is K · nr, where K is the number of worker unit. Usually, the range of nr is constrained to the memory in GPUs. Then, the batch size of real images scales linearly with the number of machines. For the special "2 GPUs" setting in experiments, we utilize two GPUs equipped at different machines to make a fair comparison, because the inner-machine communication is much faster than the inter-machine protocol.

To demonstrate the advantage of our proposed method compared to the traditional distributed deep learning framework, we implement parallel WGAN-GP through MPI interface in PyTorch. MPI follows All-Reduce policy, equivalent to parameter server framework depicted in Figure 1 (a).

Figure 2 (a) and (b) illustrate the convergence speed with regard to the number of the GPUs, where the y-axis shows the change of estimated Wasserstein distance between the real distribution and the generated distribution during the training. From the figures, we can see both of the proposed method and WGAN-GP with more GPUs always converge faster than the same approach with fewer GPUs.

We also employ the standard speedup measurement in distributed computation to demonstrate the scalability of our proposed method, and the measurement is defined as below:

In the experiments, we measure the duration from starts to the first checkpoint where the corresponding estimated Wasserstein distance is less than 5.0. We consider 5.0 as an ending point because all methods converge with less vibration when the estimated distance is less than 5.0. From Figure 2 (c), we can easily conclude that our method enjoys a better convergence speed than the traditional distributed deep learning framework, which is coincident with our expectation.

Due to limited space, we do not demonstrate the figures for CIFAR10 dataset. Actually, the behavior of two methods in CIFAR10 is quite similar to Figure 2-both methods converge faster if more GPUs, and our approach is faster than the traditional parallel framework.

4.2 Quantitative and Qualitative Evaluation

We first conduct the quantitative evaluation to demonstrate the stability of the proposed method. We calculate the Inception Score during the training and evaluate the generator performance on CIFAR-10. A higher Inception Score indicates a better ability of the generative model to produce samples with variability. Practically, we calculate the maximal Inception Score reached in 20,000 generator updates.

We also follow [Wu et al., 2018] to evaluate the model performance on LSUN dataset by Fr´echet Inception Distance (FID) [Heusel et al., 2017], which measures the difference between real and fake data distributions.

From Table 1, we observe that the performance of the proposed method becomes better when the number of worker unit increases, which mainly benefits from larger size of real data batch among multiple worker units. Moreover, our method enjoys much better performance over WGAN-GP.

We also take qualitative evaluation from the visualization performance. Figure 4 indicates that the generated images of our proposals with different number of GPUs are comparable to results in WGAN-GP. Moreover, comparable results in Figure 4 (b-d) show that our parallel method has no restriction on the number of GPUs used to generate plausible results. Figure 3 shows the performance of our proposed method is also comparable with WGAN-GP on CIFAR-10 data. Hence our method enjoys superior performance concerning the convergence speed with the comparable generation quality.

5 Conclusions

In this paper, we introduce a novel parallel architecture designed for GANs framework to speed up the costly compuwork tation in Wasserstein GANs with multiple computation units. Experimental reports demonstrate that our proposed architecture enjoys an attractive scalability performance to decrease the training time remarkably. Moreover, both quantitative scores and qualitative demonstration display the proposed loss function with the Wasserstein metric and the parallel algorithm are reasonable and suitable for the parallel environment.

Figure 3: Qualitative comparisons between the state-of-the-art WGAN-GP and the proposed method.

Acknowledgements

The work described in this paper was partially supported by the Research Grants Council of the Hong Kong Special Administrative Region, China (No. CUHK 14208815 and No. CUHK 14210717 of the General Research Fund).