Quartz 5

15. Tiny code examples

Properties1
tagsai, llm, deep-learning, transformers

2 min read

One attention head in plain NumPy

This is the attention formula from 8a. Self-attention, line for line:

import numpy as np
 
def softmax(x, axis=-1):
    x = x - x.max(axis=axis, keepdims=True)   # numerical stability
    e = np.exp(x)
    return e / e.sum(axis=axis, keepdims=True)
 
def causal_self_attention(X, W_Q, W_K, W_V):
    # X: [n, d_model]   W_*: [d_model, d_head]
    Q, K, V = X @ W_Q, X @ W_K, X @ W_V
    d_k = Q.shape[-1]
    scores = (Q @ K.T) / np.sqrt(d_k)              # [n, n]
    mask = np.triu(np.ones_like(scores), k=1)      # 1s above the diagonal = future
    scores = np.where(mask == 1, -1e9, scores)     # block looking ahead
    weights = softmax(scores, axis=-1)             # each row sums to 1
    return weights @ V                             # [n, d_head]
 
# toy run
n, d_model, d_head = 5, 16, 8
rng = np.random.default_rng(0)
X = rng.normal(size=(n, d_model))
out = causal_self_attention(
    X,
    rng.normal(size=(d_model, d_head)),
    rng.normal(size=(d_model, d_head)),
    rng.normal(size=(d_model, d_head)),
)
print(out.shape)  # (5, 8)

The training loop shape in PyTorch

import torch
import torch.nn.functional as F
 
model = MyTinyGPT(...)                       # embeddings -> blocks -> unembedding
opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.1)
 
for step, (x, y) in enumerate(loader):       # y is x shifted left by one token
    logits = model(x)                        # [batch, seq, vocab]
    loss = F.cross_entropy(
        logits.view(-1, logits.size(-1)),
        y.view(-1),
    )
    opt.zero_grad()
    loss.backward()                          # backprop through everything
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    opt.step()

Generation loop

@torch.no_grad()
def generate(model, tokens, max_new=100, temperature=0.8):
    for _ in range(max_new):
        logits = model(tokens)[:, -1, :]                 # last position only
        probs = F.softmax(logits / temperature, dim=-1)
        next_tok = torch.multinomial(probs, num_samples=1)
        tokens = torch.cat([tokens, next_tok], dim=1)
    return tokens

That’s genuinely the skeleton. Production systems add KV caching, batching, better sampling, and a mountain of engineering, but the core idea fits on a page.

← 14. Limitations and failure modes · Contents · 16. Where to learn more →

Backlinks