The feed-forward network — per-token computation

The half of the block that nobody talks about

Attention gets the attention. But roughly two-thirds of a transformer's parameters live in the feed-forward network, and the clearest way to see what it does is to contrast how information moves through each half of the block.

In self-attention the lanes cross. Three tokens means nine connections; a thousand tokens means a million. This is the only sub-layer in the entire architecture where information moves sideways between positions, and it is the sole source of the quadratic cost.

In the feed-forward network the lanes never meet. The same small network — expand from 512 to 2048, apply a nonlinearity, project back to 512 — is applied to each position separately, with identical weights and no knowledge that other positions exist. Token 1's output would be unchanged if tokens 2 and 3 were deleted.

So the block is: mix, then think. Attention gathers relevant context from elsewhere; the feed-forward network does the actual nonlinear work on what was gathered. Neither is sufficient. Attention alone is a weighted average — a linear operation dressed up, unable to compute anything a weighted average cannot. A feed-forward network alone can compute a great deal but only ever about one token in isolation. Alternating them is what makes the stack expressive.

The independence has two payoffs worth naming. It is trivially parallel — every position runs at once with no coordination, which is why the feed-forward network is rarely the latency problem despite holding most of the weights. And it is where the model's knowledge appears to live: there is a growing body of interpretability work reading the expand-then-contract structure as a key-value memory, where the first matrix detects patterns and the second writes associated information back into the residual stream. That reading is not settled, but it explains a stubborn empirical fact — making the hidden dimension wider reliably buys more factual recall, while making attention wider does not.

Note the ratio. $d_{\text{ff}} = 4 \times d_{\text{model}}$ appeared in the original paper with no derivation and has survived nearly every architecture since. Later designs changed the nonlinearity, split the projection into gated halves, and replaced dense layers with sparse experts — but the factor of four kept showing up as roughly optimal.

Every transformer block contains two sub-layers: an attention layer that mixes information across positions, and a feed-forward network (FFN) that processes each position independently. The attention layer lets tokens communicate; the FFN transforms each token's representation in isolation, using the same weights at every position.

Structure

The standard FFN is a two-layer MLP with a nonlinearity between:

The input x is a vector of dimension d_model. The first linear layer projects it up to a higher-dimensional space d_ff, applies a nonlinear activation function, and the second linear layer projects back down to d_model. The expansion ratio d_ff / d_model is typically 4 — the inner dimension is 4x the model dimension. This ratio has been used since the original transformer and persists across nearly every model family.

W_1 has shape (d_ff, d_model). W_2 has shape (d_model, d_ff). For the original transformer with $d_{\text{model}} = 512$ and $d_{\text{ff}} = 2048$, that's $512 \times 2048 + 2048 \times 512 = 2, 097, 152$ parameters per FFN — about twice the parameter count of the attention layer in the same block.

At modern scales the FFN dominates the parameter budget. Llama 3 8B uses $d_{\text{model}} = 4096$ and $d_{\text{ff}} = 14336$ (a 3.5x ratio — slightly below 4x due to the SwiGLU activation, which adds a third weight matrix). Parameters per FFN layer:

• Two-matrix FFN (ReLU/GELU): $2 \times d_{\text{model}} \times d_{\text{ff}} = 2 \times 4096 \times 14336 = 117, 440, 512$ (~117M)
• Three-matrix FFN (SwiGLU): $3 \times d_{\text{model}} \times d_{\text{ff}} = 3 \times 4096 \times 14336 = 176, 160, 768$ (~176M)

With 32 transformer blocks in Llama 3 8B, the SwiGLU FFN layers alone account for $32 \times 176M = 5.64B$ parameters — approximately 70% of the model's 8B total. The FFN is where most of the model's parameters live, and by implication, where most of the model's knowledge is stored.

What FFNs store

Research on transformer internals (Geva et al. 2021) shows that FFN layers function as key-value memories. The first linear layer W_1 acts as a collection of "keys" — each row of W_1 is a pattern detector that activates on specific input features. The activation function gates which patterns fire. The second linear layer W_2 provides the "values" — each column of W_2 specifies what information to add to the residual stream when the corresponding pattern fires.

Meng et al. (2022) demonstrated this concretely by localizing factual knowledge to specific FFN neurons. The association "The Eiffel Tower is in Paris" is stored in a small number of neurons in middle-layer FFNs. By modifying the value vectors (columns of W_2) for those neurons, they could change the model's factual recall — making it say "The Eiffel Tower is in Rome" — without affecting other knowledge. This technique, called ROME (Rank-One Model Editing), established that FFN layers are the primary site of factual knowledge storage in transformers.

The implication is that doubling the FFN width doesn't just add capacity in an abstract sense — it adds more key-value memory slots. A model with $d_{\text{ff}} = 14336$ has 14,336 pattern detectors per layer, each capable of recognizing a different input pattern and contributing a different piece of information to the output. This is why FFN width correlates strongly with the model's factual knowledge and reasoning capability.

Activation functions

The choice of activation function between W_1 and W_2 has evolved significantly. Three activations dominate the landscape:

ReLU (original transformer)

The simplest activation: zero out all negative values, pass positive values unchanged. The original transformer (Vaswani et al. 2017) used ReLU. It works, but has a known failure mode: "dead neurons." If a neuron's pre-activation is consistently negative (across all training examples), its gradient is always zero and the neuron permanently stops learning. In large models with aggressive learning rates, 10–30% of ReLU neurons can die during training (Lu et al. 2020).

GELU (GPT-2, GPT-3, BERT)

where Phi(x) is the standard Gaussian CDF. In practice this is computed as: GELU(x) ≈ 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3))). GELU was introduced by Hendrycks & Gimpel (2016).

Unlike ReLU's hard cutoff at zero, GELU smoothly gates values near zero — small negative inputs are attenuated rather than killed outright. This eliminates dead neurons and provides smoother gradient flow. GELU became the default activation for BERT (Devlin et al. 2019) and the GPT family (Radford et al. 2018, 2019; Brown et al. 2020).

SwiGLU (Llama 3, PaLM, Gemma)

SwiGLU (Shazeer 2020) restructures the FFN by introducing a gating mechanism:

where $\text{swish}(x) = x \cdot \text{sigmoid}(x)$ and ⊙ is element-wise multiplication. This uses three weight matrices — W_1, W_gate, and W_2 — instead of two. The gate swish(W_gate * x) learns to selectively pass or suppress features from W_1 * x before the output projection.

SwiGLU is a member of the GLU (Gated Linear Unit) family, where the core idea is that one linear transformation produces the content and another produces a gate that modulates it. Dauphin et al. (2017) introduced GLUs with sigmoid gating; Shazeer (2020) systematically evaluated all combinations of gating activation (sigmoid, ReLU, GELU, swish) and found SwiGLU consistently outperformed the others across model sizes from 50M to 1B parameters.

The extra weight matrix adds ~50% more parameters per FFN layer. To compensate and keep the total parameter budget comparable, models using SwiGLU often reduce d_ff relative to the 4x ratio used with ReLU/GELU. Llama 3 8B uses $d_{\text{ff}} = 14336$ instead of the $4 \times 4096 = 16384$ you'd expect, keeping the per-layer parameter count close to what a two-matrix FFN with the full 4x ratio would have.

Implementation

import torch
import torch.nn as nn
import torch.nn.functional as F


class FFN_ReLU(nn.Module):
    """Original transformer FFN: two linear layers with ReLU."""
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff)
        self.w2 = nn.Linear(d_ff, d_model)

    def forward(self, x):
        return self.w2(F.relu(self.w1(x)))


class FFN_GELU(nn.Module):
    """GPT-2/BERT FFN: two linear layers with GELU."""
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff)
        self.w2 = nn.Linear(d_ff, d_model)

    def forward(self, x):
        return self.w2(F.gelu(self.w1(x)))


class FFN_SwiGLU(nn.Module):
    """Llama/PaLM FFN: gated linear unit with swish activation."""
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff, bias=False)
        self.w_gate = nn.Linear(d_model, d_ff, bias=False)
        self.w2 = nn.Linear(d_ff, d_model, bias=False)

    def forward(self, x):
        return self.w2(self.w1(x) * F.silu(self.w_gate(x)))

Parameter counts

d_model = 4096   # Llama 3 8B
d_ff = 14336     # Llama 3 8B

for name, cls in [("ReLU", FFN_ReLU), ("GELU", FFN_GELU), ("SwiGLU", FFN_SwiGLU)]:
    ffn = cls(d_model, d_ff)
    params = sum(p.numel() for p in ffn.parameters())
    print(f"{name:6s}: {params:>12,} parameters ({params / 1e6:.1f}M)")

Output:

ReLU  :  117,464,064 parameters (117.5M)
GELU  :  117,464,064 parameters (117.5M)
SwiGLU:  176,160,768 parameters (176.2M)

ReLU and GELU have identical parameter counts — 2 × d_model × d_ff plus biases. SwiGLU is 50% larger because of the third matrix. Across 32 layers:

• ReLU/GELU: $32 \times 117.5M = 3.76B$ in FFN layers
• SwiGLU: $32 \times 176.2M = 5.64B$ in FFN layers

For Llama 3 8B's 8.03B total parameters, the SwiGLU FFNs account for 70.2%. Attention layers (with their Q, K, V, and output projections) contribute most of the rest. The embedding matrix, layer norms, and output projection are comparatively small.

The residual connection

The FFN doesn't replace the token's representation — it adds to it. The transformer uses a residual connection: output = x + FFN(LayerNorm(x)) (in pre-norm architectures like GPT-2 and Llama) or output = LayerNorm(x + FFN(x)) (in post-norm architectures like the original transformer).

The residual stream is a core concept in transformer interpretability (Elhage et al. 2021). Each token's representation starts as its embedding vector and flows through the network. Every attention layer and every FFN layer reads from this stream and writes an additive update back into it. The final representation is the embedding plus the sum of all attention and FFN contributions across all layers.

This means each FFN layer can be understood as answering the question: "Given what this token's representation currently encodes (after all prior layers), what additional information should I add?" The FFN at layer 5 might add syntactic role information ("this is a verb"). The FFN at layer 15 might add factual associations ("this refers to Paris"). The FFN at layer 30 might adjust the representation to match the output distribution for next-token prediction.

import torch
import torch.nn as nn

class TransformerBlockPreNorm(nn.Module):
    """Single transformer block with pre-norm (GPT-2/Llama style)."""
    def __init__(self, d_model, n_heads, d_ff):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
        self.ln2 = nn.LayerNorm(d_model)
        self.ffn = FFN_SwiGLU(d_model, d_ff)

    def forward(self, x, mask=None):
        # Attention: read from stream, write additive update
        h = self.ln1(x)
        attn_out, _ = self.attn(h, h, h, attn_mask=mask)
        x = x + attn_out

        # FFN: read from stream, write additive update
        x = x + self.ffn(self.ln2(x))
        return x

block = TransformerBlockPreNorm(d_model=4096, n_heads=32, d_ff=14336)
params = sum(p.numel() for p in block.parameters())
print(f"Parameters per block: {params:,} ({params/1e6:.1f}M)")
# ~243M per block (attention ~67M + SwiGLU FFN ~176M)

The FFN and attention layer serve complementary roles: attention moves information between positions (letting tokens communicate), and the FFN transforms information at each position (adding knowledge and adjusting representations). Together, one block of attention + FFN constitutes a single "reasoning step" in the transformer's computation. Stack 32 such blocks and you get a model capable of language understanding, generation, and — increasingly — multi-step reasoning.