| """Shared helpers for the M1 notebooks (notes 0101-0110). |
| |
| Everything is fp64 CPU torch. The "language model" is a deliberately tiny |
| autoregressive policy whose *entire* sequence space can be enumerated, so every |
| expectation in the notes can be computed EXACTLY (no Monte Carlo) and compared |
| against the hand-derived formula to ~1e-15. Monte Carlo checks are kept for the |
| few identities that are about sampling (and are reported against 4 standard errors). |
| |
| Policy: pi_theta(a | prev, t) = softmax(theta_tok[prev] + theta_pos[t])_a |
| - vocabulary {0,..,V-1} where token V-1 is EOS |
| - `prev` is the previous token (BOS = index V for t=0) |
| - a sequence stops at EOS or after L tokens (truncation, no EOS) |
| This is a valid autoregressive factorisation pi(o) = prod_t pi(o_t | o_<t) |
| (with a Markov-1 context) so every identity about the general case holds verbatim. |
| """ |
| import os, sys, math, itertools |
| os.environ.setdefault("CUDA_VISIBLE_DEVICES", "") |
| import torch |
| torch.set_default_dtype(torch.float64) |
| torch.set_num_threads(2) |
|
|
| CHECKS = [] |
|
|
| def check(name, err, tol, kind="exact"): |
| """Record + assert a numerical identity. kind='exact' (fp64 arithmetic) or 'mc' (4-sigma Monte Carlo).""" |
| err = float(err) |
| ok = err <= tol |
| CHECKS.append((name, err, tol, kind, ok)) |
| print(f"CHECK|{kind}|{name}|{err:.3e}|tol={tol:.1e}|{'PASS' if ok else 'FAIL'}") |
| assert ok, f"{name}: err {err:.3e} > tol {tol:.1e}" |
|
|
| def maxabs(a, b): |
| return float((torch.as_tensor(a) - torch.as_tensor(b)).abs().max()) |
|
|
| class TinyLM: |
| def __init__(self, V=4, L=4, seed=0, scale=1.0): |
| self.V, self.L = V, L |
| self.EOS = V - 1 |
| g = torch.Generator().manual_seed(seed) |
| self.n_tok = (V + 1) * V |
| self.n_pos = L * V |
| self.theta0 = scale * torch.randn(self.n_tok + self.n_pos, generator=g) |
| self.seqs, self.tokens, self.mask = self._enumerate() |
| self.N = len(self.seqs) |
| self.lengths = self.mask.sum(1) |
|
|
| |
| def _enumerate(self): |
| V, L, EOS = self.V, self.L, self.EOS |
| seqs = [] |
| def rec(prefix): |
| if len(prefix) == L or (prefix and prefix[-1] == EOS): |
| seqs.append(tuple(prefix)); return |
| for a in range(V): |
| rec(prefix + [a]) |
| rec([]) |
| tokens = torch.full((len(seqs), L), EOS, dtype=torch.long) |
| mask = torch.zeros(len(seqs), L) |
| for i, s in enumerate(seqs): |
| tokens[i, :len(s)] = torch.tensor(s) |
| mask[i, :len(s)] = 1.0 |
| return seqs, tokens, mask |
|
|
| def _logits_tables(self, theta): |
| V, L = self.V, self.L |
| tok = theta[:self.n_tok].reshape(V + 1, V) |
| pos = theta[self.n_tok:].reshape(L, V) |
| return tok, pos |
|
|
| def token_logprobs(self, theta): |
| """[N, L] log pi(o_t | o_<t) for every enumerated sequence (0 where masked).""" |
| V, L = self.V, self.L |
| tok, pos = self._logits_tables(theta) |
| prev = torch.cat([torch.full((self.N, 1), V, dtype=torch.long), self.tokens[:, :-1]], 1) |
| logits = tok[prev] + pos[None, :, :] |
| lp = torch.log_softmax(logits, -1) |
| lp = lp.gather(-1, self.tokens[..., None]).squeeze(-1) |
| return lp * self.mask |
|
|
| def seq_logprob(self, theta): |
| return self.token_logprobs(theta).sum(1) |
|
|
| def probs(self, theta): |
| return self.seq_logprob(theta).exp() |
|
|
| |
| def seq_logprob_naive(self, theta, seq): |
| V = self.V |
| tok, pos = self._logits_tables(theta) |
| lp, prev = 0.0, V |
| for t, a in enumerate(seq): |
| logits = tok[prev] + pos[t] |
| z = logits - logits.max() |
| lp = lp + (z[a] - torch.log(torch.exp(z).sum())) |
| prev = a |
| return lp |
|
|
| |
| def prefixes(self): |
| """All non-terminal prefixes (the MDP states), shortest first.""" |
| out = [()] |
| for k in range(1, self.L): |
| out += [p for p in itertools.product(range(self.V - 1), repeat=k)] |
| return out |
|
|
| def cond_probs(self, theta, prefix): |
| V = self.V |
| tok, pos = self._logits_tables(theta) |
| prev = V if len(prefix) == 0 else prefix[-1] |
| return torch.softmax(tok[prev] + pos[len(prefix)], -1) |
|
|
| def default_reward(lm): |
| """Binary verifier-style reward: success iff the sequence ends in EOS and its token sum is odd and length>=2.""" |
| R = torch.zeros(lm.N) |
| for i, s in enumerate(lm.seqs): |
| if s[-1] == lm.EOS and len(s) >= 2 and (sum(s[:-1]) % 2 == 1): |
| R[i] = 1.0 |
| return R |
|
|
| def flat_grad(scalar, theta): |
| g, = torch.autograd.grad(scalar, theta, retain_graph=True, allow_unused=False) |
| return g |
|
|
| def score_matrix(lm, theta): |
| """S[o] = grad_theta log pi(o) as an [N, P] matrix (autograd, one backward per sequence).""" |
| th = theta.clone().requires_grad_(True) |
| lp = lm.seq_logprob(th) |
| rows = [torch.autograd.grad(lp[i], th, retain_graph=True)[0] for i in range(lm.N)] |
| return torch.stack(rows) |
|
|
| def token_score_tensor(lm, theta): |
| """T[o, t] = grad_theta log pi(o_t | o_<t) as [N, L, P].""" |
| th = theta.clone().requires_grad_(True) |
| lp = lm.token_logprobs(th) |
| P = theta.numel() |
| out = torch.zeros(lm.N, lm.L, P) |
| for i in range(lm.N): |
| for t in range(int(lm.lengths[i])): |
| out[i, t] = torch.autograd.grad(lp[i, t], th, retain_graph=True)[0] |
| return out |
|
|
| def perturb(theta, delta, seed=1): |
| g = torch.Generator().manual_seed(seed) |
| d = torch.randn(theta.numel(), generator=g) |
| return theta + delta * d / d.norm() |
|
|
