Alternative Loss Functions (WGAN, WGAN-GP, LSGAN)

The standard GAN min-max game, while elegant, often suffers from training instability. The generator and discriminator might fail to converge, gradients can vanish, or the generator might only learn to produce a limited variety of outputs (mode collapse). A significant contributor to these problems is the nature of the Jensen-Shannon (JS) divergence that the original GAN objective implicitly minimizes. When the distributions of real and generated data are disjoint or have negligible overlap (which happens frequently early in training or with high-dimensional data), the JS divergence saturates, leading to near-zero gradients for the generator.

To address these shortcomings, researchers have proposed alternative loss functions that provide more stable gradient signals and better reflect the distance between the real and generated data distributions. Let's examine three prominent alternatives: Wasserstein GAN (WGAN), WGAN with Gradient Penalty (WGAN-GP), and Least Squares GAN (LSGAN).

Wasserstein GAN (WGAN)

The core idea behind WGAN is to replace the JS divergence with the Wasserstein-1 distance, also known as the Earth Mover's distance ($W_1$). Intuitively, if you imagine the real and generated distributions as piles of earth, $W_1$ measures the minimum "cost" (amount of earth multiplied by the distance moved) to transform one pile into the other. Unlike JS divergence, the Wasserstein distance provides a meaningful gradient even when the distributions don't overlap significantly, making it a more suitable metric for GAN training.

Calculating $W_1$ directly is intractable. However, the Kantorovich-Rubinstein duality provides a way to compute it:

$$W_1(P_r, P_g) = \sup_{||f||_L \le 1} \mathbb{E}_{x \sim P_r}[f(x)] - \mathbb{E}_{\tilde{x} \sim P_g}[f(\tilde{x})]$$

Here, $P_r$ is the real data distribution, $P_g$ is the generator's distribution ($\tilde{x} = G(z)$), and the supremum is taken over all 1-Lipschitz functions $f$. A function $f$ is 1-Lipschitz if $|f(x_1) - f(x_2)| \le |x_1 - x_2|$ for all $x_1, x_2$.

In the WGAN framework, the discriminator (now often called a "critic" denoted by $D$ or $f_w$) is trained to approximate this optimal function $f$. The critic outputs a scalar score (not a probability) reflecting the "realness" of the input. The WGAN objective becomes:

$$\min_G \max_{w \in \mathcal{W}} \mathbb{E}_{x \sim P_r}[D_w(x)] - \mathbb{E}_{z \sim p(z)}[D_w(G(z))]$$

Here, $w$ represents the critic's parameters. The constraint $||f||_L \le 1$ (the Lipschitz constraint) is important. The original WGAN paper proposed enforcing this by clipping the weights of the critic $w$ to lie within a small fixed range, like $[-c, c]$, after each gradient update.

Critic Update: Maximize $\mathbb{E}_{x \sim P_r}[D_w(x)] - \mathbb{E}_{z \sim p(z)}[D_w(G(z))]$. This pushes the score for real samples up and the score for fake samples down. Generator Update: Minimize $-\mathbb{E}_{z \sim p(z)}[D_w(G(z))]$. This is equivalent to maximizing the critic's score for fake samples, encouraging the generator to produce samples that the critic scores higher (i.e., considers more "real").

While WGAN with weight clipping often leads to more stable training and helps prevent mode collapse compared to the standard GAN, weight clipping is a somewhat crude way to enforce the Lipschitz constraint. It can lead to:

1. Capacity Underuse: Pushing weights towards the clipping boundaries might limit the critic's modeling capacity.
2. Exploding or Vanishing Gradients: If the clipping parameter $c$ is too large or too small, gradients can still behave poorly. Finding the right $c$ can be difficult.

WGAN with Gradient Penalty (WGAN-GP)

WGAN-GP addresses the issues of weight clipping by proposing a more direct way to enforce the Lipschitz constraint: penalizing the gradient norm of the critic with respect to its input. A differentiable function is 1-Lipschitz if and only if its gradients have a norm of at most 1 everywhere. Instead of enforcing this strictly, WGAN-GP adds a penalty term to the critic's loss that encourages this condition.

The penalty focuses on points sampled between the real and generated distributions. For a pair of real sample $x$ and generated sample $\tilde{x} = G(z)$, an interpolated sample $\hat{x}$ is created:

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

where $\epsilon$ is sampled uniformly from $U[0, 1]$. The gradient penalty term is then:

$$\lambda \mathbb{E}_{\hat{x} \sim P_{\hat{x}}}[(||\nabla_{\hat{x}} D_w(\hat{x})||_2 - 1)^2]$$

Here, $P_{\hat{x}}$ is the distribution of interpolated samples, $||\cdot||_2$ is the L2 norm (Euclidean norm), and $\lambda$ is a penalty coefficient (typically set to 10). This term penalizes deviations of the gradient norm from 1 at these interpolated points.

The WGAN-GP critic loss becomes:

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

The critic aims to minimize this loss (note the sign change compared to the WGAN maximization formulation, which is common in implementations). The generator loss remains the same as in WGAN, aiming to maximize the critic's score for generated samples:

$$L_{Generator} = - \mathbb{E}_{\tilde{x} \sim P_g}[D_w(\tilde{x})]$$

WGAN-GP typically results in significantly more stable training than the original WGAN and standard GANs, often producing higher-quality samples without requiring careful tuning of clipping parameters. It has become a widely adopted baseline for GAN training. Important implementation details include removing batch normalization layers from the critic (as the penalty is sensitive to batch statistics) and using layer normalization or other alternatives if normalization is needed.

Least Squares GAN (LSGAN)

LSGAN tackles training instability from a different angle. It observes that the sigmoid cross-entropy loss function used in the standard GAN discriminator can lead to vanishing gradients. When the discriminator successfully classifies a generated sample as fake (outputting a probability close to 0), the gradient flowing back to the generator becomes very small, slowing down learning.

LSGAN replaces the sigmoid cross-entropy loss with a least squares (mean squared error) objective. The LSGAN objectives are:

Discriminator Loss:

$$L_D^{\text{LSGAN}} = \frac{1}{2} \mathbb{E}_{x \sim P_r}[(D(x) - b)^2] + \frac{1}{2} \mathbb{E}_{z \sim p(z)}[(D(G(z)) - a)^2]$$

Generator Loss:

$$L_G^{\text{LSGAN}} = \frac{1}{2} \mathbb{E}_{z \sim p(z)}[(D(G(z)) - c)^2]$$

Here, $a$ and $b$ are the target labels for fake and real data, respectively, and $c$ is the value the generator wants the discriminator to output for fake data. A common choice is $a=0$, $b=1$, and $c=1$.

The main idea is that the least squares loss penalizes samples even if they are classified correctly but lie far from the decision boundary (defined by $D(x)=c$). By minimizing $L_G$, the generator tries to fool the discriminator by producing samples $\tilde{x}$ such that $D(\tilde{x})$ is close to the target label for real data ($b$, or $c$ in the generator's objective formulation). The quadratic loss ensures that gradients do not vanish as quickly as with the sigmoid cross-entropy, leading to a more stable learning process and potentially higher quality results. LSGAN is generally simpler to implement than WGAN-GP as it doesn't require gradient penalties or weight clipping.

Choosing an Alternative Loss

• WGAN-GP is often considered a strong choice for stabilizing GAN training, particularly for image generation tasks. It directly addresses the theoretical limitations of the JS divergence and provides smoother gradients. The main drawback is the computational cost of the gradient penalty.
• LSGAN provides a simpler alternative that also improves stability by avoiding the saturation problem of sigmoid cross-entropy. It can be effective and computationally less demanding than WGAN-GP.
• Original WGAN (with weight clipping) is less commonly used now due to the practical issues with clipping, but understanding its principles is important context for WGAN-GP.

Experimentation is often necessary, but WGAN-GP and LSGAN provide powerful tools to overcome the notorious instability of GAN training, allowing you to train more complex models and achieve better results. The hands-on practical later in this chapter will guide you through implementing WGAN-GP.