# How to calculate the gradient penalty proposed in "Improved Training of Wasserstein GANs"?

The research paper titled Improved Training of Wasserstein GANs proposed a gradient penalty in order to avoid undesired behavior due to weight clipping of the discriminator.

We now propose an alternative way to enforce the Lipschitz constraint. A differentiable function is 1-Lipschtiz if and only if it has gradients with norm at most 1 everywhere, so we consider directly constraining the gradient norm of the critic’s output with respect to its input. To circumvent tractability issues, we enforce a soft version of the constraint with a penalty on the gradient norm for random samples $$\hat{x} \sim P_\hat{x}$$. Our new objective is

$$L = \mathop{\mathbb{E}}\limits_{\tilde x \sim \mathbb{P}_g} [D(\tilde x)] - \mathop{\mathbb{E}}\limits_{ x \sim \mathbb{P}_r} [D(x)] + \mathop{\mathbb{E}}_{\hat{x} \sim P_{\hat{x}}} [ (\| \triangledown_{\hat{x}} D(\hat{x})\|_2 - 1 )^2 ]$$

The last term in the discriminator's loss function is related to the gradient penalty. It is easy to calculate the first two terms. Since discriminator, in general, gives value in range $$[0, 1]$$, the first two terms are just the average of the sequence of probability values given by discriminator on generated and real images respectively.

But, how to calculate $$\triangledown_{\hat{x}} D(\hat{x})$$ for a given image $$\hat{x}$$?

First of all, the discriminator in WGAN does not give a value in the range $$[0,1]$$. Compared to the traditional discriminator, it has a linear activation in the output layer. Therefore, the authors call it critic instead.

To calculate the penalty, we sample an image that lies on the line between the real and the generated image. This is done by sampling a real image $$x$$, generating an image $$\tilde{x}$$, and mixing these images $$\hat{x} = \alpha \widetilde{x}+(1-\alpha)x$$ with $$\alpha \sim U(0,1)$$. That is, $$\hat{x}$$ is uniformly sampled from the lines between real and fake images, which can be illustrated as follows:

We then feed $$\hat{x}$$ into the critic and calculate the gradient norm of the discriminator's output with respect to its input, $$(\| \triangledown_{\hat{x}} D(\hat{x})\|_2 - 1 )^2$$. Here is a snippet code for PyTorch:

def gradient_penalty(D, real_data, generated_data, device):
batch_size = real_data.shape[0]

# Calculate interpolation
alpha = torch.rand(batch_size, 1, 1, 1)
alpha = alpha.expand_as(real_data).to(device)
# getting x hat
interpolated = alpha * real_data + (1 - alpha) * generated_data

dis_interpolated = D(interpolated)

# Calculate gradients of probabilities with respect to examples