File size: 3,560 Bytes
f31f59b | 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 | """Add words from EXAMPLES, no gradient: a new word's frozen code = sign of the average readout direction the
trained model produces just before the word, over up to K example sentences from the TRAINING text; prior = log of the
word's training frequency. Evaluated on HELD-OUT text: total nats per character over whole passages (words merged where
they occur) vs no new words. Also reports tokens saved."""
import sys, numpy as np, torch, torch.nn.functional as F
import new_words as nw
from transformers import AutoTokenizer
ARM, CKPT = (sys.argv[1], sys.argv[2]) if len(sys.argv) > 2 else ("B", "finals/B/step_07820.pt")
NW, K, CTX = 200, 50, 128
tok = AutoTokenizer.from_pretrained("google/gemma-3-270m")
counts = np.load("data/unigram_counts.npy"); val = np.load("data/val.npy", mmap_mode="r")
tr = np.load("data/train_000.npy", mmap_mode="r")[:100_000_000]
words = nw.corpus_words(tok, val, NW, 5)
wp = [tok(" " + w, add_special_tokens=False).input_ids for w in words]
m, tag = nw.load(ARM, CKPT, counts)
T0, S0 = m.table.clone(), m.seal.detach().clone(); V = T0.shape[0]
def occ(p, arr, limit):
k = len(p); h = np.flatnonzero(arr[:len(arr) - k + 1] == p[0])
for j in range(1, k): h = h[arr[h + j] == p[j]]
return h, len(h)
codes, priors = [], []
with torch.no_grad():
for p in wp:
h, n = occ(p, tr, K)
h = [x for x in h if x >= CTX][:K]
if not h:
v = torch.zeros(512, device="cuda")
else:
ctx = torch.tensor(np.stack([tr[x - CTX:x] for x in h]).astype(np.int64), device="cuda")
with torch.autocast("cuda", dtype=torch.bfloat16):
z = m.out(m.hidden(ctx)[:, -1]).float() # readout direction just before the word
v = z.mean(0)
c = torch.sign(v); c[c == 0] = 1; codes.append(c)
priors.append(np.log((n * 20.5 + 1) / (counts.sum() + 262144))) # 100M-token shard scaled to 2.05B
codes = torch.stack(codes); priors = torch.tensor(priors, dtype=torch.float32, device="cuda")
NSEQ, L = 600, 512
seqs = [np.asarray(val[i * L:(i + 1) * L]).tolist() for i in range(NSEQ)]
chars = sum(len(tok.decode(s)) for s in seqs)
order = sorted(range(len(wp)), key=lambda i: -len(wp[i]))
def merge(s, ids):
out, i = [], 0
while i < len(s):
for j in order:
p = wp[j]
if s[i:i + len(p)] == p: out.append(ids[j]); i += len(p); break
else: out.append(s[i]); i += 1
return out
@torch.no_grad()
def total(seq_list):
t = 0.0
for s in seq_list:
x = torch.tensor([[1] + s], device="cuda")
with torch.autocast("cuda", dtype=torch.bfloat16):
lg = m.logits(m.hidden(x[:, :-1])).float()
t += float(F.cross_entropy(lg[0], x[0, 1:], reduction="sum"))
return t
m.table, m.seal = T0, S0
base = total(seqs)
ids = list(range(V, V + len(wp)))
merged = [merge(s, ids) for s in seqs]
saved = sum(map(len, seqs)) - sum(map(len, merged))
for name, cd, pr in (("examples code + frequency prior", codes, priors),
("piece-mean code + frequency prior", None, priors)):
if cd is None:
E = nw.gemma_embedding(); vv = torch.stack([E[p].mean(0) for p in wp]).cuda()
cd = torch.sign((vv - m.table_mu) @ m.table_W); cd[cd == 0] = 1
m.table = torch.cat([T0, cd]); m.seal = torch.cat([S0, pr])
t = total(merged)
print(f"{tag} | {name}: {t / chars:.4f} vs {base / chars:.4f} nats/char ({(t - base) / chars:+.4f}, {(t - base) / base * 100:+.2f}%) | "
f"{saved} tokens saved of {sum(map(len, seqs))}", flush=True)
|