Inside the Transformer Block

Introduction

In our last session, we took a bird's-eye view of the Transformer, examining its full encoder-decoder structure and the clever use of positional encodings to imbue the model with a sense of sequence order. We noted that both the encoder and decoder are composed of a stack of identical layers, or "blocks." Today, we will zoom in and perform a detailed dissection of these fundamental building blocks.

Our objective is to detail the components of a Transformer block: multi-head attention, feed-forward networks, residual connections, and layer normalization. These are not merely implementation details; they are the core engineering and theoretical innovations that enable Transformers to be trained at unprecedented depths and scales. We will analyze the specific role each component plays, the problems they solve, and the theoretical principles that justify their inclusion.

The Anatomy of a Transformer Block

A standard Transformer block, whether in an encoder or a decoder, is primarily composed of two key computational sub-layers. In an encoder block, these are:

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

However, what makes these blocks truly effective is how these sub-layers are "wired" together. Each sub-layer is wrapped by two additional components: a residual connection and a layer normalization step. It is this combination of four elements that constitutes a complete, functional Transformer block. Let's examine each in turn.

Component 1: Multi-Head Attention

We have already established the mathematical foundation of attention. Within the context of a Transformer block, its role is to act as a dynamic, content-aware information routing system. It allows every token's representation to be updated by aggregating information from all other tokens in the sequence, weighted by relevance.

The "multi-head" aspect is crucial for capturing the richness of language. A single attention mechanism might learn to track one type of relationship (e.g., subject-verb agreement). By using multiple heads in parallel, the model can simultaneously track diverse syntactic and semantic dependencies.

Each head operates in its own representational subspace, defined by its unique projection matrices ($W_i^Q,W_i^K,W_i^V$). This allows for specialization. The final output projection ($W^O$) then synthesizes these diverse perspectives into a single, coherent representation.

For a thorough refresher on the motivation and implementation details of multi-head attention, the following text is an excellent resource.

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

After the attention sub-layer has mixed information across the sequence, the position-wise feed-forward network (FFN) provides the next stage of processing. Its primary function is to introduce non-linearity and increase the model's representational capacity.

The attention operation itself is fundamentally a weighted average of the value vectors. If we were to stack only attention layers, the composition of these linear transformations (in V) would collapse into a single, equivalent linear transformation. The FFN breaks this linearity, which is essential for learning complex functions.

The FFN is a simple two-layer multilayer perceptron (MLP) that is applied identically and independently to each token's representation. The architecture follows a distinct "expand, activate, contract" pattern:

1. Expand: A linear layer projects the $d_{\text{model}}$-dimensional representation to a larger intermediate dimension, typically $4\times d_{\text{model}}$.
2. Activate: A non-linear activation function, such as GELU (Gaussian Error Linear Unit), is applied.
3. Contract: A second linear layer projects the intermediate representation back down to $d_{\text{model}}$.

This structure can be expressed as:

(using ReLU for simplicity, though GELU is more common in Transformers).

A compelling theoretical interpretation views the FFN as a form of key-value memory. The first linear layer acts as a "key" that pattern-matches against the input representation, and the second layer retrieves the corresponding "value". From this perspective, a significant portion of a Transformer's factual knowledge is stored within the weights of its FFNs.

Component 3: Residual (Skip) Connections

One of the most significant breakthroughs enabling very deep neural networks was the introduction of residual connections. Transformers adopt this principle wholeheartedly. Each sub-layer (both attention and FFN) in a Transformer block is wrapped in a residual connection. The output of a sub-layer is not a new representation, but rather an update to the input representation:

This has profound theoretical and practical implications:

1. Gradient Highways: During backpropagation, the gradient of the loss with respect to $x_{\text{input}}$ has a direct path through the addition operator. The chain rule gives $\frac{\partial L}{\partial x_{\text{input}}}=\frac{\partial L}{\partial x_{\text{output}}}\cdot (1+\frac{\partial \text{Sublayer}}{\partial x_{\text{input}}})$. The '1' term creates an unimpeded "gradient highway," allowing the loss signal to flow directly to earlier layers, which drastically mitigates the vanishing gradient problem that plagues deep sequential models.

2. Smoother Loss Landscape: Empirically, residual connections have been shown to make the optimization landscape much smoother, removing many of the sharp local minima that can trap gradient-based optimizers.

3. Default Identity Mapping: A layer's default behavior is to pass its input through unchanged (an identity mapping). The network only needs to learn the residual—the necessary modification—which is a much simpler learning problem than learning an entire transformation from scratch.

This leads to the powerful mental model of a residual stream, where the main representation flows through the network, and each layer simply adds its contribution or correction to this stream.

Component 4: Layer Normalization

The final piece of the puzzle is Layer Normalization. While residual connections help gradients flow, Layer Normalization stabilizes the training process by controlling the distribution of the activations.

Deep networks can suffer from internal covariate shift, where the distribution of each layer's inputs changes during training as the parameters of the previous layers are updated. This forces layers to constantly adapt to a moving target, slowing down training. Normalization techniques address this by standardizing the activations to have a specific mean and variance.

While Batch Normalization is common in computer vision (normalizing across the batch dimension), it is ill-suited for NLP and Transformers due to variable sequence lengths and its poor performance with small batch sizes.

Layer Normalization, by contrast, normalizes across the feature dimension for each individual example in the batch. The mean and variance are computed for each token's embedding vector independently. This makes LayerNorm independent of the batch size and sequence length.

The formula for Layer Normalization is:

where $\mu$ and $\sigma$ are the mean and standard deviation of the elements in the vector $x$, and $\gamma$ (gain) and $\beta$ (bias) are learnable parameters that allow the model to scale and shift the normalized output, giving it the flexibility to restore the representational power that pure normalization might remove.

Pre-LN vs. Post-LN: A Critical Distinction

A seemingly minor implementation detail with major consequences is the placement of the LayerNorm operation relative to the residual connection.

• Post-LN (Original Transformer): LayerNorm(x + Sublayer(x))
• Pre-LN (Modern Standard): x + Sublayer(LayerNorm(x))

In the Post-LN scheme, the output of the addition flows directly into the LayerNorm. This arrangement is notoriously unstable and requires a "learning rate warm-up" phase to prevent divergence early in training. Conversely, in the Pre-LN scheme, the residual path is a clean addition, providing an unimpeded gradient highway. The normalization is applied to the input of the sublayer. This is significantly more stable and is the standard in almost all modern Transformers (e.g., GPT-2 and beyond).

The following research paper provides the definitive theoretical analysis of this phenomenon, explaining exactly why Post-LN is unstable and Pre-LN is not. Given your background, you will find its formal treatment of gradient scaling insightful.

Assembling the Complete Block

With all four components understood, we can now assemble a complete Pre-LN Transformer encoder block. For a given input $x$, the process is as follows:

1. First Sub-layer (Multi-Head Attention):
   a. Normalize the input: $x_{\text{norm1}}=\text{LayerNorm}(x)$
   b. Compute self-attention: $x_{\text{attn}}=\text{MultiHeadAttention}(x_{\text{norm1}},x_{\text{norm1}},x_{\text{norm1}})$
   c. Add the residual: $x^{\prime }=x+x_{\text{attn}}$

2. Second Sub-layer (Feed-Forward Network):
   a. Normalize the intermediate representation: $x_{\text{norm2}}^{\prime }=\text{LayerNorm}(x^{\prime })$
   b. Compute the FFN output: $x_{\text{ffn}}=\text{FFN}(x_{\text{norm2}}^{\prime })$
   c. Add the residual: $x_{\text{output}}=x^{\prime }+x_{\text{ffn}}$

The output, $x_{\text{output}}$, has the exact same shape as the input $x$, allowing these blocks to be stacked seamlessly to form a deep network. A decoder block follows the same pattern but inserts a third sub-layer for cross-attention between the self-attention and the FFN.

Conclusion

In this lesson, we have deconstructed the Transformer block into its four essential components and analyzed their individual contributions.

Here are the key takeaways:

• Multi-Head Attention acts as the dynamic information-mixing layer, allowing tokens to communicate.
• The Position-wise Feed-Forward Network provides non-linearity and is thought to be a primary locus of factual knowledge storage.
• Residual Connections are the architectural backbone, creating "gradient highways" that enable the stable training of very deep networks by allowing layers to learn updates rather than entire transformations.
• Layer Normalization stabilizes training by controlling the statistics of activations. The modern Pre-LN configuration is crucial for this stability, eliminating the need for delicate learning rate warm-up schedules required by the original Post-LN design.

Having now mastered the intricacies of the Transformer's core architecture, we are ready to move from a static model to a dynamic learning system. In the next lesson, we will explore the dominant training methodology in modern NLP: the pre-training and fine-tuning paradigm. We will see how these massive Transformer models are first trained on vast, general-purpose datasets and then adapted to solve specific downstream tasks, forming the foundation of models like BERT and GPT.