Inside the block: residuals, LayerNorm, FFN and embeddings

The residual stream

Write one transformer layer (pre-norm, as in GPT-2 and everything since) as two updates to a running vector per token:

The vector $\mathbf{x}$ that flows straight through, untouched except for additions, is the residual stream. Each sub-layer reads from it (after normalization) and writes back by adding. This view, popularized by interpretability research (8.11 Interpretability: what is happening inside), explains a lot: gradients reach early layers through the identity path (4.6 Initialization and normalization, 5.3 Architectures: LeNet → AlexNet → VGG → Inception → ResNet → EfficientNet → ConvNeXt); layers can specialize because each only has to add a correction; and the stream is a shared communication bus of width $d$ that attention heads and FFN neurons use to pass information to later layers — the model's working memory.

The original paper used post-norm ($\operatorname{LN}(\mathbf{x} + \text{sublayer}(\mathbf{x}))$), which normalizes the residual path itself. Deep post-norm stacks are hard to train without careful warmup; pre-norm keeps the identity path clean and is now universal, with one final LayerNorm before the output head.

Normalization: LayerNorm and RMSNorm

LayerNorm standardizes each token's vector across its $d$ features (4.6 Initialization and normalization):

Unlike BatchNorm it needs no batch statistics, so it behaves identically in training and inference and works for variable-length sequences. RMSNorm (used by Llama, Gemma, Qwen) drops the mean subtraction and the bias:

It is cheaper, equally effective in practice, and keeps the stream's direction while fixing its scale — a direct application of normalizing a vector to the unit sphere (1.1 Vectors, dot products and norms).

The feed-forward network: a key–value memory

Attention moves information between positions; the FFN processes each position independently with the same weights:

The original $\phi$ was ReLU; GPT-2 switched to GELU (4.2 Activation functions). Llama uses SwiGLU, a gated variant $\mathbf{W}_2\big(\operatorname{SiLU}(\mathbf{W}_1\mathbf{x})\odot\mathbf{W}_3\mathbf{x}\big)$ with $d_{ff}\approx\tfrac{8}{3}d$ so that its three matrices hold the same $8d^2$ parameters as the two-matrix version.

Intuition. Read the rows of $\mathbf{W}_1$ as keys: each neuron fires when the token's vector matches its pattern. Read the columns of $\mathbf{W}_2$ as values: a firing neuron adds its value vector to the residual stream. With $d_{ff}=4d$ the layer stores thousands of such pattern→update pairs — this is where factual associations like "Michael Jordan → basketball" are believed to live (3Blue1Brown's chapter 7 walks through exactly this picture).

Embeddings in and out

Token ids enter through an embedding matrix $\mathbf{E}\in\mathbb{R}^{V\times d}$; after the last layer, the unembedding $\mathbf{W}_U\in\mathbb{R}^{d\times V}$ maps each position back to vocabulary logits, and a softmax gives next-token probabilities (7.6 Decoder-only: GPT and autoregressive generation). GPT-2 ties $\mathbf{W}_U = \mathbf{E}^\top$ (saving $Vd$ parameters and acting as a regularizer); Llama keeps them separate. Dropout, when used, is applied to embeddings, attention weights and sub-layer outputs; models trained for a single epoch over trillions of tokens generally use none (4.7 Regularization in deep nets).

Counting parameters

Per layer, attention contributes $4d^2$ (7.2 Scaled dot-product and multi-head attention, fully derived) and the FFN $8d^2$, so

LayerNorm gains and biases add only $O(d)$ per layer. Two worked counts:

• GPT-2 small: $L=12$, $d=768$, $V=50{,}257$, $n_{\max}=1024$. Layers: $12\times 12\times 768^2 = 84.9$M. Token embeddings (tied): $50{,}257\times 768 = 38.6$M. Positions: $1024\times 768 = 0.8$M. Total $\approx 124.3$M — the reported 124M.
• Llama 2 7B: $L=32$, $d=4096$, $d_{ff}=11{,}008$ (SwiGLU), $V=32{,}000$, untied embeddings, RoPE (no position table). Attention $4d^2 = 67.1$M; FFN $3\,d\,d_{ff} = 135.3$M; per layer $202.4$M; $\times 32 = 6.48$B; embeddings $2\times 32{,}000\times 4096 = 0.26$B. Total $\approx 6.74$B — matching the published figure.

The llm-calculator demo does this arithmetic for any configuration, and also estimates training FLOPs and inference memory (7.9 Scaling laws and compute-optimal training, 7.10 Efficient inference: KV cache, GQA, FlashAttention, speculative decoding, quantization).

Where the compute goes

A matrix–vector product with a $p$-parameter matrix costs about $2p$ FLOPs (one multiply and one add per weight). So a token passing through one layer costs roughly $2\times 12d^2 = 24d^2$ FLOPs in the weight matrices, plus about $4nd$ for the attention scores and value mixing over a context of $n$ tokens. For $n\ll 6d$ (GPT-2: $n=1024$, $6d=4608$) the weights dominate; at 128k-token contexts the attention term takes over — which is why long context is expensive (7.12 Beyond vanilla attention: state space models, long context and hybrids). Summed over the forward and backward passes this becomes the $6ND$ rule of 7.9 Scaling laws and compute-optimal training.

A block in code

import torch.nn as nn

class Block(nn.Module):
    def __init__(self, d, n_heads, d_ff=None, dropout=0.0):
        super().__init__()
        d_ff = d_ff or 4 * d
        self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.attn = MultiHeadAttention(d, n_heads)          # from lesson 7.2
        self.ffn = nn.Sequential(nn.Linear(d, d_ff), nn.GELU(),
                                 nn.Linear(d_ff, d), nn.Dropout(dropout))

    def forward(self, x):
        x = x + self.attn(self.ln1(x))     # read, write back
        x = x + self.ffn(self.ln2(x))
        return x

Stack $L$ of these between an embedding and an unembedding and you have a GPT (7.8 Build GPT from scratch (nanoGPT walkthrough)).

Where this shows up later

• Pre-norm blocks with RMSNorm/SwiGLU are the Llama recipe — 7.6 Decoder-only: GPT and autoregressive generation; replacing the FFN with routed experts gives MoE — 7.11 Mixture of Experts.
• The parameter and FLOP counts feed scaling laws (7.9 Scaling laws and compute-optimal training) and memory planning for inference (7.10 Efficient inference: KV cache, GQA, FlashAttention, speculative decoding, quantization, 4.12 Compute, hardware and scale basics).
• The residual stream, the FFN key–value view and the unembedding ("logit lens") are the vocabulary of mechanistic interpretability — 8.11 Interpretability: what is happening inside.
• LoRA adapts exactly these weight matrices with low-rank updates — 8.2 Fine-tuning and PEFT.