Transformers from Scratch

A transformer does not see text. It sees integers, which it then turns into vectors. The job of the tokenizer is to do the text-to-integers part, and the choice of tokenizer is the first place a serious decision gets made.

The naive options are bad. Character-level tokenization (one integer per Unicode code point) means a 100-character sentence becomes 100 tokens and the model spends most of its compute mixing redundancy. Word-level tokenization gives a vocabulary that explodes with every new domain and chokes on misspellings. Both have been used. Neither is what GPT-3 and Llama and Claude actually do.

The dominant choice since GPT-2 is byte-pair encoding (BPE). Start with a vocabulary of single bytes (256 entries). Look at your training corpus, find the most frequent adjacent token pair, merge it into one new token. Repeat ten thousand times. You end up with ≈50k tokens, the common ones being whole words ( the, and), the rare ones being subwords ( chronological → chrono, logical) or single bytes for things the corpus never saw.

This is the part that surprises people: the BPE vocabulary includes the leading space. cat and cat are different tokens. That is why prompts beginning with cat versus cat produce different outputs in some models. The tokenizer is part of the model.

tok = Tokenizer(BPE(unk_token="[UNK]"))
tok.pre_tokenizer = Whitespace()
tok.train(["data/tiny_shakespeare.txt"], BpeTrainer(vocab_size=10000))
ids = tok.encode("To be, or not to be").ids

def get_pair_counts(seq):
    return Counter(zip(seq, seq[1:]))
def merge(seq, pair, new_token):
    out, i = [], 0
    while i < len(seq):
        if i < len(seq)-1 and (seq[i], seq[i+1]) == pair:
            out.append(new_token); i += 2
        else:
            out.append(seq[i]); i += 1
    return out
def train_bpe(text, target_vocab):
    seq = list(text.encode("utf-8"))     # start from bytes 0-255
    merges, next_id = {}, 256
    while next_id < target_vocab:
        pairs = get_pair_counts(seq)
        if not pairs: break
        best = pairs.most_common(1)[0][0]    # most frequent adjacent pair
        seq = merge(seq, best, next_id)
        merges[best] = next_id; next_id += 1
    return seq, merges

Run this on Tiny Shakespeare with target_vocab=1000 and you get about 2.5× compression. With target_vocab=10000 you get about 4×. The library version does some additional things, mainly handling the regex-based pre-tokenization that GPT-2 uses to keep token boundaries at word edges. Look up regex.findall(GPT2_PATTERN,...) if you want the production version.

Once you have token IDs, you turn each one into a vector. This is what nn.Embedding(vocab_size, d_model) does: a lookup into a learned table of shape (vocab_size, d_model). For each integer id, you fetch row id.

The embedding table is a learned parameter. At initialization it is small random gaussians. After training, similar tokens end up at similar positions in the d_model-dimensional space, in a way that is qualitatively interpretable (the famous king − man + woman ≈ queen experiment was on word2vec embeddings, but transformer embeddings learn the same kind of structure).

tok_emb = nn.Embedding(vocab_size, d_model)
# idx: (B, T) int64  ->  tok_emb(idx): (B, T, d_model), one learned row per id

class Embedding:
    def __init__(self, vocab_size, d_model):
        # init scales variance with embedding dim
        self.weight = np.random.randn(vocab_size, d_model) * (1.0 / np.sqrt(d_model))
    def __call__(self, ids):
        # ids: (batch, seq_len) integer array
        return self.weight[ids]   # fancy indexing -> (batch, seq_len, d_model)

There is nothing more to embeddings than that. The interesting question is whether the unembedding (the final projection from d_model to vocab_size that produces logits) should share weights with the embedding. In GPT-2, it does. In some other models, it doesn't. Both work. Tying saves parameters.

This is the part where you have to be careful with shapes. Self-attention is the function that, for each position in a sequence, computes a weighted sum of all positions in the same sequence, where the weights come from how similar each position's "query" is to every position's "key".

Three vectors per token: $Q$ (query), $K$ (key), $V$ (value). Each is a linear projection of the embedding. The attention score from position $i$ to position $j$ is:

The score gets passed through softmax along the $j$ axis to make it a probability distribution. Then the output at position $i$ is the weighted sum of the $V_j$'s using those probabilities.

The whole thing in matrix form, for a batch of one sequence:

Where $Q, K, V$ are matrices of shape (seq_len, d_k). The result is (seq_len, d_k).

Three things are worth labeling explicitly because every implementation gets one of them wrong at first.

Note 1: The $\sqrt{d_k}$ scaling matters. Without it, for large $d_k$, the dot products become large in magnitude, the softmax saturates, and gradients vanish. Vaswani et al. derive this from the variance of a dot product of two unit gaussians; it shows up in every implementation because if you don't scale, your model trains slower and uses less of the head's capacity.

Note 2: The softmax is along the key axis (the position you are attending to). This means each query gets a probability distribution over keys. The other direction (softmax over queries) is a different operation and is not what attention does. In matrix code: F.softmax(scores, dim=-1) when scores is (..., seq_len_q, seq_len_k).

Note 3: $Q$, $K$, $V$ all come from the same input in self-attention. They are different linear projections of the same token embeddings. In cross-attention (encoder-decoder), $Q$ comes from one sequence and $K, V$ come from another. We only care about self-attention here.

# q, k, v: (B, n_heads, T, d_k)  -- you still do the head reshape yourself
out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
# out: (B, n_heads, T, d_k) = softmax(QK^T / sqrt(d_k) + causal_mask) @ V, fused

scores = q @ k.transpose(-2, -1) / math.sqrt(self.d_k)   # QK^T / sqrt(d_k)
scores = scores.masked_fill(self.mask[:, :, :T, :T], float('-inf'))  # causal
attn = F.softmax(scores, dim=-1)                         # over the key axis
out = attn @ v                                           # weighted sum of values
# q, k, v are already (B, n_heads, T, d_k); self.mask is
# torch.triu(torch.ones(L, L), diagonal=1).bool() registered in __init__.

Read those two implementations side by side. They are the same function. The library version uses nn.Linear (which is x @ W.T + b), the from-scratch version uses raw x @ W. The transpose on $K$ to compute scores: k.transpose(-2, -1) versus k.transpose(0, 2, 1). Same thing, different APIs. If you understand this side-by-side, you understand self-attention.

A single attention head can only learn one type of relationship at a time. Maybe it learns "the next-noun pattern". Maybe it learns "the matching-bracket pattern". If you want multiple patterns running in parallel, you give the model multiple heads.

Multi-head attention runs $h$ attention heads on the same input, in parallel, each with its own $W_q, W_k, W_v$ projections, each producing a separate $(seq\_len, d_k)$ output. The outputs are concatenated along the d_k axis (giving $(seq\_len, h \cdot d_k)$), then projected back down to $d_{model}$ by a final linear layer $W_o$.

By convention $h \cdot d_k = d_{model}$, so each head sees a $d_k = d_{model}/h$-dimensional subspace. For GPT-2 small, $d_{model}=768$ and $h=12$, so $d_k = 64$.

The clever implementation does not allocate $h$ separate nn.Linear modules. It allocates one big nn.Linear(d_model, 3 * d_model) to produce all of $Q, K, V$ across all heads at once, then reshapes the result into $(batch, h, seq\_len, d_k)$ and lets PyTorch's batched matmul handle the rest.

Library path (the production idiom):

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model: int, n_heads: int):
        super().__init__()
        assert d_model % n_heads == 0
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads
        self.qkv_proj = nn.Linear(d_model, 3 * d_model, bias=False)
        self.out_proj = nn.Linear(d_model, d_model, bias=False)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, C = x.shape   # batch, seq_len, d_model
        qkv = self.qkv_proj(x)   # (B, T, 3*C)
        q, k, v = qkv.split(self.d_model, dim=-1)   # each (B, T, C)
        # Reshape to (B, n_heads, T, d_k)
        q = q.view(B, T, self.n_heads, self.d_k).transpose(1, 2)
        k = k.view(B, T, self.n_heads, self.d_k).transpose(1, 2)
        v = v.view(B, T, self.n_heads, self.d_k).transpose(1, 2)
        # Batched matmul over heads
        scores = q @ k.transpose(-2, -1) / (self.d_k ** 0.5)   # (B, n_heads, T, T)
        attn = F.softmax(scores, dim=-1)
        out = attn @ v   # (B, n_heads, T, d_k)
        out = out.transpose(1, 2).contiguous().view(B, T, C)   # (B, T, C)
        return self.out_proj(out)

The from-scratch version is the same idea with explicit NumPy reshapes. (See the lab at the bottom of this chapter.)

If you want an autoregressive model (one that predicts the next token given previous ones), you cannot let position $i$ attend to position $j > i$. Otherwise the model just looks at the answer and is useless at inference time.

You enforce this with a mask. Before the softmax, you set $\text{score}_{ij} = -\infty$ for all $j > i$. After softmax, those positions become 0. The query at position $i$ can only mix values from positions $0$ through $i$.

In PyTorch:

# T = seq_len
mask = torch.triu(torch.ones(T, T), diagonal=1).bool()
# True above the diagonal. Mask those out.
scores = scores.masked_fill(mask, float('-inf'))
attn = F.softmax(scores, dim=-1)

In NumPy:

mask = np.triu(np.ones((T, T), dtype=bool), k=1)
scores = np.where(mask, -np.inf, scores)

Without something extra, attention is permutation-invariant: shuffling the tokens in the input shuffles the outputs the same way, but the function the model computes is the same. That is a disaster for sequence modeling because word order matters.

There are three families of fixes.

Absolute positional embeddings (GPT-2, BERT): add a learned vector to each token's embedding based on its position. pos_emb = nn.Embedding(max_seq_len, d_model), then x = token_emb + pos_emb. Simple. Limits the model to max_seq_len because positions beyond that have no learned embedding.

Sinusoidal positional encoding (the original transformer): add a sin/cos-based encoding instead of a learned one. Generalizes to longer sequences (in theory). In practice, models trained with sinusoidal positions also struggle to extrapolate, just less catastrophically.

Rotary positional embeddings (RoPE — Llama, Mistral, GPT-NeoX, most modern LLMs): instead of adding a positional vector, rotate the query and key vectors in 2D subspaces by an angle proportional to their position. Crucially, the dot product $Q_i \cdot K_j$ depends only on the relative position $i - j$. This lets the model handle relative positions naturally.

The RoPE math is the part people skip. Here it is.

For the $k$-th pair of dimensions $(2k, 2k+1)$ in a Q or K vector at position $p$, the rotation is:

where $\theta_k = 10000^{-2k/d_{model}}$ varies across pairs of dimensions ($k = 0, 1, \dots, d_{model}/2 - 1$), giving each pair a different rotation frequency. Pairs with low $k$ rotate fast, pairs with high $k$ rotate slowly. The slow-rotating pairs encode long-range positional information; the fast-rotating pairs encode short-range.

After rotation, the dot product $Q_i \cdot K_j$ ends up being a function of the dimension-pair contents and of $\cos((i-j)\theta_k)$ and $\sin((i-j)\theta_k)$ for each pair $k$. Hence "relative".

If the rotation-matrix algebra is the first place in this chapter you stall, that is expected — it goes a notch past the linear algebra in another chapter. You can use RoPE correctly by treating the code below as the contract (rotate $Q$ and $K$ by a position-dependent angle, leave $V$ alone) and come back to the derivation later; nothing downstream depends on you re-deriving it.

Library path (HuggingFace's RoPE impl, simplified):

def rotary_embed(x: torch.Tensor, pos: torch.Tensor, base: float = 10000) -> torch.Tensor:
    """x: (..., seq_len, d). pos: (seq_len,) integer positions."""
    d = x.shape[-1]
    half = d // 2
    freqs = base ** (-torch.arange(0, half, device=x.device).float() / half)
    angles = pos[:, None].float() * freqs[None, :]   # (seq_len, half)
    cos, sin = angles.cos(), angles.sin()
    x_even, x_odd = x[..., 0::2], x[..., 1::2]
    out_even = x_even * cos - x_odd * sin
    out_odd = x_even * sin + x_odd * cos
    out = torch.stack([out_even, out_odd], dim=-1).flatten(-2)
    return out

You apply this to $Q$ and $K$ before computing scores. Not to $V$.

After every attention layer there's a per-position feed-forward network (FFN) — the same MLP applied independently to each position's vector. Two linear layers with a non-linearity between them.

class FeedForward(nn.Module):
    def __init__(self, d_model: int, d_ff: int = None):
        super().__init__()
        d_ff = d_ff or 4 * d_model   # convention: FFN is 4x wider
        self.fc1 = nn.Linear(d_model, d_ff)
        self.fc2 = nn.Linear(d_ff, d_model)
    
    def forward(self, x):
        return self.fc2(F.gelu(self.fc1(x)))

That's the whole thing. The 4× expansion is convention from the original transformer; modern models sometimes use different ratios (Llama uses 2.67× with SwiGLU instead of GELU). The choice of activation matters less than people think for accuracy, more than people think for inference speed.

Why is the FFN here at all? Attention is a linear function of its values (the softmax weights are non-linear, but each output is a linear combination of values). Without an FFN, the model is too limited. The FFN provides the non-linearity. There is a mech-interp literature on what FFNs actually do in trained models (they appear to encode key-value memories: each column of the first linear layer is a "key" direction that some input pattern activates, and the matching row of the second layer is the "value" it then writes back), but for now treat the FFN as the "compute" of the architecture and attention as the "communication".

LayerNorm normalizes each token's vector to have zero mean and unit variance, then scales and shifts by learned per-dimension parameters $\gamma, \beta$. This is per token, not across the batch (that's BatchNorm).

class LayerNorm(nn.Module):
    def __init__(self, d_model: int, eps: float = 1e-5):
        super().__init__()
        self.gamma = nn.Parameter(torch.ones(d_model))
        self.beta = nn.Parameter(torch.zeros(d_model))
        self.eps = eps
    
    def forward(self, x):
        mean = x.mean(dim=-1, keepdim=True)
        var = x.var(dim=-1, keepdim=True, unbiased=False)
        return self.gamma * (x - mean) / torch.sqrt(var + self.eps) + self.beta

The interesting question is where in the transformer block you put the LayerNorm.

The original "Attention Is All You Need" paper put it after each sub-layer (attention, FFN): x = LayerNorm(x + Attention(x)). This is "post-norm". It is hard to train. The gradients explode at deep stacks unless you use careful learning-rate warmup.

GPT-2 and every modern model put LayerNorm before each sub-layer: x = x + Attention(LayerNorm(x)). This is "pre-norm". It trains easily without warmup. It is the dominant choice.

Llama and later models use RMSNorm instead of LayerNorm — same idea but without subtracting the mean (only normalizing the variance). Faster, no meaningful accuracy difference.

Every sub-layer is wrapped in a residual connection: x = x + Sublayer(x). This started as an "easier to train deep networks" trick (He et al. 2016, ResNet). Anthropic's mech-interp work reframed it as something much more useful conceptually.

Look at the architecture: each token's vector flows from the embedding through 12 (or 24, or 96) blocks, each of which adds to the vector. Attention reads from the residual stream and writes back. FFN reads and writes. The residual stream is the central "bus" through which all information flows; every component reads-then-writes-to it.

This is the "residual stream" view from 22-anthropic-recent/2021-framework-index §residual-stream. It is the most useful way to think about why mech-interp findings like "the IOI circuit" work: specific attention heads read specific features from the residual stream and write specific corrections back, and you can decompose the model's behavior into those reads and writes.

When you're writing the from-scratch code, the residual connection is just +=:

class Block(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.ln1 = LayerNorm(d_model)
        self.attn = MultiHeadAttention(d_model, n_heads)
        self.ln2 = LayerNorm(d_model)
        self.ffn = FeedForward(d_model)
    
    def forward(self, x):
        x = x + self.attn(self.ln1(x))   # read from stream, write back
        x = x + self.ffn(self.ln2(x))    # read from stream, write back
        return x

A GPT-style transformer is:

1. Token embedding + positional encoding → input vectors
2. $N$ identical transformer blocks (attention + FFN with residuals + pre-norm)
3. Final LayerNorm
4. Linear projection to vocab_size → logits

That's it. ~50 lines of PyTorch.

class GPT(nn.Module):
    def __init__(self, vocab_size, d_model=768, n_heads=12, n_layers=12, max_seq_len=1024):
        super().__init__()
        self.tok_emb = nn.Embedding(vocab_size, d_model)
        self.pos_emb = nn.Embedding(max_seq_len, d_model)
        self.blocks = nn.ModuleList([Block(d_model, n_heads) for _ in range(n_layers)])
        self.ln_f = LayerNorm(d_model)
        self.head = nn.Linear(d_model, vocab_size, bias=False)
        # Weight tying: share embedding and output head
        self.head.weight = self.tok_emb.weight
    
    def forward(self, idx):
        B, T = idx.shape
        pos = torch.arange(T, device=idx.device)
        x = self.tok_emb(idx) + self.pos_emb(pos)
        for block in self.blocks:
            x = block(x)
        x = self.ln_f(x)
        return self.head(x)   # (B, T, vocab_size) — logits

That is GPT-2. Llama-style models swap LayerNorm for RMSNorm, sinusoidal/learned position for RoPE, and the FFN's GELU for SwiGLU. The skeleton is identical.

To train a language model, you give it a sequence and ask it to predict the next token at every position. Inputs: tokens 0 through T-1. Targets: tokens 1 through T. Loss: cross-entropy between the logits at each position and the target token at that position.

def train_step(model, optimizer, idx, targets):
    logits = model(idx)   # (B, T, vocab_size)
    loss = F.cross_entropy(
        logits.view(-1, logits.size(-1)),
        targets.view(-1)
    )
    optimizer.zero_grad()
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    optimizer.step()
    return loss.item()

Three details that matter at small scale:

• AdamW is the optimizer everyone uses. Don't use Adam, use AdamW. The decoupled weight decay matters at LLM scale.
• Gradient clipping at norm 1.0 prevents the occasional huge gradient from destabilizing training.
• Learning-rate warmup over the first ~2000 steps from 0 to peak LR, then cosine decay to ~10% of peak. nanoGPT's defaults are good.

Train on Tiny Shakespeare for ~5000 steps with a GPT-2 small (12 layers, 12 heads, d_model 768) and you reach validation loss ≈1.5, generating shakespeare-flavoured nonsense. Train for ~50000 steps and you reach ≈1.1, generating actually-passable Shakespeare imitation.

At inference, you have a sequence of tokens, you run the model, you get logits at the last position, you turn those into a token, you append, you repeat.

The boring version is greedy: take argmax(logits). Produces deterministic, often repetitive, often boring output.

The standard version is temperature + top-k + top-p:

@torch.no_grad()
def generate(model, idx, max_new_tokens, temperature=1.0, top_k=None, top_p=None):
    for _ in range(max_new_tokens):
        # Crop to max_seq_len
        idx_cond = idx[:, -model.max_seq_len:]
        logits = model(idx_cond)[:, -1, :]   # (B, vocab_size) — last position only
        logits = logits / temperature
        # top-k: keep only the k highest-logit tokens
        if top_k is not None:
            v, _ = torch.topk(logits, top_k)
            logits[logits < v[:, [-1]]] = float('-inf')
        # top-p (nucleus): keep tokens whose cumulative prob ≥ p
        if top_p is not None:
            sorted_logits, sorted_idx = torch.sort(logits, descending=True)
            cumprob = F.softmax(sorted_logits, dim=-1).cumsum(dim=-1)
            sorted_idx_remove = cumprob > top_p
            sorted_idx_remove[..., 1:] = sorted_idx_remove[..., :-1].clone()
            sorted_idx_remove[..., 0] = False
            indices_to_remove = sorted_idx_remove.scatter(1, sorted_idx, sorted_idx_remove)
            logits[indices_to_remove] = float('-inf')
        probs = F.softmax(logits, dim=-1)
        next_token = torch.multinomial(probs, num_samples=1)
        idx = torch.cat([idx, next_token], dim=1)
    return idx

Defaults that work well for general text: temperature=0.8, top_p=0.9. For deterministic completions: temperature=0.0 (which becomes greedy).

Karpathy's nanoGPT (12-karpathy-code/nanoGPT-master-model + nanoGPT-master-train) is the cleanest from-scratch GPT implementation that actually trains to competitive perplexity on real corpora. ~300 lines for the model, ~300 for the training loop. Worth reading top to bottom.

What he does differently from this chapter:

• Combined QKV projection (nn.Linear(d_model, 3*d_model)) instead of three separate ones — fewer kernel launches, same math
• Uses F.scaled_dot_product_attention, PyTorch's attention API that dispatches to a FlashAttention kernel when the inputs are eligible — same math, ~3× faster, much less memory. The speedup comes from fusing the whole attention computation into one GPU kernel so the big (T, T) score matrix is computed in fast on-chip memory and never written out to slow GPU DRAM; the naive version pays for a full write-then-read of that matrix between the matmul and the softmax (covered in another chapter — Efficient Inference)
• Implements the GELU non-linearity as PyTorch's nn.GELU(approximate='tanh') to match GPT-2's exact form
• Weight-ties the embedding and unembedding layers
• Uses nn.LayerNorm with the bias term following GPT-2

Diff the Block class in nanoGPT against the Block class in the lab at the end of this chapter. They are the same module modulo those four optimizations. Internalizing this diff is the goal of the chapter.