Building a Transformer

import warnings
warnings.filterwarnings("ignore")
import math, time
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.nn.functional as F

np.random.seed(0); torch.manual_seed(0)
np.set_printoptions(precision=3, suppress=True)

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 show_matrix(ax, M, xlabels, ylabels, title, cmap="Blues", vmin=None, vmax=None):
    ax.imshow(M, cmap=cmap, vmin=vmin, vmax=vmax)
    ax.set_xticks(range(len(xlabels))); ax.set_xticklabels(xlabels, rotation=45, ha="right")
    ax.set_yticks(range(len(ylabels))); ax.set_yticklabels(ylabels)
    ax.set_title(title, fontsize=10)

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

pe = sinusoidal_pe(64, 64)
fig, axes = plt.subplots(1, 3, figsize=(18, 4))
axes[0].imshow(pe.T, cmap="RdBu", aspect="auto"); axes[0].set_xlabel("position"); axes[0].set_ylabel("dimension")
axes[0].set_title("Sinusoidal positional encoding")
for dim in [0, 10, 30]:
    axes[1].plot(pe[:, dim], label=f"dim {dim}")
axes[1].legend(); axes[1].set_xlabel("position"); axes[1].set_title("Each dimension is a wave of a different frequency")
axes[2].imshow(pe @ pe.T, cmap="viridis"); axes[2].set_title("PE · PE: similarity depends on distance only")
axes[2].set_xlabel("position"); axes[2].set_ylabel("position")
plt.tight_layout(); plt.show()

tokens = ["the", "cat", "sat", "on", "the", "mat"]
vocab = sorted(set(tokens))
d_model, d_k = 8, 4
E = np.random.randn(len(vocab), d_model)                     # embedding table
X = np.array([E[vocab.index(t)] for t in tokens]) + sinusoidal_pe(len(tokens), d_model)
W_Q, W_K, W_V = (np.random.randn(d_model, d_k) / np.sqrt(d_model) for _ in range(3))

Q, K, V = X @ W_Q, X @ W_K, X @ W_V
scores = Q @ K.T / np.sqrt(d_k)                              # (6, 6): every query against every key
A = softmax(scores, axis=-1)                                  # each row sums to 1
out = A @ V                                                   # (6, 4): new representation per token

print("X:", X.shape, " Q,K,V:", Q.shape, " scores:", scores.shape, " output:", out.shape)
print("row sums of A:", A.sum(axis=1))
fig, axes = plt.subplots(1, 3, figsize=(17, 4.5))
show_matrix(axes[0], scores, tokens, tokens, "1) scores  QKᵀ/√d_k", cmap="RdBu_r")
show_matrix(axes[1], A, tokens, tokens, "2) weights  softmax(scores)  (row = query)", vmin=0, vmax=1)
show_matrix(axes[2], out, [f"v{i}" for i in range(d_k)], tokens, "3) output  A·V", cmap="RdBu_r")
plt.tight_layout(); plt.show()

dims = [4, 16, 64, 256, 1024]
ent_raw, ent_scaled = [], []
for d in dims:
    q, k = np.random.randn(1000, d), np.random.randn(1000, 32, d)
    s = np.einsum("nd,nkd->nk", q, k)
    for store, ss in [(ent_raw, s), (ent_scaled, s / np.sqrt(d))]:
        p = softmax(ss)
        store.append(-(p * np.log(p + 1e-12)).sum(1).mean())
plt.figure(figsize=(7, 4))
plt.semilogx(dims, ent_raw, "o-", label="raw  q·k"); plt.semilogx(dims, ent_scaled, "o-", label="scaled  q·k/√d")
plt.axhline(np.log(32), color="gray", ls="--", label="uniform over 32 keys")
plt.xlabel("d_k"); plt.ylabel("entropy of attention weights"); plt.legend(); plt.grid(alpha=0.3)
plt.title("Without scaling, attention collapses to one-hot as d grows"); plt.show()

def attention(Xin):
    q, k, v = Xin @ W_Q, Xin @ W_K, Xin @ W_V
    return softmax(q @ k.T / np.sqrt(d_k)) @ v

X_nopos = np.array([E[vocab.index(t)] for t in tokens])
perm = np.random.permutation(len(tokens))
print("permutation:", perm)
print("attention(shuffled X) == shuffled attention(X):", np.allclose(attention(X_nopos[perm]), attention(X_nopos)[perm]))

n = len(tokens)
causal = np.tril(np.ones((n, n), dtype=bool))
A_causal = softmax(np.where(causal, scores, -np.inf))

pad = np.array([True, True, True, True, False, False])      # pretend the last two tokens are padding
A_pad = softmax(np.where(pad[None, :], scores, -np.inf))

fig, axes = plt.subplots(1, 3, figsize=(17, 4.5))
show_matrix(axes[0], causal.astype(float), tokens, tokens, "causal mask (1 = allowed)", cmap="Greys")
show_matrix(axes[1], A_causal, tokens, tokens, "causal attention: lower-triangular", vmin=0, vmax=1)
show_matrix(axes[2], A_pad, tokens, tokens, "padding mask: last 2 columns ignored", vmin=0, vmax=1)
plt.tight_layout(); plt.show()

def multi_head_attention(X, Wq, Wk, Wv, Wo, n_heads, mask=None):
    n, d = X.shape
    dh = d // n_heads
    def split(M):                                   # (n, d) -> (heads, n, dh)
        return M.reshape(n, n_heads, dh).transpose(1, 0, 2)
    Q, K, V = split(X @ Wq), split(X @ Wk), split(X @ Wv)
    scores = Q @ K.transpose(0, 2, 1) / np.sqrt(dh)  # (heads, n, n)
    if mask is not None:
        scores = np.where(mask, scores, -np.inf)
    A = softmax(scores)
    heads = A @ V                                   # (heads, n, dh)
    concat = heads.transpose(1, 0, 2).reshape(n, d) # back to (n, d)
    return concat @ Wo, A

d_model, n_heads = 8, 4
Wq, Wk, Wv, Wo = (np.random.randn(d_model, d_model) / np.sqrt(d_model) for _ in range(4))
Y, A_heads = multi_head_attention(X, Wq, Wk, Wv, Wo, n_heads)
print("input", X.shape, "-> output", Y.shape, "| attention maps", A_heads.shape, "(heads, query, key)")
fig, axes = plt.subplots(1, n_heads, figsize=(18, 4))
for h, ax in enumerate(axes):
    show_matrix(ax, A_heads[h], tokens, tokens, f"head {h}", vmin=0, vmax=1)
plt.tight_layout(); plt.show()

def gelu(x):
    return 0.5 * x * (1 + np.tanh(np.sqrt(2 / np.pi) * (x + 0.044715 * x ** 3)))   # GPT-2's approximation

d_ff = 4 * d_model
W1, b1 = np.random.randn(d_model, d_ff) / np.sqrt(d_model), np.zeros(d_ff)
W2, b2 = np.random.randn(d_ff, d_model) / np.sqrt(d_ff), np.zeros(d_model)
ffn = lambda x: gelu(x @ W1 + b1) @ W2 + b2

whole = ffn(X)
one_by_one = np.stack([ffn(X[i:i + 1])[0] for i in range(len(X))])
print("FFN applied to the sequence == FFN applied to each token separately:", np.allclose(whole, one_by_one))
print(f"FFN params: {W1.size + b1.size + W2.size + b2.size}  vs  attention params: {4 * d_model * d_model}")

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

def transformer_block(x, p, mask=None):
    a, attn = multi_head_attention(layer_norm(x, *p["ln1"]), *p["attn"], n_heads=p["heads"], mask=mask)
    x = x + a                                         # residual 1: attention adds to the stream
    x = x + p["ffn"](layer_norm(x, *p["ln2"]))        # residual 2: MLP adds to the stream
    return x, attn

params = {"ln1": (np.ones(d_model), np.zeros(d_model)), "ln2": (np.ones(d_model), np.zeros(d_model)),
          "attn": (Wq, Wk, Wv, Wo), "heads": n_heads, "ffn": ffn}
h = X
for layer in range(3):                               # weight-shared for brevity; real models use different weights per layer
    h_new, _ = transformer_block(h, params, mask=causal)
    print(f"block {layer}: shape {h_new.shape}, size of update ‖Δ‖ = {np.linalg.norm(h_new - h):.2f}")
    h = h_new

# Check our LayerNorm against PyTorch's
x_t = torch.randn(5, 8)
ours = layer_norm(x_t.numpy(), np.ones(8), np.zeros(8))
print("LayerNorm matches PyTorch:", np.allclose(ours, F.layer_norm(x_t, (8,)).numpy(), atol=1e-5))

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        assert d_model % n_heads == 0
        self.h, self.dh = n_heads, d_model // n_heads
        self.Wq, self.Wk, self.Wv, self.Wo = (nn.Linear(d_model, d_model) for _ in range(4))
        self.last_attn = None
    def forward(self, x_q, x_kv, mask=None):
        B, Tq, D = x_q.shape
        Tk = x_kv.shape[1]
        q = self.Wq(x_q).view(B, Tq, self.h, self.dh).transpose(1, 2)    # (B, h, Tq, dh)
        k = self.Wk(x_kv).view(B, Tk, self.h, self.dh).transpose(1, 2)
        v = self.Wv(x_kv).view(B, Tk, self.h, self.dh).transpose(1, 2)
        scores = q @ k.transpose(-2, -1) / math.sqrt(self.dh)            # (B, h, Tq, Tk)
        if mask is not None:
            scores = scores.masked_fill(~mask, float("-inf"))
        attn = scores.softmax(-1)
        self.last_attn = attn.detach()
        out = (attn @ v).transpose(1, 2).reshape(B, Tq, D)
        return self.Wo(out)

class FeedForward(nn.Module):
    def __init__(self, d_model, mult=4):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(d_model, mult * d_model), nn.GELU(), nn.Linear(mult * d_model, d_model))
    def forward(self, x):
        return self.net(x)

class EncoderBlock(nn.Module):
    def __init__(self, d, h):
        super().__init__()
        self.ln1, self.attn, self.ln2, self.ffn = nn.LayerNorm(d), MultiHeadAttention(d, h), nn.LayerNorm(d), FeedForward(d)
    def forward(self, x):
        y = self.ln1(x)
        x = x + self.attn(y, y)
        return x + self.ffn(self.ln2(x))

class DecoderBlock(nn.Module):
    def __init__(self, d, h):
        super().__init__()
        self.ln1, self.self_attn = nn.LayerNorm(d), MultiHeadAttention(d, h)
        self.ln2, self.cross_attn = nn.LayerNorm(d), MultiHeadAttention(d, h)
        self.ln3, self.ffn = nn.LayerNorm(d), FeedForward(d)
    def forward(self, x, memory, causal_mask):
        y = self.ln1(x)
        x = x + self.self_attn(y, y, causal_mask)                  # look at previous outputs only
        x = x + self.cross_attn(self.ln2(x), memory)               # look at the encoded source
        return x + self.ffn(self.ln3(x))

class Transformer(nn.Module):
    def __init__(self, vocab, d=64, heads=4, n_enc=2, n_dec=2, max_len=32):
        super().__init__()
        self.src_emb, self.tgt_emb = nn.Embedding(vocab, d), nn.Embedding(vocab, d)
        self.register_buffer("pe", torch.tensor(sinusoidal_pe(max_len, d), dtype=torch.float32))
        self.enc = nn.ModuleList(EncoderBlock(d, heads) for _ in range(n_enc))
        self.dec = nn.ModuleList(DecoderBlock(d, heads) for _ in range(n_dec))
        self.ln_enc, self.ln_dec = nn.LayerNorm(d), nn.LayerNorm(d)
        self.out = nn.Linear(d, vocab)
        self.d = d
    def encode(self, src):
        x = self.src_emb(src) * math.sqrt(self.d) + self.pe[:src.shape[1]]
        for blk in self.enc:
            x = blk(x)
        return self.ln_enc(x)
    def decode(self, tgt_in, memory):
        T = tgt_in.shape[1]
        causal = torch.tril(torch.ones(T, T, dtype=torch.bool, device=tgt_in.device))
        x = self.tgt_emb(tgt_in) * math.sqrt(self.d) + self.pe[:T]
        for blk in self.dec:
            x = blk(x, memory, causal)
        return self.out(self.ln_dec(x))
    def forward(self, src, tgt_in):
        return self.decode(tgt_in, self.encode(src))

N_DIGITS, SEQ, SOS = 10, 10, 10
VOCAB = N_DIGITS + 1

def make_batch(B):
    src = torch.randint(0, N_DIGITS, (B, SEQ))
    tgt = src.sort(dim=1).values
    tgt_in = torch.cat([torch.full((B, 1), SOS), tgt[:, :-1]], dim=1)     # shift right: teacher forcing
    return src, tgt_in, tgt

torch.manual_seed(0)
model = Transformer(VOCAB)
print(f"parameters: {sum(p.numel() for p in model.parameters()):,}")
src, tgt_in, tgt = make_batch(2)
print("example source :", src[0].tolist())
print("decoder input  :", tgt_in[0].tolist(), "(10 = <sos>)")
print("target         :", tgt[0].tolist())

opt = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)
steps = 1200
sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=1e-3, total_steps=steps, pct_start=0.1)   # warm-up then decay
t0, losses = time.time(), []
for step in range(steps):
    src, tgt_in, tgt = make_batch(128)
    logits = model(src, tgt_in)
    loss = F.cross_entropy(logits.reshape(-1, VOCAB), tgt.reshape(-1))
    opt.zero_grad(); loss.backward(); nn.utils.clip_grad_norm_(model.parameters(), 1.0); opt.step(); sched.step()
    losses.append(loss.item())
    if step % 300 == 0 or step == steps - 1:
        print(f"step {step:4d}  loss {loss.item():.4f}  ({time.time() - t0:.0f}s)")
plt.figure(figsize=(7, 3.5)); plt.semilogy(losses); plt.xlabel("step"); plt.ylabel("loss"); plt.grid(alpha=0.3)
plt.title("Training the transformer to sort"); plt.show()

@torch.no_grad()
def greedy_sort(model, src):
    memory = model.encode(src)
    out = torch.full((src.shape[0], 1), SOS)
    for _ in range(SEQ):                                   # autoregressive: one digit at a time
        nxt = model.decode(out, memory)[:, -1].argmax(-1, keepdim=True)
        out = torch.cat([out, nxt], dim=1)
    return out[:, 1:]

model.eval()
torch.manual_seed(123)
src, _, tgt = make_batch(2000)
pred = greedy_sort(model, src)
print(f"fully correct sequences: {(pred == tgt).all(1).float().mean():.1%}   per-digit accuracy: {(pred == tgt).float().mean():.2%}")
for i in range(4):
    print("in ", src[i].tolist(), "\nout", pred[i].tolist(), "✓" if (pred[i] == tgt[i]).all() else "✗")

example = torch.tensor([[7, 2, 9, 0, 4, 4, 1, 8, 3, 6]])
with torch.no_grad():
    memory = model.encode(example)
    res = greedy_sort(model, example)
    tgt_in = torch.cat([torch.tensor([[SOS]]), res[:, :-1]], dim=1)
    model.decode(tgt_in, memory)                         # fills last_attn for the full decoded sequence
src_lbl = [f"{d}" for d in example[0].tolist()]
out_lbl = [f"→{d}" for d in res[0].tolist()]
print("input :", example[0].tolist(), "\noutput:", res[0].tolist())

H = model.enc[0].attn.h
fig, axes = plt.subplots(2, H, figsize=(4 * H, 8))
for l, blk in enumerate(model.enc):
    for h in range(H):
        show_matrix(axes[l, h], blk.attn.last_attn[0, h], src_lbl, src_lbl, f"encoder L{l} head {h}", vmin=0, vmax=1)
fig.suptitle("Encoder self-attention (rows = query digit, cols = key digit)"); plt.tight_layout(); plt.show()

fig, axes = plt.subplots(2, H, figsize=(4 * H, 8))
for l, blk in enumerate(model.dec):
    for h in range(H):
        show_matrix(axes[l, h], blk.cross_attn.last_attn[0, h], src_lbl, out_lbl, f"decoder L{l} cross-attn head {h}", cmap="Greens", vmin=0, vmax=1)
fig.suptitle("Decoder cross-attention (rows = output step, cols = input position)"); plt.tight_layout(); plt.show()

def rollout(attn_maps):
    R = torch.eye(attn_maps[0].shape[-1])
    for A in attn_maps:
        A_bar = A.mean(0)                                       # average heads
        A_res = 0.5 * A_bar + 0.5 * torch.eye(A_bar.shape[0])   # account for the residual path
        A_res = A_res / A_res.sum(-1, keepdim=True)
        R = A_res @ R
    return R

enc_maps = [blk.attn.last_attn[0] for blk in model.enc]
fig, axes = plt.subplots(1, 3, figsize=(16, 4.8))
show_matrix(axes[0], enc_maps[0].mean(0), src_lbl, src_lbl, "layer 0 (mean over heads)", cmap="Purples")
show_matrix(axes[1], enc_maps[1].mean(0), src_lbl, src_lbl, "layer 1 (mean over heads)", cmap="Purples")
show_matrix(axes[2], rollout(enc_maps), src_lbl, src_lbl, "rollout through both layers", cmap="Purples")
plt.tight_layout(); plt.show()

model.zero_grad()
emb = (model.src_emb(example) * math.sqrt(model.d)).detach().requires_grad_(True)
x = emb + model.pe[:SEQ]
for blk in model.enc:
    x = blk(x)
memory = model.ln_enc(x)
logits = model.decode(tgt_in, memory)                        # (1, SEQ, VOCAB)

saliency = torch.zeros(SEQ, SEQ)
for k in range(SEQ):
    chosen = logits[0, k, res[0, k]]                         # the logit of the digit it actually wrote
    g, = torch.autograd.grad(chosen, emb, retain_graph=True)
    saliency[k] = g[0].norm(dim=-1)
saliency = saliency / saliency.sum(1, keepdim=True)

fig, axes = plt.subplots(1, 2, figsize=(12, 4.8))
show_matrix(axes[0], saliency, src_lbl, out_lbl, "gradient saliency (normalised per row)", cmap="Oranges")
show_matrix(axes[1], model.dec[-1].cross_attn.last_attn[0].mean(0), src_lbl, out_lbl, "last-layer cross-attention (mean of heads)", cmap="Greens")
plt.tight_layout(); plt.show()

def count(m): return sum(p.numel() for p in m.parameters())
blk = model.enc[0]
attn_p, ffn_p, ln_p = count(blk.attn), count(blk.ffn), count(blk.ln1) + count(blk.ln2)
print(f"encoder block (d=64): attention {attn_p:,}  FFN {ffn_p:,}  LayerNorm {ln_p:,}")
print(f"FFN share: {ffn_p / (attn_p + ffn_p + ln_p):.0%},  12·d² = {12 * 64 ** 2:,}")
for name, d, L in [("GPT-2 small", 768, 12), ("GPT-2 XL", 1600, 48), ("GPT-3", 12288, 96)]:
    print(f"{name:12s} d={d:5d} layers={L:3d}  ->  12·L·d² ≈ {12 * L * d * d / 1e9:6.2f} B parameters in the blocks")