| """Model-implied input second moments, built from the weights only and streamed layer by layer. |
| |
| Residual stream: Sig_0 = E_{t~pi}[e_t e_t^T] (+ positional embeddings), pi = Zipf prior on token id |
| (BPE ids are assigned in merge order, i.e. by frequency); then Sig += W^T W for every writer |
| (attention output and MLP output) in order. |
| Norm outputs: LayerNorm: C = P Sig P / (tr/n), cov = g C g, mean = beta, M2 = cov + beta beta^T. |
| RMSNorm: C = Sig / (tr/n), M2 = g C g, mean = g m / s, cov = M2 - mean mean^T. |
| MLP 2nd input: Monte Carlo of the block's own nonlinearity on x ~ N(mean, cov). |
| Attention output: E[o o^T] = mu_o mu_o^T + kappa * blockdiag_heads(C_v of the head's KV group). |
| Every S is normalised to trace n. Design H = I + lam * S. |
| """ |
| import math |
| import torch |
|
|
| F64 = torch.float64 |
|
|
|
|
| class ResidualModel: |
| def __init__(self, arch, reader, device, ctx=1024, zipf_shift=10.0, chunk=8192): |
| self.a, self.r, self.dev = arch, reader, device |
| E = reader.get(arch.embed, device=device, dtype=torch.float32) |
| V = E.shape[0] |
| pi = 1.0 / (torch.arange(V, device=device, dtype=F64) + zipf_shift); pi /= pi.sum() |
| Sig = torch.zeros((arch.d, arch.d), dtype=F64, device=device); m = torch.zeros(arch.d, dtype=F64, device=device) |
| for s in range(0, V, chunk): |
| e = E[s:s + chunk].to(F64); w = pi[s:s + chunk] |
| Sig += (e * w[:, None]).T @ e; m += w @ e |
| del E |
| if arch.pos_embed is not None and reader.has(arch.pos_embed): |
| P = reader.get(arch.pos_embed, device=device, dtype=F64) |
| T = min(ctx, P.shape[0] - arch.pos_offset) |
| P = P[arch.pos_offset: arch.pos_offset + T] |
| pm = P.mean(0) |
| Sig += P.T @ P / T + torch.outer(m, pm) + torch.outer(pm, m); m = m + pm |
| self.Sig, self.m = Sig, m |
|
|
| def add_writer(self, W_out_in): |
| """W in torch layout (out=d, in); the writer adds W h with isotropic h.""" |
| W = W_out_in.to(device=self.dev, dtype=F64) |
| self.Sig += W @ W.T |
|
|
| def norm_moments(self, L, which): |
| a = self.a; n = a.d |
| gname, bname = a.norm_names(L, which) |
| g = self.r.get(gname, device=self.dev, dtype=F64) |
| if a.norm == "ln": |
| Sg = self.Sig - self.Sig.mean(0, keepdim=True) |
| Sg = Sg - Sg.mean(1, keepdim=True) |
| C = Sg / (torch.trace(Sg) / n) |
| cov = g[:, None] * C * g[None, :] |
| b = self.r.get(bname, device=self.dev, dtype=F64) if bname and self.r.has(bname) else torch.zeros(n, dtype=F64, device=self.dev) |
| return b, cov, cov + torch.outer(b, b) |
| s2 = torch.trace(self.Sig) / n |
| M2 = g[:, None] * (self.Sig / s2) * g[None, :] |
| mu = g * self.m / torch.sqrt(s2) |
| w, Q = torch.linalg.eigh(M2 - torch.outer(mu, mu)) |
| cov = (Q * w.clamp(min=0)) @ Q.T |
| return mu, cov, M2 |
|
|
|
|
| def normalise(M): |
| return M.mul_(M.shape[0] / torch.trace(M)) |
|
|
|
|
| def S_norm(moments): |
| return normalise(moments[2].clone()) |
|
|
|
|
| def S_act(arch, moments, W_up, b_up, W_gate, n_samples=16384, batch=2048, seed=0, device="cpu"): |
| """W_up/W_gate: torch layout (inter, d). Returns (inter, inter) float32.""" |
| mu, cov, _ = moments |
| w, Q = torch.linalg.eigh(cov) |
| R = (Q * w.clamp(min=0).sqrt()).to(torch.float32) |
| mu = mu.to(torch.float32) |
| Wu = W_up.to(device=device, dtype=torch.float32).T |
| Wg = None if W_gate is None else W_gate.to(device=device, dtype=torch.float32).T |
| bu = None if b_up is None else b_up.to(device=device, dtype=torch.float32) |
| gen = torch.Generator(device=device).manual_seed(seed) |
| M = torch.zeros((Wu.shape[1], Wu.shape[1]), dtype=torch.float32, device=device); done = 0 |
| while done < n_samples: |
| b = min(batch, n_samples - done) |
| x = mu + torch.randn((b, R.shape[0]), generator=gen, device=device) @ R.T |
| u = x @ Wu |
| if bu is not None: |
| u = u + bu |
| if arch.act == "swiglu": |
| z = x @ Wg; h = z * torch.sigmoid(z) * u |
| elif arch.act == "relu": |
| h = torch.relu(u) |
| elif arch.act == "gelu_new": |
| h = 0.5 * u * (1 + torch.tanh(math.sqrt(2 / math.pi) * (u + 0.044715 * u ** 3))) |
| elif arch.act == "gelu": |
| h = torch.nn.functional.gelu(u) |
| else: |
| raise NotImplementedError(arch.act) |
| M += h.T @ h; done += b |
| return normalise(M.div_(done)) |
|
|
|
|
| def S_attn(arch, moments, W_v, b_v, kappa=1.0): |
| """W_v: torch layout (n_kv*head_dim, d). Returns (n_head*head_dim)^2 float64.""" |
| mu, cov, _ = moments |
| Wv = W_v.to(device=cov.device, dtype=F64) |
| mv = Wv @ mu + (0 if b_v is None else b_v.to(device=cov.device, dtype=F64)) |
| Cv = Wv @ cov @ Wv.T |
| H, dh = arch.n_head, arch.head_dim; per = H // arch.n_kv |
| n = H * dh |
| mo = torch.zeros(n, dtype=F64, device=cov.device); M = torch.zeros((n, n), dtype=F64, device=cov.device) |
| for h in range(H): |
| gk = h // per; s, t = slice(h * dh, (h + 1) * dh), slice(gk * dh, (gk + 1) * dh) |
| mo[s] = mv[t]; M[s, s] = kappa * Cv[t, t] |
| return normalise(M + torch.outer(mo, mo)) |
|
|
