Hands-on Practical: Implementing WGAN-GP

Wasserstein GANs (WGAN) and the Gradient Penalty (GP) technique address the weight clipping limitations of the original WGAN, stabilizing training. A practical implementation of WGAN-GP follows. This approach is widely regarded as a significant improvement for GAN training stability and sample quality.

This practical guide assumes you are comfortable with implementing basic GANs in PyTorch or TensorFlow. We will focus on the specific modifications required for WGAN-GP.

Core Components of WGAN-GP

Implementing WGAN-GP involves adjustments primarily to the critic (discriminator) network, the loss functions, and the training loop.

1. Critic Architecture: The critic in WGAN-GP does not have a final sigmoid activation function. Its role is to output a scalar score (representing the "realness" according to the Wasserstein distance approximation) rather than a probability. The output layer should be linear.
2. Loss Functions: We replace the standard GAN's log-loss with losses derived from the Wasserstein distance estimate and the gradient penalty.
3. Gradient Penalty: This is the defining feature. Instead of weight clipping, we add a penalty term to the critic's loss that encourages the norm of the critic's gradient with respect to its input to be close to 1. This enforces the Lipschitz constraint more effectively.
4. Training Loop: Typically, the critic is updated more frequently than the generator within each training iteration (e.g., 5 critic updates per generator update).

Implementing the Critic Loss

The critic aims to maximize the difference between its scores for real samples and generated samples, while also incorporating the gradient penalty. The loss function for the critic ($D$) is:

$$L_D = \mathbb{E}_{\tilde{x} \sim P_g}[D(\tilde{x})] - \mathbb{E}_{x \sim P_r}[D(x)] + \lambda \mathbb{E}_{\hat{x} \sim P_{\hat{x}}}[(\|\nabla_{\hat{x}} D(\hat{x})\|_2 - 1)^2]$$

Where:

• $P_g$ is the generator's distribution (fake samples $\tilde{x}$).
• $P_r$ is the real data distribution (real samples $x$).
• $P_{\hat{x}}$ is the distribution of interpolated samples $\hat{x}$.
• $\lambda$ is the gradient penalty coefficient (often set to 10).

Let's break down the implementation, particularly the gradient penalty term.

Implementing the Gradient Penalty

Calculating the gradient penalty involves several steps:

1. Sample Interpolated Points: For each real sample $x$ and generated sample $\tilde{x}$ in a batch, create an interpolated sample $\hat{x}$.

   $$\hat{x} = \epsilon x + (1 - \epsilon) \tilde{x}$$

   Here, $\epsilon$ is a random number sampled uniformly from $U[0, 1]$. This needs to be done element-wise for batches.

2. Calculate Critic Output for Interpolated Points: Pass these interpolated samples $\hat{x}$ through the critic network to get their scores $D(\hat{x})$.

3. Compute Gradients: Calculate the gradients of the critic's outputs $D(\hat{x})$ with respect to the interpolated inputs $\hat{x}$. This requires using your deep learning framework's automatic differentiation capabilities (e.g., torch.autograd.grad in PyTorch or tf.GradientTape in TensorFlow). It's important to ensure gradients are computed for the inputs (create_graph=True in PyTorch is needed as the gradient penalty itself forms part of the loss graph).

4. Calculate Gradient Norm: Compute the L2 norm (Euclidean norm) of these gradients for each interpolated sample.

5. Compute Penalty: Calculate the penalty for each sample as $( \| \nabla_{\hat{x}} D(\hat{x}) \|_2 - 1 )^2$.

6. Average and Scale: Average the penalties across the batch and multiply by the coefficient $\lambda$.

Here's a PyTorch implementation snippet for the gradient penalty function:

import torch
import torch.autograd as autograd

def compute_gradient_penalty(critic, real_samples, fake_samples, device):
    """Calculates the gradient penalty loss for WGAN GP"""
    batch_size = real_samples.size(0)
    # Random weight term for interpolation between real and fake samples
    alpha = torch.rand(batch_size, 1, 1, 1, device=device) # Assuming 4D tensor (B, C, H, W)
    # Expand alpha to match the image dimensions
    alpha = alpha.expand_as(real_samples)

    # Get random interpolation between real and fake samples
    interpolates = (alpha * real_samples + ((1 - alpha) * fake_samples)).requires_grad_(True)

    # Get critic scores for the interpolates
    d_interpolates = critic(interpolates)

    # Use fake ones tensor as target for gradient computation
    fake = torch.ones(batch_size, 1, device=device, requires_grad=False) # Use size matching critic output

    # Get gradient w.r.t. interpolates
    gradients = autograd.grad(
        outputs=d_interpolates,
        inputs=interpolates,
        grad_outputs=fake, # Gradient outputs must match d_interpolates shape
        create_graph=True, # Create graph for second derivative (part of GP loss)
        retain_graph=True, # Retain graph for further computations (critic loss)
        only_inputs=True,
    )[0]

    # Reshape gradients to easily compute the norm per sample
    gradients = gradients.view(gradients.size(0), -1)
    # Compute the L2 norm and the penalty
    gradient_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean()
    return gradient_penalty

# --- Inside the training loop ---
# Assuming 'critic', 'real_imgs', 'fake_imgs' are defined
# LAMBDA_GP = 10 # Gradient penalty coefficient

# gradient_penalty = compute_gradient_penalty(critic, real_imgs.data, fake_imgs.data, device)
# critic_loss = torch.mean(critic_fake) - torch.mean(critic_real) + LAMBDA_GP * gradient_penalty
# critic_loss.backward()
# optimizer_D.step()

Note: Ensure the shapes used for alpha, fake, and gradient calculations match your specific data and critic output dimensions. The requires_grad_(True) on interpolates and create_graph=True, retain_graph=True in autograd.grad are essential for correctly computing the penalty.

Implementing the Generator Loss

The generator ($G$) aims to produce samples that the critic scores highly (i.e., makes the critic think they are real). Its loss is simpler:

$$L_G = - \mathbb{E}_{\tilde{x} \sim P_g}[D(\tilde{x})]$$

In practice, this means generating a batch of fake samples, passing them through the critic, and minimizing the negative mean of the resulting scores.

# --- Inside the training loop, during generator update ---
# Generate fake images
# z = torch.randn(batch_size, latent_dim, 1, 1, device=device)
# gen_imgs = generator(z)

# Calculate generator loss
# fake_scores = critic(gen_imgs)
# generator_loss = -torch.mean(fake_scores)

# generator_loss.backward()
# optimizer_G.step()

Training Procedure

The typical WGAN-GP training loop involves alternating updates between the critic and the generator. A common practice is to perform multiple critic updates for each generator update.

Training loop structure for WGAN-GP, emphasizing multiple critic updates per generator update.

Considerations:

• Optimizers: Adam is commonly used, often with specific hyperparameters like $\beta_1=0.0$ or $\beta_1=0.5$ and $\beta_2=0.9$. Standard Adam settings ($\beta_1=0.9, \beta_2=0.999$) can also work but might require more tuning. Use separate optimizer instances for the generator and the critic.
• Learning Rates: Similar learning rates for both (e.g., 1e-4 or 2e-4) are often a good starting point, unlike TTUR which explicitly uses different rates.
• Critic Updates ($n_{critic}$): Values like 5 are common, but this can be tuned. It ensures the critic provides reliable gradients to the generator.
• Gradient Penalty Coefficient ($\lambda$): Typically set to 10, but can be adjusted if training is unstable or gradients vanish/explode.
• Batch Normalization: Generally avoided in the critic for WGAN-GP, as it can introduce dependencies between samples in a batch, interfering with the gradient penalty calculation. Layer Normalization or Instance Normalization might be alternatives if normalization is needed. For the generator, Batch Normalization is often still used.

By implementing these components, particularly the gradient penalty calculation and the adjusted loss functions, you can leverage WGAN-GP to train more stable GANs capable of generating higher quality synthetic data compared to the standard GAN formulation or the original WGAN with weight clipping. Remember to monitor the critic loss, generator loss, and the magnitude of the gradient penalty during training to diagnose potential issues.