attn.py
One attention head, and the induction-head result
349 lines. This is the file the build actually runs, copied verbatim at build time.
"""One attention head, written out by hand.
The model on /learn/ flattens a fixed window into one vector, so every position
is wired to the output separately and the only way to use a character is to have
learned a rule for that character *at that offset*. Attention replaces the
flattening with a lookup by content: build a query from the current position,
compare it against a key at every earlier position, and take a weighted average
of the values there. Which position matters is decided at run time from the
content, not baked into the weights.
x_i = C[token_i] + P[i] embedding plus position
q = x_last @ Wq one query, from where we are now
k_i = x_i @ Wk a key at every position
v_i = x_i @ Wv a value at every position
score_i = (q . k_i) / sqrt(A)
w = softmax(score) how much to look at each position
context = sum_i w_i v_i
logits = context @ Wo + b
This is a single head at a single position — the smallest thing that shows the
mechanism. A transformer does this at every position at once, several times in
parallel, stacked in layers, with a feed-forward network between.
No autodiff: every derivative is written out, and gradcheck() compares them
against finite differences. Mirrored by attention/attn.js.
"""
import math
WINDOW = 16
EMBED = 16
ATTN = 16
class Rng:
"""mulberry32, same as everywhere else here."""
def __init__(self, seed):
self.state = seed & 0xFFFFFFFF
def next(self):
self.state = (self.state + 0x6D2B79F5) & 0xFFFFFFFF
t = self.state
t = ((t ^ (t >> 15)) * (t | 1)) & 0xFFFFFFFF
t = (t ^ (t + ((t ^ (t >> 7)) * (t | 61) & 0xFFFFFFFF))) & 0xFFFFFFFF
return ((t ^ (t >> 14)) & 0xFFFFFFFF) / 4294967296.0
def normal(self, scale):
u1 = max(self.next(), 1e-12)
u2 = self.next()
return scale * math.sqrt(-2 * math.log(u1)) * math.cos(2 * math.pi * u2)
def build_vocab(text):
vocab, seen = ["\n"], {"\n"}
for ch in text:
if ch not in seen:
seen.add(ch)
vocab.append(ch)
return vocab
def make_examples(text, vocab, window=WINDOW):
index = {c: i for i, c in enumerate(vocab)}
padded = "\n" * window + text
xs, ys = [], []
for i in range(window, len(padded)):
xs.append([index[padded[i - window + j]] for j in range(window)])
ys.append(index[padded[i]])
return xs, ys
class Model:
def __init__(self, vocab, seed=1, window=WINDOW, embed=EMBED, attn=ATTN,
shift_values=False):
self.vocab = vocab
self.V = len(vocab)
self.L = window
self.D = embed
self.A = attn
rng = Rng(seed)
s = 1.0 / math.sqrt(self.D)
self.C = [[rng.normal(1.0) for _ in range(self.D)] for _ in range(self.V)]
self.P = [[rng.normal(0.3) for _ in range(self.D)] for _ in range(self.L)]
self.Wq = [[rng.normal(s) for _ in range(self.A)] for _ in range(self.D)]
self.Wk = [[rng.normal(s) for _ in range(self.A)] for _ in range(self.D)]
self.Wv = [[rng.normal(s) for _ in range(self.A)] for _ in range(self.D)]
so = 1.0 / math.sqrt(self.A)
self.Wo = [[rng.normal(so) for _ in range(self.V)] for _ in range(self.A)]
self.b = [0.0] * self.V
# When true, position i's value carries the NEXT character rather
# than its own. A real transformer arranges this with an earlier
# layer; supplying it lets one head show the matching half alone.
self.shift_values = shift_values
def parameter_count(self):
return (self.V * self.D + self.L * self.D + 3 * self.D * self.A
+ self.A * self.V + self.V)
# ---- forward ----
def forward(self, ctx):
L, D, A = self.L, self.D, self.A
x = []
for i in range(L):
emb = self.C[ctx[i]]
pos = self.P[i]
x.append([emb[d] + pos[d] for d in range(D)])
def project(vec, W):
out = [0.0] * A
for d in range(D):
val = vec[d]
if val == 0.0:
continue
row = W[d]
for a in range(A):
out[a] += val * row[a]
return out
q = project(x[L - 1], self.Wq)
k = [project(x[i], self.Wk) for i in range(L)]
if self.shift_values:
zero = [0.0] * D
xv = [x[i + 1] if i + 1 < L else zero for i in range(L)]
else:
xv = x
v = [project(xv[i], self.Wv) for i in range(L)]
inv = 1.0 / math.sqrt(A)
scores = []
for i in range(L):
dot = 0.0
ki = k[i]
for a in range(A):
dot += q[a] * ki[a]
scores.append(dot * inv)
m = max(scores)
exps = [math.exp(s - m) for s in scores]
total = 0.0
for e in exps:
total += e
w = [e / total for e in exps]
context = [0.0] * A
for i in range(L):
wi = w[i]
vi = v[i]
for a in range(A):
context[a] += wi * vi[a]
logits = list(self.b)
for a in range(A):
ca = context[a]
row = self.Wo[a]
for c in range(self.V):
logits[c] += ca * row[c]
mm = max(logits)
le = [math.exp(z - mm) for z in logits]
lt = 0.0
for e in le:
lt += e
probs = [e / lt for e in le]
return probs, {"x": x, "xv": xv, "q": q, "k": k, "v": v, "w": w,
"context": context, "ctx": ctx}
def loss(self, xs, ys):
total = 0.0
for ctx, y in zip(xs, ys):
probs, _ = self.forward(ctx)
total += -math.log(max(probs[y], 1e-12))
return total / len(xs)
def attention(self, ctx):
"""The weights themselves — what the model looked at."""
_, cache = self.forward(ctx)
return cache["w"]
# ---- backward ----
def zero_grads(self):
return {
"C": [[0.0] * self.D for _ in range(self.V)],
"P": [[0.0] * self.D for _ in range(self.L)],
"Wq": [[0.0] * self.A for _ in range(self.D)],
"Wk": [[0.0] * self.A for _ in range(self.D)],
"Wv": [[0.0] * self.A for _ in range(self.D)],
"Wo": [[0.0] * self.V for _ in range(self.A)],
"b": [0.0] * self.V,
}
def backward(self, xs, ys, grads):
L, D, A, V = self.L, self.D, self.A, self.V
inv = 1.0 / math.sqrt(A)
n = len(xs)
total_loss = 0.0
for ctx, y in zip(xs, ys):
probs, c = self.forward(ctx)
total_loss += -math.log(max(probs[y], 1e-12))
x, xv, q, k, v, w, context = (c["x"], c["xv"], c["q"], c["k"],
c["v"], c["w"], c["context"])
dlogits = [p / n for p in probs]
dlogits[y] -= 1.0 / n
dcontext = [0.0] * A
for a in range(A):
row = self.Wo[a]
grow = grads["Wo"][a]
ca = context[a]
acc = 0.0
for cidx in range(V):
dc = dlogits[cidx]
grow[cidx] += ca * dc
acc += row[cidx] * dc
dcontext[a] = acc
for cidx in range(V):
grads["b"][cidx] += dlogits[cidx]
# context = sum_i w_i v_i
dw = [0.0] * L
dv = [[0.0] * A for _ in range(L)]
for i in range(L):
vi = v[i]
wi = w[i]
acc = 0.0
dvi = dv[i]
for a in range(A):
acc += dcontext[a] * vi[a]
dvi[a] = wi * dcontext[a]
dw[i] = acc
# through the softmax over scores
dot = 0.0
for i in range(L):
dot += w[i] * dw[i]
dscore = [w[i] * (dw[i] - dot) for i in range(L)]
# score_i = (q . k_i) * inv
dq = [0.0] * A
dk = [[0.0] * A for _ in range(L)]
for i in range(L):
s = dscore[i] * inv
ki = k[i]
dki = dk[i]
for a in range(A):
dq[a] += s * ki[a]
dki[a] = s * q[a]
dx = [[0.0] * D for _ in range(L)]
def backprop_projection(vec, W, gW, dout, dvec):
for d in range(D):
val = vec[d]
row = W[d]
grow = gW[d]
acc = 0.0
for a in range(A):
da = dout[a]
grow[a] += val * da
acc += row[a] * da
dvec[d] += acc
backprop_projection(x[L - 1], self.Wq, grads["Wq"], dq, dx[L - 1])
for i in range(L):
backprop_projection(x[i], self.Wk, grads["Wk"], dk[i], dx[i])
# values may read the next position, so the gradient goes there
target = i + 1 if self.shift_values else i
if target < L:
backprop_projection(xv[i], self.Wv, grads["Wv"], dv[i],
dx[target])
for i in range(L):
gc = grads["C"][ctx[i]]
gp = grads["P"][i]
dxi = dx[i]
for d in range(D):
gc[d] += dxi[d]
gp[d] += dxi[d]
return total_loss / n
def step(self, grads, lr):
for name in ("C", "P", "Wq", "Wk", "Wv", "Wo"):
tensor = getattr(self, name)
grad = grads[name]
for i in range(len(tensor)):
row, g = tensor[i], grad[i]
for j in range(len(row)):
row[j] -= lr * g[j]
for k in range(self.V):
self.b[k] -= lr * grads["b"][k]
def train(self, xs, ys, steps, lr, batch=32, seed=7):
rng = Rng(seed)
history = []
for _ in range(steps):
idx = [int(rng.next() * len(xs)) % len(xs) for _ in range(batch)]
bx = [xs[i] for i in idx]
by = [ys[i] for i in idx]
grads = self.zero_grads()
history.append(self.backward(bx, by, grads))
self.step(grads, lr)
return history
def gradcheck(model, xs, ys, eps=1e-5, tol=1e-4, checks=40, seed=3):
grads = model.zero_grads()
model.backward(xs, ys, grads)
rng = Rng(seed)
names = ["C", "P", "Wq", "Wk", "Wv", "Wo"]
worst, failures = 0.0, []
for _ in range(checks):
if rng.next() < 0.12:
i = int(rng.next() * model.V)
original, analytic = model.b[i], grads["b"][i]
model.b[i] = original + eps
plus = model.loss(xs, ys)
model.b[i] = original - eps
minus = model.loss(xs, ys)
model.b[i] = original
label = "b"
else:
name = names[int(rng.next() * len(names))]
tensor, grad = getattr(model, name), grads[name]
i = int(rng.next() * len(tensor))
j = int(rng.next() * len(tensor[i]))
original, analytic = tensor[i][j], grad[i][j]
tensor[i][j] = original + eps
plus = model.loss(xs, ys)
tensor[i][j] = original - eps
minus = model.loss(xs, ys)
tensor[i][j] = original
label = name
numeric = (plus - minus) / (2 * eps)
scale = max(abs(analytic), abs(numeric), 1e-8)
rel = abs(analytic - numeric) / scale
worst = max(worst, rel)
if rel > tol:
failures.append((label, analytic, numeric, rel))
return worst, failures