Mapping Network in StyleGAN
This article discusses the structural components of StyleGAN, such as structure charts, and demonstrates the principles of mapping networks through the example of StyleGAN’s own mapping network.
Let’s first review the structure charts of traditional GANs and DCGANs.
In traditional GAN architectures, several fully connected (FC) layers are utilized to diversify the features. These features progress without being constrained by spatial dimensions and are directly applicable to the image domain. If the image is small, this model can still operate effectively, considering all elements.
DCGANs enhance the traditional GAN architecture by introducing convolutional layers, which enable spatial dependencies within features to be captured. Unlike traditional GANs where features are directly applied to the image domain, in DCGANs, the convolution uses weights to apply kernels that capture spatial dependencies, allowing feature maps to develop more complex structures as they progress through the network.
As an aside, while traditional GANs tend to increase the number of in_features during training, DCGANs start with a higher number of in_channels and progressively decrease it, training the model with channel counts diminishing in sequence, like [1024, 512, 256, …]. This approach contrasts with the method of traditional GANs, which expand their feature set for training.
Now let’s apply a mapping network to our current model and then proceed to train the DCGAN.
In the traditional DCGAN framework, image generation begins with a completely random distribution, known as z_latent. However, in this modified structure, the process starts with a w_latent vector that is derived from the z_latent through fully connected (FC) layers. This w_latent vector contains high-dimensional features, which are initially random but structured in a way that is more conducive to image generation.
In StyleGAN, this w_latent is further processed into a vector referred to as “style,” which is then applied across various feature map sizes. However, for the purpose of this discussion, let’s omit that part and instead focus on analyzing the impact that the mapping network has on model training.
The results vary depending on the depth of the mapping network applied to the training of the model. Below, the generated images demonstrate the differences in results.
As you can see, when the learning rate (lr) is high or the fc (fully connected) depth is great, the balance of training can be disrupted. In other words, naively applying a Mapping network could result in a non-functional model.
Get Sisung Kim’s stories in your inbox
Join Medium for free to get updates from this writer.
Let’s compare the loss graphs between times when training is successful and when it is not.
Despite the generator becoming stronger with the introduction of a mapping network, observe how the balance of training is compromised by the discriminator driving its loss towards zero. If the discriminator’s loss seems excessively low, overwhelming the generator, it may actually indicate that the generator is overperforming. Therefore, when tuning hyperparameters in GANs, even if it appears that the discriminator is dominant, consider dialing back the generator’s strength. This approach can be essential for restoring balance to the training process.
We still want to benefit from the feature diversification offered by the mapping layer, but can we do so without the balancing issues? NVIDIA has addressed this by introducing equalized layers and learning rate multipliers. Here, we will focus solely on the equalized layer.
class EqualizedLayer(nn.Module):
def __init__(self, module):
super().__init__()
self.module = module
self.gain = 1/ np.sqrt(np.prod(module.weight.shape[1:]))
def forward(self, *input):
self.module.weight.data *= self.gain
output = self.module(*input)
self.module.weight.data /= self.gain
return output
equalized_fc = EqualizedLayer(nn.Linear(64, 128))NVIDIA has created a streamlined approach to training GANs that effectively addresses the balance issues encountered when adding layers. This implementation allows us to continue enjoying the benefits of feature diversification through the mapping layer. The Equalized Layer ensures the input and output’s influence remains balanced and does not become disproportionate. Below we delve into the technical aspects of this approach.
This formula illustrates the concept of equalizing the weights, which is achieved by dividing the weights by the constant c, which is the square root of the number of inputs n. This ensures that V(z′) is a normalized version of V(z), allowing for balanced weight influence during the training process.
The provided code snippet is a simplified adaptation of NVIDIA’s original implementation for equalized layers. This implementation uses scaled weights of w′=w/c for the forward pass. However, during the backward pass, the weight gradients are computed with respect to w, not w′. By doing this, the updates to the weights w are kept small relative to their influence on the training process, maintaining stability. This subtle yet crucial technique ensures that the weights are updated in a controlled manner, contributing to the model’s overall stability during training.
Now, let’s take a look at the results when these adjustments have been applied.
Observe how the models achieve stable training at different FC depths when equalized weights are applied. While it is possible to maintain stable training by increasing the depth of the fully connected layers, applying the equalized layer to the mapping network allows for diversification of features without collapsing the model. Instability is reduced, but increasing the generator’s strength too much can lead to mode collapse, where the generator starts to overlook certain outputs, limiting the diversity of generated images.
You can find more detailed examples of the code used in this guide on my GitHub page.
https://github.com/crimson206/example_archive/tree/main/deeplearning/gans
