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
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
|
"""Post-run analysis: final metrics, [101,200] probe with sieve diagnostic, grokking signature.
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
import json
import os
import numpy as np
import torch
from src.config import Config
from src.data import build_examples, decode_tokens, encode_int, get_splits, is_prime_n, next_prime, sieve_primes
from src.model_api import build_model, greedy_decode
from src.train import evaluate
# Composites with no prime factor <= 7 inside the probe's candidate window [102, 211].
# A model that learned only the {2,3,5,7} sieve predicts THESE as "next primes" (errors on
# n = 113..120, 139..142, 167..168, 181..186, 199..200 — 22 errors total). Includes 209 = 11*19,
# which the original {121,143,169,187} set (composites <= 200) missed — see Addendum 3.
FLAGGED_SIEVE_PREDS = {121, 143, 169, 187, 209}
def probe_report(model, cfg: Config, lo: int = 101, hi: int = 200) -> dict:
primes = sieve_primes(hi + 200)
correct = 0
errors = []
easy_total = 0
easy_wrong = 0
is_prime_task = cfg.task_mode == "is_prime"
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)
if is_prime_task:
target = 1 if is_prime_n(n) else 0
ok = (pred == 1) == (target == 1) # any non-"1" output reads as "composite"
else:
target = next_prime(n, primes)
ok = pred == target
is_easy = (n % 2 == 0) or (n % 5 == 0) # trivial composites (skip-evens / skip-5s)
if is_easy:
easy_total += 1
if ok:
correct += 1
else:
errors.append({"n": n, "target": target, "pred": pred})
if is_easy:
easy_wrong += 1
total = hi - lo + 1
acc = correct / total
if is_prime_task:
flagged = [e for e in errors if e["n"] in FLAGGED_SIEVE_PREDS] # input IS the classified number
distinct_flagged = len({e["n"] for e in flagged})
else:
flagged = [e for e in errors if e["pred"] in FLAGGED_SIEVE_PREDS]
distinct_flagged = len({e["pred"] for e in flagged})
flagged_frac = len(flagged) / len(errors) if errors else 0.0
# classification per prereg + Addendum 3/5 operationalization
if easy_total and easy_wrong / easy_total > 0.5:
code = "P4" # fails trivial evens/5-multiples -> pure memorization
elif is_prime_task and len(errors) >= 3 and distinct_flagged >= 3 and flagged_frac >= 0.8:
code = "P1" # is-prime: errors concentrated on no-small-factor composites -> learned sieve
elif acc >= 0.85:
code = "P3" # surprising success beyond expectation (next_prime semantics)
elif len(errors) >= 3 and distinct_flagged >= 3 and flagged_frac >= 0.8:
code = "P1" # next_prime: errors = sieves predicting no-small-factor composites
else:
code = "P2" # scattered errors -> memorization / non-transferable heuristics
return {
"code": code, "acc": acc, "correct": correct, "total": total,
"errors": errors, "flagged_errors": flagged,
"flagged_pred_fraction": flagged_frac, "distinct_flagged_preds": distinct_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 (Addendum 2 operationalization)."""
with open(metrics_path) as fh:
rows = list(csv.DictReader(fh))
if not rows:
return {"code": "O4", "note": "no eval rows"}
train = [float(r["train_em"]) for r in rows]
val = [float(r["val_em"]) for r in rows]
n = len(rows)
# 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)
# 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:
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 = "O-PARTIAL" # val ended in (0.3, 0.9) with no 0.9 reach — prereg has no boundary
return {
"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 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):
x = torch.tensor(encode_int(n, cfg), dtype=torch.long).unsqueeze(0)
h = model._encode(x)
_, s = model._run_compute(h)
gaps.append(next_prime(n, primes) - n)
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
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" # 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,
"gaps": gaps, "steps_by_n": 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("--runs-dir", default="runs")
a = ap.parse_args()
out_dir = os.path.join(a.runs_dir, a.model, f"seed{a.seed}")
cfg = Config.load(os.path.join(out_dir, "config.json"))
_, val_in = get_splits(cfg)
val_ex = build_examples(val_in, cfg)
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(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__":
main()
|
