GANs Part3: WGAN
Techniques to reduce GAN failure due to imbalances between the generator & discriminator. Use WGAN to mitigate unstable training and mode collapse using W-Loss and Lipschitz Continuity enforcement.
Outlines.
- What is mode collapse?
- The Problems with Binary Cross-Entropy (BCE) loss.
- The Objectives of MinMax Game.
- Earth Mover’s Distance (EMD)
- Wasserstein Loss (W-Loss).
- The Critics.
- W-Loss Vs BCE Loss
- Condition on Wasserstein Critic
- 1-Lipschitz Continuity Enforcement.
- Weight (W) clipping
- Gradient Penalty.
Implementation Code: Implementing-GANs-from-Scratch/7-WGAN_GP.ipynb at main · Mustafaadel7/Implementing-GANs-from-Scratch
What is Mode Collapse?
Mode collapse sounds like things collapsing to one mode or fewer modes and some of the modes disappearing.
A mode in a probability density distribution (PDD) of data is just an area with a high concentration of observations. I.e., the mean value in a normal distribution is the single mode of that distribution, and there’re distributions with multiple modes where the mean doesn’t have to be one of them.
More intuitively, any peak on the PDD over features is a mode of that distribution. Many real-life distributions used for training GANs are multimodal.
Concrete Example.
Take a discriminator that has learned to be good at identifying which handwritten digits are fakes, except for cases where the generated images look like 1 and 7.
This could mean that the discriminator is at local minima of its cost function. The discriminator classifies most of the digits correctly, except for the ones that resembled 1 and 7, then this information is passed to the generator.
The generator looks at the feedback from the discriminator and gets a good idea of how to fool the discriminator in the next round.
It sees that all the images were misclassified by the discriminator, resemble either 1 or 7, so it generates a lot of pictures that resemble either of those numbers.
Then these generated images are then passed on to the discriminator in the next round who then misclassifies every picture except for maybe the 1 felt looks more like 7.
Generator gets that feedback and sees that the discriminator’s weakness is with the pictures that resembled a handwritten 1.
So, this time all the pictures it produces resembled that digit, collapsing to a single mode.
Eventually the discriminator will probably catch on and learn to catch the generator’s fake handwritten (1s) by getting out of that local minimum. But the generator could also migrate to another mode of the distribution and again would collapse again to a different mode. Or the generator would not be able to figure out where else to diversify.
Problems with Binary Cross-Entropy (BCE) loss.
With BCE loss, GANs are prone to mode collapse and are susceptible to vanishing gradient problems. Remember the BCE loss function, it’s just an average of the cost for misclassifying real and fake observations. Where the 1st term is for reals and the 2nd term is for the fakes.
The higher the cost value is, the worse the discriminator is doing.
- The generator wants to maximize this cost because that means the discriminator is fooled and classifying fakes into reals.
- The discriminator wants to minimize this cost function, which means it’s classifying things correctly.
This maximization & minimization is called a minimax game.
The Objectives of MinMax Game.
At the end of this minimax game, the generator and discriminator interaction translate to a more general objective for the whole GAN architecture to make the real and generated data distributions of features very similar. Trying to get the generated distribution to be as close as possible to the reals.
During the training process, the discriminator is trying to delineate real and fake distribution as much as possible, whereas the generator is trying to make the generated distribution look more like the reals.
The discriminator has a much easier task, as it needs to output just a single value prediction within [0,1]. Whereas the generator actually needs to produce a pretty complex output composed of multiple features (image) to try and fool the discriminator.
During training it’s possible for the discriminator to outperform the generator. But at the beginning of training, this isn’t such a big problem, as the discriminator isn’t that good. It has trouble distinguishing the generated and real distributions. There’s some overlap and it’s not quite sure.
As a result, it’s able to give useful feedback in the form of a non-zero gradient back to the generator. However, as it gets better at training, it starts to delineate the generated and real distributions a little bit more I.e., it starts distinguishing them much more.
Where the real distribution will be centered around 1 and the generated distribution will start to approach 0.
As a result, when the discriminator is starting to get better, it’ll start giving less informative feedback. In fact, it might give gradients closer to 0, and that becomes unhelpful for the generator because the generator doesn’t know how to improve.
Earth Mover’s Distance (EMD)
EMD measures the distance between 2 distributions and generally outperforms the BCE loss for training GANs. As, it also solves vanishing gradient problem of the BCE loss.
For illustrations, take this generated and real distributions with the same variance but different means and assume they are normal distributions.
What the EMD does, is it measures how different these 2 distributions are, by estimating the amount of effort it takes to make the generated distribution equal to the real.
So intuitively, the generate distribution was a pile of dirt, how difficult would it be to move that pile of dirt and mold it into the shape and location of the real distribution? So that’s what EMD means.
EMD depends on both the distance and the amount that the generated distribution needs to be moved. So the problem with BCE loss is that as a discriminator improves, it would start giving more extreme values between 0 and 1, so values closer to 1 and closer to 0. As a result, this became less helpful feedback back to the generator.
Get AI & Data Science 604’s stories in your inbox
Join Medium for free to get updates from this writer.
The generator would stop learning due to vanishing gradient problems. With EDM, there’s no such ceiling to the 0 and 1. So the cost function continues to grow regardless of how far apart these distributions are. The gradient of this measure won’t approach 0 and as a result, GANs are less prone to vanishing gradient problems and from vanishing gradient problems, mode collapse.
Wasserstein Loss (W-Loss).
W-loss is an alternative loss function that approximates the EMD. It is simpler and more compact than BCE loss.
The very long equation for BCE Loss can be simplified as follows.
The sum and division over examples M, is nothing but a mean or expected value. In the 1st part, inside the sum, measures how bad the discriminator classifies real observations, where y=1. And the 2nd part measures how bad it classifies fake observations produced by the generator, where y of 1 means real, but 1 — y, y of 0, means fake.
W-Loss approximates the EMD between the real and generated distributions, but it has nicer properties than BCE. However, it does look very similar to the simplified form for the BCE Loss, and in this case the function calculates the difference between the expected values of the predictions of the discriminator.
Here the discriminator is called the critic (C ), so I’m going to represent it with C. And this is C(real example) X C (fake example) which is g(z).
So, the discriminator looks at these 2 things, and it wants to maximize the distance between its thoughts on the reals versus its thoughts on the fakes. It’s trying to push away these 2 distributions to be as far apart as possible. Meanwhile, the generator wants to minimize this difference, as it wants the discriminator to think that its fake images are as close as possible to the reals.
The Critics.
in contrast with BCE there are no logs in this function, since the critics outputs are no longer bounded to be between 0 and 1. So, for the BCE Loss to make sense, the output of the discriminator needs to be a prediction between 0 and 1. And the discriminator’s NN has a sigmoid activation function in the output layer to then squash the values between 0 and 1.
W-Loss doesn’t have that requirement at all, so you can actually have a linear layer at the end of the discriminator’s NN, and that could produce any real value output.
We can interpret that output as, how real an image is considered by the critic. It’s no longer discriminating between these 2 classes. And so, as a result, it wouldn’t make that much sense to call that NN a discriminator. And so, for W-Loss, the equivalent to a discriminator is called a critic, and what it tries to do is, maximize the distance between its evaluation on a fake, and its evaluation on a real.
W-Loss Vs BCE Loss
The discriminator under BCE Loss outputs a value between 0 and 1, while the critic in W-Loss will output any number. Additionally, the forms of the cost functions is very similar, but W-Loss doesn’t have any logarithms within it, and that’s because it’s a measure of how far the prediction of the critic for the real is from its prediction on the fake.
Meanwhile, BCE Loss does measure that distance between fake or a real, but to a ground truth of 1 or 0.
The discriminator is bounded between 0 and 1, whereas the critic is no longer bounded ,and just trying to separate the 2 distributions as much as possible.
And as a result, because it’s not bounded, the critic is allowed to improve without degrading its feedback back to the generator. And this is because, it doesn’t have a vanishing gradient problem, and this will mitigate against mode collapse, because the generator will always get useful feedback back.
Condition on Wasserstein Critic
W-Loss solves some problems faced by GANs, e.g., mode collapse and vanishing gradients. But there’s a special condition for the critic to work well, which is called 1-Lipschitz (1-L) Continuous
W-Loss is a simple expression that computes the difference between the expected values of the critics output for the real examples x and its predictions on the fake examples g(z).
Critic needs to be 1-L Continuous.
For a function to be 1-Lipschitz Continuous, the norm of its gradient needs to be at most 1. The slope (its gradient) can’t be greater than 1 at any point. I.e., this function above, we want to go along every point in this function and make sure its slope ≤ 1. What we can do is actually draw 2 lines at a certain point where you’re evaluating the function:
one with slope = 1 and the other with slope = -1.
You want to make sure that the growth of this function never goes out of bounds from these lines because staying within these lines means that the function is growing linearly.
This function is not Lipschitz Continuous because it’s coming out in all sections.
It’s not staying within this green area, which suggests that it’s growing more than linearly.
Look at another example here. This is a smooth curve functions.
You want to check every single point on this function before you can determine whether or not that this is 1-L Continuous.
Let’s say you take every single value, and the function never grows more than linearly. This function is 1-Lipschitz Continuous.
This condition on the critics NN is important for W-Loss because it assures that the W-Loss function is not only continuous and differentiable, but also that it doesn’t grow too much and maintain some stability during training. This is what makes the underlying EMD valid, which is what W-Loss is founded on.
1-Lipschitz Continuity Enforcement
2 different methods to enforce 1-L continuity on the critic, namely weight clipping and gradient penalty.
Weight (W) clipping
W of the critics NN are forced to take values between a fixed interval. After updating W during gradient descent, we actually will clip any W outside of the desired interval. Which means is W over that interval (i.e., too high/ too low) will be set to the (min, max) amount allowed.
But forcing W of the critic to a limited range, could limit the critics ability to learn, because if the critic’s W can’t take on many different values, it might not be able to improve easily or find the optimal.
On the other hand, it might actually limit the critic too little if you don’t clip the W enough. There’s a lot of hyperparameter tuning involved.
Gradient Penalty.
Gradient penalty is a much softer way to enforce the critic to be 1-l continuous. All you need to do is add a regularization term to the loss function to penalize the critic when its gradient norm is higher than 1.
Lambda is just a hyperparameter value of how much to weigh this regularization term against the main loss function.
It’s not practical to check the critics gradient at every possible point of the feature space. Instead with gradient penalty, all we do is sample some points by interpolating between real and fake examples.
Epsilon could be a weight of 0.3, and here it would evaluate 1– epsilon would be 0.7. That would get you this random interpolated image (X^) that’s in-between these 2 images.
The critic looks at X^, and then we take the norm of that gradient, and we want the norm to be 1. Because this in fact is penalizing any value outside of 1.
The 2 here is just saying, "I want the squared distance as opposed to perhaps the absolute value between them, penalizing values much more when they’re further away from 1.”
Specifically, that X^ is an intermediate image where it’s weighted against the real and a fake using epsilon. With this method, you’re not strictly enforcing 1-L continuity, but you’re just encouraging it. The complete expression, the loss function used for training a GAN with W loss in gradient penalty now has these 2 components.
The 1st part, we approximate EMD with this main W loss component. This makes GAN less parental mode collapse and managed ingredients.
The 2nd part of this loss function is a regularization term that meets the condition for what the critic desires in order to make this main term valid.
This is a soft constraint on making the critic 1-Lipchitz continuous, but it has been shown to be very effective.
