summaryrefslogtreecommitdiff
path: root/src/eval.py
diff options
context:
space:
mode:
Diffstat (limited to 'src/eval.py')
-rw-r--r--src/eval.py171
1 files changed, 107 insertions, 64 deletions
diff --git a/src/eval.py b/src/eval.py
index 4485523..fdc2dec 100644
--- a/src/eval.py
+++ b/src/eval.py
@@ -1,7 +1,9 @@
"""Post-run analysis: final metrics, [101,200] probe with sieve diagnostic, grokking signature.
-Interpretation codes are locked in design/preregistration.md — this module only MEASURES
-and classifies against those definitions (O1-O4, H1-H4, P1-P4).
+Interpretation codes are locked in design/preregistration.md (with operationalization
+thresholds in Addendum 2). This module MEASURES and classifies against those definitions.
+Where the prereg prose supplies no testable boundary, the code emits the measurements
+plus "unclassified" rather than inventing a label.
"""
import argparse
import csv
@@ -23,38 +25,44 @@ def probe_report(model, cfg: Config, lo: int = 101, hi: int = 200) -> dict:
primes = sieve_primes(hi + 200)
correct = 0
errors = []
- easy_misses = 0
+ easy_total = 0
+ easy_wrong = 0
for n in range(lo, hi + 1):
x = torch.tensor(encode_int(n, cfg), dtype=torch.long).unsqueeze(0)
gen = greedy_decode(model, x, cfg)[0].tolist()
pred = decode_tokens(gen, cfg)
target = next_prime(n, primes)
+ is_easy = (n % 2 == 0) or (n % 5 == 0) # trivial composites (skip-evens / skip-5s)
+ if is_easy:
+ easy_total += 1
if pred == target:
correct += 1
else:
errors.append({"n": n, "target": target, "pred": pred})
- if n % 2 == 0 or n % 5 == 0:
- easy_misses += 1
+ if is_easy:
+ easy_wrong += 1
total = hi - lo + 1
acc = correct / total
flagged = [e for e in errors if e["n"] in FLAGGED_COMPOSITES]
- # P-code classification (see preregistration.md)
- if acc >= 0.85:
- code = "P3"
- elif easy_misses >= 5:
- code = "P4"
- elif errors and all(e["n"] in FLAGGED_COMPOSITES for e in errors) and len(errors) <= 6:
- code = "P1"
+ # classification per prereg + Addendum 2 operationalization; P4 checked first
+ if easy_total and easy_wrong / easy_total > 0.5:
+ code = "P4" # fails trivial evens/5-multiples -> pure memorization
+ elif acc >= 0.85:
+ code = "P3" # surprising success beyond expectation
+ elif len(flagged) >= 3 and len(errors) <= 6 and all(e["n"] in FLAGGED_COMPOSITES for e in errors):
+ code = "P1" # errors concentrated on composites needing divisors 11,13 -> learned sieve
else:
- code = "P2"
+ code = "P2" # scattered errors -> memorization / non-transferable heuristics
return {
"code": code, "acc": acc, "correct": correct, "total": total,
- "errors": errors, "flagged_errors": flagged, "easy_misses": easy_misses,
+ "errors": errors, "flagged_errors": flagged,
+ "easy_total": easy_total, "easy_wrong": easy_wrong,
+ "easy_err_rate": (easy_wrong / easy_total) if easy_total else None,
}
def grokking_signature(metrics_path: str) -> dict:
- """Classify the training curve against preregistered codes O1-O4."""
+ """Classify the training curve against preregistered codes O1-O4 (Addendum 2 operationalization)."""
with open(metrics_path) as fh:
rows = list(csv.DictReader(fh))
if not rows:
@@ -62,31 +70,41 @@ def grokking_signature(metrics_path: str) -> dict:
train = [float(r["train_em"]) for r in rows]
val = [float(r["val_em"]) for r in rows]
n = len(rows)
- saturated = any(all(t >= 0.95 for t in train[i:i + 10]) for i in range(n - 9)) if n >= 10 else False
+ # first index where train EM >= 0.95 for 10 CONSECUTIVE evals
+ sat_start = next((i for i in range(n - 9) if all(t >= 0.95 for t in train[i:i + 10])), None)
hi = next((i for i, v in enumerate(val) if v >= 0.9), None)
- trans = None
- if hi is not None:
- lo_cands = [i for i in range(hi) if val[i] <= 0.2]
+ # transition: last eval with val <= 0.2 strictly before hi, and after saturation window
+ lo = None
+ if hi is not None and sat_start is not None:
+ lo_cands = [i for i in range(sat_start + 10, hi) if val[i] <= 0.2]
if lo_cands:
- trans = hi - max(lo_cands)
- if saturated and hi is not None and trans is not None and trans <= 5:
- code = "O1"
- elif max(train) >= 0.95 and hi is not None and (trans is None or trans > 5):
- code = "O3"
- elif max(train) >= 0.95 and hi is None:
- code = "O2"
+ lo = max(lo_cands)
+ width = (hi - lo) if (hi is not None and lo is not None) else None
+
+ max_train = max(train)
+ if max_train < 0.95:
+ code = "O4" # train never reached 0.95 -> setup/optimization failure
+ elif sat_start is None:
+ code = "O4" # reached 0.95 but never sustained 10 evals within budget
+ elif hi is not None and width is not None and width <= 5:
+ code = "O1" # sharp transition AFTER sustained train saturation
+ elif hi is not None:
+ code = "O3" # reached 0.9+ but not via the sharp O1 pattern (gradual)
+ elif val[-1] <= 0.3:
+ code = "O2" # train memorized, val stayed low
else:
- code = "O4"
+ code = "O-PARTIAL" # val ended in (0.3, 0.9) with no 0.9 reach — prereg has no boundary
return {
- "code": code, "train_saturated": saturated, "val_hi_eval_idx": hi,
- "transition_width_evals": trans, "n_evals": n,
+ "code": code, "sat_start_eval": sat_start, "val_hi_eval_idx": hi,
+ "transition_width_evals": width, "n_evals": n,
"final_train_em": train[-1], "final_val_em": val[-1],
+ "max_train_em": max_train,
}
@torch.no_grad()
def halting_report(model, cfg: Config) -> dict:
- """RNN halting structure: mean steps + correlation with gap-to-next-prime (H1-H4)."""
+ """RNN halting structure: mean steps at run end + correlation with gap-to-next-prime (H1-H4)."""
primes = sieve_primes(300)
gaps, steps = [], []
for n in range(cfg.range_start, cfg.range_end + 1):
@@ -97,54 +115,79 @@ def halting_report(model, cfg: Config) -> dict:
steps.append(float(s.mean()))
mean = float(np.mean(steps))
rho = float(np.corrcoef(gaps, steps)[0, 1]) if len(set(gaps)) > 1 else 0.0
- if mean <= cfg.min_steps + 0.5:
- code = "H1"
- elif mean >= cfg.max_steps - 0.5:
- code = "H2"
- elif abs(rho) >= 0.3:
- code = "H3"
+ lo, hi = cfg.min_steps + 0.5, cfg.max_steps - 0.5
+ if mean <= lo:
+ code = "H1" # collapse to the min_steps floor
+ elif mean >= hi:
+ code = "H2" # pinned at K — never learned to halt
+ elif rho >= 0.3:
+ code = "H3" # intermediate + positive gap correlation -> computation budget
else:
- code = "H4"
- return {"code": code, "mean_steps": mean, "corr_gap": rho, "min_steps": cfg.min_steps, "max_steps": cfg.max_steps}
+ code = "H4" # intermediate but no positive gap correlation — noisy/unused
+ return {"code": code, "mean_steps": mean, "corr_gap": rho,
+ "min_steps": cfg.min_steps, "max_steps": cfg.max_steps}
+
+
+def _load(out_dir: str, ckpt: str, cfg: Config):
+ model = build_model(cfg)
+ model.load_state_dict(torch.load(os.path.join(out_dir, ckpt), map_location="cpu"))
+ model.eval()
+ return model
def main() -> None:
ap = argparse.ArgumentParser(description="prime-grokking eval")
ap.add_argument("model")
ap.add_argument("seed")
- ap.add_argument("--ckpt", default="best.pt")
a = ap.parse_args()
out_dir = os.path.join("runs", a.model, f"seed{a.seed}")
cfg = Config.load(os.path.join(out_dir, "config.json"))
- model = build_model(cfg)
- model.load_state_dict(torch.load(os.path.join(out_dir, a.ckpt), map_location="cpu"))
- model.eval()
_, val_in = get_splits(cfg)
val_ex = build_examples(val_in, cfg)
- tok, em, per = evaluate(model, val_ex, cfg)
-
- probe = probe_report(model, cfg)
- sig = grokking_signature(os.path.join(out_dir, "metrics.csv"))
- hlt = halting_report(model, cfg) if cfg.model == "rnn" else None
-
- results = {
- "model": cfg.model, "seed": cfg.seed, "ckpt": a.ckpt,
- "params": model.param_count(),
- "val_token_acc": tok, "val_exact_match": em,
- "per_example": [{"n": "".join(map(str, x)), "target": t, "pred": p, "ok": ok}
- for x, t, p, ok in per],
- "probe": probe,
- "signature": sig,
- "halting": hlt,
- }
+
+ report = {"model": cfg.model, "seed": cfg.seed, "params": None, "val_selected": {}, "final": {}}
+ for ckpt in ("best.pt", "last.pt"):
+ path = os.path.join(out_dir, ckpt)
+ if not os.path.exists(path):
+ continue
+ model = _load(out_dir, ckpt, cfg)
+ tok, em, per = evaluate(model, val_ex, cfg)
+ entry = {
+ "ckpt": ckpt,
+ "params": model.param_count(),
+ "val_token_acc": tok,
+ "val_exact_match": em,
+ "note": ("val-selected checkpoint (early stop / best-on-val per prereg — "
+ "val EM here is selection-holed, not an untouched test estimate" if ckpt == "best.pt"
+ else "final checkpoint, no val selection"),
+ }
+ if ckpt == "last.pt":
+ entry["probe"] = probe_report(model, cfg)
+ if cfg.model == "rnn":
+ entry["halting"] = halting_report(model, cfg)
+ entry["per_example"] = [{"n": "".join(map(str, x)), "target": t, "pred": p, "ok": ok}
+ for x, t, p, ok in per]
+ if ckpt == "best.pt":
+ report["val_selected"] = entry
+ else:
+ report["final"] = entry
+
+ report["signature"] = grokking_signature(os.path.join(out_dir, "metrics.csv"))
with open(os.path.join(out_dir, "results.json"), "w") as fh:
- json.dump(results, fh, indent=2)
- print(json.dumps({"val_token_acc": round(tok, 4), "val_em": round(em, 4),
- "probe": probe["code"], "probe_acc": round(probe["acc"], 3),
- "flagged_errors": [e["n"] for e in probe["flagged_errors"]],
- "signature": sig["code"], "halting": hlt["code"] if hlt else None,
- "results": os.path.join(out_dir, "results.json")}, indent=2))
+ json.dump(report, fh, indent=2)
+
+ f = report["final"]
+ p = f.get("probe", {})
+ print(json.dumps({
+ "val_em_best": round(report["val_selected"].get("val_exact_match", float("nan")), 4),
+ "val_em_last": round(f.get("val_exact_match", float("nan")), 4),
+ "probe_code": p.get("code"), "probe_acc": round(p.get("acc", float("nan")), 3),
+ "flagged_errors": [e["n"] for e in p.get("flagged_errors", [])],
+ "signature": report["signature"]["code"],
+ "halting": (f.get("halting") or {}).get("code"),
+ "results": os.path.join(out_dir, "results.json"),
+ }, indent=2))
if __name__ == "__main__":
generated by cgit v1.2.3 (git 2.46.0) at 2026-08-31 21:11:08 +0000