ATOM (atom-ai): code snapshot (100fae0 + batched-train), README with measured numbers, curated checkpoints (part 3)
969c687 verified | """Is the name still IN the field at the answer position? Frozen recall-bench checkpoint, eval mode, | |
| stream of recall episodes (held-out filler half); at the first answer byte collect α (32x32 field) and | |
| fit a ridge one-vs-rest linear probe (name, 16-way) on 1200 episodes, test on 600, per distance bucket. | |
| Chance = 6.25 %. Separates 'memory retains the information' from 'readout learned to use it'. | |
| usage: probe_recall_alpha.py OUT.json CKPT [CKPT ...]""" | |
| import json, random, sys | |
| from pathlib import Path | |
| import numpy as np | |
| import torch | |
| HERE = Path(__file__).resolve().parent | |
| sys.path.insert(0, str(HERE)) | |
| import longmem_bench as LB # noqa: E402 | |
| sys.path.insert(0, str(LB.ROOT / "tools")) | |
| from chat_atom_native import load_model # noqa: E402 | |
| def collect(m, eps): | |
| m.eval(); m.reset_state(reset_atomizer=True); X, y, d = [], [], []; tick = 0 | |
| for text, ans, dist, name in eps: | |
| pk = m.atomizer.encode_bytes(text, reset=False) | |
| for i in range(len(pk)): | |
| m.forward_packet(pk[i]); tick += 1 | |
| if tick % 64 == 0: LB.MB.flush(m) | |
| if i + 1 == ans: | |
| X.append(m.core.state.alpha.detach().reshape(-1).numpy().copy()); y.append(LB.NAMES.index(name)); d.append(dist) | |
| return np.array(X), np.array(y), np.array(d) | |
| def ridge_probe(Xtr, ytr, Xte, lam=1.0): | |
| mu, sd = Xtr.mean(0), Xtr.std(0) + 1e-6 | |
| A, B = (Xtr - mu) / sd, (Xte - mu) / sd | |
| Y = np.eye(16)[ytr] | |
| W = np.linalg.solve(A.T @ A + lam * len(A) * np.eye(A.shape[1]), A.T @ Y) | |
| return (B @ W).argmax(1) | |
| def main(): | |
| torch.set_num_threads(1) | |
| out, cks = sys.argv[1], sys.argv[2:] | |
| src = LB.filler_source(); ev = src[len(src) // 2:] | |
| rng = random.Random(4242) | |
| eps = [LB.episode(rng, ev) for _ in range(1800)] | |
| res = {} | |
| for ck in cks: | |
| if ck.startswith("fresh"): # untrained controls: "fresh" (λ .05) / "fresh_slow" (+ slow block .25) | |
| import argparse | |
| a = argparse.Namespace(seed=20260913, leak=0.05, slow_block=0.25 if ck == "fresh_slow" else 0.0, | |
| slow_cap=0.02, slow_gain="ema", k=0, ladrop=0.5) | |
| m = LB.build(a) | |
| else: | |
| m = load_model(ck) | |
| X, y, d = collect(m, eps) | |
| pred = ridge_probe(X[:1200], y[:1200], X[1200:], lam=1e-2) | |
| yt, dt = y[1200:], d[1200:] | |
| r = {"overall_acc": float((pred == yt).mean())} | |
| for lo, hi, b in ((0, 32, "d<32"), (32, 64, "d32-63"), (64, 128, "d64-127"), (128, 999, "d128+")): | |
| sel = (dt >= lo) & (dt < hi) | |
| r[b] = {"n": int(sel.sum()), "acc": float((pred[sel] == yt[sel]).mean()) if sel.any() else None} | |
| res[ck] = r; print(ck, json.dumps(r), flush=True) | |
| Path(out).write_text(json.dumps(res, indent=1)) | |
| if __name__ == "__main__": | |
| main() | |