Anatomy of a Modern Transformer Block: From RMSNorm to SwiGLU

In deep neural networks, the primary challenge is the degradation of the signal as it passes through many successive transformations. When we stack dozens of layers, the gradient used for optimization can either explode or vanish entirely before it reaches the earlier parts of the model. Residual connections solve this by providing a shortcut for information to flow through the network.

A residual connection works by adding the original input of a block to its transformed output. This architecture ensures that the layer only needs to learn the difference or the residual between the input and the ideal output. If the transformation is unnecessary, the model can simply learn to push the weights of the layer toward zero, effectively allowing the identity to pass through.

From a mathematical perspective, this creates a gradient highway that bypasses the complex non-linearities of the attention or feed-forward modules. This bypass mechanism is critical for training models with hundreds of billions of parameters. Without these skip connections, the optimization landscape would become too fractured and chaotic for standard backpropagation to converge effectively.

1import torch
2import torch.nn as nn
3
4class ResidualBlock(nn.Module):
5    def __init__(self, module: nn.Module):
6        super().__init__()
7        self.module = module
8
9    def forward(self, x: torch.Tensor) -> torch.Tensor:
10        # The addition of the input x back to the output of the module
11        # ensures that gradients can bypass the module during backpropagation.
12        return x + self.module(x)
13
14# Usage in an LLM context
15# layer_output = residual_block(input_tensor)

Residual connections do not just help with training; they fundamentally change the loss landscape by making it significantly smoother and more convex, which allows for higher learning rates.

Stability in large models is also heavily dependent on normalization layers that control the distribution of activations. As data flows through the model, the variance of the values can shift significantly, leading to internal covariate shift. Normalization layers recenter and rescale these values to keep them within a predictable range for the next layer.

In the early days of transformers, LayerNorm was the standard choice for stabilizing the training process. LayerNorm calculates the mean and variance across the feature dimension for each token and uses them to normalize the activations. While effective, the calculation of the mean adds a small amount of computational overhead that scales with the model size.

Modern architectures like Llama and Mistral have transitioned to Root Mean Square Layer Normalization, or RMSNorm. RMSNorm simplifies the process by only scaling the activations based on their root mean square, skipping the mean subtraction step. This reduction in operations leads to faster training times on modern GPU hardware without sacrificing the stability of the model.

• LayerNorm: Normalizes using both mean and variance, requiring more computation.
• RMSNorm: Normalizes using only the root mean square, providing better efficiency.
• Batch Normalization: Rarely used in LLMs because it depends on batch size and is unstable for variable sequence lengths.
• Weight Tying: A related technique where embedding and output layers share parameters to reduce memory footprint.

1class RMSNorm(nn.Module):
2    def __init__(self, dim: int, eps: float = 1e-6):
3        super().__init__()
4        self.eps = eps
5        self.weight = nn.Parameter(torch.ones(dim))
6
7    def _norm(self, x: torch.Tensor):
8        # We calculate the root mean square across the last dimension
9        # then multiply by the learnable weight parameter.
10        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
11
12    def forward(self, x: torch.Tensor):
13        output = self._norm(x.float()).type_as(x)
14        return output * self.weight

While the attention mechanism handles the relationships between tokens, the position-wise feed-forward network is where the model stores most of its factual knowledge. This component processes each token independently, applying a series of linear transformations and non-linear activations. It is often referred to as the knowledge engine of the transformer block.

In a standard transformer, the FFN consists of two linear layers with an expansion factor of four. The first layer projects the embedding into a higher-dimensional space, and the second layer projects it back down to the original size. This expansion allows the model to form complex representations that would be impossible in a lower-dimensional space.

State-of-the-art models have largely moved away from the standard ReLU activation function in favor of Gated Linear Units like SwiGLU. SwiGLU uses a gating mechanism where one linear projection is multiplied by the output of a non-linear activation of another projection. This element-wise multiplication allows for more fine-grained control over the information flow through the FFN.

The placement of normalization layers within the transformer block is a critical design decision that affects training stability. In the original transformer design, normalization was applied after the residual addition, a configuration known as Post-Norm. While Post-Norm can lead to better performance, it is notoriously difficult to train because the gradients near the output are much larger than those near the input.

To address this, most modern models utilize the Pre-Norm architecture, where normalization is applied before the attention and feed-forward modules. Pre-Norm makes the model much more stable from the very beginning of the training process. This stability allows researchers to use higher learning rates and skip the complex learning rate warm-up schedules required by Post-Norm designs.

However, Pre-Norm is not without its drawbacks, as it can lead to a phenomenon known as representation collapse in extremely deep models. Because the identity path is never normalized, the scale of the hidden states can grow linearly with the number of layers. Engineers must carefully monitor the magnitude of these states to ensure the model continues to learn effectively at great depths.