Prb-GAN: A Probabilistic Framework for GAN Modelling

Recall that the original GAN formulation features a generator and a discriminator with a single parameter setting $\theta_{g}\text{ and }\theta_{d}$ respectively rather than a parameter distribution. The optimal parameter point estimate is found by solving

Here $z$ is the latent noise vector. We show in the next sections the simple changes required to convert the standard GAN into a probabilistic GAN.

Probabilistic Generators: Assume a gaussian prior distribution over the weights of generator $\theta_{g}$, where $n_{\theta_{g}}$ represents number of generator parameters,

Let the probability that the data point produced by the generator $G(z;\theta_{g})$ being marked real by the discriminator be given by

Here $\theta_{d}$ are the parameters of discriminator network. Then we seek to find the posterior distribution of the weights $p(\theta_{g}|y_{g}=1,z)$ assuming that the generated point was marked real. Direct computation is infeasible and we infer using variational inference.

We setup our variational distribution for the generator weights using Bernoulli dropout as

Here $W_{g}$ serves as the variational parameter, $K_{i}$ is the number of neurons in the $i^{th}$ layer, $L$ is the number of layers and $p$ is the dropout probability. We seek out the optimal variational parameters for the generator by maximizing the evidence lower bound (ELBO) function associated with the above. Thus

The first term is an expectation and can be approximated using Monte-Carlo(MC) [40] integration for some arbitrary MC iteration, $M_{g}$, and the second term corresponds to the $L_{2}$ norm regularization term which forces the parameters to stay close to 0. Thus the loss function we would like to minimize becomes

Probabilistic Discriminators: The setup for probabilistic discriminators are nearly the same. We place a gaussian prior over the discriminator parameters as well. The likelihood term however is slightly different. The probability that the discriminator marks the real generated points as real and fake respectively appropriately is:

Then, the posterior distribution of the discriminator weights is sought out using a similar bernouilli dropout based variational distribution $q(\theta_{d};W_{d})$. The loss function that follows is given by

$M_{d}$ is the number of sampled discriminator.

Training scheme: In practice, the updated loss function for the generator is equivalent to actually sampling the probabilistic generator $M_{g}$ many times and back-propagating the averaged loss for a single generator training step and similarly sampling the discriminator $M_{d}$ many times and averaging the loss during a single discriminator training step. For each discriminator training step, a single random generator is sampled and vice versa for a single generator training step. Generally, we set the number of MC samples for generator and discriminator to be equal $M_{d}=M_{g}=N$. Figure 1 and  2 describes how the training steps would proceed each case. For a generator training step, while the variational parameters $\theta_{g}^{(l)}$ are the same for all generators, unique sampling of dropout layers ${\beta^{\prime}}_{n}^{(l)}$ create a unique generator. The case holds similarly for the discriminator training step. The algorithm is presented in  1. We term our formulation of GAN as Prb-GAN.

Inspired by the idea of determining aleatoric uncertainty for neural nets, we allow each dropout sampled discriminator to decide on how certain or uncertain it wants to be about its prediction on its input data-point. Every randomly sampled discriminator corresponds to instantiating a decision boundary for separating real and fake data points. When input data-points are closer to the captured mode of one of these sampled discriminators, if that discriminator makes a prediction with high certainty, its score is weighted higher. This encourages discriminators to capture diverse modes of the data by allowing it to be sure of what it knows and unsure of what it does not know. Figure 3 illustrates that since the input datapoint is closer to discriminator 2’s captured modes, we would like the discriminator 2 to be more sure about its prediction as compared to discriminator 1.

Following this, we create a new loss function that essentially weighs the prediction scores of the discriminators based on their uncertainty scores. Uncertainty scores are estimated by adding another branch from the penultimate layer of the discriminator. Greater the uncertainty measure, the closer the predicted logit score of that discriminator is to 0 (correspondingly the actual prediction being close to 0.5). Consider the logits predicted by a discriminator $D(x;\theta_{d})$. Suppose we allow the discriminator to specify how certain it is about its prediction - uncertainty represented as $u(x)$, then the modification is

Here $b_{1}$ is a hyperparameter that stands for a bias term which in we set a small number in our experiments. This is to prevent division by zero in case the uncertainty predicted is 0. The effect of distorting the predicted logits as above is to increase or decrease the prediction magnitude thus making the logit prediction closer to 0 in case $u(x;\theta_{d})>1$ or further away from 0 otherwise.

To prevent the discriminator from assigning high uncertainty to all the points we add the uncertainty estimate itself as an additional penalty term. Thus the loss function for the discriminator factoring in this input data uncertainty is changed Eq. 12 to Eq. 13, BCE is the binary cross-entropy loss.

BCE is the binary cross-entropy loss.

Our second modification is based on the idea that we could measure the variance of the set of discriminator scores and use this information to improve gradients for the generator. Essentially, if the discriminator scores exhibit high variance, then the sampled discriminators must disagree with each other’s scores and some discriminators must have assigned high logit scores while others must have scored the point with low logit scores. This would happen presumably when the data point is closer to modes captured by some of the discriminators while lying outside the high probability regions of the remaining discriminators which would indicate that the discriminators have captured diverse modes. When the generator generates such data points we’d like to reward it based on the magnitude of the score set variance.

Consider a concrete example wherein two separate cases, the discriminators assign logit scores to the same generated point as $[7,-1,7,-1]$ and $[3,3,3,3]$. Now the mean score in both cases is equal to 3; however, the first case though, exhibits high variance among the scores while the second set does not. We argue that the first case might be more desirable due to the fact that one-half of the discriminators recognize the point as being close to their captured modes while the other half marks them as unsure or outside their captured modes thus again incentivizing diverse mode capture.

To implement this idea, we introduce an additional term in the loss function of the generators that rewards generated points that have high variance in the discriminator logit score set. The updated loss function for the generator is given as

In the above $D_{n}$ represents the $n^{th}$ sampled discriminator. The term features $\lambda$ - a hyperparameter, $b_{2}$ - a fixed bias term to prevent division by 0, and the mean and variance of the set of discriminator scores. If the mean of the prediction is closer to zero, then the term has higher weightage. This is done to emphasize that only when the mean of the discriminator scores is near 0, is the term relevant. Greater the variance of the discriminator scores, the greater the reward.

Initially, we evaluated our model on a set of two simple datasets, a mixture of 1-dimensional Gaussians and the MNIST dataset. We show that on both the datasets, whereas the standard GAN exhibits mode loss behavior, the probabilistic formulation, as described above, is able to significantly prevent mode loss. We use a simple, fully connected 4 layer neural network with a net input size of $z_{d}=100$ where $z_{d}$ is the dimensionality of the random latent code space and hidden layer size of 600 neurons with the leaky Relu activation function. Xavier’s initialization of weights is used. We set the probability of dropout $p$= 0.4 and $N$ = 20.

The 1D Gaussian dataset is comprised of data-points in $R$ sampled from a mixture of 5 Gaussians. The means of the Gaussians are located at $[10,20,60,80,110]$ with standard deviations as $[3,3,2,2,1]$. The histogram representing the real dataset is shown in red, while the generated data histogram is shown in green. This is the same dataset used in [14]. The dataset consists of 600,000 data-points sampled from the mixture distribution. An ideal generative model would be able to recreate the dataset perfectly, and the histogram associated would resemble the true distribution. The vanilla GAN and Prb-GAN are trained for 5 epochs over the entire dataset.

Figure 4(a) and 4(b) shows the histogram plots of the data generated by both the models. It is observed that while the vanilla GAN seems to capture only one mode well, the dropout GAN is able to capture almost all the modes. In fact, by iteration 50, we observed that Prb-GAN was able to capture four of the most important modes, while vanilla GAN was only able to capture one mode. Also, we noted that the values for dropout probability and $N$ (number of MC sample iterations) used can change the results. In general, the greater the dropout probability used, the greater the $N$ value that would have to be in order to produce good samples.

Following this experiment, we tested our model on real images, specifically the MNIST dataset. We use the same GANs but make a change to the input layer size to match the input latent code size, which we set to be 200. We run the models over the entire dataset for 100 epochs and plot the images generated by each model at the end of the training. Figure 5(a) and 5(b) shows how while Vanilla GAN has collapsed the entire latent code mapping to the digit 1, Prb-GAN is able to produce diverse images. We note that the visual quality of the Prb-GAN may appear slightly less appealing, but what we gain in return is larger diversity. Through the use of Progressive GANs [42], one could conceive of using a Prb-GAN backend to allow the network to produce diverse images while having a standard Progressive GAN front end to make the produced images more crisp.

Further, in Figure 6, we see that using the exact same random latent code, Prb-GAN in separate Bernoulli dropout sampling instances is able to produce very different images from each other. Thus separate instances of the generators actually capture different sets of modes. This observation shows that through the use of a single GAN model and the dropout technique, we’re able to achieve performance similar to peers like MADGAN while having the advantage of not having to know in advance the number of modes that could be present in the data.

To test our GAN performance more rigorously, we turn towards the work done by [43]. Their work compares all the standard GANs on various datasets on a variety of hyperparameters and measures them using the FID score. As opposed to the earlier used Inception score [44], the FID score [45] is the current best GAN evaluation metric since it is more robust to noise and sensitive to mode loss compared to Inception score. The surprising conclusion reached is that it is possible that all the tested GANs perform nearly equal. We demonstrate that while this might be the case, the use of our probabilistic approach in GANs almost always improves performance.

Our architectural setup is as used in CompareGAN [46]. We compare three of the best variations of GANs - Non-Saturating GANs (NSGANs, the original GAN) [1] and Least-squares GANs (LSGANs) [47]. We make modifications to the networks by introducing dropout and training through MC integration based averaged loss. The comparison is made using the MNIST [48], CIFAR 10 [49] and the CelebA datasets [50]. We trained the models using a batch-size of 64, iteration steps of 18750 for the MNIST dataset, 200,000 for the CIFAR-10 dataset, and 80,000 iteration steps on the CelebA.

The FID score comparison for the various GAN models is given in Table I. Lower the FID score, the better the performance of the GAN. We noted that for the NSGAN and the LSGAN, our probabilistic approach improves on the FID score but not on the WGAN. We used the number of discriminators instances sampled $N$ = 20 and dropout probability $p$ = 0.4. Dropout was applied to both the generator and the discriminator side and the averaged loss was back-propagated. Qualitative results are presented in Figure 7.

Using the two modifications as described in section IV, we now implement our uncertainty measure equipped probabilistic GAN. In this section, we present the tabulation of the FID scores. The FID scores and the Inception scores are mentioned; however, note that the FID score is superior to the Inception score, and hence we emphasize the former’s results. We note that in almost all cases, the use of weighted discriminator scores improves the FID score. Prb-GAN represents the probabilistic GAN based on the setup described in the previous chapter. Prb-GAN (v1) describes the probabilistic GAN setup that uses the idea of uncertainty estimation and weighted discriminator scoring. Prb-GAN (v2) is the probabilistic GAN that makes additional use of the variance of the set of discriminator logit scores. Table II shows the FID score and Table III the Inception scores. We found amongst all our experiments Prb-GANs equipped with uncertainty measures perform the best. We used the values $N$ = 10, and p = 0.4, where $N$ is the number of discriminators instances sampled, and p is the dropout probability. Note that we apply dropout to only the discriminator and not to the generator.

In this section, we compare the performance of GANs that use only probabilistic discriminators against GANs that use probabilistic generators and discriminators. We find empirically that the use of probabilistic discriminators alone is sufficient to improve performance. In Table IV, Prb-GAN(v1) represents that Probabilistic GAN version that uses probabilistic generators and probabilistic discriminators and Prb-GAN (v2) which uses probabilistic discriminators and deterministic generators.