mlp.py
The neural language model behind /learn/, by hand
311 lines. This is the file the build actually runs, copied verbatim at build time.
"""A small neural language model, written out by hand.
The n-gram on /predict/ cannot answer a question it was never asked: given a
context it has not seen, it backs off to a shorter one and eventually to noise.
This model can, because it does not store contexts — it stores a vector for each
character and learns a function of those vectors. Characters that behave alike
end up near each other, and a context made of familiar parts is answerable even
if that exact context never occurred.
Architecture, deliberately the smallest thing that shows the point:
context of K characters
-> embedding lookup (V x D), concatenated (K*D)
-> linear W1 (K*D x H) + b1, tanh (H)
-> linear W2 (H x V) + b2 (V)
-> softmax probabilities
No autodiff and no matrix library: the forward pass and every gradient are
written out so they can be read. verify_gradients() checks the analytic
gradients against finite differences, which is the real correctness argument —
float arithmetic will not reproduce bit-for-bit across languages, but a wrong
derivative shows up immediately.
Mirrored by learn/mlp.js.
"""
import math
CONTEXT = 3 # characters of context
EMBED = 8 # dimensions per character
HIDDEN = 64
class Rng:
"""mulberry32 again, so initialisation is reproducible in both languages."""
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):
"""Box-Muller, so weights start small and centred."""
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):
"""Characters in first-appearance order, plus a boundary marker at index 0."""
vocab = ["\n"]
seen = {"\n"}
for ch in text:
if ch not in seen:
seen.add(ch)
vocab.append(ch)
return vocab
def make_examples(text, vocab, context=CONTEXT):
"""Every position becomes (context indices -> next index)."""
index = {c: i for i, c in enumerate(vocab)}
padded = "\n" * context + text
xs, ys = [], []
for i in range(context, len(padded)):
xs.append([index[padded[i - context + j]] for j in range(context)])
ys.append(index[padded[i]])
return xs, ys
class Model:
def __init__(self, vocab, seed=1, context=CONTEXT, embed=EMBED, hidden=HIDDEN):
self.vocab = vocab
self.V = len(vocab)
self.K = context
self.D = embed
self.H = hidden
rng = Rng(seed)
self.C = [[rng.normal(1.0) for _ in range(self.D)] for _ in range(self.V)]
fan_in = self.K * self.D
s1 = 1.0 / math.sqrt(fan_in)
self.W1 = [[rng.normal(s1) for _ in range(self.H)] for _ in range(fan_in)]
self.b1 = [0.0] * self.H
s2 = 1.0 / math.sqrt(self.H)
self.W2 = [[rng.normal(s2) for _ in range(self.V)] for _ in range(self.H)]
self.b2 = [0.0] * self.V
# ---- forward ----
def forward(self, ctx):
"""Returns (probabilities, cache) for one example."""
emb = []
for idx in ctx:
emb.extend(self.C[idx])
h_pre = list(self.b1)
for i, e in enumerate(emb):
if e == 0.0:
continue
row = self.W1[i]
for j in range(self.H):
h_pre[j] += e * row[j]
h = [math.tanh(v) for v in h_pre]
logits = list(self.b2)
for j in range(self.H):
hj = h[j]
row = self.W2[j]
for k in range(self.V):
logits[k] += hj * row[k]
m = max(logits)
exps = [math.exp(v - m) for v in logits]
total = 0.0
for e in exps:
total += e
probs = [e / total for e in exps]
return probs, {"emb": emb, "h": h, "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)
# ---- backward ----
def zero_grads(self):
return {
"C": [[0.0] * self.D for _ in range(self.V)],
"W1": [[0.0] * self.H for _ in range(self.K * self.D)],
"b1": [0.0] * self.H,
"W2": [[0.0] * self.V for _ in range(self.H)],
"b2": [0.0] * self.V,
}
def backward(self, xs, ys, grads):
"""Accumulate gradients of mean cross-entropy over the batch."""
n = len(xs)
total_loss = 0.0
for ctx, y in zip(xs, ys):
probs, cache = self.forward(ctx)
total_loss += -math.log(max(probs[y], 1e-12))
# dL/dlogits for softmax + cross-entropy is (p - onehot)/n
dlogits = [p / n for p in probs]
dlogits[y] -= 1.0 / n
h = cache["h"]
emb = cache["emb"]
dh = [0.0] * self.H
for j in range(self.H):
row = self.W2[j]
grow = grads["W2"][j]
hj = h[j]
acc = 0.0
for k in range(self.V):
dk = dlogits[k]
grow[k] += hj * dk
acc += row[k] * dk
dh[j] = acc
for k in range(self.V):
grads["b2"][k] += dlogits[k]
# through tanh
dh_pre = [dh[j] * (1.0 - h[j] * h[j]) for j in range(self.H)]
demb = [0.0] * (self.K * self.D)
for i in range(self.K * self.D):
row = self.W1[i]
grow = grads["W1"][i]
e = emb[i]
acc = 0.0
for j in range(self.H):
dj = dh_pre[j]
grow[j] += e * dj
acc += row[j] * dj
demb[i] = acc
for j in range(self.H):
grads["b1"][j] += dh_pre[j]
for slot, idx in enumerate(cache["ctx"]):
base = slot * self.D
grow = grads["C"][idx]
for d in range(self.D):
grow[d] += demb[base + d]
return total_loss / n
def step(self, grads, lr):
for i in range(self.V):
row, g = self.C[i], grads["C"][i]
for d in range(self.D):
row[d] -= lr * g[d]
for i in range(self.K * self.D):
row, g = self.W1[i], grads["W1"][i]
for j in range(self.H):
row[j] -= lr * g[j]
for j in range(self.H):
self.b1[j] -= lr * grads["b1"][j]
row, g = self.W2[j], grads["W2"][j]
for k in range(self.V):
row[k] -= lr * g[k]
for k in range(self.V):
self.b2[k] -= lr * grads["b2"][k]
def train(self, xs, ys, steps, lr, batch=32, seed=7, on_step=None):
rng = Rng(seed)
history = []
for s 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()
loss = self.backward(bx, by, grads)
self.step(grads, lr)
history.append(loss)
if on_step:
on_step(s, loss)
return history
def verify_gradients(model, xs, ys, eps=1e-5, tol=1e-4, checks=40, seed=3):
"""Analytic gradients vs finite differences.
This is the correctness argument for the whole page. Float arithmetic will
not agree bit-for-bit between Python and JavaScript, so "both implementations
match exactly" is not available here — but a wrong derivative fails this
immediately, in either language.
"""
grads = model.zero_grads()
model.backward(xs, ys, grads)
rng = Rng(seed)
params = [
("C", model.C, grads["C"], True),
("W1", model.W1, grads["W1"], True),
("W2", model.W2, grads["W2"], True),
("b1", model.b1, grads["b1"], False),
("b2", model.b2, grads["b2"], False),
]
worst = 0.0
failures = []
for _ in range(checks):
name, tensor, grad, two_d = params[int(rng.next() * len(params))]
if two_d:
i = int(rng.next() * len(tensor))
j = int(rng.next() * len(tensor[i]))
original = tensor[i][j]
analytic = 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
else:
i = int(rng.next() * len(tensor))
original = tensor[i]
analytic = grad[i]
tensor[i] = original + eps
plus = model.loss(xs, ys)
tensor[i] = original - eps
minus = model.loss(xs, ys)
tensor[i] = original
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((name, analytic, numeric, rel))
return worst, failures
def generate(model, prompt, n, temp=1.0, seed=11):
"""Sample n characters. Shares the sampling arithmetic from /predict/."""
rng = Rng(seed)
index = {c: i for i, c in enumerate(model.vocab)}
ctx = [index.get(c, 0) for c in ("\n" * model.K + prompt)[-model.K:]]
out = []
for _ in range(n):
probs, _ = model.forward(ctx)
if temp <= 0:
pick = max(range(len(probs)), key=lambda i: probs[i])
else:
scaled = [p ** (1.0 / temp) for p in probs]
total = 0.0
for v in scaled:
total += v
scaled = [v / total for v in scaled]
r = rng.next()
acc = 0.0
pick = len(scaled) - 1
for i, v in enumerate(scaled):
acc += v
if r < acc:
pick = i
break
out.append(model.vocab[pick])
ctx = ctx[1:] + [pick]
return "".join(out)