Transformer From Scratch in NumPy: A Walkthrough You Can Run

If you can write a transformer forward pass in NumPy from memory, most "explain attention" interview questions stop being scary, because you can answer them by pointing at a line of code. This post gives you a complete decoder-only transformer of 4,320 parameters in about 130 lines. I ran it on 2026-10-09 and pasted its real output below, so you can check your own run against mine.

Two honest limits up front. The weights are random and seeded, so the model predicts nothing useful; the goal is the mechanism. And I'm citing Attention Is All You Need (Vaswani et al., 2017) only for details I read in the paper text on 2026-10-09: the model dimensions, sinusoidal positions, the residual connection with layer normalisation, the shared embedding matrix scaled by the square root of the model dimension, and masking future positions.

If you want the memory and serving side of this topic afterwards, read LLM interview questions on attention and the KV cache. For the broader question bank see AI engineer interview questions on RAG, agents and evals.

What will you build, and what does the paper say?

A decoder-only stack: token embedding, position encoding, then N layers of masked self-attention plus a feed-forward block, then a projection back to the vocabulary. The original paper's base model uses N = 6 layers, d_model = 512, d_ff = 2048 and h = 8 heads with d_k = 64 per head. The toy below shrinks all of that:

The paper's architecture is an encoder-decoder built for translation. Most LLMs you'll be asked about are decoder-only, so the toy keeps only the masked self-attention part. It follows the paper on: sinusoidal positions, a residual connection around each sub-layer followed by layer normalisation, ReLU in the feed-forward network, and one embedding matrix shared between input and output, multiplied by the square root of d_model on the way in. Dropout is omitted because nothing is being trained, and layer norm has no learned scale or shift.

The full script

Save this as tiny_transformer.py and run it with Python 3 and NumPy (I used NumPy 2.5.3). Every number in the next section comes from this exact file.

import numpy as np

rng = np.random.default_rng(0)
np.set_printoptions(precision=3, suppress=True, linewidth=120)

# ---- config
vocab = ["<pad>", "the", "cat", "sat", "on", "mat", "dog", "ran"]
V = len(vocab)
d_model, n_heads, d_ff, n_layers = 16, 2, 32, 2
d_k = d_model // n_heads

def softmax(x, axis=-1):
    x = x - x.max(axis=axis, keepdims=True)
    e = np.exp(x)
    return e / e.sum(axis=axis, keepdims=True)

def layer_norm(x, eps=1e-5):
    mu = x.mean(-1, keepdims=True)
    var = x.var(-1, keepdims=True)
    return (x - mu) / np.sqrt(var + eps)

def sinusoidal(T, d):
    pos = np.arange(T)[:, None]
    i = np.arange(d // 2)[None, :]
    ang = pos / (10000 ** (2 * i / d))
    pe = np.zeros((T, d))
    pe[:, 0::2] = np.sin(ang)
    pe[:, 1::2] = np.cos(ang)
    return pe

def init(shape, scale):
    return rng.normal(0, scale, shape)

emb = init((V, d_model), 0.5)
layers = []
for _ in range(n_layers):
    layers.append(dict(
        Wq=init((d_model, d_model), d_model ** -0.5),
        Wk=init((d_model, d_model), d_model ** -0.5),
        Wv=init((d_model, d_model), d_model ** -0.5),
        Wo=init((d_model, d_model), d_model ** -0.5),
        W1=init((d_model, d_ff), d_model ** -0.5),
        b1=np.zeros(d_ff),
        W2=init((d_ff, d_model), d_ff ** -0.5),
        b2=np.zeros(d_model),
    ))

n_params = emb.size + sum(v.size for L in layers for v in L.values())
print("params (embeddings tied to output, no LayerNorm params):", n_params)

def attention(x, L, return_w=False):
    T = x.shape[0]
    Q, K, Vv = x @ L["Wq"], x @ L["Wk"], x @ L["Wv"]
    # split heads: (T, d_model) -> (heads, T, d_k)
    sp = lambda a: a.reshape(T, n_heads, d_k).transpose(1, 0, 2)
    Qh, Kh, Vh = sp(Q), sp(K), sp(Vv)
    scores = Qh @ Kh.transpose(0, 2, 1) / np.sqrt(d_k)      # (heads, T, T)
    mask = np.triu(np.ones((T, T), dtype=bool), k=1)         # True above diagonal = future
    scores = np.where(mask, -1e9, scores)
    w = softmax(scores)
    out = (w @ Vh).transpose(1, 0, 2).reshape(T, d_model)    # concat heads
    return out @ L["Wo"], w, dict(Q=Qh, K=Kh, V=Vh, scores=scores)

def forward(ids, verbose=False):
    T = len(ids)
    x = emb[ids] * np.sqrt(d_model) + sinusoidal(T, d_model)
    if verbose:
        print("tokens:", [vocab[i] for i in ids], "ids:", ids)
        print("input x shape:", x.shape)
    ws = []
    for li, L in enumerate(layers):
        a, w, dbg = attention(x, L)
        ws.append(w)
        if verbose and li == 0:
            print("layer0 Q/K/V per-head shape:", dbg["Q"].shape, dbg["K"].shape, dbg["V"].shape)
            print("layer0 scores shape (heads,T,T):", dbg["scores"].shape)
            print("layer0 attention output shape:", a.shape)
        x = layer_norm(x + a)                                # post-LN, as in the paper
        h = np.maximum(0, x @ L["W1"] + L["b1"]) @ L["W2"] + L["b2"]
        x = layer_norm(x + h)
    logits = x @ emb.T                                       # tied output embedding
    return logits, ws

ids = [1, 2, 3, 4, 1, 5]   # the cat sat on the mat
logits, ws = forward(ids, verbose=True)
print("logits shape:", logits.shape)

print("\nLayer 0, head 0 attention weights (rows = query token, cols = key token):")
print("cols:", [vocab[i] for i in ids])
for t, row in enumerate(ws[0][0]):
    print(f"{vocab[ids[t]]:>4s}", row)
print("row sums:", ws[0][0].sum(-1))
print("max weight above diagonal:", np.triu(ws[0][0], k=1).max())

probs = softmax(logits[-1])
print("\nNext-token distribution after 'the cat sat on the mat' (random weights, meaningless):")
for i in np.argsort(-probs):
    print(f"  {vocab[i]:>6s} {probs[i]:.3f}")
print("sum:", probs.sum(), " uniform would be", 1 / V)

# causality check: change last token, earlier logits must not move
ids2 = ids[:-1] + [6]
logits2, _ = forward(ids2)
print("\nCausality: max |logit diff| at positions 0..4 after changing token 5:",
      np.abs(logits[:-1] - logits2[:-1]).max())
print("Same for position 5:", np.abs(logits[-1] - logits2[-1]).max())

# scaling check
q = rng.normal(size=(20000, 64)); k = rng.normal(size=(20000, 64))
dots = (q * k).sum(1)
print("\nstd of q.k at d_k=64:", dots.std().round(2), " after /sqrt(64):", (dots / 8).std().round(2))

# KV cache: incremental decode matches full forward, layer by layer
def decode_with_cache(ids):
    cache = [dict(K=np.zeros((0, d_model)), V=np.zeros((0, d_model))) for _ in layers]
    outs = []
    for t, tok in enumerate(ids):
        x = (emb[tok] * np.sqrt(d_model) + sinusoidal(t + 1, d_model)[t])[None, :]
        for li, L in enumerate(layers):
            c = cache[li]
            q, k, v = x @ L["Wq"], x @ L["Wk"], x @ L["Wv"]
            c["K"] = np.vstack([c["K"], k]); c["V"] = np.vstack([c["V"], v])
            n = c["K"].shape[0]
            qh = q.reshape(1, n_heads, d_k).transpose(1, 0, 2)
            Kh = c["K"].reshape(n, n_heads, d_k).transpose(1, 0, 2)
            Vh = c["V"].reshape(n, n_heads, d_k).transpose(1, 0, 2)
            w = softmax(qh @ Kh.transpose(0, 2, 1) / np.sqrt(d_k))
            a = (w @ Vh).transpose(1, 0, 2).reshape(1, d_model) @ L["Wo"]
            x = layer_norm(x + a)
            h = np.maximum(0, x @ L["W1"] + L["b1"]) @ L["W2"] + L["b2"]
            x = layer_norm(x + h)
        outs.append((x @ emb.T)[0])
    return np.array(outs), cache

lc, cache = decode_with_cache(ids)
print("\nKV cache decode vs full forward, max abs diff:", np.abs(lc - logits).max())
print("cache K shape per layer:", cache[0]["K"].shape, " cache floats total:",
      sum(c["K"].size + c["V"].size for c in cache),
      " formula 2*layers*T*d_model =", 2 * n_layers * len(ids) * d_model)

What does the output show?

This is the real output from python3 tiny_transformer.py with the seed above:

params (embeddings tied to output, no LayerNorm params): 4320
tokens: ['the', 'cat', 'sat', 'on', 'the', 'mat'] ids: [1, 2, 3, 4, 1, 5]
input x shape: (6, 16)
layer0 Q/K/V per-head shape: (2, 6, 8) (2, 6, 8) (2, 6, 8)
layer0 scores shape (heads,T,T): (2, 6, 6)
layer0 attention output shape: (6, 16)
logits shape: (6, 8)

Layer 0, head 0 attention weights (rows = query token, cols = key token):
cols: ['the', 'cat', 'sat', 'on', 'the', 'mat']
 the [1. 0. 0. 0. 0. 0.]
 cat [0.991 0.009 0.    0.    0.    0.   ]
 sat [0.25  0.627 0.123 0.    0.    0.   ]
  on [0.066 0.005 0.216 0.712 0.    0.   ]
 the [0.222 0.489 0.207 0.001 0.081 0.   ]
 mat [0.186 0.    0.001 0.215 0.595 0.003]
row sums: [1. 1. 1. 1. 1. 1.]
max weight above diagonal: 0.0

Read the shapes first, because interviewers often ask for them. A sequence of 6 tokens with d_model = 16 enters as (6, 16). After the head split, queries, keys and values are (2, 6, 8): 2 heads, 6 positions, 8 dimensions each. Scores are (2, 6, 6), one square matrix per head. The head outputs are concatenated back to (6, 16) and multiplied by the output projection. The logits are (6, 8): one distribution over the 8-word vocabulary for every position, which is also how training gets a loss at every position in one pass.

Then the weights. Each row sums to 1, which is the softmax. Everything above the diagonal is exactly zero, because the mask turns future scores into a huge negative number before the softmax. The first token can only attend to itself, so its row is [1, 0, 0, ...] whatever the weights are. Random weights give arbitrary patterns (the second "the" puts 0.489 on "cat"), so don't read linguistic meaning into them.

Parameter count, so you can derive it on a whiteboard: the embedding is 8 x 16 = 128. Each layer has four 16 x 16 attention matrices (1,024), plus 16 x 32 + 32 + 32 x 16 + 16 = 1,072 in the feed-forward block, so 2,096 per layer and 4,192 for two. Total 4,320. The output projection costs nothing extra because it reuses the embedding.

Next-token distribution after 'the cat sat on the mat' (random weights, meaningless):
     mat 0.735
     cat 0.242
     dog 0.009
     the 0.009
     sat 0.002
   <pad> 0.002
     ran 0.001
      on 0.001
sum: 0.9999999999999998  uniform would be 0.125

The model "predicts" mat, the token it just saw. That's an artefact of tying the output matrix to the input embedding with random weights, where a hidden state stays close to its own input embedding. A trained model learns to move away from that. It's a good reminder that a plausible-looking distribution proves nothing about correctness.

How do you prove the causal mask works?

Don't argue that it works; test it. The script changes only the last token (from "mat" to "dog") and compares logits at every earlier position:

Causality: max |logit diff| at positions 0..4 after changing token 5: 0.0
Same for position 5: 5.165342478492287

Positions 0 to 4 are bit-for-bit unchanged. Only the final position moves. This is the unit test I'd want in any from-scratch implementation: if a future token can change an earlier output, your mask is wrong, and a model trained that way will cheat on next-token prediction by reading the answer.

Why divide by sqrt(d_k)?

The paper's argument is a variance one. If query and key components are independent with mean 0 and variance 1, their dot product has variance d_k, so the scores grow with head size and the softmax saturates. The script checks it with random normal vectors:

std of q.k at d_k=64: 8.04  after /sqrt(64): 1.0

A standard deviation of 8.04 before scaling, 1.0 after. That matches sqrt(64) = 8. The caveat to say out loud in an interview: this holds for unit-variance inputs, such as at initialisation, not as a law about trained weights.

Does the KV cache give the same answer?

A KV cache is just an optimisation, so it must not change the output. The script runs generation token by token, appending each new key and value to a per-layer cache and attending over it, then compares with the full forward pass:

KV cache decode vs full forward, max abs diff: 1.7763568394002505e-15
cache K shape per layer: (6, 16)  cache floats total: 384  formula 2*layers*T*d_model = 384

The difference is floating-point rounding, about 1e-15. The cache holds 384 floats, which matches 2 * layers * T * d_model: K and V, two layers, six tokens, sixteen dimensions. In a real model with grouped-query attention the d_model term becomes kv_heads * head_dim, and the post linked earlier turns that into GiB for real configs.

One detail worth noticing in decode_with_cache: for each new token the query is a single row, so the score matrix is (heads, 1, n). No mask is needed, since the cache only contains the past. That is why decoding cost per token is linear in the cached length.

What does this toy leave out?

Be ready to list these, because "what's different from a real LLM?" is a common follow-up.

• Training. There are no gradients or optimiser. Adding them in NumPy means writing backpropagation for every operation, which is a good exercise and a bad use of a weekend if your goal is a job. A framework does it for you.
• Tokenisation. The vocabulary is 8 whole words. Real models use subword tokenisers; see tokenization and embeddings interview questions.
• Scale and kernels. Real attention is computed with fused, memory-efficient kernels, not by building the full score matrix as done here.
• Design drift. This follows the 2017 paper (post-layer-norm, sinusoidal positions, ReLU). I haven't checked which variants any particular current model uses, so don't claim "modern LLMs do X" from this post.
• Batching and padding. One sequence only. A batch needs padding masks on top of the causal mask.

Which interview questions does this answer?

Short answers, each tied to a line in the script:

For a practice drill, try these four changes and predict the result before running: remove the mask and re-run the causality test; delete the / np.sqrt(d_k) and compare the max attention weight; set n_heads = 4 and check shapes; replace sinusoidal positions with zeros and ask whether "the cat sat" and "sat cat the" still differ (hint: look at what else carries order information in this model).

If you're preparing for a loop, pair this with the mock format in mock interview script for AI engineers, and build one thing you can show, as suggested in AI portfolio projects that get interviews.

What I did not verify

The output above is from the script exactly as printed (I extracted the code block from this page and re-ran it; output identical), on NumPy 2.5.3 with Python 3. I did not run it on other NumPy versions, and floating-point digits in the last places may differ across machines. I read the paper's text for the architecture details listed above and did not check any claim about specific modern models. The "tied embedding predicts its own input" explanation is my interpretation of this toy's behaviour, not something I tested by ablation.