Weight Initialization Techniques

The stability and performance of Generative Adversarial Networks, especially the complex architectures discussed in this course, are highly sensitive to the initial values assigned to the network weights. While modern deep learning frameworks provide default initializations (often He or Glorot/Xavier), the adversarial nature of GAN training introduces unique challenges that may require more careful consideration of how weights are set at the beginning of training. Poor initialization can lead to immediate problems like vanishing or exploding gradients, mode collapse, or one network overwhelming the other before learning can effectively begin.

The Importance of Controlled Initialization

Standard initialization techniques like Xavier/Glorot initialization (often used with tanh or sigmoid activations) and He initialization (designed for ReLU activations) aim to maintain the variance of activations and gradients as they propagate through the network layers. This helps prevent gradients from becoming too small or too large, facilitating smoother training in typical supervised learning scenarios.

However, GAN training involves a delicate balance between the generator ($G$) and the discriminator ($D$). If the initial weights cause the discriminator to be too strong, its gradients might vanish early on, providing no useful signal to the generator. Conversely, if the generator produces easily distinguishable outputs initially, the discriminator might learn too quickly and saturate, again stalling the generator's learning. Therefore, the choice of initialization can significantly impact the early training dynamics and overall convergence.

Common Initialization Strategies in GANs

While standard initializations are often a reasonable starting point, several strategies have been found particularly useful in the context of GANs:

1. Gaussian Initialization with Small Standard Deviation

Early GAN papers, including DCGAN, often recommended initializing weights from a zero-centered Normal (Gaussian) distribution with a small standard deviation, such as $\mathcal{N}(0, 0.02^2)$.

$$W \sim \mathcal{N}(0, \sigma^2) \quad \text{where } \sigma = 0.02$$

The rationale is to keep initial activations and gradients small, potentially preventing one network from overwhelming the other immediately. This can be particularly relevant for convolutional and fully connected layers. While simple, this approach might not scale well to very deep networks where variance control becomes more significant.

2. Orthogonal Initialization

Orthogonal initialization aims to set weight matrices $W$ such that $W^T W = I$ (or $W W^T = I$). This property helps preserve the norm of vectors (and gradients) during both forward and backward propagation. By preventing gradients from exploding or vanishing due to repeated matrix multiplications, orthogonal initialization can contribute to more stable training, especially in deeper networks or RNNs.

It's often used in conjunction with techniques like Spectral Normalization, which also constrains the Lipschitz constant of the discriminator's layers. Implementing orthogonal initialization typically involves techniques like Singular Value Decomposition (SVD) or QR decomposition applied to an initial random matrix. Many deep learning libraries provide functions for this (e.g., torch.nn.init.orthogonal_ in PyTorch).

Illustration of gradient norm preservation with good initialization versus potential explosion or vanishing with poor initialization during backpropagation.

3. Re-evaluating Xavier/Glorot and He Initialization

Despite the specialized techniques, He initialization (for layers followed by ReLU or LeakyReLU) and Xavier/Glorot initialization (for layers followed by tanh or similar bounded activations) remain common and often effective choices for GANs. They provide a principled way to scale initial weights based on layer dimensions. When using standard convolutional or linear layers, these are often the default and a sensible first choice.

It's important to apply them correctly based on the activation function used in subsequent layers. If your generator uses ReLUs extensively, He initialization is generally preferred. If it uses tanh for the output layer, Xavier might be considered for the final layer's weights.

Initialization in Specific Architectures

Advanced GAN architectures sometimes employ specific initialization schemes tailored to their unique components:

• StyleGAN/StyleGAN2: These models often use He initialization for standard convolutional layers but might use customized initialization (e.g., scaling weights to 1.0) for layers involved in the AdaIN (Adaptive Instance Normalization) modulation or specific parts of the mapping network. The goal is often to ensure that learned style vectors have a predictable initial effect.
• BigGAN: Known for large-scale training, BigGAN often utilizes orthogonal initialization, which complements its use of Spectral Normalization and helps manage training dynamics at large batch sizes.

Consulting the original papers or reliable implementations for these specific architectures is recommended if you are aiming for faithful reproduction.

Practical Recommendations

• Start Standard: Unless training proves unstable, begin with He initialization for ReLU/LeakyReLU layers and Xavier/Glorot for tanh/sigmoid layers. These are often the defaults in frameworks like PyTorch and TensorFlow.

• Consider $\mathcal{N}(0, 0.02^2)$: If you encounter instability early in training, especially with architectures similar to DCGAN, try the small-variance Gaussian initialization as an alternative.

• Experiment with Orthogonal: For deeper networks or if using techniques like Spectral Normalization, orthogonal initialization is a strong candidate worth experimenting with.

• Initialize Biases: Biases are typically initialized to zero. This is usually a safe and standard practice.

• Consistency: Apply your chosen initialization scheme consistently across similar layer types within both the generator and discriminator.

• Code Example (PyTorch):

  import torch
  import torch.nn as nn
  import torch.nn.init as init
  
  def weights_init_normal(m):
      classname = m.__class__.__name__
      if classname.find('Conv') != -1:
          init.normal_(m.weight.data, 0.0, 0.02) # Gaussian init
          if hasattr(m, 'bias') and m.bias is not None:
              init.constant_(m.bias.data, 0.0)
      elif classname.find('BatchNorm2d') != -1:
          init.normal_(m.weight.data, 1.0, 0.02) # Batch norm weights often init near 1
          init.constant_(m.bias.data, 0.0)
      elif classname.find('Linear') != -1:
           init.normal_(m.weight.data, 0.0, 0.02) # Gaussian init for linear
           if hasattr(m, 'bias') and m.bias is not None:
               init.constant_(m.bias.data, 0.0)
  
  # Example Usage:
  # model = Generator(...) or Discriminator(...)
  # model.apply(weights_init_normal)
  

  Note: You might use init.xavier_uniform_, init.kaiming_normal_ (He), or init.orthogonal_ similarly.

• Interaction with Normalization: Remember that normalization techniques (Batch Norm, Instance Norm, Layer Norm, Spectral Norm) interact significantly with weight initialization. Spectral Normalization, for instance, directly constrains the spectral norm of weight matrices, potentially overriding some effects of the initial scale but not the initial direction/structure (like orthogonality).

Choosing the right weight initialization is not always straightforward and can sometimes require experimentation based on the specific architecture, dataset, and training configuration. However, understanding these common strategies provides a solid foundation for tackling the practical challenges of implementing and stabilizing advanced GANs.