diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/config.py | 18 | ||||
| -rw-r--r-- | src/data.py | 28 | ||||
| -rw-r--r-- | src/eval.py | 171 | ||||
| -rw-r--r-- | src/model_api.py | 7 | ||||
| -rw-r--r-- | src/models/rnn.py | 8 | ||||
| -rw-r--r-- | src/train.py | 20 |
6 files changed, 172 insertions, 80 deletions
diff --git a/src/config.py b/src/config.py index c0532bc..0b161e7 100644 --- a/src/config.py +++ b/src/config.py @@ -17,11 +17,10 @@ class Config: d_model: int = 128 max_steps: int = 20 # K tied iterations (RNN) halting: bool = True # False -> fixed-K ablation - halt_eps: float = 0.05 # ACT cumulative threshold (documented; K small so no early break) halt_penalty: float = 0.01 # lambda on mean steps halt_warmup_steps: int = 1000 # penalty = 0 before this (anti-collapse) halt_ramp_end_steps: int = 5000 # penalty ramps linearly 0 -> halt_penalty between warmup and here - min_steps: int = 2 # ACT: halt prob forced to 0 before this step + min_steps: int = 2 # ACT: halt prob forced to 0 for the first min_steps-1 steps n_layers: int = 2 n_heads: int = 4 # training @@ -38,10 +37,19 @@ class Config: @property def vocab(self) -> int: - """Token count. Integers mode: 0..range_end+1 (next prime can exceed range_end).""" + """Token count. Integers mode: value tokens 0..next_prime(range_end), plus EOS (+1) and pad (+1).""" if self.vocab_mode == "integers": - return self.range_end + 2 - return 11 # digits 0-9 + EOS + # inline mini-sieve: next prime above range_end bounds the target domain + limit = self.range_end + 100 + is_prime = [True] * (limit + 1) + is_prime[0] = is_prime[1] = False + for p in range(2, int(limit ** 0.5) + 1): + if is_prime[p]: + for m in range(p * p, limit + 1, p): + is_prime[m] = False + np_ = next(i for i in range(self.range_end + 1, limit + 1) if is_prime[i]) + return np_ + 2 # values 0..np_, EOS=np_+1, pad=np_+2 + return 11 # digits 0-9 + EOS @property def eos_id(self) -> int: diff --git a/src/data.py b/src/data.py index 35682d3..e6ea57b 100644 --- a/src/data.py +++ b/src/data.py @@ -63,15 +63,33 @@ def build_examples(inputs: list[int], cfg: Config) -> list[tuple[list[int], list return out +def _global_lengths(cfg: Config) -> tuple[int, int]: + """(in_max, out_max): fixed global lengths so batch layout == singleton layout (codex BLOCKER fix).""" + primes = sieve_primes(cfg.range_end + 100) + max_target = next_prime(cfg.range_end, primes) + if cfg.vocab_mode == "integers": + return 1, 2 # [value], [value, EOS] + return len(str(cfg.range_end)), len(str(max_target)) + 1 # digits + EOS + + +def pad_inputs(x: torch.Tensor, cfg: Config) -> torch.Tensor: + """LEFT-pad inputs to the global in_max so absolute positions are layout-invariant.""" + in_max, _ = _global_lengths(cfg) + if x.shape[1] < in_max: + pad = torch.full((x.shape[0], in_max - x.shape[1]), cfg.pad_id, dtype=x.dtype, device=x.device) + x = torch.cat([pad, x], dim=1) + return x + + def make_batch(examples, cfg: Config) -> dict[str, torch.Tensor]: + """Fixed global layout (not batch-max): x left-padded to in_max, y right-padded to out_max.""" xs, ys = zip(*examples) - T_in = max(len(x) for x in xs) - T_out = max(len(y) for y in ys) + in_max, out_max = _global_lengths(cfg) B = len(examples) - x = torch.full((B, T_in), cfg.pad_id, dtype=torch.long) - y = torch.full((B, T_out), cfg.pad_id, dtype=torch.long) + x = torch.full((B, in_max), cfg.pad_id, dtype=torch.long) + y = torch.full((B, out_max), cfg.pad_id, dtype=torch.long) for i, (xi, yi) in enumerate(examples): - x[i, : len(xi)] = torch.tensor(xi, dtype=torch.long) + x[i, in_max - len(xi):] = torch.tensor(xi, dtype=torch.long) y[i, : len(yi)] = torch.tensor(yi, dtype=torch.long) # teacher-forced decoder input: BOS(=EOS reuse) then shifted y y_in = torch.cat([torch.full((B, 1), cfg.eos_id, dtype=torch.long), y[:, :-1]], dim=1) 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__": diff --git a/src/model_api.py b/src/model_api.py index 66c2ace..31ffe31 100644 --- a/src/model_api.py +++ b/src/model_api.py @@ -3,6 +3,7 @@ import torch import torch.nn as nn from src.config import Config +from src.data import pad_inputs class PrimeModel(nn.Module): @@ -27,7 +28,11 @@ def build_model(cfg: Config) -> PrimeModel: @torch.no_grad() def greedy_decode(model: PrimeModel, x: torch.Tensor, cfg: Config, max_len: int | None = None) -> torch.Tensor: - """Autoregressive greedy decode of output digits. Returns (B, max_len) tokens (BOS stripped).""" + """Autoregressive greedy decode of output digits. Returns (B, max_len) tokens (BOS stripped). + + Inputs are LEFT-padded to the global layout so positions match training exactly + (batch/singleton invariance — codex BLOCKER fix).""" + x = pad_inputs(x, cfg) max_len = max_len or cfg.max_out_len B = x.shape[0] y_in = torch.full((B, 1), cfg.eos_id, dtype=torch.long, device=x.device) diff --git a/src/models/rnn.py b/src/models/rnn.py index 19123b7..44f0a0d 100644 --- a/src/models/rnn.py +++ b/src/models/rnn.py @@ -45,7 +45,8 @@ class TiedRNN(PrimeModel): return pe def _encode(self, x: torch.Tensor) -> torch.Tensor: - """Read input digits through the tied cell (pad positions inject nothing).""" + """Read input digits through the tied cell. Pad positions are exact no-ops + (state update masked) so batch composition cannot change an example's state.""" B, T = x.shape d = self.cfg.d_model pos = self._sinusoidal(T, d).to(x.device) # (T,d) @@ -53,7 +54,8 @@ class TiedRNN(PrimeModel): mask = (x != self.cfg.pad_id).float().unsqueeze(-1) # (B,T,1) h = torch.zeros(B, d, device=x.device) for t in range(T): - h = self._cell_step(h + (e[:, t] + pos[t]) * mask[:, t]) + h_new = self._cell_step(h + (e[:, t] + pos[t]) * mask[:, t]) + h = mask[:, t] * h_new + (1 - mask[:, t]) * h # pad step = no-op return h def _run_compute(self, h0: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: @@ -72,7 +74,7 @@ class TiedRNN(PrimeModel): for t in range(cfg.max_steps): h = self._cell_step(h) p = torch.sigmoid(self.halt_head(self.ln1(h))).squeeze(-1) # (B,) - if t < cfg.min_steps: + if t < cfg.min_steps - 1: p = p * 0.0 h_list.append(h) p_list.append(p) diff --git a/src/train.py b/src/train.py index 975598d..46bd8e7 100644 --- a/src/train.py +++ b/src/train.py @@ -1,8 +1,10 @@ """Training loop. Usage: python -m src.train [model] [seed] [--flag ...]""" import csv +import json import math import os import random +import sys import numpy as np import torch @@ -62,8 +64,21 @@ def main() -> None: cfg = parse_args() set_seed(cfg.seed) out_dir = os.path.join(cfg.out_dir, cfg.model, f"seed{cfg.seed}") + csv_path = os.path.join(out_dir, "metrics.csv") + if os.path.exists(csv_path): + raise SystemExit(f"REFUSING to rerun in place: {csv_path} exists. Use a fresh --out_dir " + f"(reruns would corrupt the CSV and checkpoint provenance).") os.makedirs(out_dir, exist_ok=True) cfg.save(os.path.join(out_dir, "config.json")) + with open(os.path.join(out_dir, "run_meta.json"), "w") as fh: + json.dump({ + "python": sys.version.split()[0], + "torch": torch.__version__, + "numpy": np.__version__, + "device": "cpu", + "torch_threads": torch.get_num_threads(), + "cmd": sys.argv, + }, fh, indent=2) train_in, val_in = get_splits(cfg) train_ex = build_examples(train_in, cfg) @@ -75,9 +90,8 @@ def main() -> None: opt = torch.optim.AdamW(model.parameters(), lr=cfg.lr, weight_decay=cfg.weight_decay) ce = nn.CrossEntropyLoss(reduction="none") - log_examples = sorted(val_ex, key=lambda x: (len(str(x)), x))[: cfg.log_n_examples] + log_examples = sorted(val_ex, key=lambda ex: (len(ex[0]), ex[0]))[: cfg.log_n_examples] - csv_path = os.path.join(out_dir, "metrics.csv") fieldnames = ["step", "train_loss", "train_token_acc", "train_em", "val_token_acc", "val_em", "mean_halt_steps", "log_examples", "param_count"] best_val_em = -1.0 @@ -152,6 +166,8 @@ def main() -> None: _eval_pass(loss.detach(), halt.detach() if halt is not None else None) if done: break + if step >= cfg.max_train_steps: + break torch.save(model.state_dict(), os.path.join(out_dir, "last.pt")) print(f"DONE steps={step} best_val_em={best_val_em:.4f}") |