Decoding the Feed-Forward Network in Transformers

Introduction

In our previous lessons, we assembled the core components for understanding a sequence. We implemented multi-head attention to allow tokens to gather contextual information from across the entire input, and we added positional encodings to make the model aware of token order.

At this point, the output from our attention sub-layer is a set of context-rich, position-aware vectors. However, the attention mechanism itself, being a weighted sum of value vectors, is fundamentally a linear operation. To build a truly powerful model capable of learning complex patterns, we need to introduce non-linearity. This is the crucial role of the component we'll study today.

Our learning outcome is to explain the purpose of the position-wise feed-forward network (FFN) in each transformer block. We will break down its architecture, understand why it operates "position-wise," and see how it provides the computational depth necessary for the Transformer to model complex functions.

The Need for Non-Linearity

The self-attention mechanism is exceptional at routing and mixing information between tokens. However, after the weighted sums are calculated, the transformation from input values to output is linear. Stacking multiple layers of purely linear operations is mathematically equivalent to a single, more complex linear operation. To give the model the ability to approximate any function, we need to introduce non-linear transformations.

The following resource provides an excellent explanation for why attention alone is not enough and why a component like the FFN is essential.

As the article explains, the Transformer elegantly divides its labor:

• Self-Attention: Handles inter-token communication (routing information across the sequence).
• Feed-Forward Network: Handles intra-token computation (transforming the information at each position).

The "Position-Wise" Architecture

The "position-wise" aspect of the FFN is a key design choice. It means that the exact same feed-forward network (i.e., the same set of weights) is applied independently to each token's representation in the sequence.

This approach has two major benefits:

1. Parallelization: Since the computation for each position is independent, this entire step can be massively parallelized on a GPU, making it very efficient.
2. Parameter Efficiency: Instead of learning a different transformation for every possible position in a sequence, the model learns a single, general-purpose transformation that can be applied anywhere.

The Expand-and-Contract Structure

The FFN itself is a simple two-layer fully connected neural network. However, it has a distinct "expand-and-contract" structure. For a single token representation vector $x$:

1. Expand: The first linear layer projects the vector from its original dimension $d_{\text{model}}$ to a larger, intermediate dimension $d_{\text{ff}}$. Typically, $d_{\text{ff}}=4\times d_{\text{model}}$.
2. Transform: A non-linear activation function (like ReLU or GELU) is applied element-wise.
3. Contract: The second linear layer projects the vector back down from $d_{\text{ff}}$ to the original $d_{\text{model}}$.

The formula for the FFN using a ReLU activation is:

Where:

• $x$ is the input vector for one position (shape: $d_{\text{model}}$).
• $W_1$ is the first weight matrix (shape: $d_{\text{model}}\times d_{\text{ff}}$).
• $b_1$ is the first bias vector (shape: $d_{\text{ff}}$).
• $W_2$ is the second weight matrix (shape: $d_{\text{ff}}\times d_{\text{model}}$).
• $b_2$ is the second bias vector (shape: $d_{\text{model}}$).

The expansion into a higher-dimensional space allows the network to learn more complex relationships between features before compressing them back down. The output dimension must match the input dimension to facilitate the residual connection that follows.

To dive deeper into this architecture, let's return to the Brenndörfer article.

A Look at the Matrix Math

Your background in ML engineering and data science will make you appreciate seeing this in action with matrix operations. While we've discussed the FFN conceptually, seeing the matrix multiplications makes it concrete.

The following video from StatQuest provides a very clear, step-by-step walkthrough of all the matrix math in a Transformer. We will focus only on the part relevant to the FFN.

The FFN's Place in the Transformer Block

Now, let's put everything together. Within a single Transformer encoder or decoder block, the FFN is the second main sub-layer, appearing after the multi-head attention sub-layer.

The data flow through an encoder block is as follows:

1. The input vectors (embeddings + positional encodings) enter the block.
2. They are processed by the Multi-Head Attention sub-layer.
3. A residual connection adds the original input to the attention output, and the result is passed through Layer Normalization.
4. This normalized output becomes the input to the Position-Wise Feed-Forward Network.
5. A second residual connection adds the FFN's input to its output, and the result is again passed through Layer Normalization.
6. This final vector is the output of the block, ready to be fed into the next Transformer block.

This video gives a quick visual summary of this structure.

Conclusion

You've now dissected the second major computational unit within a Transformer block. The FFN may seem simple compared to the attention mechanism, but it is just as critical to the model's success.

Key Takeaways:

• Primary Purpose: The FFN's main role is to introduce non-linearity into the Transformer block, allowing the model to learn far more complex functions than attention alone could.
• Position-Wise Operation: The FFN applies the same learned transformation to each token's representation independently, enabling massive parallelization and parameter efficiency.
• Architecture: It consists of two linear layers in an "expand-and-contract" pattern. It projects the input to a higher dimension (d_ff), applies a non-linear activation (like ReLU or GELU), and projects it back to the original dimension (d_model).
• Role in the Block: It serves as the second sub-layer, processing the context-aware representations generated by the multi-head attention mechanism.

What's Next?

Throughout this module, we've repeatedly mentioned the "Add & Norm" step that follows both the attention and FFN sub-layers. In our next lesson, we will finally demystify this component by focusing on Layer Normalization. You will learn what it does, why it's absolutely essential for stabilizing the training of deep Transformers, and how to implement it.