Position-wise Feed-Forward Networks in Transformers

Introduction

In our last lesson, we dissected the multi-head attention mechanism, establishing it as the Transformer's communication layer. We saw how it allows tokens to look at each other and mix information, creating contextually rich representations. However, communication is only half the story. Once a token has gathered information from its neighbors, it needs to process that information. This is where the second major component of a Transformer block comes into play: the position-wise feed-forward network (FFN).

This lesson will focus on the crucial role of the FFN. We'll explore why it's an indispensable part of the architecture and what function it serves. You'll learn that the FFN has two primary roles. First, it serves a fundamental computational purpose by introducing non-linearity, a requirement for any deep neural network. Second, and more profoundly, research has revealed that it acts as a form of distributed memory, storing the vast factual and linguistic knowledge the model learns during training.

The Computational Role: Non-linearity and Independent Processing

From your background in deep learning, you know that stacking linear layers without a non-linear activation function in between is futile; the entire stack can be collapsed into a single, equivalent linear layer, defeating the purpose of depth. The attention mechanism, at its core, is a linear operation—it computes a weighted average of the Value vectors.

If we were to stack multiple attention layers on top of each other, the model's expressive power would be severely limited. The position-wise feed-forward network solves this problem by introducing a non-linear transformation after the attention sub-layer.

This FFN has two defining characteristics:

It is a simple two-layer multilayer perceptron (MLP). It consists of two linear transformations with a non-linear activation function in between.
It is "position-wise." This means the exact same network (i.e., the same weight matrices) is applied independently to the representation of each token in the sequence. The FFN processes each token's vector in isolation, without any information sharing between positions.

This creates a clear separation of concerns within the Transformer block:

Multi-Head Attention: Mixes information across different token positions. (Communication)
Feed-Forward Network: Transforms the representation at each token position. (Computation/Thinking)

To get a clear, code-grounded intuition for this, let's turn to Andrej Karpathy's "Let's build GPT" series.

Let's build GPT: from scratch, in code, spelled out.

In this clip, Karpathy explains the FFN's role as the 'thinking' part that follows the 'communication' of self-attention. He frames it as a necessary step for tokens to process the information they have gathered.

Watch the segment from this point. Notice his distinction between self-attention (where tokens look at each other) and the FFN (where tokens 'think' on the data individually). This is a powerful mental model for understanding the FFN's purpose.

The FFN Architecture: Expand and Contract

The structure of the FFN is simple yet highly effective. It takes the output from the attention sub-layer (a vector of dimension $d_{model}$ ), projects it to a much higher-dimensional space, applies a non-linearity, and then projects it back down to the original $d_{model}$ dimension.

This diagram illustrates the 'expand and contract' pattern of the FFN. The input $x \in \mathbb{R}^{d_{model}}$ is expanded to an intermediate representation $y \in \mathbb{R}^{d_{ff}}$, and then contracted back to the original dimension in the output $\hat{x} \in \mathbb{R}^{d_{model}}$.

The standard equation for the FFN, using a ReLU activation function, is:

FFN (x) = max (0, x W_{1} + b_{1}) W_{2} + b_{2}

Where:

$x$ is the input vector for a single token position (dimension $d_{model}$ ).
$W_{1}$ and $b_{1}$ are the weights and biases of the first linear layer, which expands the dimension from $d_{model}$ to $d_{ff}$ .
$W_{2}$ and $b_{2}$ are the weights and biases of the second linear layer, which contracts the dimension from $d_{ff}$ back to $d_{model}$ .

A common convention, established in the original "Attention Is All You Need" paper, is to set the inner-layer dimension $d_{ff}$ to be four times the model dimension $d_{model}$ . For a model with $d_{model} = 512$ , the intermediate dimension would be $d_{ff} = 2048$ . This expanded "thinking space" is critical; it's where a significant portion of the model's parameters and representational capacity reside. In fact, the FFN layers typically account for about two-thirds of a Transformer's total parameters.

The following diagram provides a more detailed view of this process applied across an entire sequence of token embeddings.

This diagram visualizes how the same FFN parameters are applied independently to each position in the sequence of embeddings. Each embedding is passed through an expansion layer, a non-linearity (implied), and a contraction layer.

To solidify these architectural details, the following reading provides a concise and clear explanation.

The Transformer Architecture: A Deep Dive - Imad Dabbura

This article by Imad Dabbura explains the necessity of the FFN and details its 'expand, activate, contract' architecture. It provides a superb summary of the FFN's role and structure.

Please read the sections titled Why Is It Needed? and Architecture Details. Pay close attention to the argument for non-linearity and the explanation of the 4x expansion factor.

The Empirical Role: The FFN as a Knowledge Store

While the computational role of the FFN is clear from a network design perspective, recent research has uncovered a deeper, more fascinating function. The FFN appears to be the primary location where the model stores factual and associative knowledge.

A groundbreaking interpretation by Geva et al. (2021) frames the FFN as a distributed key-value memory. In this view:

The first linear layer ( $W_{1}$ ) acts as the keys. The rows of this matrix learn to detect specific patterns or concepts in the input token's representation.
The second linear layer ( $W_{2}$ ) acts as the values. Each column in this matrix corresponds to a "value" to be retrieved when its associated "key" in $W_{1}$ is activated. This value is an instruction to modify the output distribution, for example, by promoting certain vocabulary tokens.
The activation function acts as a gate, determining which key-value pairs are activated by the input and by how much.

This leads to the powerful aphorism: Attention is the routing system; the FFN is the knowledge store. Attention determines which tokens should send information to which other tokens, but the FFN is where the substance of that information—the learned associations like "Paris is the capital of France"—is actually stored.

The following video from Mor Geva, the lead author of this research, provides a direct explanation of this concept.

Mor Geva: Transformer Feed Forward Layers are Key-Value Memories, and Build Predictions

This seminar talk introduces the idea that FFNs were a relatively underexplored part of Transformers and then presents the core hypothesis that they function as key-value memories.

First, watch the introduction from this section to understand the motivation for this research. Then, watch the detailed explanation of the key-value memory hypothesis from this segment. Focus on how the dot product with the first matrix acts as a pattern-matching step (finding the key) and how the output is a linear combination of the columns of the second matrix (retrieving the values).

This interpretation explains how an LLM can answer a factual question. When the model processes "The capital of France is ___", the representation of "France" at the attention output activates specific "key" patterns in an FFN layer. These keys, in turn, retrieve corresponding "value" vectors that promote tokens like "Paris" in the model's final prediction. This process happens in a distributed manner across many neurons and many layers, gradually building the final prediction.

Conclusion

In this lesson, we have established the dual role of the position-wise feed-forward network. It is not merely a standard component plugged in after the attention layer; it is a thoughtfully designed and essential part of the Transformer block.

Let's summarize the key takeaways:

Computational Role: The FFN provides the essential non-linearity that allows Transformer models to have depth. Without it, stacked attention layers would collapse into a single linear transformation.
Processing Role: It acts on each token's representation independently, providing a "computation" or "thinking" step after the "communication" of the attention layer.
Architectural Signature: It follows an "expand-and-contract" pattern, typically projecting the $d_{model}$ representation to a $4 \times d_{model}$ intermediate space before contracting it back down. This intermediate space holds most of the model's parameters.
Empirical Function: Research strongly suggests that FFNs function as a distributed key-value memory, storing the factual and linguistic knowledge learned during pre-training.

You have now examined the two core sub-layers of the Transformer block: multi-head attention and the position-wise feed-forward network. In the upcoming lessons, we will explore the finer details, starting with the GELU activation function commonly used within the FFN. Afterward, we will see how these two sub-layers are combined using residual connections and layer normalization to form a complete, stable, and trainable Transformer block.

Recommended resources

Introduction

The Computational Role: Non-linearity and Independent Processing

The FFN Architecture: Expand and Contract

The Empirical Role: The FFN as a Knowledge Store

Conclusion