Implementing WGAN-GP: Practice

Having explored the theoretical underpinnings of Wasserstein GANs and the limitations of weight clipping, we now turn to a more robust method for enforcing the Lipschitz constraint: the Gradient Penalty (WGAN-GP). This technique directly penalizes the critic (discriminator) if its gradient norm deviates significantly from 1 along the paths between real and generated data points. Implementing WGAN-GP is a standard practice for achieving stable GAN training, particularly for complex architectures and datasets. This section provides practical guidance and code examples to help you incorporate it into your projects.

Understanding the Gradient Penalty Calculation

Recall that the Wasserstein distance requires the critic to be 1-Lipschitz. WGAN-GP enforces this by adding a penalty term to the critic's loss function. This penalty discourages the gradient norm of the critic's output with respect to its input from deviating far from 1.

The penalty is specifically calculated on points sampled uniformly along the straight lines connecting pairs of points from the real data distribution ($p_{data}$) and the generator distribution ($p_g$). Let $x$ be a real sample and $\tilde{x}$ be a generated sample. An interpolated sample $\hat{x}$ is defined as:

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

where $\epsilon$ is a random number sampled uniformly from $U[0, 1]$.

The gradient penalty term added to the critic loss is:

$$L_{GP} = \lambda \mathbb{E}_{\hat{x} \sim p_{\hat{x}}}[(\| \nabla_{\hat{x}} D(\hat{x}) \|_2 - 1)^2]$$

Here:

• $D(\hat{x})$ is the critic's output for the interpolated sample $\hat{x}$.
• $\nabla_{\hat{x}} D(\hat{x})$ is the gradient of the critic's output with respect to the interpolated input $\hat{x}$.
• $\| \cdot \|_2$ denotes the L2 norm (Euclidean norm).
• $\lambda$ is the penalty coefficient, a hyperparameter controlling the strength of the penalty (commonly set to 10).
• $p_{\hat{x}}$ represents the distribution of these interpolated samples.

The expectation $\mathbb{E}_{\hat{x} \sim p_{\hat{x}}}$ is approximated using a batch of interpolated samples created from the current batch of real and fake data.

Implementing the Gradient Penalty Function

Let's outline the steps to compute this penalty within a typical deep learning framework like PyTorch or TensorFlow. Assume you have a batch of real images (real_samples) and a batch of fake images (fake_samples) generated by the generator, both of the same shape (e.g., [batch_size, channels, height, width]).

1. Generate Epsilon: Create a tensor epsilon with shape [batch_size, 1, 1, 1] (or appropriate shape to broadcast with your image tensors) containing random numbers sampled uniformly between 0 and 1.
2. Create Interpolated Samples: Calculate $\hat{x}$ using broadcasting: interpolated_samples = epsilon * real_samples + (1 - epsilon) * fake_samples. Ensure these samples require gradients for the subsequent step.
3. Calculate Critic Scores for Interpolated Samples: Pass interpolated_samples through the critic network: interpolated_scores = critic(interpolated_samples).
4. Calculate Gradients: Compute the gradients of interpolated_scores with respect to interpolated_samples. Most frameworks provide a function for this (e.g., torch.autograd.grad in PyTorch, tf.GradientTape context in TensorFlow). It's important to set create_graph=True (PyTorch) or ensure the gradient computation is within the tape's context (TensorFlow) if these gradients will be part of a graph used for optimizing the critic (which they are). You need the gradients themselves, not just their contribution to a final loss.
5. Calculate Gradient Norms: Reshape the gradients if necessary to have shape [batch_size, -1] and compute the L2 norm for each sample's gradient across all features. Add a small value (e.g., 1e-8) under the square root for numerical stability when calculating the norm: gradient_norms = sqrt(sum(gradients**2, axis=1) + 1e-8).
6. Compute the Penalty: Apply the WGAN-GP formula: gradient_penalty = lambda * mean((gradient_norms - 1)**2).

The following diagram visualizes the sampling process for $\hat{x}$:

Interpolated samples ($\hat{x}$, green) are chosen along lines connecting real samples ($x$, blue) and generated samples ($\tilde{x}$, red). The critic's gradient norm is penalized at these interpolated points.

Below is an implementation snippet using PyTorch syntax:

import torch
import torch.autograd as autograd

def compute_gradient_penalty(critic, real_samples, fake_samples, lambda_gp, device):
    """Calculates the gradient penalty loss for WGAN GP"""
    # Random weight term for interpolation between real and fake samples
    alpha = torch.rand(real_samples.size(0), 1, 1, 1, device=device)
    # Get random interpolation between real and fake samples
    interpolates = (alpha * real_samples + ((1 - alpha) * fake_samples)).requires_grad_(True)

    critic_interpolates = critic(interpolates)

    # Use autograd to compute gradients
    gradients = autograd.grad(
        outputs=critic_interpolates,
        inputs=interpolates,
        grad_outputs=torch.ones(critic_interpolates.size(), device=device), # Ensure gradients flow for all outputs
        create_graph=True, # Create graph for second derivative during critic update
        retain_graph=True, # Retain graph for generator update
        only_inputs=True,
    )[0] # Get grads w.r.t. inputs

    # Reshape gradients and compute norm
    gradients = gradients.view(gradients.size(0), -1)
    # Add small epsilon for numerical stability
    gradient_norm = torch.sqrt(torch.sum(gradients ** 2, dim=1) + 1e-12)
    # Compute penalty
    gradient_penalty = lambda_gp * ((gradient_norm - 1) ** 2).mean()
    return gradient_penalty

# Example usage within a training step:
# lambda_gp = 10
# Assuming critic, real_batch, fake_batch, device are defined
# gp = compute_gradient_penalty(critic, real_batch, fake_batch, lambda_gp, device)
# critic_loss = ... + gp # Add to the main critic loss

Modifying the Training Loop

Integrating WGAN-GP involves two main changes compared to the original GAN or WGAN with weight clipping:

1. Critic Loss Calculation: The critic's objective function becomes:

   $$L_{Critic} = \mathbb{E}_{\tilde{x} \sim p_g}[D(\tilde{x})] - \mathbb{E}_{x \sim p_{data}}[D(x)] + \lambda \mathbb{E}_{\hat{x} \sim p_{\hat{x}}}[(\| \nabla_{\hat{x}} D(\hat{x}) \|_2 - 1)^2]$$

   Note that we aim to minimize this loss. The first two terms approximate the negative Wasserstein distance, and the last term is the gradient penalty. Some implementations might flip the signs of the first two terms if the goal is maximization. Ensure your optimizer step aligns with minimizing or maximizing the objective.

2. Optimizer and Architecture:

   • Remove weight clipping from the critic optimizer.
   • Adam is often a suitable optimizer, potentially with adjusted hyperparameters (e.g., beta1=0.0, beta2=0.9 as used in some WGAN-GP papers, though standard beta1=0.5, beta2=0.999 also works).
   • It's common practice to update the critic several times (e.g., 5 times) for each update of the generator, especially early in training, to ensure the critic provides meaningful gradients.

A typical WGAN-GP training step looks like this (pseudo-code):

# Hyperparameters
lambda_gp = 10
critic_iterations = 5
learning_rate = 0.0001
beta1 = 0.0
beta2 = 0.9

# Optimizers (Adam often used)
optimizer_G = Adam(generator.parameters(), lr=learning_rate, betas=(beta1, beta2))
optimizer_C = Adam(critic.parameters(), lr=learning_rate, betas=(beta1, beta2))

for epoch in range(num_epochs):
    for i, real_batch in enumerate(data_loader):
        # ---------------------
        # Train Critic
        # ---------------------
        optimizer_C.zero_grad()

        real_samples = real_batch.to(device)
        batch_size = real_samples.size(0)

        # Sample noise and generate fake samples
        z = torch.randn(batch_size, latent_dim, 1, 1, device=device)
        fake_samples = generator(z).detach() # Detach to avoid training G through C

        # Get critic scores
        real_scores = critic(real_samples)
        fake_scores = critic(fake_samples)

        # Calculate gradient penalty
        gradient_penalty = compute_gradient_penalty(
            critic, real_samples.data, fake_samples.data, lambda_gp, device
        )

        # Calculate critic loss: -(Wasserstein Loss) + Gradient Penalty
        # We minimize this loss, equivalent to maximizing (Real Scores - Fake Scores - GP)
        critic_loss = -torch.mean(real_scores) + torch.mean(fake_scores) + gradient_penalty
        
        # Backpropagate and update critic
        critic_loss.backward()
        optimizer_C.step()

        # Train the generator only every 'critic_iterations' steps
        if i % critic_iterations == 0:
            # -----------------
            # Train Generator
            # -----------------
            optimizer_G.zero_grad()

            # Generate a new batch of fake samples (with gradients)
            z = torch.randn(batch_size, latent_dim, 1, 1, device=device)
            gen_samples = generator(z)

            # Get critic scores for generated samples
            gen_scores = critic(gen_samples)

            # Calculate generator loss (maximize score of fake samples)
            # We minimize the negative score
            generator_loss = -torch.mean(gen_scores)

            # Backpropagate and update generator
            generator_loss.backward()
            optimizer_G.step()

        # Logging, saving models, etc.
        # ...

Implementation Considerations

• Lambda ($\lambda$) Value: The choice of $\lambda=10$ is empirically derived and works well for many problems. However, it is a hyperparameter. If training is unstable, or if the gradient norm consistently deviates far from 1, you might experiment with adjusting $\lambda$, though 10 is a standard starting point.
• Numerical Stability: The small epsilon (e.g., 1e-8 or 1e-12) added before taking the square root when calculating the L2 norm is important to prevent NaN values if the gradient vector happens to be zero.
• Framework Differences: While the concept is identical, the exact syntax for computing gradients with respect to intermediate variables differs slightly between frameworks. Ensure you use the correct functions (torch.autograd.grad, tf.GradientTape.gradient) and manage gradient calculation contexts appropriately. Pay attention to arguments like create_graph=True (PyTorch) or nesting GradientTape (TensorFlow) to allow gradients to flow back through the penalty calculation during the critic's optimization step.
• Debugging: Monitor the components of the critic loss: the Wasserstein estimate (mean(fake_scores) - mean(real_scores)) and the gradient penalty term. Also, monitor the average gradient norm (mean(gradient_norms)) itself. Ideally, the average gradient norm should hover around 1 during stable training. If it consistently stays much higher or lower, it might indicate issues with learning rates, network capacity, or the $\lambda$ value.

By replacing weight clipping with the gradient penalty, you provide a smoother, more reliable way to enforce the Lipschitz constraint, often leading to significantly improved training stability and sample quality compared to the original WGAN formulation. This makes WGAN-GP a valuable technique for your advanced GAN toolkit.