The Feed-Forward Block

What is the Feed Forward Network (FFN) Block?

Every Transformer layer contains two main components:

1. Multi Head Attention Sub Layer: Allows tokens to communicate and share context across sequence positions.
2. Feed Forward Network (FFN) Sub Layer: Processes each token vector independently.

While Attention routes information between tokens, the FFN processes and transforms information within each individual token.

  Token Embedding Vector x [1 x d_model]
             │
             ▼
  [ Linear Layer 1 (W1, b1) ] ──► Expands dimension to d_ff = 4 * d_model
             │
             ▼
  [ Non Linear Activation (GELU / Swish) ]
             │
             ▼
  [ Linear Layer 2 (W2, b2) ] ──► Shrinks dimension back to d_model
             │
             ▼
  Output Vector [1 x d_model]

Standard FFN Math

In classic Transformers (BERT, GPT 2):

$$\text{FFN}(x) = \text{GELU}(x W_1 + b_1) W_2 + b_2$$

• $W_1 \in \mathbb{R}^{d_{\text{model}} \times d_{\text{ff}}}$: Expands dimension by 4x ($d_{\text{ff}} = 4 \cdot d_{\text{model}}$).
• $W_2 \in \mathbb{R}^{d_{\text{ff}} \times d_{\text{model}}}$: Projects dimension back to $d_{\text{model}}$.

The FFN block contains roughly two thirds of the total parameters in a Transformer model!

Modern Variant: SwiGLU (LLaMA Standard)

Modern LLMs (LLaMA 2, LLaMA 3, PaLM) replace standard FFNs with SwiGLU (Swish Gated Linear Unit):

$$\text{SwiGLU}(x) = \left( \text{Swish}(x W_{\text{gate}}) \odot x W_{\text{up}} \right) W_{\text{down}}$$

SwiGLU uses a gating mechanism where one branch acts as a smooth gate for another branch, improving representation quality and training stability.

FFNs as Key Value Factual Memories

Interpretability research (Geva et al., 2021) revealed what FFN layers actually do:

1. First Linear Layer ($W_1$): Acts as Keys. Neurons fire when they recognize specific concepts or patterns in the input token (for example, detecting "Eiffel Tower").
2. Second Linear Layer ($W_2$): Acts as Values. When a key fires, the second layer injects corresponding factual predictions (for example, adding distribution weight for "Paris") into the token embedding.

FFN layers are the primary memory storage where Large Language Models store world knowledge learned during pretraining.

Say this out loud

The Feed Forward Network is a point wise two layer dense block applied to each token independently. It expands model dimension by 4x before projecting back. While attention routes context between tokens, FFN layers transform features within each token and act as key value factual memory stores holding pretraining world knowledge. Modern LLMs use SwiGLU gated activations for improved stability.

Followups to expect

1. How does Mixture of Experts (MoE) modify the FFN block? MoE replaces the single dense FFN block with multiple parallel Expert FFN blocks (for example, 8 experts), using a Router network to direct each token to the top 2 most relevant experts.
2. Why are FFN parameters called Point Wise? Because the weight matrices W1 and W2 operate on single token vectors at position t independently, without mixing information across neighboring sequence positions.