Anatomy of a Transformer Decoder Block

Hello! Welcome to the fifth lesson in our module on LLM and Transformer Fundamentals.

In our last lesson, we tackled a critical piece of the puzzle: positional encodings. We now understand how to inject information about sequence order into our token embeddings, creating position-aware input vectors.

This lesson is about assembling all the pieces. We'll move from individual components to a complete, functional Transformer decoder block, which is the fundamental repeating unit of powerful generative models like GPT. Our learning goal is to diagram a full transformer decoder block, explaining the role of each component (attention, FFN, layer norm, and residuals).

By the end of this lesson, you'll understand how these components work in concert to process information and transform an input sequence into a representation ready for generating the next token.

1. The Anatomy of a Decoder Block

A decoder-only Transformer, like the one used in GPT, is essentially a stack of identical blocks. The output of one block feeds into the input of the next. Let's start by looking at the overall structure of a single block.

As you can see, a decoder block consists of two main sub-layers:

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

These core computational layers are wrapped with two other crucial components that enable stable training of deep networks: Residual Connections and Layer Normalization. Let's break each of these down.

2. Component 1: Masked Multi-Head Self-Attention

We have covered this mechanism in detail in previous lessons. Its role within the decoder block is to allow each token to look at all preceding tokens (and itself) and decide which ones are most important for understanding its own meaning in context. The key aspects to remember here are:

• Self-Attention: Queries, Keys, and Values are all derived from the same input sequence.
• Multi-Head: The attention mechanism is run multiple times in parallel with different learned weights, allowing the model to focus on different types of relationships (e.g., syntactic, semantic) simultaneously.
• Masked (Causal): This is the defining feature of a decoder. The attention mask ensures that a token at position t can only attend to tokens at positions 0 to t. It cannot "see into the future." This is essential for the autoregressive, next-token-prediction task.

Andrej Karpathy provides an excellent, code-first explanation of how the attention mechanism is implemented, including the crucial scaling factor and the causal mask.

3. Component 2: Position-wise Feed-Forward Network (FFN)

After the attention sub-layer has aggregated relevant information from across the sequence, the output is passed to a Feed-Forward Network (FFN).

The FFN's role is to provide additional computational depth. You can think of it as giving the model "time to think" about the information it just gathered via attention. Unlike the attention mechanism, which mixes information across the sequence, the FFN processes each token's vector independently and identically.

It's typically a simple two-layer multilayer perceptron (MLP):

1. A linear layer that expands the embedding dimension (e.g., from d_model to 4 * d_model).
2. A non-linear activation function (like GELU or SwiGLU).
3. A linear layer that projects the dimension back down (e.g., from 4 * d_model to d_model).

The article "Decoder-Only Transformers: The Workhorse of Generative LLMs" also provides a concise description and a clean PyTorch implementation of this component.

4. Components 3 & 4: Residual Connections and Layer Normalization

Simply stacking attention and FFN layers would make a network that is very difficult to train. Two key innovations, typically used together, solve this: Residual Connections and Layer Normalization.

Residual Connections (Add)

A residual connection is a "shortcut" that adds the input of a sub-layer to its output. The operation is simply x + SubLayer(x).

Its primary purpose is to combat the vanishing gradient problem in very deep networks. By providing a direct path for gradients to flow from the final output back to the initial input, it acts as a "gradient superhighway," dramatically improving training stability.

Layer Normalization (Norm)

Layer Normalization (LayerNorm) stabilizes the network by ensuring that the inputs to each layer have a consistent distribution (e.g., zero mean, unit variance). Unlike Batch Normalization, which normalizes across the batch dimension, LayerNorm normalizes across the feature (embedding) dimension for each individual sequence in the batch. This makes it independent of batch size and better suited for variable-length NLP tasks.

5. Putting It All Together: The Pre-Norm Architecture

Now we can assemble the full block. The arrangement of these components matters. While the original Transformer paper used a "Post-Norm" architecture (LayerNorm(x + SubLayer(x))), modern GPT-style models almost universally use a "Pre-Norm" architecture: x + SubLayer(LayerNorm(x)). This has been found to be more stable for training very deep models.

The data flow for a pre-norm decoder block is as follows:

1. The input x is first passed through Layer Normalization.
2. The normalized output is processed by the Masked Multi-Head Attention sub-layer.
3. The output of the attention sub-layer is added to the original input x (the residual connection).
4. The result of this addition is then passed through a second Layer Normalization.
5. This normalized output is processed by the Feed-Forward Network.
6. The output of the FFN is added to the result from step 3 (the second residual connection).
7. The final result is the output of the block.

The article "Meet GPT, The Decoder-Only Transformer" provides an excellent visual and code-based comparison of the Post-Norm (used in GPT-1) and Pre-Norm (used in GPT-2/3) architectures.

Finally, let's look at a complete, clean implementation of a Block class that combines all these elements using the pre-norm structure.

Conclusion

In this lesson, we have constructed a complete Transformer decoder block from its constituent parts. We saw that it isn't just a random assortment of layers, but a carefully engineered unit where each component plays a vital role.

Key Takeaways:

• A decoder block is composed of two main computational sub-layers: Masked Multi-Head Self-Attention (for communication across tokens) and a Position-wise Feed-Forward Network (for independent computation on each token).
• Residual Connections (Add) act as gradient shortcuts, enabling stable training of very deep networks by adding a layer's input to its output.
• Layer Normalization (Norm) stabilizes training by normalizing the activations for each token across its embedding dimension.
• Modern architectures like GPT use a Pre-Norm structure (x + SubLayer(LayerNorm(x))), which applies normalization before the main sub-layers, as it generally leads to more stable training.
• A full decoder-only Transformer is simply a stack of these identical blocks, allowing the model to build progressively more abstract and complex representations of the input sequence.

Next Lesson Preview

We now have a complete, functional decoder block. But how do we use it to actually generate language? In our next lesson, we will explore the training process: autoregressive language modeling. We'll see how the model is trained to predict the next token in a sequence and how this simple objective, when combined with the powerful architecture we've just dissected, leads to the emergent capabilities of Large Language Models.