#!/usr/bin/env python3 """Curriculum-strategy comparison figures (validation set). Produces four panels comparing the three L10 curriculum-management strategies (backtracking vs. two no-backtracking baselines): (A) Stage-wise learning curve: validation accuracy vs training epoch, with the curriculum stage drawn as a step line (shows *how* each method climbs the 10 stages, and where the no-repair baseline deadlocks). (B) Per-hop latent depth-identification accuracy on the validation set (frontier_acc: does latent slot m decode to a correct depth-m node?). (C) Per-hop decoy-arm leakage on the validation set (fraction of latents that decode into the UNREACHABLE arm -- the interpretability failure signature). (D) Per-hop exact-depth accuracy (confusion-matrix diagonal / N): the latent decodes to a node whose TRUE reachable depth == its slot index. Interpretability metrics are computed via logit-lens on the final checkpoint of each arm (node id == token id, so argmax of the LM head at a latent position is a predicted graph node). Usage: CUDA_VISIBLE_DEVICES=7 PYTHONPATH=. python scripts/plot_curriculum_analysis.py """ import argparse import json import os import re import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import torch from transformers import AutoModelForCausalLM, AutoConfig from stokenizer import STokenizer from coconut import Coconut from scripts.probe_latents import build_prefix_tokens, node_depth_maps ARMS = [ # (label, ckpt-dir slug, log slug, color) ("Backtracking (CE-gated)", "backtrack-ce", "backtrack_ce", "#e08214"), ("Backtracking (superposition-gated)", "backtrack-superpos", "backtrack_superpos", "#762a83"), ("Backtracking (frontier-gated)", "backtrack", "backtrack", "#1b7837"), ("Current-stage-only", "curstage", "curstage", "#2166ac"), ("Retention-gated (no repair)", "accstage-nobt", "accstage_nobt", "#b2182b"), ] def latest_ckpt(slug): d = f"ckpts/star-coconut-L10-bfs-{slug}" cks = [f for f in os.listdir(d) if f.startswith("checkpoint_")] cks.sort(key=lambda x: int(x.split("_")[1])) return os.path.join(d, cks[-1]) def parse_log(log_slug): """Return (epochs, val_acc, stage_at_eval, promote_epochs). We align each 'Accuracy on validation set' line to the most recent 'train epoch N/.. stage=S' line preceding it. """ path = f"logs/star_coconut_L10_bfs_{log_slug}.log" ep_re = re.compile(r"train epoch (\d+)/\d+ stage=(\d+)") acc_re = re.compile(r"Accuracy on validation set: \d+ / \d+ = ([0-9.]+)") prom_re = re.compile(r"PROMOTE stage (\d+) -> (\d+)") cur_ep, cur_stage = 0, 0 epochs, accs, stages, promotes = [], [], [], [] with open(path, errors="ignore") as fh: for line in fh: m = ep_re.search(line) if m: cur_ep, cur_stage = int(m.group(1)), int(m.group(2)) continue m = acc_re.search(line) if m: epochs.append(cur_ep) accs.append(float(m.group(1))) stages.append(cur_stage) continue m = prom_re.search(line) if m: promotes.append((cur_ep, int(m.group(2)))) return epochs, accs, stages, promotes @torch.no_grad() def probe(ckpt, val_path, model_id, L, device, batch_size=64): tok = STokenizer() latent_id = tok.convert_tokens_to_ids("<|latent|>") base = AutoModelForCausalLM.from_config(AutoConfig.from_pretrained(model_id)) model = Coconut(base, latent_id, tok.convert_tokens_to_ids("<|start-latent|>"), tok.convert_tokens_to_ids("<|end-latent|>"), tok.eos_token_id) sd = torch.load(ckpt, map_location="cpu") model.load_state_dict(sd, strict=False) model.to(device).eval() data = json.load(open(val_path)) frontier = [0] * (L + 1) # decoded node on correct depth-m frontier (either arm) exact = [0] * (L + 1) # decoded node's TRUE reachable depth == m neg = [0] * (L + 1) # decoded node in the UNREACHABLE arm superpos = [0] * (L + 1) # top-|F| logits == the full depth-m frontier SET total = 0 # confusion: row = latent slot m (1..L), col = TRUE role of decoded node # cols 0..L = reachable depth, L+1 = decoy(neg) arm, L+2 = off-graph NEG, OFF = L + 1, L + 2 conf = [[0] * (L + 3) for _ in range(L + 1)] for i in range(0, len(data), batch_size): batch = data[i:i + batch_size] seqs = [build_prefix_tokens(s, tok) + [latent_id] * L for s in batch] maxlen = max(len(x) for x in seqs) seqs = [x for x in seqs if len(x) == maxlen] # fixed L => equal already input_ids = torch.tensor([build_prefix_tokens(s, tok) + [latent_id] * L for s in batch], device=device) attn = torch.ones_like(input_ids) pos = torch.arange(input_ids.shape[1], device=device).unsqueeze(0).expand(len(batch), -1) logits = model.forward(input_ids, attn, input_ids.clone(), pos).logits for bi, s in enumerate(batch): total += 1 role = node_depth_maps(s, L) root_pos = len(build_prefix_tokens(s, tok)) - 1 for m in range(1, L + 1): slot_logits = logits[bi, root_pos + (m - 1)] pred = int(torch.argmax(slot_logits).item()) F = {int(n) for n in s["neighbor_k"].get(str(m), [])} if F: top = torch.topk(slot_logits, k=len(F)).indices.tolist() if {int(t) for t in top} == F: superpos[m] += 1 if pred in s["neighbor_k"].get(str(m), []): frontier[m] += 1 r = role.get(pred) if r is None: conf[m][OFF] += 1 elif r[0] == "neg": neg[m] += 1 conf[m][NEG] += 1 else: conf[m][r[1]] += 1 if r[1] == m: exact[m] += 1 hops = list(range(1, L + 1)) return { "hops": hops, "total": total, "frontier": [frontier[m] / total for m in hops], "exact": [exact[m] / total for m in hops], "neg": [neg[m] / total for m in hops], "superposition": [superpos[m] / total for m in hops], "conf": conf, } def main(): ap = argparse.ArgumentParser() ap.add_argument("--val", default="data/star_2arm_L10_valid_fo_bfs.json") ap.add_argument("--model_id", default="configs/symbol-2layer-8head-768dim-L20.json") ap.add_argument("--L", type=int, default=10) ap.add_argument("--device", default="cuda:0") ap.add_argument("--out", default="figs/curriculum_analysis.png") ap.add_argument("--cache", default="figs/curriculum_metrics.json") args = ap.parse_args() os.makedirs(os.path.dirname(args.out), exist_ok=True) results = {} for label, ck_slug, log_slug, color in ARMS: ckpt = latest_ckpt(ck_slug) print(f"[{label}] log={log_slug} ckpt={ckpt}") pr = probe(ckpt, args.val, args.model_id, args.L, args.device) ep, acc, stg, prom = parse_log(log_slug) results[label] = {"color": color, "ckpt": ckpt, "probe": pr, "epochs": ep, "val_acc": acc, "stage": stg, "promotes": prom} json.dump(results, open(args.cache, "w"), indent=2) fig, axes = plt.subplots(2, 2, figsize=(15, 11)) axA, axB, axC, axD = axes[0, 0], axes[0, 1], axes[1, 0], axes[1, 1] # (A) stage-wise learning curve for label, *_ , color in ARMS: r = results[label] axA.plot(r["epochs"], r["val_acc"], color=color, lw=1.6, label=label) axA.set_xlabel("training epoch"); axA.set_ylabel("validation accuracy") axA.set_title("(A) Stage-wise learning: val accuracy vs epoch") axA.set_ylim(0, 1); axA.grid(alpha=0.3); axA.legend(fontsize=8, loc="upper left") axA2 = axA.twinx() for label, *_, color in ARMS: r = results[label] axA2.step(r["epochs"], r["stage"], color=color, lw=1.0, ls=":", alpha=0.6, where="post") axA2.set_ylabel("curriculum stage (dotted)") axA2.set_ylim(0, args.L + 0.5) hops = list(range(1, args.L + 1)) # (B) per-hop depth-identification (frontier) accuracy for label, *_, color in ARMS: axB.plot(hops, results[label]["probe"]["frontier"], "-o", color=color, label=label, ms=4) axB.set_xlabel("latent slot / hop m"); axB.set_ylabel("frontier accuracy") axB.set_title("(B) Latent depth-ID accuracy per hop (val)") axB.set_ylim(0, 1.02); axB.grid(alpha=0.3); axB.legend(fontsize=8, loc="lower left") # (C) per-hop decoy-arm leakage for label, *_, color in ARMS: axC.plot(hops, results[label]["probe"]["neg"], "-o", color=color, label=label, ms=4) axC.set_xlabel("latent slot / hop m"); axC.set_ylabel("fraction decoding to decoy arm") axC.set_title("(C) Decoy-arm leakage per hop (val)") axC.grid(alpha=0.3); axC.legend(fontsize=8, loc="upper left") # (D) per-hop SUPERPOSITION accuracy: top-|F| logits == full frontier set for label, *_, color in ARMS: axD.plot(hops, results[label]["probe"]["superposition"], "-o", color=color, label=label, ms=4) axD.set_xlabel("latent slot / hop m"); axD.set_ylabel("superposition accuracy") axD.set_title("(D) Superposition per hop: top-2 == BOTH frontier nodes (val)") axD.set_ylim(0, 1.02); axD.grid(alpha=0.3); axD.legend(fontsize=8, loc="lower left") fig.suptitle("L10 2-arm star: curriculum strategy comparison (validation set)", fontsize=13) fig.tight_layout(rect=[0, 0, 1, 0.98]) fig.savefig(args.out, dpi=150) print(f"\nsaved combined figure -> {args.out}") outdir = os.path.dirname(args.out) # ---- individual panels ---- def save_line(fname, ykey, ylabel, title, ylim=None, loc="lower left"): f, ax = plt.subplots(figsize=(7, 5)) for label, *_, color in ARMS: ax.plot(hops, results[label]["probe"][ykey], "-o", color=color, label=label, ms=5, lw=1.8) ax.set_xlabel("latent slot / hop m"); ax.set_ylabel(ylabel); ax.set_title(title) if ylim: ax.set_ylim(*ylim) ax.grid(alpha=0.3); ax.set_xticks(hops); ax.legend(fontsize=9, loc=loc) f.tight_layout(); p = os.path.join(outdir, fname); f.savefig(p, dpi=150); plt.close(f) print(f"saved -> {p}") # (A) stage-wise learning as its own figure fA, ax = plt.subplots(figsize=(9, 5.5)) for label, *_, color in ARMS: r = results[label] ax.plot(r["epochs"], r["val_acc"], color=color, lw=1.7, label=label) ax.set_xlabel("training epoch"); ax.set_ylabel("validation accuracy") ax.set_title("Stage-wise learning: val accuracy vs epoch (dotted = curriculum stage)") ax.set_ylim(0, 1); ax.grid(alpha=0.3); ax.legend(fontsize=9, loc="upper left") ax2 = ax.twinx() for label, *_, color in ARMS: r = results[label] ax2.step(r["epochs"], r["stage"], color=color, lw=1.1, ls=":", alpha=0.65, where="post") ax2.set_ylabel("curriculum stage (dotted)"); ax2.set_ylim(0, args.L + 0.5) fA.tight_layout(); pA = os.path.join(outdir, "panelA_stagewise_learning.png") fA.savefig(pA, dpi=150); plt.close(fA); print(f"saved -> {pA}") save_line("panelB_depth_id_accuracy.png", "frontier", "frontier accuracy", "Latent depth-ID accuracy per hop (validation)", ylim=(0, 1.02)) save_line("panelC_decoy_leakage.png", "neg", "fraction decoding to decoy arm", "Decoy-arm leakage per hop (validation)", loc="upper left") save_line("panelD_exact_depth.png", "exact", "exact-depth accuracy", "Exact BFS-depth identification per hop (validation)", ylim=(0, 1.02)) save_line("panelE_superposition.png", "superposition", "superposition accuracy", "Superposition per hop: top-2 == BOTH frontier nodes (validation)", ylim=(0, 1.02)) # ---- confusion-matrix heatmaps (one per arm) ---- fh, axes_h = plt.subplots(1, len(ARMS), figsize=(6.3 * len(ARMS), 6)) for ax, (label, *_, color) in zip(axes_h, ARMS): conf = results[label]["probe"]["conf"] L = args.L mat = [[conf[m][c] for c in range(L + 3)] for m in range(1, L + 1)] im = ax.imshow(mat, aspect="auto", cmap="magma") ax.set_xticks(range(L + 3)) ax.set_xticklabels([str(d) for d in range(L + 1)] + ["neg", "off"], fontsize=8) ax.set_yticks(range(L)); ax.set_yticklabels(range(1, L + 1)) ax.set_xlabel("TRUE role of decoded node (reachable depth / decoy / off-graph)") ax.set_ylabel("latent slot m") ax.set_title(label) for mi in range(L): for c in range(L + 3): v = mat[mi][c] if v: ax.text(c, mi, str(v), ha="center", va="center", color="white" if v < results[label]["probe"]["total"] * 0.5 else "black", fontsize=6) fh.colorbar(im, ax=ax, fraction=0.046, pad=0.04) fh.suptitle("Latent confusion matrices: slot m vs. TRUE decoded-node depth (validation, N=256)", fontsize=13) fh.tight_layout(rect=[0, 0, 1, 0.96]) ph = os.path.join(outdir, "confusion_heatmaps.png") fh.savefig(ph, dpi=150); plt.close(fh); print(f"saved -> {ph}") print(f"\nsaved metrics -> {args.cache}") if __name__ == "__main__": main()