Attention to Output: The Transformer Block's Final Transformation

The Role of the Position-wise Feed-Forward Network

After the Multi-Head Attention sub-layer has mixed information across different token positions, the Transformer needs a component to process this new information within each position. This is the job of the Position-wise Feed-Forward Network (FFN).

At first glance, it might seem like just another set of layers. However, its role is critical for two main reasons:

1. Introducing Non-linearity: The attention mechanism, while powerful, is fundamentally a sophisticated weighted averaging operation. In terms of its Value inputs, it is a linear transformation. If we were to stack multiple attention layers without any non-linear processing in between, the entire stack would be mathematically equivalent to a single, more complex linear transformation. The FFN introduces essential non-linearity, typically via a GELU or ReLU activation function, which dramatically increases the model's expressive power and allows it to learn far more complex functions.

2. Per-Position Transformation: The FFN acts as a small, two-layer Multi-Layer Perceptron (MLP) that is applied independently to each token's vector representation. While the attention layer's job is to handle communication between tokens, the FFN's job is to "think" or "digest" the information gathered for each token individually. You can conceptualize this as a two-step process within each block: first, gather context (attention), then process that context (FFN).

Recent research suggests that the FFN is more than just a source of non-linearity; it may also function as a key-value memory or knowledge store, where the model's factual knowledge is encoded.

To dive deeper into the motivation and architecture of the FFN, let's turn to an excellent explanation.

Tracing the Data Flow: The Second Sub-layer

Having established the FFN's purpose, let's now trace the data flow from the output of the attention sub-layer, which we'll call $O$, through this second component. Just as with the attention sub-layer, we'll examine both the modern Pre-LN and original Post-LN patterns.

The Pre-LN FFN Sub-layer (Modern Standard)

The flow is consistent with the first sub-layer: Output = Input + Sublayer(LayerNorm(Input)).

1. Input: We begin with the tensor $O$ produced by the MHA sub-layer. It has the shape [batch_size, seq_len, d_model].
2. Pre-Normalization: The tensor $O$ is passed through the block's second LayerNorm module. $$O_{\text{norm}}={\text{LayerNorm}}_2(O)$$
3. FFN Processing: The normalized tensor is processed by the FFN. As described in the reading, this involves an expansion, a GELU activation, and a contraction. $$F=\text{FFN}(O_{\text{norm}})$$ The output $F$ has the same shape as the input, [batch_size, seq_len, d_model].
4. Residual Connection: The FFN's output $F$ is added back to its un-normalized input $O$. A dropout layer is typically applied to $F$ for regularization. $$H=O+\text{Dropout}(F)$$

The resulting tensor $H$ is the final output of the entire Transformer block. Its dimensions are unchanged, allowing it to be fed directly into the next identical block in the stack.

The Post-LN FFN Sub-layer (Original Architecture)

The original paper applied normalization after the residual connection: Output = LayerNorm(Input + Sublayer(Input)).

1. Input: The tensor $O$ from the MHA sub-layer.
2. FFN Processing: The FFN is applied directly to the un-normalized tensor $O$. $$F=\text{FFN}(O)$$
3. Residual Connection: The output $F$ is added to its input $O$. $$H_{\text{sum}}=O+\text{Dropout}(F)$$
4. Post-Normalization: The result of the sum is passed through the second LayerNorm module. $$H={\text{LayerNorm}}_2(H_{\text{sum}})$$

Putting It All Together: The Complete Transformer Block

We have now traced the data flow through both sub-layers. Let's consolidate this into a complete picture of a single Transformer block.

The diagram below shows the FFN sub-layer in context. The input comes from the MHA sub-layer (which we covered last time), gets processed by the FFN, and is then integrated via a residual connection and a final LayerNorm. This specific diagram illustrates the Post-LN flow.

The following diagram shows a complete encoder block, integrating both the Multi-Head Attention and Feed-Forward Network sub-layers.

For a rigorous, mathematical summary of the entire block's data flow, the Jurafsky and Martin textbook provides a step-by-step breakdown.

Finally, let's see how this entire block is implemented in code. Andrej Karpathy's "Let's build GPT" video provides a hands-on demonstration where he assembles these components into a Block class.

Your DevOps experience with CI/CD pipelines provides a useful mental model here. A Transformer model is a data processing pipeline composed of stacked, identical blocks. Each block is a mini-pipeline with two stages (Attention and FFN). The use of residual connections ensures a "main branch" of information flows cleanly through, while each stage contributes modifications. The Pre-LN architecture adds verification steps (normalization) before each critical stage, ensuring the stability and predictability of the entire pipeline, much like running linters and tests before a build or deployment step.

Conclusion

In this lesson, we have completed our detailed examination of the Transformer block by tracing the data flow through its second major component, the Position-wise Feed-Forward Network.

Here are the key takeaways:

• The FFN provides essential non-linearity and acts as a per-token computation/digestion step after the attention mechanism has mixed information across the sequence.
• The FFN typically follows an "expand-activate-contract" architecture, projecting the d_model representation to a much larger intermediate dimension (4 * d_model) before projecting it back down.
• The data flow through the FFN sub-layer mirrors the attention sub-layer, incorporating Layer Normalization and a residual connection to ensure stable and effective training.
• A complete Transformer block is a self-contained unit that takes a tensor of shape [batch_size, seq_len, d_model] and produces an output tensor of the exact same shape, making the blocks stackable to create deep networks.

You now possess a complete, end-to-end conceptual and mathematical map of a single Transformer block. With this foundational knowledge, you are well-prepared to understand how these blocks are arranged into the large-scale architectures that power modern LLMs.

In our next lesson, we will zoom out to the architectural level, comparing the original encoder-decoder model from "Attention Is All You Need" with the decoder-only (GPT-style) and encoder-only (BERT-style) variants that have become standards in the field.