Build a Transformer from Scratch: Attention, Masking & LayerNorm

June 12, 2017. Eight researchers at Google Brain and Google Research upload a 15-page paper titled Attention Is All You Need. They discard recurrence, they discard convolutions, and they build an entire sequence model out of one operation: attention. Every major language model of the 2020s — GPT-4, Claude, Gemini, Llama, Mistral — is a direct descendant of that one paper. The core 70 lines of Python haven't meaningfully changed in 8 years.

By the end of this lab you will have built a real transformer from four basic operations (matrix multiply, softmax, layernorm, addition), trained it on GPU, and used it to generate text. No nn.Transformer, no F.scaled_dot_product_attention — you'll write every piece.

Canonical references (read alongside)

• Attention Is All You Need (Vaswani et al., 2017) — the paper. Figure 1 is the whole architecture. Section 3.2.1 is scaled dot-product attention. Read those two first.
• Andrej Karpathy — Let's Build GPT from Scratch — 2-hour video that builds up to what this lab does, with live coding. The best single educational resource on transformers ever produced.
• Karpathy's nanoGPT — a 300-line PyTorch GPT-2 reimplementation. Production-ready but readable. Our lab tracks its structure closely.
• The Illustrated Transformer (Jay Alammar) — the architectural intuitions in pictures.
• The Annotated Transformer (Harvard NLP) — the original paper with runnable PyTorch inline.

What you'll build

A decoder-only (GPT-style) transformer, piece by piece:

1. Scaled dot-product attention with a causal mask — the one operation that started it all.
2. Multi-head attention — attention in parallel, concatenated.
3. The Transformer Block — attention + feed-forward + LayerNorm + residual connections.
4. Stack blocks into a tiny GPT, train it, generate from it.

Scale context

The architecture is the same at every scale — you multiply d_model, layers, and heads by constants. Understanding 60K params means understanding 175B.

---

Attention answers: "For each position in my sequence, how much should I pay attention to every other position?" It outputs, per position, a weighted mix of the values at all other positions.

The formula (it fits on one line)

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right) V$$

Three matrices per input:

• Q (queries) — "what am I looking for?" — one vector per position
• K (keys) — "what do I contain?" — one vector per position
• V (values) — "what do I contribute?" — one vector per position

QK^T gives a T × T score matrix: score[i, j] = how much position i cares about position j. Softmax turns the scores into probabilities. Multiply by V to get the weighted mix.

Why the √d_k scaling?

Without it, dot products grow with the embedding dim and softmax gets too peaky (one position gets ~100% of the weight, everything else ~0%). Gradients die. The /√d_k keeps the variance of QK^T ~ 1 regardless of embedding size. Skipping this is one of the top-5 bugs people hit when implementing attention from scratch.

The causal mask

For a decoder (generative) transformer, position i must only see positions ≤ i. Seeing the future is cheating — the model would learn to copy the next token instead of predicting it. We enforce this by setting future-position scores to −∞ before softmax, so they become 0 after softmax.

Scaled dot-product attention — one function, five steps

The core operation of every transformer. Five steps, each numbered in the comments:

1. Score each query against each key (dot product).
2. Scale by √d to keep softmax sharpness bounded as d grows.
3. Mask future positions to -∞ (causal attention — token t can only see tokens ≤ t).
4. Softmax over keys turns scores into a probability distribution per query.
5. Weighted sum of value vectors using those probabilities.

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

device = torch.device('cuda')
torch.manual_seed(42)

def scaled_dot_product_attention(Q, K, V, causal=True):
    """
    Q, K, V: (B, T, D) tensors.
    Returns: output (B, T, D) and attention weights (B, T, T).
    """
    B, T, D = Q.shape
    # 1. Score each query against each key
    scores = Q @ K.transpose(-2, -1)              # (B, T, T)
    # 2. Scale by √d_k to keep softmax sharpness bounded
    scores = scores / (D ** 0.5)
    # 3. Apply causal mask — set future positions to -inf so softmax zeroes them
    if causal:
        mask = torch.triu(torch.ones(T, T, device=Q.device), diagonal=1).bool()
        scores = scores.masked_fill(mask, float('-inf'))
    # 4. Softmax over the key dimension → probability distribution per row
    weights = F.softmax(scores, dim=-1)          # (B, T, T)
    # 5. Weighted sum of values
    output = weights @ V                          # (B, T, D)
    return output, weights

Sanity check — shapes, row sums, causality

Three invariants every attention implementation should pass:

• Output shape matches input shape (B, T, D in → B, T, D out).
• Each row of the attention weights sums to 1.0 (softmax normalization).
• With causal masking, the top-right triangle of the weights should be zero.

# Verify it runs and the shapes match expectation
B, T, D = 2, 8, 16
Q = torch.randn(B, T, D, device=device)
K = torch.randn(B, T, D, device=device)
V = torch.randn(B, T, D, device=device)
attention_output, attention_weights = scaled_dot_product_attention(Q, K, V, causal=True)
print(f'output shape:     {tuple(attention_output.shape)}')
print(f'weights shape:    {tuple(attention_weights.shape)}')
print(f'row sum (should be 1.0 everywhere): {attention_weights.sum(dim=-1).mean().item():.4f}')

# Visual sanity: show the causal mask in action — upper triangle should be exactly 0.
print('Attention weights for the first example (rows=query, cols=key):')
print(attention_weights[0].cpu().numpy().round(2))
print()
print('Notice the upper triangle is all zeros — that is the causal mask at work.')