GANs Deep Dive

1. GAN Theory

1.1 Minimax Formulation

The Generative Adversarial Network (Goodfellow et al., 2014) frames generative modelling as a two-player minimax game:

where:

• is the generator mapping latent codes to samples
• is the discriminator (critic) classifying real vs fake
• is the true data distribution
• is the prior (typically )

1.2 Optimal Discriminator

For fixed , the optimal discriminator is:

where is the distribution induced by .

Proof: For any , the objective is:

For each , we maximise where , .

Setting gives .

1.3 Global Optimality

Substituting back into :

where is the Jensen-Shannon divergence:

Theorem (Goodfellow et al., 2014): The global minimum of is achieved if and only if , at which point .

1.4 Jensen-Shannon Divergence Properties

The JS divergence has several important properties:

1. Symmetry:
2. Boundedness:
3. Non-zero gradients: Unlike KL divergence, JS divergence has non-vanishing gradients when distributions are disjoint

For disjoint supports:

This is the vanishing gradient problem that WGAN addresses.

2. Training Dynamics

2.1 Alternating Optimisation

The training alternates between:

1. Discriminator step: Maximise for fixed
2. Generator step: Minimise for fixed

In practice, we maximise instead of minimising to avoid saturation.

2.2 Training Instability

GAN training is notoriously unstable due to:

1. Mode collapse: maps different to the same output
2. Oscillation: and oscillate without converging
3. Vanishing gradients: becomes too good, providing no useful gradients to
4. Non-convergence: The game has no Nash equilibrium in function space

2.3 Spectral Normalization

To stabilise training, spectral normalisation constrains the Lipschitz constant of :

where is the spectral norm (largest singular value).

This ensures:

with controlled by the spectral norms of each layer.

3. Wasserstein GAN (WGAN)

3.1 Wasserstein Distance

The Wasserstein distance (Earth Mover's distance) measures the minimum "work" to transform one distribution into another:

where is the set of all joint distributions with marginals and .

3.2 Kantorovich-Rubinstein Duality

By the Kantorovich-Rubinstein duality:

where the supremum is over all 1-Lipschitz functions .

3.3 WGAN Objective

where is the set of 1-Lipschitz functions.

Advantages:

1. Meaningful loss metric (correlates with sample quality)
2. No mode collapse
3. Stable training

3.4 Gradient Penalty (WGAN-GP)

Instead of weight clipping, enforce the Lipschitz constraint via gradient penalty:

where for .

Total objective:

4. StyleGAN Architecture

4.1 Mapping Network

StyleGAN uses a mapping network to map the latent code to an intermediate space:

The mapping network is an 8-layer MLP that learns a more disentangled representation.

4.2 Adaptive Instance Normalisation (AdaIN)

Styles are injected via adaptive instance normalisation:

where are learned affine transformations from .

4.3 Style Mixing

During training, styles from different latent codes are mixed at different layers:

where can come from different latent codes, enabling disentanglement of coarse (pose) and fine (texture) styles.

4.4 Truncation Trick

To trade off diversity for quality:

where is the mean latent code and controls the truncation.

5. Mode Collapse

5.1 Types of Mode Collapse

1. Intra-batch collapse: Generator produces limited variety within a batch
2. Temporal collapse: Generator produces same outputs across iterations
3. Manifold collapse: Generator maps all inputs to a lower-dimensional manifold

5.2 Detection

Mode collapse can be detected by:

• Inception Score (IS): Low diversity
• Fréchet Inception Distance (FID): Poor coverage
• Precision vs Recall: High precision, low recall

5.3 Mitigation Strategies

6. Code Examples

6.1 WGAN-GP Implementation

import torch
import torch.nn as nn
import torch.autograd as autograd

class Generator(nn.Module):
    """WGAN-GP Generator."""
    
    def __init__(self, latent_dim=128, channels=3, features=64):
        super().__init__()
        
        self.net = nn.Sequential(
            # Input: latent_dim x 1 x 1
            nn.ConvTranspose2d(latent_dim, features * 8, 4, 1, 0, bias=False),
            nn.BatchNorm2d(features * 8),
            nn.ReLU(True),
            
            nn.ConvTranspose2d(features * 8, features * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(features * 4),
            nn.ReLU(True),
            
            nn.ConvTranspose2d(features * 4, features * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(features * 2),
            nn.ReLU(True),
            
            nn.ConvTranspose2d(features * 2, features, 4, 2, 1, bias=False),
            nn.BatchNorm2d(features),
            nn.ReLU(True),
            
            nn.ConvTranspose2d(features, channels, 4, 2, 1, bias=False),
            nn.Tanh()
        )
    
    def forward(self, z):
        return self.net(z.view(z.size(0), -1, 1, 1))


class Critic(nn.Module):
    """WGAN-GP Critic (not sigmoid output)."""
    
    def __init__(self, channels=3, features=64):
        super().__init__()
        
        self.net = nn.Sequential(
            nn.Conv2d(channels, features, 4, 2, 1),
            nn.LeakyReLU(0.2, inplace=True),
            
            nn.Conv2d(features, features * 2, 4, 2, 1, bias=False),
            nn.InstanceNorm2d(features * 2, affine=True),
            nn.LeakyReLU(0.2, inplace=True),
            
            nn.Conv2d(features * 2, features * 4, 4, 2, 1, bias=False),
            nn.InstanceNorm2d(features * 4, affine=True),
            nn.LeakyReLU(0.2, inplace=True),
            
            nn.Conv2d(features * 4, features * 8, 4, 2, 1, bias=False),
            nn.InstanceNorm2d(features * 8, affine=True),
            nn.LeakyReLU(0.2, inplace=True),
            
            nn.Conv2d(features * 8, 1, 4, 1, 0),
        )
    
    def forward(self, x):
        return self.net(x).view(-1)


def compute_gradient_penalty(critic, real_data, fake_data, device, lambda_gp=10):
    """
    Compute gradient penalty for WGAN-GP.
    
    GP = λ * E[(||∇D(û)||₂ - 1)²]
    where û = εx + (1-ε)G(z)
    """
    batch_size = real_data.size(0)
    epsilon = torch.rand(batch_size, 1, 1, 1, device=device)
    epsilon = epsilon.expand_as(real_data)
    
    # Interpolated samples
    interpolated = (epsilon * real_data + (1 - epsilon) * fake_data).requires_grad_(True)
    
    # Critic output for interpolated samples
    d_interpolated = critic(interpolated)
    
    # Compute gradients
    gradients = autograd.grad(
        outputs=d_interpolated,
        inputs=interpolated,
        grad_outputs=torch.ones_like(d_interpolated),
        create_graph=True,
        retain_graph=True
    )[0]
    
    gradients = gradients.view(batch_size, -1)
    gradient_penalty = lambda_gp * ((gradients.norm(2, dim=1) - 1) ** 2).mean()
    
    return gradient_penalty


def train_wgan_gp(generator, critic, dataloader, latent_dim=128, 
                   n_epochs=200, lr=0.0001, n_critic=5, device='cuda'):
    """Train WGAN-GP."""
    
    optimizer_G = torch.optim.Adam(generator.parameters(), lr=lr, betas=(0.0, 0.9))
    optimizer_C = torch.optim.Adam(critic.parameters(), lr=lr, betas=(0.0, 0.9))
    
    for epoch in range(n_epochs):
        for i, (real_images, _) in enumerate(dataloader):
            real_images = real_images.to(device)
            batch_size = real_images.size(0)
            
            # ---------------------
            #  Train Critic
            # ---------------------
            optimizer_C.zero_grad()
            
            # Generate fake images
            z = torch.randn(batch_size, latent_dim, device=device)
            fake_images = generator(z).detach()
            
            # Critic losses
            loss_real = critic(real_images).mean()
            loss_fake = critic(fake_images).mean()
            gp = compute_gradient_penalty(critic, real_images, fake_images, device)
            
            loss_C = loss_fake - loss_real + gp
            loss_C.backward()
            optimizer_C.step()
            
            # -----------------
            #  Train Generator
            # -----------------
            if (i + 1) % n_critic == 0:
                optimizer_G.zero_grad()
                
                z = torch.randn(batch_size, latent_dim, device=device)
                fake_images = generator(z)
                
                loss_G = -critic(fake_images).mean()
                loss_G.backward()
                optimizer_G.step()
        
        # Print progress
        if (epoch + 1) % 10 == 0:
            print(f"Epoch [{epoch+1}/{n_epochs}] "
                  f"C_loss: {loss_C.item():.4f} "
                  f"G_loss: {loss_G.item():.4f} "
                  f"Wasserstein: {(loss_real - loss_fake).item():.4f}")

6.2 StyleGAN Generator

class MappingNetwork(nn.Module):
    """StyleGAN mapping network: z -> w."""
    
    def __init__(self, latent_dim=512, w_dim=512, num_layers=8):
        super().__init__()
        
        layers = []
        for i in range(num_layers):
            in_dim = latent_dim if i == 0 else w_dim
            layers.append(nn.Linear(in_dim, w_dim))
            layers.append(nn.LeakyReLU(0.2))
        
        self.net = nn.Sequential(*layers)
    
    def forward(self, z):
        return self.net(z)


class AdaIN(nn.Module):
    """Adaptive Instance Normalisation."""
    
    def __init__(self, style_dim, num_features):
        super().__init__()
        self.norm = nn.InstanceNorm2d(num_features, affine=False)
        self.style_scale = nn.Linear(style_dim, num_features)
        self.style_bias = nn.Linear(style_dim, num_features)
    
    def forward(self, x, style):
        # Normalise
        x = self.norm(x)
        
        # Apply style
        scale = self.style_scale(style).unsqueeze(2).unsqueeze(3)
        bias = self.style_bias(style).unsqueeze(2).unsqueeze(3)
        
        return x * (1 + scale) + bias


class StyleBlock(nn.Module):
    """Single StyleGAN block with style injection."""
    
    def __init__(self, in_channels, out_channels, style_dim):
        super().__init__()
        
        self.conv = nn.Conv2d(in_channels, out_channels, 3, padding=1)
        self.adain = AdaIN(style_dim, out_channels)
        self.noise_weight = nn.Parameter(torch.zeros(1, out_channels, 1, 1))
        self.activation = nn.LeakyReLU(0.2)
    
    def forward(self, x, style, noise=None):
        # Convolution
        x = self.conv(x)
        
        # Add noise
        if noise is None:
            noise = torch.randn(x.size(0), 1, x.size(2), x.size(3), device=x.device)
        x = x + self.noise_weight * noise
        
        # Apply style
        x = self.adain(x, style)
        x = self.activation(x)
        
        return x


class StyleGANGenerator(nn.Module):
    """Simplified StyleGAN generator."""
    
    def __init__(self, latent_dim=512, style_dim=512, channels=3):
        super().__init__()
        
        self.mapping = MappingNetwork(latent_dim, style_dim)
        
        # Initial constant input
        self.constant = nn.Parameter(torch.randn(1, 512, 4, 4))
        
        # Style blocks
        self.blocks = nn.ModuleList([
            StyleBlock(512, 512, style_dim),  # 4x4
            StyleBlock(512, 512, style_dim),  # 8x8
            StyleBlock(512, 256, style_dim),  # 16x16
            StyleBlock(256, 128, style_dim),  # 32x32
            StyleBlock(128, 64, style_dim),   # 64x64
            StyleBlock(64, 32, style_dim),    # 128x128
            StyleBlock(32, 16, style_dim),    # 256x256
        ])
        
        # To RGB
        self.to_rgb = nn.Conv2d(16, channels, 1)
    
    def forward(self, z):
        # Map z to w
        w = self.mapping(z)
        
        # Start from constant
        x = self.constant.expand(z.size(0), -1, -1, -1)
        
        # Apply style blocks
        for block in self.blocks:
            x = F.interpolate(x, scale_factor=2, mode='bilinear')
            x = block(x, w)
        
        # To RGB
        x = self.to_rgb(x)
        x = torch.tanh(x)
        
        return x


# Example: Create StyleGAN
stylegan = StyleGANGenerator(latent_dim=512, style_dim=512)

z = torch.randn(4, 512)
fake_images = stylegan(z)

print(f"Input shape: {z.shape}")
print(f"Output shape: {fake_images.shape}")
print(f"Parameters: {sum(p.numel() for p in stylegan.parameters()):,}")

6.3 Evaluation Metrics

import torch
from torch.nn import functional as F

def compute_inception_score(images, inception_model, batch_size=32, splits=10):
    """
    Compute Inception Score.
    
    IS = exp(E[KL(p(y|x) || p(y))])
    
    High IS: diverse and realistic samples
    """
    inception_model.eval()
    
    # Get predictions
    preds = []
    with torch.no_grad():
        for i in range(0, len(images), batch_size):
            batch = images[i:i+batch_size]
            pred = F.softmax(inception_model(batch), dim=1)
            preds.append(pred)
    
    preds = torch.cat(preds, dim=0)
    
    # Compute marginal distribution
    py = preds.mean(dim=0)
    
    # Compute KL divergence
    scores = []
    for i in range(splits):
        chunk = preds[i * len(preds) // splits:(i + 1) * len(preds) // splits]
        kl = (chunk * (torch.log(chunk + 1e-10) - torch.log(py + 1e-10))).sum(dim=1).mean()
        scores.append(torch.exp(kl).item())
    
    return sum(scores) / len(scores)


def compute_fid(real_features, fake_features):
    """
    Compute Fréchet Inception Distance.
    
    FID = ||μ_r - μ_f||² + Tr(Σ_r + Σ_f - 2(Σ_r Σ_f)^{1/2})
    
    Low FID: similar distributions
    """
    mu_real = real_features.mean(dim=0)
    mu_fake = fake_features.mean(dim=0)
    
    sigma_real = torch.cov(real_features.T)
    sigma_fake = torch.cov(fake_features.T)
    
    # Compute FID
    diff = mu_real - mu_fake
    
    # Product of covariances
    covmean = torch.mm(sigma_real, sigma_fake).clamp(min=0)
    
    # Matrix square root (approximation)
    covmean = torch.linalg.matrix_sqrt(covmean + 1e-6 * torch.eye(len(covmean)))
    
    fid = diff @ diff + torch.trace(sigma_real + sigma_fake - 2 * covmean)
    
    return fid.item()

7. Summary

1. GANs frame generative modelling as a minimax game, with the global optimum at .

2. WGAN uses Wasserstein distance for more stable training and meaningful loss metrics.

3. Mode collapse is a fundamental challenge; WGAN and architectural innovations help mitigate it.

4. StyleGAN achieves state-of-the-art image generation through style-based synthesis.

5. Evaluation metrics (IS, FID) provide quantitative measures of sample quality and diversity.

References

• Goodfellow, I., et al. (2014). Generative adversarial nets. NeurIPS.
• Arjovsky, M., et al. (2017). Wasserstein GAN. arXiv.
• Gulrajani, I., et al. (2017). Improved training of Wasserstein GANs. NeurIPS.
• Karras, T., et al. (2019). A style-based generator architecture for generative adversarial networks. CVPR.
• Karras, T., et al. (2020). Analyzing and improving the image quality of StyleGAN. CVPR.
• Salimans, T., et al. (2016). Improved techniques for training GANs. NeurIPS.
• Metz, L., et al. (2016). Unrolled generative adversarial networks. ICLR.