File size: 5,204 Bytes
c675b8d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 | """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) # P Sig P
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))
|