Anatomy of a Transformer Encoder Block

Introduction

In our previous lesson, we established a crucial foundation for the Transformer architecture by addressing its inherent lack of sequential awareness. You learned that self-attention is permutation-equivariant and requires an explicit mechanism—sinusoidal positional encodings—to inject information about token order. By adding these encodings to the input embeddings, we created a representation that is rich in both semantic content and positional context.

Now, we will assemble all the pieces into the fundamental processing unit of the Transformer's encoder side: the Transformer Encoder Block. This lesson directly addresses the learning outcome of tracing the complete data flow through this block. We will move beyond the individual components of multi-head attention and positional encodings to see how they are integrated with two other critical architectural elements: residual connections and layer normalization, plus a final position-wise feed-forward network. By the end of this lesson, you will be able to describe precisely how a sequence of input vectors is transformed into a sequence of contextually-aware output vectors, ready to be fed into the next identical block.

The Anatomy of an Encoder Block

At a high level, a single encoder block consists of two main sub-layers. The input embeddings (with positional information added) first pass through a multi-head self-attention layer. The output of this layer then passes through a position-wise feed-forward network.

However, a key architectural detail, crucial for enabling deep stacks of these blocks, is the use of residual connections and layer normalization around each of these two sub-layers.

A detailed view of a single Transformer Encoder block. Notice the two primary components—Multihead Attention and the Feed Forward network—each followed by an "Add & Norm" step, which represents a residual connection and layer normalization.

This "Add & Norm" step is the glue that holds the architecture together. Let's break down its components and why they are so vital.

Architectural Glue: Residuals and Normalization

1. Residual Connections

You might recall from our study of ResNets (Module 8) that as networks get deeper, they can suffer from a "degradation" problem where performance gets worse, not better. The solution was the residual, or "skip," connection. The Transformer architecture leverages the same principle.

A residual connection takes the input to a layer and adds it to the layer's output:

$$\text{Output}=\text{Layer}(\mathbf{x})+\mathbf{x}$$

In the context of the Transformer encoder, the input $\mathbf{x}$ to the multi-head attention sub-layer is added to the output of that sub-layer.

Socratic Question: In our previous lesson, we established that the self-attention mechanism is permutation-equivariant. We solved this by adding positional encodings to the initial input embeddings. What might happen to this carefully added positional information after it passes through the self-attention layer, which essentially re-expresses each vector as a weighted average of all other vectors? Why would a residual connection be particularly important here?

Answer

The self-attention mechanism computes each output vector as a weighted sum of the value vectors from the entire sequence. This process of averaging could "wash out" or dilute the unique positional signal associated with each specific token. By adding the original input vector $\mathbf{x}$ (which contains the untouched positional encoding) back to the attention output, the residual connection creates a direct path for the positional information to be preserved. It ensures that the output of the sub-layer still has a strong, unmodified signal about its own position, even after incorporating contextual information from other tokens.

This preservation of positional information is a critical, and often overlooked, function of the residual connection in Transformers. It allows the model to maintain a sense of order even as it is stacked dozens of layers deep.

2. Layer Normalization

After the residual connection, the result is passed through a Layer Normalization step. This is another crucial component for stabilizing the training of deep networks.

Layer Normalization operates by computing the mean and variance for each training example independently across all its features. For a token embedding vector of size $d_{\text{model}}$, LayerNorm calculates the mean and variance of those 512 values and uses them to normalize that single vector. This differs from Batch Normalization, which normalizes across the batch dimension for each feature.

The primary role of LayerNorm is to keep the activations within each layer in a consistent range, preventing the explosion or vanishing of gradients during backpropagation. This ensures a smoother, more stable training process.

To see a clear, code-oriented explanation of how Layer Normalization is implemented, and how its learnable parameters gamma and beta function, let's turn to a practical resource.

Pre-LN vs. Post-LN

The original "Attention Is All You Need" paper placed the layer normalization after the residual addition (Post-LN). However, subsequent research found that placing it before the sub-layer and applying the residual connection to the un-normalized input leads to more stable training. This configuration is known as Pre-LN and is standard in most modern Transformer implementations.

A comparison of the standard (Post-LN, left) and Pre-LayerNorm (Pre-LN, right) Transformer encoder blocks. In this lesson, we will trace the data flow through the more modern Pre-LN architecture. Image from "An even more annotated Transformer".

Tracing the Data Flow: A Step-by-Step Walkthrough

Let's now trace the journey of a tensor X with shape [batch_size, sequence_length, d_model] through a single Pre-LN encoder block.

Input: Tensor X representing the position-aware token embeddings.

Step 1: First Sub-layer (Multi-Head Self-Attention)

1. Normalization: The input tensor X is first passed through a Layer Normalization layer. $$X_{\text{norm1}}=\text{LayerNorm}(X)$$
2. Self-Attention: The normalized tensor $X_{\text{norm1}}$ is fed into the Multi-Head Attention layer. Since this is self-attention, the Query, Key, and Value inputs are all derived from $X_{\text{norm1}}$. $$\text{AttentionOutput}=\text{MultiHeadAttention}(Q=X_{\text{norm1}},K=X_{\text{norm1}},V=X_{\text{norm1}})$$
3. Dropout & Residual Connection: A dropout is applied to the attention output for regularization, and then the original input X is added. This creates the residual connection. $$X_{\text{sub1}}=X+\text{Dropout}(\text{AttentionOutput})$$

The output of this first stage, $X_{\text{sub1}}$, now holds a representation where each token has gathered context from all other tokens, while still strongly retaining its original positional identity.

Step 2: Second Sub-layer (Position-wise Feed-Forward Network)

This sub-layer provides additional processing capacity and non-linearity. The "position-wise" part of its name simply means it is applied independently to each token's vector in the sequence.

The network itself is a simple two-layer Multi-Layer Perceptron (MLP), typically with a ReLU activation in between. It first projects the $d_{\text{model}}$-dimensional vector to a much larger dimension (e.g., $4\times d_{\text{model}}$) and then projects it back down to $d_{\text{model}}$.

1. Normalization: The output from the first sub-layer, $X_{\text{sub1}}$, is normalized. $$X_{\text{norm2}}=\text{LayerNorm}(X_{\text{sub1}})$$
2. Feed-Forward Network: The normalized tensor is passed through the FFN. $$\text{FFNOutput}=\text{FFN}(X_{\text{norm2}})={\text{Linear}}_2(\text{ReLU}({\text{Linear}}_1(X_{\text{norm2}})))$$
3. Dropout & Residual Connection: Dropout is applied, and the residual connection is formed by adding the input to this stage, $X_{\text{sub1}}$. $$X_{\text{output}}=X_{\text{sub1}}+\text{Dropout}(\text{FFNOutput})$$

Output: The tensor $X_{\text{output}}$ is the final output of the encoder block. It has the exact same shape as the input X ([batch_size, sequence_length, d_model]), but each vector is now a deep, context-aware representation of its corresponding token. This output can be directly fed as input to the next encoder block in the stack.

To solidify this entire process, let's walk through it with a detailed step-by-step example focused on the matrix operations and tensor transformations.

As a final, comprehensive overview, Jay Alammar's blog provides a beautifully illustrated summary of the entire encoder block, which you can use to tie all these concepts together visually.

Conclusion

In this lesson, we have successfully assembled the Transformer encoder block, the workhorse of models like BERT and GPT. You can now trace the complete path of a data tensor as it is transformed from a simple, position-aware embedding into a deeply contextualized representation.

Key Takeaways:

• Encoder Block Structure: An encoder block consists of two sub-layers: a multi-head self-attention mechanism and a position-wise feed-forward network.
• Residual Connections: These skip connections are vital for training deep Transformers. They combat vanishing gradients and, critically, ensure that positional information is preserved as data flows through the layers.
• Layer Normalization (Pre-LN): Applied before each sub-layer, LayerNorm stabilizes training by keeping activations in a controlled range, allowing for faster and more reliable convergence.
• Feed-Forward Network (FFN): This component introduces essential non-linearity and provides further per-token processing after contextual information has been aggregated by the attention mechanism.
• Data Flow: The process for each sub-layer in the modern Pre-LN architecture is: Norm -> Sub-layer -> Dropout -> Add Residual. The output tensor of an encoder block maintains the exact same dimensions as the input, allowing for easy stacking.

Preview of the Next Lesson:
Having now understood the complete architecture of the encoder block from a conceptual and data-flow perspective, you are perfectly positioned for the next step. In our upcoming lesson, we will move from theory to practice and begin to implement scaled dot-product attention and a single self-attention head from scratch in Python. This will be the first hands-on step in building your own Transformer.