Transformer Block Architecture

Introduction

In the previous module, we built the attention mechanism from the ground up. We explored scaled dot-product attention, saw how multi-head attention allows the model to focus on different aspects of the input, and finally, addressed the "order" problem by injecting sinusoidal positional encodings. We now have all the necessary ingredients.

In this lesson, we will assemble these pieces into a complete Transformer block. Our learning outcome is to describe the structure of a transformer block, including its four key components: self-attention, a feed-forward network, residual connections, and layer normalization.

This block is the fundamental, repeating unit that forms the backbone of modern large language models like GPT. Understanding its structure is the critical next step in our journey toward deconstructing state-of-the-art architectures.

1. A High-Level View of the Transformer Block

At its core, a Transformer block is a computational unit designed to take a sequence of context-aware token embeddings and refine them further. It consists of two main sub-layers:

1. A Multi-Head Self-Attention layer.
2. A Position-wise Feed-Forward Network (FFN).

Crucially, the block also includes two other components that are essential for training very deep networks: residual connections and layer normalization. These are applied around each of the two sub-layers.

Let's start with a clear, high-level diagram.

This "Add & Norm" step is where the magic of residual connections and layer normalization happens. Let's break down each of these four components.

2. The Two Computational Sub-layers

Self-Attention Layer (A Quick Review)

As we've covered, the self-attention layer is the heart of the Transformer. Its purpose is to update each token's representation by aggregating information from the entire input sequence. It computes Query, Key, and Value vectors for each token, allowing it to dynamically weigh the importance of other tokens when creating a new, context-aware representation for the current token. In a Transformer block, this is typically a multi-head self-attention layer, allowing it to perform this process in parallel across different "representation subspaces."

Position-wise Feed-Forward Network (FFN)

After the self-attention layer has produced its output (a context-aware embedding for each token), these embeddings are passed to the second sub-layer: a feed-forward network.

Why is this necessary? The self-attention layer itself is primarily composed of linear transformations (matrix multiplications) and weighted sums. While the softmax provides some non-linearity, the FFN adds significant computational depth and non-linearity, which is characteristic of deep learning models.

A "position-wise" FFN has two key properties:

1. It is a simple two-layer Multi-Layer Perceptron (MLP), usually with a ReLU or similar activation in between.
2. It processes each token's representation independently. The same FFN weights are applied to every token in the sequence, but it doesn't mix information between tokens. That job was already done by the self-attention layer.

The following video explains the role of the FFN in adding this necessary non-linearity and processing the output from the attention mechanism.

3. The "Glue": Residual Connections and Layer Normalization

Stacking dozens of self-attention and FFN layers would be impossible to train effectively without two critical techniques borrowed from computer vision and adapted for NLP: residual connections and layer normalization.

You'll almost always see these paired together in diagrams as "Add & Norm".

Residual Connections ("Add")

A residual connection, or "shortcut connection," is a simple but powerful idea. The input to a sub-layer is added directly to the output of that sub-layer.

$\text{Sub-layer Output}=\text{Input}+\text{Sub-layer}(\text{Input})$

The primary benefit is that it makes training very deep networks feasible. By providing a "shortcut" path, it allows gradients to flow directly through the network during backpropagation. This helps mitigate the vanishing gradient problem, where gradients become too small to provide a useful learning signal in deep architectures.

The following video segment provides a fantastic intuition for why residual connections are so helpful for the optimization process.

Layer Normalization ("Norm")

After the residual connection adds the layer's input and output, the result is passed through a Layer Normalization step.

Layer Normalization (LayerNorm) is a technique that stabilizes training by normalizing the activations. For each token's vector individually, it recalculates the features to have a mean of 0 and a standard deviation of 1.

Where:

• $\mu$ and ${\sigma }^2$ are the mean and variance of the elements within that single vector $\mathbf{x}$.
• $ϵ$ is a small constant for numerical stability.
• $\gamma$ (gamma) and $\beta$ (beta) are learnable parameters that allow the model to scale and shift the normalized output if needed.

The key difference from the more common Batch Normalization is that LayerNorm's calculations are completely independent of other examples in the batch. It normalizes across the feature dimension for each token, which has been found to work much better for variable-length sequences in NLP.

4. Assembling the Full Transformer Block

Now, let's put all four components together to trace the path of data through one complete Transformer block.

A sequence of token embeddings, $\mathbf{X}$, enters the block.

1. First "Add & Norm":

   • The embeddings $\mathbf{X}$ are passed into a Multi-Head Self-Attention layer to produce $\mathbf{Z}$.
   • A residual connection adds the original input to the attention output: $\mathbf{X}+\mathbf{Z}$.
   • The result is Layer Normalized: $\text{LayerNorm}(\mathbf{X}+\mathbf{Z})$. Let's call this intermediate result $\mathbf{Y}$.

2. Second "Add & Norm":

   • The intermediate result $\mathbf{Y}$ is passed into the Position-wise Feed-Forward Network to produce $\text{FFN}(\mathbf{Y})$.
   • A second residual connection adds the input of the FFN to its output: $\mathbf{Y}+\text{FFN}(\mathbf{Y})$.
   • The result is Layer Normalized again: $\text{LayerNorm}(\mathbf{Y}+\text{FFN}(\mathbf{Y}))$.

The output of this final normalization is the output of the Transformer block, which is then fed as input to the next identical block in the stack.

For a programming-oriented perspective, the TransformerEncoderBlock class described in the "Dive into Deep Learning" resource shows how these components are structured in code.

Conclusion

In this lesson, we have successfully assembled the Transformer block, the fundamental repeating unit of modern language models.

Key Takeaways:

• A Transformer block is composed of two main sub-layers: a multi-head self-attention layer and a position-wise feed-forward network.
• The self-attention layer creates context-aware representations by allowing tokens to communicate with each other.
• The feed-forward network provides additional non-linear processing for each token independently.
• Residual connections ("Add") are used around each sub-layer to enable effective gradient flow and the training of very deep networks.
• Layer normalization ("Norm") is applied after each residual connection to stabilize training by normalizing the activations for each token.

Preview of the Next Lesson:
We've now seen all the components of the standard Transformer block. In the next lesson, "Explain LayerNorm and how it stabilizes training in deep networks," we will take a deeper dive into Layer Normalization. We'll explore its mathematical properties more formally and build a stronger intuition for exactly why it is so crucial for achieving stable training in deep Transformer models.