Feed-Forward Networks › The Position-Wise Feed-Forward Network

To understand why transformers need the feed-forward network, consider what attention alone provides. Self-attention computes weighted averages of value vectors, where the weights come from query-key similarities. This is powerful for gathering contextual information: a token can attend to any position in the sequence and assemble a representation that reflects its surrounding context. But no matter how sophisticated the attention patterns, the operation remains a convex combination of value vectors. Each output is literally a weighted sum of inputs. This is a linear function.

Linear functions have severe limitations for language understanding. They can only rotate and scale the input space, then translate it. They cannot learn the curved decision boundaries that distinguish "bank" (financial institution) from "bank" (riverbank). They cannot capture the complex feature interactions that tell you whether "not bad" is negative or positive in context. They cannot represent the conditional relationships that make "the cat sat on the mat" grammatically fine but "the cat sat on the cat" slightly odd. For a model to approximate arbitrary functions of language, it needs the ability to carve up representation space in nonlinear ways.

The key insight is that any sufficiently wide network with a nonlinear activation function can approximate arbitrary continuous functions. This is the universal approximation theorem, a foundational result in neural network theory. The FFN uses this result by introducing a nonlinear activation between two linear layers. The first layer projects the input into a higher-dimensional space. The activation function applies a nonlinearity. The second layer projects back. The result is a function that can approximate far more complex relationships than any linear operation could.

There is a critical design choice embedded in how the FFN is applied: instead of mixing information across positions like attention does, the FFN processes each position independently. If your sequence has 100 tokens, the FFN applies the same transformation to each of those 100 representations separately, using identical weights for all positions. This is called a position-wise or position-independent transformation, and it fundamentally shapes the FFN's role in the transformer.

This division of labor is elegant. Attention handles inter-position communication: it routes information between tokens, allowing the model to build representations that depend on context. The FFN handles intra-position computation: it transforms each position's representation using a learned nonlinear function. By separating these concerns, the transformer achieves both contextual awareness (from attention) and expressive power (from the FFN).

Why share the same network across all positions? Two reasons justify this choice. First, language exhibits translation invariance in a deep sense: the grammatical patterns that help understand "the cat" are equally useful whether those words appear at the beginning, middle, or end of a sentence. A transformation that converts a noun phrase's representation into something more useful should work wherever that noun phrase appears. Second, parameter sharing dramatically reduces model size. Instead of learning separate networks for each of the potentially thousands of positions in a sequence, we learn one network that generalizes across all positions. This is a major efficiency win.

Think of the position-wise FFN as a shared "dictionary of transformations." After attention has looked up what each position means in context, the FFN looks up how to transform that meaning into a richer representation, using the same dictionary regardless of where the position appears in the sequence.

With this motivation in place, let's examine what the FFN computes. The architecture is surprisingly simple: two linear transformations with a nonlinear activation function sandwiched between them. But this simplicity is deceptive. The expand-transform-contract structure enables remarkably expressive transformations, and the specific choice of expansion factor shapes the model's capacity in ways that affect both learning and efficiency.

Consider a single token's representation: a vector $x$ with $d_{\text{model}}$ dimensions. This vector encodes everything the model currently knows about that token in its context, having passed through however many transformer layers preceded this one. The FFN transforms this vector in three stages:

1. Expand: Project $x$ into a higher-dimensional space
2. Transform: Apply nonlinearity to enable learning of curved boundaries
3. Contract: Project back to the original dimension

The complete formula captures all three stages in one expression. Given an input vector $x \in \mathbb{R}^{d_{\text{model}}}$, the FFN computes:

$$\text{FFN}(x) = \sigma(xW_1 + b_1)W_2 + b_2$$

where:

• $x \in \mathbb{R}^{d_{\text{model}}}$: the input vector for a single position (the token's representation from the previous layer)
• $W_1 \in \mathbb{R}^{d_{\text{model}} \times d_{ff}}$: the first projection matrix, which expands the dimension from $d_{\text{model}}$ to $d_{ff}$
• $b_1 \in \mathbb{R}^{d_{ff}}$: the first bias vector, added after the first linear transformation
• $\sigma$: a nonlinear activation function (e.g., ReLU, GELU) applied element-wise to the hidden representation
• $W_2 \in \mathbb{R}^{d_{ff} \times d_{\text{model}}}$: the second projection matrix, which contracts the dimension back from $d_{ff}$ to $d_{\text{model}}$
• $b_2 \in \mathbb{R}^{d_{\text{model}}}$: the second bias vector, added to produce the final output
• $d_{ff}$: the hidden (or intermediate) dimension, typically set to $4 \times d_{\text{model}}$

To understand this formula, let's read it from the inside out, following the order of operations.

Stage 1: The Expansion ($xW_1 + b_1$)

The input vector $x$ is multiplied by a weight matrix $W_1$ and shifted by a bias vector $b_1$. This is a standard linear transformation, but with a twist: the output dimension is larger than the input dimension. If $x$ has 512 dimensions, the result might have 2048 dimensions. We're projecting into a higher-dimensional space where the data is easier to manipulate. In this expanded space, the model has more "room" to organize features before applying nonlinearity.

Why does projecting to higher dimensions help? Think of it this way: in the original $d_{\text{model}}$-dimensional space, the token's representation may conflate many different features in a compressed form. By expanding to $d_{ff}$ dimensions, the first layer can disentangle these features, spreading them across a larger space where they are more separable. The nonlinearity can then act on these individual features independently, creating the curved boundaries that linear operations cannot.

Stage 2: The Nonlinearity ($\sigma(\cdot)$)

The activation function $\sigma$ is applied element-wise to the expanded representation. Common choices include ReLU (which zeros out negative values) and GELU (a smoother alternative that we'll examine in detail in the next chapter). This is where the magic happens: the nonlinearity allows the network to learn curved decision boundaries that would be impossible with linear transformations alone.

The element-wise nature of the activation matters. Because it operates on each hidden dimension independently, it creates a kind of gating: some hidden dimensions activate strongly and contribute to the output, while others are suppressed to near zero. Different input patterns activate different subsets of hidden dimensions, which is exactly how the network implements different transformations for different types of inputs.

Stage 3: The Contraction ($(\cdot)W_2 + b_2$)

Finally, the transformed representation is projected back to the original dimension via weight matrix $W_2$ and bias $b_2$. The output has the same dimensionality as the input, which is essential for the residual connection that adds the FFN output back to its input. The second linear layer combines the activated hidden dimensions into a compressed $d_{\text{model}}$-dimensional output, essentially summarizing the computation done in the expanded space.

Why does this formula make sense? Notice that without the nonlinearity $\sigma$, the composition of two linear maps is itself a linear map: $x W_1 W_2 + \text{const}$ is just another linear function of $x$. The nonlinearity is the only thing that prevents the two layers from collapsing into a single linear layer. The biases also allow the network to shift the activation function's operating range, making it easier to learn the right threshold for each hidden dimension.

The expansion ratio $d_{ff} / d_{\text{model}}$ is a key hyperparameter. The original transformer used a 4x expansion: for $d_{\text{model}} = 512$, the hidden dimension was $d_{ff} = 2048$. This ratio has become a de facto standard, though modern architectures sometimes use different values, especially when combined with gated variants, as we'll explore later.

With the formula understood, translating it to code reveals its simplicity. The entire FFN is just two matrix multiplications with a nonlinearity in between: