"""Main manuscript tables and headline values, recomputed from the artifact. Run this and diff against the draft. Validation-checkpoint epochs and epoch-matched shares are emitted by the companion ``checkpoint_selection_audit.py``. python scripts/paper_numbers.py # human-readable python scripts/paper_numbers.py --json # machine-readable Sources (all under ../drift_run unless noted): drift_out/drift_by_stage.csv per (chain, layer, stage) drift, 3 inner products drift_out/drift_per_feature.csv per-latent drift, density, persistence drift_out/base_feature_z.csv per-latent logit-lens peakedness at the base SAE drift_out/null_replicates.csv null-step replicates + real first transitions drift_out/null_trajectory_l18.csv the originally reported 7-step null trajectory drift_out/null_trajectory_l18_seeds.csv the same, 4 seeds (supersedes the above) null_chain_all/layer{L}_seed{S}/ per-layer null trajectories, 3 seeds each null_budget/ epoch-budget sweep + budget-matched trajectories arch_null/arch_null_trajectory.csv TopK / BatchTopK / JumpReLU at matched recipe saes_local/, null_ctrl/, null_seeds/ SAE weights (for cosines and hashes) ../sae_rl/results/sae_checkpoint_metrics.csv reconstruction fidelity ../sae_rl/results/task_performance.csv GSM8K accuracy ../sae_rl/data/activations*/ activation-set sizes Sections needing SAE weights (direction, budget, trajectories, alignment, persistent null) run on CPU by default and are skipped with a warning if the weights are absent, so the script still works from the CSVs alone. """ import argparse import glob import json import os import pathlib import re import pandas as pd from scipy.stats import spearmanr DEV = "cpu" HERE = pathlib.Path(__file__).resolve().parent def _find(*parts): """Locate a data directory whether this paper tree sits inside sae_rl/ or beside it. The paper moved out of sae_rl/paper/ on 2026-08-20, so the data root is one level further away than it used to be. Try both rather than pinning either. """ for base in (HERE.parent.parent, HERE.parent.parent / "sae_rl"): cand = base.joinpath(*parts) if cand.exists(): return cand return HERE.parent.parent.joinpath(*parts) DRIFT = _find("drift_run", "drift_out") RESULTS = _find("sae_rl", "results") ACC = "val-core/openai/gsm8k/acc/mean@1" # Chains sharing a bit-identical base SAE, so cross-chain latent indices are # comparable. Verified by hashing W_dec; see base_alignment(). ALIGNED = ["sae_sft", "sae_flexible", "sae_strict/k64"] def _round(o, n=4): if isinstance(o, dict): return {k: _round(v, n) for k, v in o.items()} if isinstance(o, list): return [_round(v, n) for v in o] return round(o, n) if isinstance(o, float) else o def artifact(drift): """Size of the artifact: SAEs, chains, layers.""" return { "n_saes": int(len(drift)), "n_chains": int(drift.chain.nunique()), "layers": sorted(drift.layer.unique().tolist()), "chain_n_checkpoints": {c: int(g.stage.nunique()) for c, g in drift.groupby("chain")}, "d_sae": int(drift.n_features.iloc[0]), } def identity_vs_fidelity(drift, metrics): """Inherited-index alignment turns over while reconstruction stays flat.""" out = {} fin = drift[(drift.chain == "sae_sft") & (drift.step == 348)] for _, r in fin.iterrows(): out[f"L{int(r.layer)}"] = { "dec_cos_geom": r.dec_geom_all_cos_matched_mean, "dec_cos_readout": r.dec_readout_all_cos_matched_mean, "dec_cos_lens": r.dec_lens_all_cos_matched_mean, "tok_jaccard": r.tok_jaccard_matched_all, "frac_top1_identical": r.frac_identical_top1_all, "frac_top1_changed": 1 - r.frac_identical_top1_all, } f = metrics[(metrics.chain == "sft") & (metrics.layer == 18)] out["L18_fidelity_over_sft"] = { "nmse_min": f.nmse_recomputed.min(), "nmse_max": f.nmse_recomputed.max(), "frac_rec_min": f.frac_rec.min(), "frac_rec_max": f.frac_rec.max(), "dead_pct_max": f.dead_latents_pct.max(), } out["dead_pct_max_all_chains"] = metrics.dead_latents_pct.max() return out def metric_agreement(drift): """The three inner products must agree, else the geometry story is metric-dependent.""" cols = ["dec_geom_all_drift_norm", "dec_readout_all_drift_norm", "dec_lens_all_drift_norm"] d = drift[drift.step > 0].copy() d["spread"] = d[cols].max(axis=1) - d[cols].min(axis=1) per_cell = d.groupby(["chain", "layer"]).spread.max() return { "n_cells": int(len(per_cell)), "max_spread": per_cell.max(), "median_spread": per_cell.median(), "argmax_cell": "/".join(str(x) for x in per_cell.idxmax()), } def null_step(reps): """How much of a real first transition a warm-started retraining reproduces.""" nulls, out = {}, {} for L, g in reps[reps.kind == "nullstep"].groupby("layer"): nulls[L] = {"n": int(len(g)), "drift_mean": g.drift.mean(), "drift_min": g.drift.min(), "drift_max": g.drift.max(), "tokJ_mean": g.tok_jaccard.mean()} shares = [] for _, r in reps[reps.kind == "real"].iterrows(): n = nulls[r.layer] shares.append({ "layer": int(r.layer), "chain": r.label, "real_drift": r.drift, "null_drift_mean": n["drift_mean"], "pct_mean": 100 * n["drift_mean"] / r.drift, "pct_lo": 100 * n["drift_min"] / r.drift, "pct_hi": 100 * n["drift_max"] / r.drift, # top-token turnover share: (1 - J_null) / (1 - J_real) "tok_pct": 100 * (1 - n["tokJ_mean"]) / (1 - r.tok_jaccard), }) out["null_replicates"] = {f"L{k}": v for k, v in nulls.items()} out["share_of_real"] = shares # Headline range, excluding L12 (null replicates there are unstable: see # drift_min/drift_max) and excluding cells where the real transition drifts # *less* than the null step (pct > 100). core = [s for s in shares if s["layer"] != 12 and s["pct_mean"] <= 110] out["headline_pct_range"] = [min(s["pct_mean"] for s in core), max(s["pct_mean"] for s in core)] out["headline_tok_pct_range"] = [min(s["tok_pct"] for s in core), max(s["tok_pct"] for s in core)] out["exceeds_100"] = [f"{s['chain']}@L{s['layer']}={s['pct_mean']:.0f}%" for s in shares if s["pct_mean"] > 110] return out def null_trajectory(traj, drift): """A whole trajectory of retrainings, with the model held fixed throughout.""" real = (drift[(drift.chain == "sae_sft") & (drift.layer == 18)] .sort_values("step").reset_index(drop=True)) rows = [] for i, r in traj.sort_values("null_step").iterrows(): rr = real.iloc[i] if i < len(real) else None rows.append({ "step_index": int(r.null_step), "null_dec_cos": r.dec_cos, "null_tok_jaccard": r.tok_jaccard, "real_step": int(rr.step) if rr is not None else None, "real_dec_cos": rr.dec_geom_all_cos_matched_mean if rr is not None else None, "real_tok_jaccard": rr.tok_jaccard_matched_all if rr is not None else None, }) end = rows[-1] return { "per_step": rows, "final_null_cos": end["null_dec_cos"], "final_real_cos": end["real_dec_cos"], "final_null_tokJ": end["null_tok_jaccard"], "final_real_tokJ": end["real_tok_jaccard"], # Fraction of the real end-to-end drift reproduced with the model frozen. "pct_drift_reproduced": 100 * (1 - end["null_dec_cos"]) / (1 - end["real_dec_cos"]), "pct_tok_turnover_reproduced": 100 * (1 - end["null_tok_jaccard"]) / (1 - end["real_tok_jaccard"]), } def shuffled_control(drift, task): """A reward carrying no task information, at a matched optimiser step.""" out = {"accuracy": {}} for c, g in task.groupby("chain"): g = g.sort_values("step") out["accuracy"][c] = {"start": g[ACC].iloc[0], "final": g[ACC].iloc[-1], "max": g[ACC].max()} for step in (10, 30): p = (drift[drift.step == step] .pivot_table(index="chain", columns="layer", values="dec_geom_all_drift_norm")) cell = {} for ref in ("sae_flexible", "sae_strict/k64"): ratios = (p.loc["sae_shuffled"] / p.loc[ref]) cell[ref] = {f"L{int(k)}": v for k, v in ratios.items()} out[f"step{step}_shuffled_over_informative"] = cell out[f"step{step}_drift_norm"] = { ch: {f"L{int(k)}": v for k, v in row.items()} for ch, row in p.iterrows()} return out def heterogeneity(per_feat, z): """Stability is not uniform, and our selection heuristic rides the tail.""" d = per_feat.merge(z, on=["chain", "layer", "feature"]) out = {"spearman_final_sft": {}, "selection_bias_final_sft": {}} for L in (6, 12, 18, 23): s = d[(d.chain == "sae_sft") & (d.layer == L) & (d.step == 348)] out["spearman_final_sft"][f"L{L}"] = { "cos_vs_density": spearmanr(s.dec_geom_cos, s.density_base).statistic, "cos_vs_peakedness": spearmanr(s.dec_geom_cos, s.logit_z).statistic, "tokJ_vs_density": spearmanr(s.tok_jaccard, s.density_base).statistic, "tokJ_vs_peakedness": spearmanr(s.tok_jaccard, s.logit_z).statistic, } # The sweep's own heuristic: frequency band, then top-n by peakedness. pool = s[(s.density_base > 0.05) & (s.density_base < 0.5)] sel = pool.nlargest(12, "logit_z") out["selection_bias_final_sft"][f"L{L}"] = { "n_pool": int(len(pool)), "sel_tokJ": sel.tok_jaccard.mean(), "pop_tokJ": s.tok_jaccard.mean(), "tokJ_ratio": sel.tok_jaccard.mean() / s.tok_jaccard.mean(), "sel_cos": sel.dec_geom_cos.mean(), "pop_cos": s.dec_geom_cos.mean(), } r = [v["tokJ_ratio"] for v in out["selection_bias_final_sft"].values()] out["tokJ_ratio_range"] = [min(r), max(r)] return out def _torch(): import torch return torch def _wdec(path): """Decoder directions [d_sae, d_model] from any of the on-disk SAE formats.""" import sys sys.path.insert(0, str(_find("sae_tuned_lens").parent)) from sae_tuned_lens.sae_io import load_sae_full return load_sae_full(str(path))["W_dec"].float() def _cos(A, B): return _torch().nn.functional.cosine_similarity(A, B, dim=-1) ROOT = _find("drift_run") SFT_SAES = _find("sae_rl", "checkpoints", "saes_sft") def direction_of_drift(): """Does the null step move the basis the same WAY, or only as far? (§4.5) A share of drift is a magnitude claim. If retraining and fine-tuning rotated latents by similar amounts into unrelated subspaces, the shares would hold while retraining explained none of the real chain's particular movement. """ torch = _torch() base = _wdec(ROOT / "saes_local/sae_sft/sae_instruct_base_layer18.pt") out = {} cases = [ ("one_retraining", ROOT / "null_ctrl/sae_instruct_base_layer18_nullstep.pt", ROOT / "saes_local/sae_sft/sae_sft_step29_layer18.pt"), ("seven_retrainings", ROOT / "null_chain/sae_nullchain7_layer18.pt", ROOT / "saes_local/sae_sft/sae_sft_step348_layer18.pt"), ] for label, npath, rpath in cases: if not (npath.exists() and rpath.exists()): out[label] = "weights absent" continue nul, real = _wdec(npath), _wdec(rpath) c_bn = _cos(base, nul).mean().item() c_br = _cos(base, real).mean().item() c_nr = _cos(nul, real).mean().item() n = base.shape[0] g = torch.Generator().manual_seed(0) perm = torch.randperm(n, generator=g) idx = torch.arange(n) perm[perm == idx] = (perm[perm == idx] + 1) % n out[label] = { "cos_base_null": c_bn, "cos_base_real": c_br, "cos_null_real_observed": c_nr, "cos_null_real_if_independent": c_bn * c_br, "random_pairing_floor": _cos(nul, real[perm]).mean().item(), "excess_over_independent": c_nr - c_bn * c_br, } return out def activation_budgets(): """Optimisation budget of the null step vs each chain's own transitions (§3.3).""" torch = _torch() act = _find("sae_rl", "data") out = {} for sub, stem in [("activations", "instruct_base_layer18_train.pt"), ("activations", "ppo_step10_layer18_train.pt"), ("activations_sft", "sft_step29_layer18_train.pt"), ("activations_shuffled", "ppo_step10_layer18_train.pt"), ("activations_kl0p025", "ppo_step10_layer18_train.pt")]: p = act / sub / stem if not p.exists(): continue n = torch.load(p, map_location="cpu", mmap=True).shape[0] out[f"{sub}/{stem}"] = {"tokens": int(n), "batches_at_512": n // 512, "adam_steps_20_epochs": (n // 512) * 20} base = out.get("activations/instruct_base_layer18_train.pt") sft = out.get("activations_sft/sft_step29_layer18_train.pt") if base and sft: out["null_step_shortfall_vs_sft_pct"] = 100 * ( 1 - base["adam_steps_20_epochs"] / sft["adam_steps_20_epochs"]) out["epochs_to_match_sft"] = 20 * sft["batches_at_512"] / base["batches_at_512"] return out def budget_sweep(drift): """Nominal budget sweep, interpreted as LR-horizon sensitivity (§4.4).""" real = drift[(drift.chain == "sae_sft") & (drift.layer == 18)].sort_values("step") real_drift = 1 - real.dec_geom_all_cos_matched_mean.iloc[1] base = ROOT / "saes_local/sae_sft/sae_instruct_base_layer18.pt" if not base.exists(): return "weights absent" W0 = _wdec(base) rows = {} for p in sorted(glob.glob(str(ROOT / "null_budget/sae_instruct_base_layer18_e*_s*.pt"))): m = re.search(r"_e(\d+)_s(\d+)", p) rows.setdefault(int(m.group(1)), []).append(1 - _cos(W0, _wdec(p)).mean().item()) out = {"real_first_transition_drift": real_drift, "by_epochs": {}} for E in sorted(rows): v = rows[E] out["by_epochs"][E] = { "adam_steps": E * 696, "n_seeds": len(v), "drift_mean": sum(v) / len(v), "drift_min": min(v), "drift_max": max(v), "pct_of_real": 100 * (sum(v) / len(v)) / real_drift, } return out def trajectories_per_layer(drift): """Null trajectories against the real SFT chain, per layer, over seeds (§4.4).""" out = {} for L in (6, 12, 18, 23): base = ROOT / f"saes_local/sae_sft/sae_instruct_base_layer{L}.pt" if not base.exists(): continue dirs = [d for d in sorted(glob.glob(str(ROOT / f"null_chain_all/layer{L}_seed*"))) if os.path.exists(os.path.join(d, "DONE"))] if L == 18: # the originally reported chain belongs to this population if (ROOT / "null_chain").exists(): dirs = [str(ROOT / "null_chain")] + dirs if not dirs: continue W0 = _wdec(base) real = (drift[(drift.chain == "sae_sft") & (drift.layer == L)] .sort_values("step").reset_index(drop=True)) rf = real.dec_geom_all_cos_matched_mean.iloc[-1] ends = [] for d in dirs: ch = sorted(glob.glob(f"{d}/sae_nullchain*_layer{L}.pt"), key=lambda q: int(re.search(r"nullchain(\d+)", q).group(1))) if len(ch) < len(real): continue ends.append(_cos(W0, _wdec(ch[len(real) - 1])).mean().item()) if not ends: continue pct = [100 * (1 - e) / (1 - rf) for e in ends] out[f"L{L}"] = { "n_chains": len(ends), "real_end_cos": rf, "null_end_cos_mean": sum(ends) / len(ends), "null_end_cos_min": min(ends), "null_end_cos_max": max(ends), "pct_reproduced_mean": sum(pct) / len(pct), "pct_reproduced_min": min(pct), "pct_reproduced_max": max(pct), } return out def persistent_null(drift, per_feat): """The null step restricted to evaluation-persistent latents (§4.3). The mask is based on firing density over shared frozen-base evaluation activations; it does not record whether a latent was resampled in training. Without this, a persistent-subset real drift gets divided by an all-latent null and the share is wrong. """ torch = _torch() first = {"sae_sft": 29, "sae_flexible": 10, "sae_strict/k64": 10} out = {} for L in (6, 12, 18, 23): base = ROOT / f"saes_local/sae_sft/sae_instruct_base_layer{L}.pt" if not base.exists(): continue W0 = _wdec(base) # Every seed on disk, not a fixed subset: the seed count is what makes the # published shares reproducible, and s10/s11 sort after s9 only numerically. paths = [ROOT / f"null_ctrl/sae_instruct_base_layer{L}_nullstep.pt"] paths += [pathlib.Path(p) for p in sorted(glob.glob(str(ROOT / f"null_seeds/sae_instruct_base_layer{L}_null_s*.pt")), key=lambda q: int(re.search(r"_null_s(\d+)", q).group(1)))] cs = [_cos(W0, _wdec(p)) for p in paths if p.exists()] if not cs: continue for ch, st in first.items(): row = drift[(drift.chain == ch) & (drift.layer == L) & (drift.step == st)] m = per_feat[(per_feat.chain == ch) & (per_feat.layer == L) & (per_feat.step == st)] if not len(row) or not len(m): continue row = row.iloc[0] idx = torch.tensor(m.sort_values("feature").persistent.values.astype(bool)) na = sum(1 - c.mean().item() for c in cs) / len(cs) npr = sum(1 - c[idx].mean().item() for c in cs) / len(cs) ra = 1 - row.dec_geom_all_cos_matched_mean rp = 1 - row.dec_geom_persist_cos_matched_mean out[f"L{L}/{ch}"] = { "n_persist": int(idx.sum()), "null_all": na, "null_persist": npr, "real_all": ra, "real_persist": rp, "pct_all": 100 * na / ra, "pct_persist": 100 * npr / rp, } return out def architectures(): """Does the confound depend on the SAE parameterisation? (Jacob's grid, L12.)""" p = ROOT / "arch_null/arch_null_trajectory.csv" if not p.exists(): return "not run" a = pd.read_csv(p) last = a.null_step.max() return {r.arm: {"arch": r.arch, "final_dec_cos": r.dec_cos, "final_tok_jaccard": r.tok_jaccard, "mean_l0": r.mean_l0, "val_mse": r.val_mse, "steps": int(last)} for _, r in a[a.null_step == last].iterrows()} def fidelity_detail(metrics, drift): """The Table 1 / §4.1 numbers the earlier version of this script did not emit.""" fin = metrics[(metrics.chain == "sft")].copy() out = {"sft_final_nmse_by_layer": {}, "sft_mse_base_to_final": {}, "persist_counts": {}, "layer23_nmse_all_runs": {}} for L in (6, 12, 18, 23): s = fin[fin.layer == L].sort_values("step") if not len(s): continue out["sft_final_nmse_by_layer"][f"L{L}"] = s.nmse_recomputed.iloc[-1] out["sft_mse_base_to_final"][f"L{L}"] = [s.mse.iloc[0], s.mse.iloc[-1]] n = drift.n_persist out["persist_counts"] = {"min": int(n.min()), "max": int(n.max()), "d_sae": int(drift.n_features.iloc[0])} l23 = metrics[metrics.layer == 23] out["layer23_nmse_all_runs"] = {"min": l23.nmse_recomputed.min(), "max": l23.nmse_recomputed.max()} return out def base_alignment(): """Cross-chain latent indices are only comparable if the base SAE is shared.""" import hashlib import sys sys.path.insert(0, str(_find("sae_tuned_lens").parent)) from sae_tuned_lens.sae_io import load_sae_full root = _find("drift_run", "saes_local") chains = ["sae_sft", "sae_flexible", "sae_strict/k64", "sae_shuffled", "sae_kl0p025"] out = {} for L in (6, 12, 18, 23): hashes = {} for c in chains: p = root / c / f"sae_instruct_base_layer{L}.pt" if not p.exists(): hashes[c] = None continue W = load_sae_full(str(p))["W_dec"].float().cpu().numpy() hashes[c] = hashlib.md5(W.tobytes()).hexdigest()[:12] ref = hashes["sae_sft"] out[f"L{L}"] = {"hashes": hashes, "shares_base": [c for c, h in hashes.items() if h == ref], "own_init": [c for c, h in hashes.items() if h is not None and h != ref]} always_own = [c for c in chains if all(c in out[f"L{L}"]["own_init"] for L in (6, 12, 18, 23))] out["own_init_at_every_layer"] = always_own out["n_chains_own_init"] = len(always_own) return out def main(): ap = argparse.ArgumentParser() ap.add_argument("--json", action="store_true") ap.add_argument("--skip-weights", "--skip-alignment", dest="skip_weights", action="store_true", help="skip every section that loads SAE weights (hashes, drift " "direction, budget sweep, per-layer trajectories, persistent null)") args = ap.parse_args() drift = pd.read_csv(DRIFT / "drift_by_stage.csv") per_feat = pd.read_csv(DRIFT / "drift_per_feature.csv") z = pd.read_csv(DRIFT / "base_feature_z.csv") reps = pd.read_csv(DRIFT / "null_replicates.csv") traj = pd.read_csv(DRIFT / "null_trajectory_l18.csv") metrics = pd.read_csv(RESULTS / "sae_checkpoint_metrics.csv") task = pd.read_csv(RESULTS / "task_performance.csv") report = { "artifact": artifact(drift), "identity_vs_fidelity": identity_vs_fidelity(drift, metrics), "fidelity_detail": fidelity_detail(metrics, drift), "metric_agreement": metric_agreement(drift), "null_step": null_step(reps), "null_trajectory_l18": null_trajectory(traj, drift), "shuffled_control": shuffled_control(drift, task), "heterogeneity": heterogeneity(per_feat, z), "activation_budgets": activation_budgets(), "architectures": architectures(), } if not args.skip_weights: report["direction_of_drift"] = direction_of_drift() report["budget_sweep"] = budget_sweep(drift) report["trajectories_per_layer"] = trajectories_per_layer(drift) report["persistent_null"] = persistent_null(drift, per_feat) report["base_alignment"] = base_alignment() report = _round(report) if args.json: print(json.dumps(report, indent=2, default=str)) else: for section, body in report.items(): print(f"\n{'=' * 70}\n{section}\n{'=' * 70}") print(json.dumps(body, indent=2, default=str)) if __name__ == "__main__": main()