The Feed-Forward Network: Position-Wise Compute

After attention mixes information across positions, every token passes through a position-wise feed-forward network (FFN) — the same small MLP applied independently to each position. The FFN expands d_model up to d_ff (the "intermediate" dimension, typically 4× d_model), applies a non-linear activation, and projects back down.

FFN(x) = Activation(x · W₁ + b₁) · W₂ + b₂

The expand-then-contract pattern is the point. The intermediate representation has 4x as many dimensions as d_model, giving the FFN room to compute non-linear functions of features the attention layer extracted. The projection back down reuses learned reductions to compress the result back into d_model.

How much of the parameter budget lives here

For a typical decoder block with d_model=4096 and d_ff=14336 (Llama 3 8B):

Attention: 4 × d_model² = 67M parameters per layer.
FFN (SwiGLU, 3 matrices): 3 × d_model × d_ff ≈ 176M parameters per layer.

The FFN holds about 70% of per-layer parameters. The same ratio holds at almost every model size. When people say "MoE shifts compute around the FFN," they mean the FFN is most of the model.

Code

Standard FFN with GELU·python

class FeedForward(nn.Module):
    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)
        self.act = nn.GELU()
    def forward(self, x):
        return self.W2(self.act(self.W1(x)))

Parameter math for one block·python

def block_params(d_model, d_ff, n_kv_heads=None, n_q_heads=None):
    n_q_heads = n_q_heads or 32
    n_kv_heads = n_kv_heads or n_q_heads
    d_head = d_model // n_q_heads
    attn = (n_q_heads + 2 * n_kv_heads) * d_model * d_head + d_model * d_model
    ffn  = 3 * d_model * d_ff      # SwiGLU has 3 matrices
    norm = 2 * d_model              # 2 RMSNorms per block
    return attn, ffn, norm

a, f, n = block_params(4096, 14336, n_kv_heads=8, n_q_heads=32)
print(f"attn {a/1e6:.0f}M  ffn {f/1e6:.0f}M  norm {n}  per block")

Exercise

Compute the parameter breakdown (attention vs FFN vs norm) for: GPT-2 (d=768, d_ff=3072, 12 heads, 12 layers); Llama 3 8B (d=4096, d_ff=14336, 32 Q heads, 8 KV heads, 32 layers); Mixtral 8x22B (d=6144, d_ff per expert ≈16384, 48 Q heads, 8 KV heads, 56 layers). What fraction of parameters lives in each component?

The Feed-Forward Network: Position-Wise Compute

How much of the parameter budget lives here

Code

External links

Exercise

Progress

Comments 0