Feed-forward network sublayers

What is the Feed-Forward Network (FFN)?

Architecture Deep Dive

Why the Expansion-Contraction Pattern?

Let's derive why we expand then contract, starting from what we need:

Goal: Learn a non-linear function $f: \mathbb{R}^{d_{\text{model}}} \to \mathbb{R}^{d_{\text{model}}}$ that can represent complex patterns.

Step 1 — Universal Approximation

By the universal approximation theorem, a two-layer network with enough hidden units can approximate any continuous function. But how many units?

Step 2 — Information Bottleneck

If $d_{ff} = d_{\text{model}}$ , we have a bottleneck—limited capacity. The hidden layer needs to represent more features than the input to capture complex combinations.

Step 3 — The4x Rule

Empirically, $d_{ff} = 4 \times d_{\text{model}}$ works well. Why?

Expansion: 4x more neurons means we can compute $\sim 4 \times d_{\text{model}}$ different feature combinations

Sparsity: ReLU activations are sparse (~50% neurons active), so effective capacity is $\sim 2 \times d_{\text{model}}$

Compression: Second layer projects back, forcing the network to learn useful high-dimensional features

Formula Complete FFN Equations

For input $x \in \mathbb{R}^{d_{\text{model}}}$ : $h = \sigma(xW_1 + b_1)$ $\text{FFN}(x) = hW_2 + b_2$

Where:

$W_1 \in \mathbb{R}^{d_{\text{model}} \times d_{ff}}$ — expansion matrix
$b_1 \in \mathbb{R}^{d_{ff}}$ — expansion bias
$\sigma$ — activation function (ReLU, GELU, SiLU)
$W_2 \in \mathbb{R}^{d_{ff} \times d_{\text{model}}}$ — projection matrix
$b_2 \in \mathbb{R}^{d_{\text{model}}}$ — projection bias Parameter count: $2\times d_{\text{model}} \times d_{ff} + d_{ff} + d_{\text{model}} \approx 8d_{\text{model}}^2$ (dominates Transformer parameters!)

Activation Functions: ReLU → GELU → SiLU

Why not just ReLU?

Starting with ReLU (Rectified Linear Unit): $\text{ReLU}(x) = \max(0, x)$

Problem: Hard cutoff at zero creates dead neurons (gradient = 0 for $x < 0$ ).

Evolution to GELU: $\text{GELU}(x) = x \cdot \Phi(x) = x \cdot P(X \leq x), \quad X \sim \mathcal{N}(0, 1)$

Intuition: Multiply input by probability it's greater than a random Gaussian. Small negative values get small but non-zero gradients (smooth!).

Approximation: $\text{GELU}(x) \approx 0.5x\left(1 + \tanh\left[\sqrt{\frac{2}{\pi}}\left(x + 0.044715x^3\right)\right]\right)$

Modern SiLU/Swish: $\text{SiLU}(x) = x \cdot \sigma(x) = \frac{x}{1 + e^{-x}}$

Why better? Smooth everywhere, self-gated (output depends on input magnitude), empirically stronger for large models.

Worked example Dimension Flow Example

Setup: BERT-base with $d_{\text{model}} = 768$ , $d_{ff} = 3072$

Input token embedding: $x = [0.2, -0.5, 0.8, \ldots] \in \mathbb{R}^{768}$

Step 1 — Expansion: $h = \text{GELU}(xW_1 + b_1)$

Matrix multiply: $768 \to 3072$ dimensions
Why this step? Create a richer feature space. E.g., $h_{100}$ might fire for "positive sentiment + past tense" combinations
Result: $h \in \mathbb{R}^{3072}$ , typically ~50% entries near-zero (sparsity)

Step 2 — Projection: $y = hW_2 + b_2$

Matrix multiply: $3072 \to 768$ dimensions
Why this step? Compress learned features back to residual stream dimension for addition
Result: $y \in \mathbb{R}^{768}$

Residual connection: $\text{output} = \text{LayerNorm}(x + y)$

Worked example Concrete Calculation (Tiny Transformer)

Micro-example: $d_{\text{model}} = 4$ , $d_{ff} = 8$ , single position

Given:

Input: $x = [1, 0, -1, 2]$
$W_1 = \begin{bmatrix} 0.5 & -0.3 & \cdots \ 0.2 & 0.4 & \cdots \\ \vdots & \vdots & \ddots \end{bmatrix}_{4 \times 8}$ (showing first 2 cols)
$b_1 = [0.1, -0.2, 0.3, -0.4, 0.5, -0.6, 0.7, -0.8]$ Computation: First hidden unit: $h_1 = \text{ReLU}(x \cdot W_1[:,1] + b_1[1])$ $= \text{ReLU}(1(0.5) + 0(0.2) + (-1)(\cdots) + 2(\cdots) + 0.1)$ $= \text{ReLU}(\text{some value}) = \max(0, \text{value})$

Why this step? The linear combination $x \cdot W_1[:,1]$ detects a pattern (e.g., "first dim positive AND third dim negative"), ReLU decides if that pattern is present.

Repeat for all8 hidden units → $h \in \mathbb{R}^8$

Then project back: $y = hW_2 + b_2 \in \mathbb{R}^4$

Position-Wise = No Cross-Position Mixing

Mathematical statement: For sequence $X = [x_1, x_2, \ldots, x_n]$ where $x_i \in \mathbb{R}^{d_{\text{model}}}$ : $\text{FFN}(X) = [\text{FFN}(x_1), \text{FFN}(x_2), \ldots, \text{FFN}(x_n)]$

Each $\text{FFN}(x_i)$ uses the same weights $W_1, W_2, b_1, b_2$ but different input data.

Role in the Transformer Block

The complete encoder block:

\begin{align} x' &= \text{LayerNorm}(x + \text{MultiHeadAttention}(x)) \\ x'' &= \text{LayerNorm}(x' + \text{FFN}(x')) \end{align}

Why this order?

Attention first: Aggregate context (linear mixing)
FFN second: Extract non-linear patterns from aggregated context
Repeat 12-24 times: Each layer refines representations further

Common Mistakes

Common mistake Mistake 1: "FFN is just a regular neural network layer"

Why it feels right: It's literally two linear layers with activation—looks like MLP!

What's wrong:

Regular MLP processes entire input together (e.g., flatten image, process as vector)
FFN processes each position independently with shared weights
This is actually a 1x1 convolution in sequence modeling terms!

Fix: Think "position-wise" = broadcast the same network across sequence dimension. In code:

# Wrong mental model
output = mlp(x.view(-1))  # Flatten sequence
# Correct
output = torch.stack([mlp(x[i]) for i in range(seq_len)])
# Or efficiently: mlp(x)  # broadcasts automatically

Common mistake Mistake 2: "The 4x expansion is arbitrary"

Why it feels right: Seems like a magic number; why not 2x or 8x? What's wrong: There's deep reasoning:

Too small (e.g., 1x): Becomes an autoencoder—compresses then reconstructs, limited capacity
Too large (e.g., 8x): More parameters → slower, overfitting risk, minimal accuracy gain
4x sweet spot: Empirically tested across thousands of experiments. Balances:
- Expressiveness (enough hidden units to learn complex functions)
- Efficiency (parameters scale as $O(d^2)$ , so 8x would double params for<5% gain)
- Gradient flow (very wide layers can have vanishing gradient issues)

Fix: Understand it as an empirical result backed by theory (universal approximation + sparsity), not a random choice. Modern models sometimes use 3x or 5.33x (MoE models), but 4x is the default baseline.

Modern Variations

GLU Variants (Gated Linear Units)

Recent models (PaLM, LaMA) replace standard FFN with gated variants:

$\text{FFN}_{\text{GLU}}(x) = (\sigma(xW_1) \odot xW_3)W_2$

Where $\odot$ is element-wise multiplication.

Why gating?

$\sigma(xW_1)$ acts as a learned gate (which features to pass)
$xW_3$ computes the features
Gate × Features = selective amplification
Benefit: 15-20% fewer parameters for same performance (gates provide implicit regularization)

Mixture of Experts (MoE)

Replace single FFN with multiple expert FFNs + router:

$\text{FFN}_{\text{MoE}}(x) = \sum_{i=1}^N \text{Router}(x)_i \cdot \text{Expert}_i(x)$

Why?

Each expert specializes (e.g., Expert 1 = syntax, Expert 2 = semantics)
Sparse activation: only top-k experts compute per token
Benefit: 10x parameters, 1.5x compute (most experts dormant per token)

Connections

Multi-Head Attention: Attention mixes information, FFN processes it
Residual Connections: FFN output added to input, enables deep networks
Layer Normalization: Applied after residual, stabilizes FFN training
Positional Encoding: FFN doesn't encode position—relies on positional encoding from input
Transformer Encoder Block: FFN is second sub-layer after attention
Activation Functions: Choice (ReLU/GELU/SiLU) significantly impacts performance
Universal Approximation Theorem: Justifies two-layer FFN structure
Batch Normalization vs Layer Normalization: LayerNorm used instead of BatchNorm for sequences

Flashcards

#flashcards/ai-ml

What is the feed-forward network sublayer in Transformers? :: A position-wise fully connected network applied independently to each token, consisting of two linear layers with a non-linear activation between them, typically expanding dimension by 4x then projecting back.

Why does the FFN expand to 4x the model dimension?

To create a richer feature space for learning complex non-linear patterns. The expansion provides enough capacity for universal approximation, while ReLU sparsity (~50% active) makes the effective capacity ~2x model dimension.

What is the formula for the standard Transformer FFN?

FFN(x) = max(0, xW₁ + b₁)W₂ + b₂, where W₁ expands dimensions (d_model → d_ff), activation is applied, and W₂ projects back (d_ff → d_model).

Why is the FFN applied position-wise?

Attention already mixes information across positions. FFN specializes in position-independent feature extraction on the already-mixed representations, which paralelizes perfectly and prevents redundant cross-position modeling.

What percentage of Transformer parameters are in the FFN?

Approximately 2/3 of each layer's parameters. With d_ff = 4×d_model, the FFN has ~8d²_model parameters vs ~4d²_model for attention.

Why did models move from ReLU to GELU to SiLU?

ReLU has hard cutoff (dead neurons for x<0). GELU smooths the activation with Gaussian probability weighting. SiLU (x·σ(x)) is self-gated, smooth everywhere, and empirically performs best for large models.

What is the difference between FFN and a regular MLP?

FFN applies the same network independently to each sequence position (like 1x1 convolution), while MLP processes the entire input as a single vector. FFN has position-wise weight sharing.

What is a GLU variant and why use it?

Gated Linear Unit: FFN_GLU(x) = (σ(xW₁) ⊙ xW₃)W₂, where one path gates the other. Provides 15-20% parameter savings for same performance through learned selective amplification.

Recall Explain to a 12-year-old

Imagine you're in a class and the teacher asks everyone to share one interesting fact. That's like attention—everyone hears what others said and combines those ideas in their head.

Now, each student goes home and thinks deeply about all the facts they heard, connecting them to what they already know. Maybe you heard "dogs are loyal" and "wolves hunt in packs," and at home you realize "oh, dogs evolved from wolves, that's why they're social!" That deep thinking alone at home is the feed-forward network.

You don't call your classmates while thinking (no mixing with others). You just process what you learned. Then tomorrow, you come back to class with better understanding, and the cycle repeats.

The "4x expansion" is like having 4x more brain space when you think at home—you can consider way more connections and ideas before summarizing your final thought to bring back to class.

Concept Map

Hinglish (regional understanding)

Intuition Hinglish mein samjho

Feed-forward network (FFN) transformer ka wo component hai jo har token ko independently process karta hai. Socho jaise attention ne sab tokens se information mix kar di, ab FFN ka kaam hai us mixed information se patterns nikalna.

Architecture simple hai: pehle input ko4x zyada dimensions mein expand karo (agar model dimension 768 hai to 3072 ban jata hai), phir ek non-linear activation (ReLU ya GELU) apply karo, aur phir wapas original dimension mein compress karo. Ye expansion zaruri hai kyunki zyada neurons ka matlabzyada complex patterns seekh sakte ho. ReLU se GELU tak evolutionhua kyunki GELU smooth hai—negative values ke liye bhi thoda gradient milta hai, matlab training better hoti hai.

Ek important baat: FFN har position ko alag-alag process karta hai, kisi position ke data ko dusre position ke sath mix nahi karta. Ye deliberate design choice hai kyunki attention layer already mixing kar chuka hota hai. FFN ka kaam sirf feature extraction hai, not information routing. Aur parameters ki baat karein to FFN mein transformer ki2/3 parameters hoti hain (attention sezyada!), isliye ye model ki real learning capacity hai. Modern variants jaise GLU ya Mixture-of-Experts (MoE) FFN ko aur efficient banate hain by using gating mechanisms ya specialized expert networks.

What is the Feed-Forward Network (FFN)?

Architecture Deep Dive

Why the Expansion-Contraction Pattern?

Activation Functions: ReLU → GELU → SiLU

Position-Wise = No Cross-Position Mixing

Role in the Transformer Block

Common Mistakes

Modern Variations

GLU Variants (Gated Linear Units)

Mixture of Experts (MoE)

Connections

Flashcards

Concept Map

Hinglish (regional understanding)

Go deeper — visual, from zero

Test yourself — Transformer Architecture

Connections