File size: 5,056 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 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 | """Validation loss split by how often the target word appeared in TRAINING.
Question: does the frozen table win on rare words, where a learned table is starved of examples, even if the average
loss favours the learned table (common words dominate the average)?
python freq_eval.py --arm A --ckpt A/step_07820.pt [--tag A] [--seqs 512]
python freq_eval.py --compare reports/freq_A.json reports/freq_B.json reports/freq_C.json reports/freq_B2.json
Buckets by the target token's count in the 2.05B training tokens (data/unigram_counts.npy):
never seen (0) | 1-9 | 10-99 | 100-999 | 1k-9.9k | 10k-99k | 100k-999k | 1M+
Saved per validation sequence (loss sum and count per bucket), so --compare gives paired 95% ranges (bootstrap over
sequences, 10,000 re-draws) for every arm minus the first file.
"""
import argparse, json, pathlib, time
import numpy as np
import torch
import torch.nn.functional as F
SEQ = 1024
EDGES = [0, 1, 10, 100, 1_000, 10_000, 100_000, 1_000_000, np.inf]
NAMES = ["never seen", "1-9", "10-99", "100-999", "1k-9.9k", "10k-99k", "100k-999k", "1M+"]
def bucket_of(counts):
b = np.searchsorted(np.array(EDGES[1:-1]), counts, side="right") # 0 -> bucket 0 (count 0), 1..9 -> 1, ...
b[counts == 0] = 0
b[(counts > 0) & (b == 0)] = 1
return b
@torch.no_grad()
def run(a):
import train
counts = np.load(pathlib.Path(a.data) / "unigram_counts.npy")
bk = torch.from_numpy(bucket_of(counts)).cuda()
m = train.Model(a.arm, counts).cuda().eval()
st = torch.load(a.ckpt, map_location="cuda", weights_only=False)
m.load_state_dict(st["model"])
val = np.load(pathlib.Path(a.data) / "val.npy", mmap_mode="r")
x = torch.from_numpy(val[:a.seqs * (SEQ + 1)].astype(np.int64).reshape(a.seqs, SEQ + 1))
nb = len(NAMES)
sums = np.zeros((a.seqs, nb)); cnts = np.zeros((a.seqs, nb))
for i in range(a.seqs):
ids = x[i:i + 1].cuda()
with torch.autocast("cuda", dtype=torch.bfloat16):
lg = m.logits(m.hidden(ids[:, :-1])).float()
tgt = ids[0, 1:]
l = F.cross_entropy(lg[0], tgt, reduction="none")
b = bk[tgt]
sums[i] = torch.zeros(nb, device="cuda").scatter_add_(0, b, l).cpu().numpy()
cnts[i] = torch.bincount(b, minlength=nb).float().cpu().numpy()
tag = a.tag or st.get("tag", a.arm)
out = {"tag": tag, "ckpt": a.ckpt, "buckets": NAMES, "sums": sums.tolist(), "counts": cnts.tolist(),
"vocab_per_bucket": np.bincount(bucket_of(counts), minlength=nb).tolist(), "date": time.strftime("%F %H:%M")}
pathlib.Path("reports").mkdir(exist_ok=True)
p = pathlib.Path("reports") / f"freq_{tag}.json"; json.dump(out, open(p, "w"))
print(f"{tag}:", " | ".join(f"{n} {sums[:, j].sum() / max(cnts[:, j].sum(), 1):.3f}" for j, n in enumerate(NAMES)))
print("saved", p)
def compare(files):
rng = np.random.default_rng(0)
R = [json.load(open(f)) for f in files]
ref = R[0]; S0 = np.array(ref["sums"]); C0 = np.array(ref["counts"])
nseq = S0.shape[0]; idx = rng.integers(0, nseq, (10_000, nseq))
L = ["# Validation loss by training frequency of the target word", "",
f"{nseq} held-out sequences x 1,024 tokens. Loss in nats/token, lower is better. Differences vs {ref['tag']}, paired,"
" 95% bootstrap range over sequences.", "",
"| bucket (count in training) | vocab words | share of val tokens | " + " | ".join(r["tag"] for r in R) + " | "
+ " | ".join(f"{r['tag']} - {ref['tag']}" for r in R[1:]) + " |",
"|---|---|---|" + "---|" * (len(R) + len(R) - 1)]
tot = C0.sum()
for j, name in enumerate(ref["buckets"]):
if C0[:, j].sum() == 0:
continue
cells = []
for r in R:
S = np.array(r["sums"]); C = np.array(r["counts"])
cells.append(f"{S[:, j].sum() / C[:, j].sum():.3f}")
diffs = []
for r in R[1:]:
S = np.array(r["sums"]); C = np.array(r["counts"])
d = S[:, j].sum() / C[:, j].sum() - S0[:, j].sum() / C0[:, j].sum()
bs = S[idx, j].sum(1) / np.maximum(C[idx, j].sum(1), 1) - S0[idx, j].sum(1) / np.maximum(C0[idx, j].sum(1), 1)
lo, hi = np.percentile(bs, [2.5, 97.5])
verdict = "tied" if lo <= 0 <= hi else ("better" if d < 0 else "worse")
diffs.append(f"{d:+.3f} ({lo:+.3f} to {hi:+.3f}) {verdict}")
L.append(f"| {name} | {ref['vocab_per_bucket'][j]:,} | {C0[:, j].sum() / tot * 100:.2f}% | " + " | ".join(cells) + " | "
+ " | ".join(diffs) + " |")
open("reports/FREQ_SPLIT.md", "w").write("\n".join(L) + "\n")
print("\n".join(L))
if __name__ == "__main__":
ap = argparse.ArgumentParser()
ap.add_argument("--arm", choices=["A", "B", "C", "D"])
ap.add_argument("--ckpt"); ap.add_argument("--tag", default="")
ap.add_argument("--data", default="data"); ap.add_argument("--seqs", type=int, default=512)
ap.add_argument("--compare", nargs="+")
a = ap.parse_args()
compare(a.compare) if a.compare else run(a)
|