Transformer from scratch in PyTorch, block by block

Attention is the famous part. It is also less than half of a transformer, and most of the design decisions that make a deep stack trainable live in the boring parts around it.

This is a complete decoder-only transformer — the GPT architecture — in about 120 lines, with the reasoning for each block. It trains.

The shape of the thing

tokens ──► embedding + positions
             │
             ├─► [ Block 1 ]  ─┐
             ├─► [ Block 2 ]   │  N identical blocks
             ├─►    ...        │
             └─► [ Block N ]  ─┘
             │
           final LayerNorm
             │
           linear ──► logits over vocabulary


each Block:
    x = x + attention(norm(x))     <- residual
    x = x + mlp(norm(x))           <- residual

Two sub-layers, each wrapped in a residual connection with normalisation applied before the sub-layer. That pattern is the whole architecture, repeated.

Multi-head causal self-attention

import math
import torch
import torch.nn as nn
import torch.nn.functional as F


class CausalSelfAttention(nn.Module):
    def __init__(self, d_model, n_heads, dropout=0.1):
        super().__init__()
        assert d_model % n_heads == 0
        self.n_heads = n_heads
        self.d_head = d_model // n_heads

        # One projection producing Q, K and V together — a single fused matmul
        # is meaningfully faster than three separate ones.
        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
        self.proj = nn.Linear(d_model, d_model, bias=False)
        self.drop = nn.Dropout(dropout)
        self.dropout_p = dropout

    def forward(self, x):
        B, T, C = x.shape                       # batch, time, channels

        q, k, v = self.qkv(x).split(C, dim=2)
        # (B, T, C) -> (B, n_heads, T, d_head)
        q = q.view(B, T, self.n_heads, self.d_head).transpose(1, 2)
        k = k.view(B, T, self.n_heads, self.d_head).transpose(1, 2)
        v = v.view(B, T, self.n_heads, self.d_head).transpose(1, 2)

        # is_causal builds the mask internally and uses a fused kernel
        # (FlashAttention where available) — same maths, far less memory.
        out = F.scaled_dot_product_attention(
            q, k, v,
            dropout_p=self.dropout_p if self.training else 0.0,
            is_causal=True,
        )

        out = out.transpose(1, 2).contiguous().view(B, T, C)
        return self.drop(self.proj(out))

Three decisions worth naming.

Fused QKV. Three separate nn.Linear calls produce the same numbers, but one (d_model, 3*d_model) matmul is a single kernel launch on a bigger matrix — measurably faster.

scaled_dot_product_attention instead of hand-rolled. It dispatches to FlashAttention when the shapes allow. Identical output, dramatically lower memory, because the (T, T) matrix is never materialised.

is_causal=True rather than a manual mask. Same result, no mask tensor to allocate or move to the right device.

The MLP, and why it is 4× wide

class MLP(nn.Module):
    def __init__(self, d_model, expansion=4, dropout=0.1):
        super().__init__()
        hidden = expansion * d_model
        self.fc1 = nn.Linear(d_model, hidden)
        self.fc2 = nn.Linear(hidden, d_model)
        self.drop = nn.Dropout(dropout)

    def forward(self, x):
        return self.drop(self.fc2(F.gelu(self.fc1(x))))

Attention moves information between positions. It applies no per-position nonlinearity worth speaking of — the output is a weighted average of values. The MLP is where each position does actual computation on what it gathered.

Expanding to 4 × d_model and back gives the nonlinearity room to work in. The expansion factor of 4 is empirical convention rather than theory, and this block is where roughly two thirds of the parameters live — at d_model = 768, the MLP holds 4.7M parameters against attention's 2.4M.

GELU rather than ReLU: smooth, non-zero gradient for small negative inputs, and consistently slightly better in transformers.

Pre-norm vs post-norm: the change that made deep stacks trainable

class Block(nn.Module):
    def __init__(self, d_model, n_heads, dropout=0.1):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = CausalSelfAttention(d_model, n_heads, dropout)
        self.ln2 = nn.LayerNorm(d_model)
        self.mlp = MLP(d_model, dropout=dropout)

    def forward(self, x):
        x = x + self.attn(self.ln1(x))     # PRE-norm: normalise, then sub-layer
        x = x + self.mlp(self.ln2(x))
        return x

The original 2017 paper put the norm after the residual add:

post-norm (original):  x = LayerNorm(x + attn(x))
pre-norm  (modern):    x = x + attn(LayerNorm(x))

The difference looks cosmetic and is not. In pre-norm, the residual path from input to output is completely unobstructed — no normalisation sits on it. Gradients flow from the loss to layer 1 without passing through 48 LayerNorms.

post-norm:  x --> [+] --> LN --> [+] --> LN --> ...   gradient crosses every LN
                   ^             ^
                  attn          mlp

pre-norm:   x --> [+] --------> [+] --------> ...     clean residual highway
                   ^             ^
                LN,attn       LN,mlp

Post-norm transformers deeper than about 12 layers need careful learning-rate warmup or they diverge. Pre-norm trains stably at 100+ layers. Essentially every model since GPT-2 uses pre-norm, which is why the final LayerNorm before the output head is necessary — without it the residual stream leaves the last block unnormalised.

The residual stream

Worth stating explicitly, because it reframes the architecture usefully.

x0 = embeddings
x1 = x0 + attn_1(...) + mlp_1(...)
x2 = x1 + attn_2(...) + mlp_2(...)
...
xN = x0 + sum of every sub-layer's contribution

Every block adds to a shared stream rather than transforming it. The residual stream is a communication channel that all layers read from and write to. Early layers write positional and lexical information; later layers read it and write semantic information.

Two practical consequences follow. Gradients reach early layers directly through the addition chain, which is why depth works at all. And layers can be partially independent — this is why techniques like layer dropout and early exit work without the model collapsing.

Positional information

Attention is permutation-invariant: shuffle the tokens and the score matrix permutes identically. Order must be injected.

class Transformer(nn.Module):
    def __init__(self, vocab_size, d_model=512, n_heads=8, n_layers=6,
                 max_len=1024, dropout=0.1):
        super().__init__()
        self.tok_emb = nn.Embedding(vocab_size, d_model)
        self.pos_emb = nn.Embedding(max_len, d_model)     # learned positions
        self.drop = nn.Dropout(dropout)

        self.blocks = nn.ModuleList(
            Block(d_model, n_heads, dropout) for _ in range(n_layers)
        )
        self.ln_f = nn.LayerNorm(d_model)
        self.head = nn.Linear(d_model, vocab_size, bias=False)

        # Weight tying: input embedding and output projection share weights.
        # Saves vocab_size * d_model parameters and usually improves quality.
        self.head.weight = self.tok_emb.weight

        self.apply(self._init)
        # Scale residual-path output projections by 1/sqrt(2*n_layers) so the
        # residual stream's variance does not grow with depth.
        for name, p in self.named_parameters():
            if name.endswith("proj.weight") or name.endswith("fc2.weight"):
                nn.init.normal_(p, mean=0.0,
                                std=0.02 / math.sqrt(2 * n_layers))

    def _init(self, module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        pos = torch.arange(T, device=idx.device)

        x = self.drop(self.tok_emb(idx) + self.pos_emb(pos))
        for block in self.blocks:
            x = block(x)
        x = self.ln_f(x)
        logits = self.head(x)

        loss = None
        if targets is not None:
            loss = F.cross_entropy(
                logits.view(-1, logits.size(-1)), targets.view(-1)
            )
        return logits, loss

Three details in that constructor do real work.

Weight tying (self.head.weight = self.tok_emb.weight). Input and output embeddings represent the same relationship in opposite directions. Sharing them saves vocab_size × d_model parameters — 38M for a 50k vocabulary at 768 dims — and typically improves quality slightly.

Scaled residual initialisation. Each block adds to the residual stream, so variance accumulates with depth. Scaling the output projections by 1/sqrt(2·n_layers) — two residual writes per block — keeps the stream's scale stable from layer 1 to layer 48. Without this, deep models start with an exploding activation scale and need a long warmup to recover.

Learned position embeddings are the simplest option. They do not extrapolate beyond max_len at all. Modern models use RoPERotary Position Embedding: rotates query and key vectors by an angle proportional to position, so attention scores depend on relative distance instead, which encodes relative position directly into the attention computation and extrapolates far better.

Generation

@torch.no_grad()
def generate(model, idx, max_new_tokens, temperature=1.0, top_k=None):
    model.eval()
    for _ in range(max_new_tokens):
        idx_cond = idx[:, -model.pos_emb.num_embeddings:]   # crop to context
        logits, _ = model(idx_cond)
        logits = logits[:, -1, :] / temperature             # last position only

        if top_k is not None:
            v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
            logits[logits < v[:, [-1]]] = -float("inf")

        probs = F.softmax(logits, dim=-1)
        idx = torch.cat([idx, torch.multinomial(probs, 1)], dim=1)
    return idx

Note the inefficiency: every step re-runs the whole sequence. Real inference uses a KV cache — keys and values for previous positions do not change, so you compute them once and append. That turns generation from O(T^2) per token into O(T), and it is the single most important inference optimisation.

Training loop

model = Transformer(vocab_size=50257, d_model=512, n_heads=8, n_layers=6)
opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.1,
                        betas=(0.9, 0.95))

# Warmup then cosine decay — transformers are unstable at full LR from step 0
warmup, total = 2000, 100_000
sched = torch.optim.lr_scheduler.LambdaLR(
    opt,
    lambda s: s / warmup if s < warmup
    else 0.5 * (1 + math.cos(math.pi * (s - warmup) / (total - warmup))),
)

for step, (x, y) in enumerate(loader):
    _, loss = model(x, y)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)   # essential
    opt.step()
    sched.step()
    opt.zero_grad(set_to_none=True)

Two non-negotiables in there. Gradient clipping at 1.0 — transformers produce occasional huge gradients, and one unclipped spike destroys a run that took days. Learning-rate warmup — Adam's second-moment estimate is unreliable in the first few hundred steps, and a full learning rate on top of that unreliability diverges.

betas=(0.9, 0.95) rather than the PyTorch default 0.999 is the standard transformer setting: a shorter second-moment window adapts faster to the changing loss landscape.

Parameter budget

For d_model = 768, n_layers = 12, vocab = 50257:

Two things stand out. The MLP is twice the attention parameters per layer, which surprises people who think of transformers as "attention models". And at this scale the embedding table alone is nearly a third of the model — which is why weight tying matters and why vocabulary size is an architectural decision, not a preprocessing detail.

Frequently asked questions

Why pre-norm instead of post-norm?

Pre-norm leaves the residual path completely clean, so gradients reach early layers without crossing every normalisation layer. Post-norm models deeper than about 12 layers need careful warmup to train at all; pre-norm is stable at 100+ layers. Every modern model uses pre-norm.

Why is the MLP four times wider than the model dimension?

Attention moves information between positions but applies almost no per-position nonlinearity. The MLP is where each position computes on what it gathered, and the 4× expansion gives that computation room. The factor is empirical convention, and this block holds about two thirds of the parameters.

What is weight tying and should I use it?

Sharing the input embedding matrix with the output projection. It saves vocab_size × d_model parameters and usually improves quality slightly, since both matrices encode the same token-to-vector relationship. Standard practice for language models.

Why do transformers need learning-rate warmup?

Adam's second-moment estimate is unreliable in the first hundreds of steps, and applying a full learning rate on top of unreliable statistics diverges. Warm up linearly over 1–5% of training, then decay. Pair it with gradient clipping at 1.0 — one unclipped gradient spike can destroy a multi-day run.

Should I use learned positional embeddings?

They are the simplest thing that works, and they do not extrapolate past the trained maximum length at all. If long-context behaviour matters, use RoPE, which encodes relative position inside the attention computation and generalises much better beyond the training length.

Next reading on Sythra Articles

• The attention mechanism, worked out by hand
• CNNs from scratch in NumPy
• Embeddings explained mathematically