Feed-Forward Networks & Normalization

Every Transformer block has two sub-layers: attention and FFN. This chapter covers the FFN side — the “expand-then-compress” network that processes each token independently — plus the normalization and residual connections that make deep networks trainable.

1. The FFN: Expand and Compress

Standard FFN

The original Transformer FFN:

\text{FFN}(\mathbf{x}) = \text{ReLU}(\mathbf{x}\mathbf{W}_1 + \mathbf{b}_1)\mathbf{W}_2 + \mathbf{b}_2

Shape flow: $(B, S, d) \to (B, S, d_{ff}) \to (B, S, d)$

The “expand-then-compress” design: up-project to $d_{ff} = 4d$ (expand the representational space), apply non-linearity, then down-project back to $d$ (compress). This bottleneck structure gives each token a chance to process information in a higher-dimensional space.

Activation Functions

Function	Formula	Properties
ReLU	$\text{ReLU}(x) = \max(0, x)$	Simple, but “dead neurons” (zero gradient for $x < 0$ )
GELU	$\text{GELU}(x) = x \cdot \Phi(x) \approx x \cdot \sigma(1.702x)$	Smooth approximation of ReLU, used in BERT/GPT-2
Swish/SiLU	$\text{Swish}(x) = x \cdot \sigma(x) = \frac{x}{1 + e^{-x}}$	Smooth, non-monotonic, used in SwiGLU

Where $\Phi(x)$ is the Gaussian CDF and $\sigma(x) = 1/(1+e^{-x})$ is the sigmoid function.

SwiGLU: The Modern Standard

SwiGLU Data Flow

Most modern LLMs (LLaMA, Qwen, Mistral, DeepSeek) use SwiGLU:

\text{SwiGLU}(\mathbf{x}) = \left(\text{Swish}(\mathbf{x}\mathbf{W}_{\text{gate}})\right) \odot \left(\mathbf{x}\mathbf{W}_{\text{up}}\right)

\text{FFN}(\mathbf{x}) = \text{SwiGLU}(\mathbf{x}) \cdot \mathbf{W}_{\text{down}}

The gating mechanism: instead of applying the same activation to all dimensions, the model learns a per-dimension gate that controls which information passes through. $\mathbf{W}_{\text{gate}}$ produces gate signals, $\mathbf{W}_{\text{up}}$ produces candidate values, element-wise multiply selects what to keep.

Impact on parameter count: SwiGLU has 3 matrices (gate, up, down) instead of 2. To maintain similar total parameters, $d_{ff}$ is reduced from $4d$ to approximately $(8/3)d$ .

FFN Type	Matrices	$d_{ff}$ (for 7B, $d=4096$ )	Params per block
Standard (ReLU/GELU)	2	$4d = 16384$	$2 \times d \times d_{ff} = 134M$
SwiGLU	3	$(8/3)d \approx 11008$	$3 \times d \times d_{ff} = 135M$

Activation Memory Cost

SwiGLU produces large intermediate tensors that must be saved for the backward pass:

Gate output: $(B, S, d_{ff})$
Up output: $(B, S, d_{ff})$
Gate activation output: $(B, S, d_{ff})$

Total activation memory per block FFN: $\sim 3 imes B imes S imes d_{ff} imes 2$ bytes (BF16).

For LLaMA-2-7B ( $d_{ff}=11008$ , $B=1, S=4096$ ): $3 imes 4096 imes 11008 imes 2 pprox 258$ MB per block, × 32 blocks = 8.3 GB of activation memory. This is why activation checkpointing is essential for training.

FFN as a Knowledge Store

Recent research (“Transformer Feed-Forward Layers Are Key-Value Memories”, Geva et al. 2021) shows that FFN layers function as learned key-value memories: the up-projection keys match input patterns, and the down-projection values store associated information. This is why:

MoE experts can specialize in different domains (each expert stores different knowledge)
Larger FFN dimensions improve factual recall
Knowledge editing techniques target FFN weights specifically

2. Residual Connections

Every sub-layer in a Transformer is wrapped with a residual (skip) connection:

\mathbf{h}_{out} = \mathbf{h}_{in} + \text{SubLayer}(\mathbf{h}_{in})

Why? In a deep network, gradients must flow through many layers. Without residuals, the gradient of layer 1 is a product of many Jacobians — likely to vanish or explode. The residual creates a gradient highway:

\frac{\partial \mathbf{h}_{out}}{\partial \mathbf{h}_{in}} = \mathbf{I} + \frac{\partial \text{SubLayer}}{\partial \mathbf{h}_{in}}

The identity matrix $\mathbf{I}$ ensures gradients always have a direct path, even if the sub-layer’s gradient is small. This is why Transformers can be 100+ layers deep.

3. Layer Normalization

LayerNorm

Normalizes across the hidden dimension for each token independently:

\text{LN}(\mathbf{x}) = \gamma \odot \frac{\mathbf{x} - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta

Where $\mu = \frac{1}{d}\sum_i x_i$ and $\sigma^2 = \frac{1}{d}\sum_i (x_i - \mu)^2$ . Learnable parameters: scale $\gamma$ and shift $\beta$ (both size $d$ ).

RMSNorm

Used by LLaMA, Qwen, Mistral — skips mean centering:

\text{RMSNorm}(\mathbf{x}) = \gamma \odot \frac{\mathbf{x}}{\sqrt{\frac{1}{d}\sum_i x_i^2 + \epsilon}}

Only $d$ learnable parameters (just the scale $\gamma$ , no shift $\beta$ ).

Why RMSNorm won: ~10-15% faster than LayerNorm (fewer operations — no mean computation, no shift) with comparable model quality. The mean centering in LayerNorm turns out to be unnecessary for most LLM architectures.

Pre-Norm vs Post-Norm

Two ways to arrange normalization in a Transformer block:

Post-Norm (original Transformer):

x → Attention → Add(x, ·) → LayerNorm → FFN → Add(·, ·) → LayerNorm → out

Pre-Norm (modern LLMs):

x → LayerNorm → Attention → Add(x, ·) → LayerNorm → FFN → Add(·, ·) → out

Pre-Norm is now standard because:

More stable training — gradients are normalized before each sub-layer
Enables training deeper models without careful learning rate tuning
Used by GPT-2/3, LLaMA, Qwen, Mistral, and virtually all modern LLMs

The tradeoff: Post-Norm can achieve slightly better final quality with careful training, but Pre-Norm is much easier to train at scale.

4. Putting It Together: One Decoder Block

class DecoderBlock(nn.Module):
    def __init__(self, d_model, num_heads, d_ff):
        super().__init__()
        self.norm1 = RMSNorm(d_model)
        self.attn = MultiHeadSelfAttention(d_model, num_heads)
        self.norm2 = RMSNorm(d_model)
        self.ffn = SwiGLU_FFN(d_model, d_ff)

    def forward(self, x, mask=None):
        # Pre-Norm + Attention + Residual
        x = x + self.attn(self.norm1(x), mask)
        # Pre-Norm + FFN + Residual
        x = x + self.ffn(self.norm2(x))
        return x

Data flow with shapes ( $d=4096, d_{ff}=11008$ ):

Input x: (B, S, 4096)
  → RMSNorm → (B, S, 4096)
  → Attention → (B, S, 4096)
  → + x (residual) → (B, S, 4096)
  → RMSNorm → (B, S, 4096)
  → SwiGLU FFN:
      → W_gate: (B, S, 4096) × (4096, 11008) → (B, S, 11008)
      → W_up:   (B, S, 4096) × (4096, 11008) → (B, S, 11008)
      → Swish(gate) ⊙ up → (B, S, 11008)
      → W_down: (B, S, 11008) × (11008, 4096) → (B, S, 4096)
  → + x (residual) → (B, S, 4096)
Output: (B, S, 4096)

Self-Check

Explain the “expand-then-compress” design of FFN and why $d_{ff} > d$
Explain how SwiGLU gating works and why it has 3 weight matrices
Explain why $d_{ff} \approx (8/3)d$ with SwiGLU instead of $4d$
Explain why residual connections prevent gradient vanishing
Explain the difference between LayerNorm and RMSNorm
Explain Pre-Norm vs Post-Norm and why Pre-Norm is standard