diff --git a/README.md b/README.md index 47b43dd..89da933 100644 --- a/README.md +++ b/README.md @@ -12,7 +12,8 @@ frozen forward → capture per-layer residual stream → [optional intervention] - **Exp4** — train per-layer decoders (the logit-lens tool itself; the one asset that costs compute). - **Exp5** — looping / repeating layers, quantified with the logit lens. -- **Exp6** — self-repair: skip a layer, quantify how downstream layers compensate. +- **Exp6** — ablate a layer, decode every depth (the paper's self-repair figure). + Whether downstream actually compensates is a separate question — see `evaluation/path_patching.py`. Model-specific code lives only in `adapters/`; capture and interventions use plain PyTorch forward hooks (no model-source surgery, no third-party lens libs). diff --git a/analysis/de-te-figures/README.md b/analysis/de-te-figures/README.md index 5ea3730..c0d848d 100644 --- a/analysis/de-te-figures/README.md +++ b/analysis/de-te-figures/README.md @@ -38,13 +38,13 @@ uv run --group viz python scripts/plot_de_te_perrow.py \ --in-dir out/perrow --coord margin --hist-only --out out/fig2_perrow_cehist_margin.png # E2 — per-layer CE~DE law across depth (Hydra Fig-4d) -uv run --group viz python scripts/plot_ce_de_law.py \ +uv run --group viz python scripts/plot_compensation_fit.py \ --in-dir out/perrow --coord margin --agg --out out/layer_law_agg_margin.png # E2 — DE vs CE at the apex layer (Hydra Fig-4b) -uv run --group viz python scripts/plot_ce_de_law.py \ +uv run --group viz python scripts/plot_compensation_fit.py \ --in-dir out/perrow --coord margin --agg --apex-scatter --out out/apex_scatter_agg_margin.png ``` -- `dip_recover_fig8.png` = a copy of `out/self_repair_combined_fig8.png` (exp6 tuned-decoder trajectory), from an uncommitted local plot script. +- `dip_recover_fig8.png` = a copy of `out/balef_exp6_combined_fig8.png` (exp6 tuned-decoder trajectory), from an uncommitted local plot script. - `hydra_ref_*.png` are screenshots from McGrath et al. 2023 — not reproducible from this repo. diff --git a/docs/self-repair.md b/docs/self-repair.md index 6f6a825..ade087c 100644 --- a/docs/self-repair.md +++ b/docs/self-repair.md @@ -16,7 +16,7 @@ Context for anyone picking up a self-repair issue (#34–#37). Defines the terms ## 3. Reproduction & the open question -We reproduced Figure 8 across **LimiX-2M, TabICLv2, Mitra, TabFM** (`evaluation/self_repair.py`, `scripts/plot_self_repair*.py`). +We reproduced Figure 8 across **LimiX-2M, TabICLv2, Mitra, TabFM** (`evaluation/balef_exp6.py`, `scripts/plot_balef_exp6_*.py`). - **Observation.** Removing a mid-late layer changes the final output by ~0 → **total effect ≈ 0**. - **The trap.** "Removed, output unchanged" has two explanations that look identical: diff --git a/docs/vastai-gpu-setup.md b/docs/vastai-gpu-setup.md index 0d54475..49c4068 100644 --- a/docs/vastai-gpu-setup.md +++ b/docs/vastai-gpu-setup.md @@ -45,7 +45,7 @@ git checkout # a fresh clone is on main — see n Notes: - **Branch matters — a fresh clone defaults to `main`.** Features under active development live on their branch, not `main`: e.g. the GT-logit sweep metric + - `plot_self_repair.py --metric gt_logit` are on `feat/logit-metric`. Running the + `plot_balef_exp6_trajectory.py --metric gt_logit` are on `feat/logit-metric`. Running the sweep from `main` silently gives you only the AUC trajectory. `git checkout` the branch that has the code you actually want **before** installing / running. - Finetune only needs core + `tabicl` (synthetic prior); it does **not** need the @@ -119,7 +119,7 @@ Exit-code triage: Inside tmux, and **`mkdir -p out` first** (see the redirect trap below): ```bash mkdir -p out -tmux new -d -s sweep "/venv/main/bin/python scripts/run_self_repair_sweep.py \ +tmux new -d -s sweep "/venv/main/bin/python scripts/run_balef_exp6_sweep.py \ --model --subsample-train 1000 --subsample-test 500 \ --skip-diffs --device cuda --out out/.json > out/.log 2>&1" ``` @@ -135,7 +135,7 @@ tmux new -d -s sweep "/venv/main/bin/python scripts/run_self_repair_sweep.py \ `mkdir -p out` *before* launching tmux, not inside it. - High-feature tables may OOM on VRAM: rerun those alone with `--tasks --subsample-train 3000`, then merge the json. -- Plot: `python scripts/plot_self_repair.py --model ` +- Plot: `python scripts/plot_balef_exp6_trajectory.py --model ` ## 5. Copy files over a flaky / session-limited link (both directions) Same root cause as §2's "false success": the vast box drops long connections, so diff --git a/scripts/validate_resample_donor_health.py b/scripts/check_donor_delta_health.py similarity index 98% rename from scripts/validate_resample_donor_health.py rename to scripts/check_donor_delta_health.py index 35e5983..c9e2640 100644 --- a/scripts/validate_resample_donor_health.py +++ b/scripts/check_donor_delta_health.py @@ -14,7 +14,7 @@ the intervention (upstream is frozen), so one clean target forward + one forward per donor suffices — no per-layer re-forward, no decoders. - uv run --group eval --group viz python scripts/validate_resample_donor_health.py \ + uv run --group eval --group viz python scripts/check_donor_delta_health.py \ --models limix_2m mitra tabicl_v2 --tasks 363621,363671,363696 """ @@ -27,7 +27,7 @@ import torch from tfm_lens.core.capture import capture_layers -from tfm_lens.core.resample_ablation import build_donor_delta, donor_deltas +from tfm_lens.core.donor_delta import build_donor_delta, donor_deltas from tfm_lens.evaluation.datasets import TABARENA_BINARY_TASK_IDS, load_tabarena_task from tfm_lens.evaluation.preprocess import ( limix_preprocess, diff --git a/scripts/validate_resample_self_repair.py b/scripts/compare_ablation_stability.py similarity index 85% rename from scripts/validate_resample_self_repair.py rename to scripts/compare_ablation_stability.py index a5e6a17..c4fb015 100644 --- a/scripts/validate_resample_self_repair.py +++ b/scripts/compare_ablation_stability.py @@ -1,32 +1,29 @@ -"""V2 self-repair cross-check: does the *faithful* (resample) ablation independently -show self-repair, and more stably than zero? (#35) - -V1 (``validate_resample_donor_health.py``) showed resample is on-manifold + -norm-preserving. V2 reruns the self-repair sweep under **both** ablations and -asks, on the decision **margin**: - -- **Claim 1 — self-repair holds:** under resample each ablated layer dips - (immediate hit) then climbs back toward baseline -> a property of the model, - not an artifact of zero's off-manifold shock. Zero is shown **for reference**, - not ground truth. -- **Claim 2 — stability:** resample perturbs *consistently* (small per-layer +"""Which ablation perturbs better — resample or zero? (#35, V2) + +V1 (``check_donor_delta_health.py``) showed resample is on-manifold + +norm-preserving. V2 reruns the Exp6 sweep under **both** ablations and compares +them on the decision **margin**: + +- **Same trajectory shape:** under resample each ablated layer still dips + (immediate hit) then climbs back toward baseline -> the dip-recover shape is a + property of the model, not an artifact of zero's off-manifold shock. +- **Stability (the criterion):** resample perturbs *consistently* (small per-layer spread, no overshoot); zero, off-manifold, occasionally sends the readout off (big swings + overshoots = ablation makes the margin *better* than baseline). -We deliberately do **not** frame this as "resample reproduces zero" — resample is -the better method that happens to land on the same self-repair, so no zero-vs- -resample correlation stat. +Zero is shown **for reference, not as ground truth** — so there is deliberately no +zero-vs-resample correlation stat. The question is which intervention to trust. -Inputs: ``out/v2_{model}_{zero,resample}.json`` (``run_self_repair_sweep.py`` +Inputs: ``out/v2_{model}_{zero,resample}.json`` (``run_balef_exp6_sweep.py`` with ``--ablation zero|resample``). Each stores baseline + per-ablated-layer trajectories under ``sweep["skip"]`` (the generic *ablated* condition, whatever the ablation) for metrics {auc, gt_logit, margin}. -Per ablated layer ``m`` (metric-normalized, averaged over tasks): +The only quantity scored, per ablated layer ``m`` (metric-normalized, averaged over tasks): + +- **immediate** ``imm(m) = baseline[m+1] - ablated_m[m+1]`` — the direct hit at depth m+1. -- **immediate** ``imm(m) = baseline[m+1] - ablated_m[m+1]`` — the direct hit. -- **total** ``TE(m) = baseline[-1] - ablated_m[-1]`` — what survives to the final layer. -- **self-repair** ``SR(m) = imm(m) - TE(m)`` — recovered downstream. +Its per-layer spread and overshoot count are what rank the two ablations. Figures: @@ -34,7 +31,7 @@ - ``v2_pair_gtlogit_{model}.png`` same, gt_logit (off-manifold blow-up in raw-logit space). - ``v2_stability.png`` per-layer std(imm) + overshoot count, all models. - uv run --group viz python scripts/validate_resample_self_repair.py \ + uv run --group viz python scripts/compare_ablation_stability.py \ --models limix_2m mitra tabicl_v2 tabfm --in-dir out """ @@ -154,7 +151,7 @@ def _pair_fig(in_dir, model, metric, out): _draw_panel(axes[0], zb, zs, "zero (reference)", metric, ylim) _draw_panel(axes[1], rb, rs, "resample (faithful)", metric, ylim) axes[0].set_ylabel(_METRIC_YLABEL[metric]) - fig.suptitle(f"{_MODEL_LABELS[model]} · self-repair · {metric}", fontweight="bold") + fig.suptitle(f"{_MODEL_LABELS[model]} · ablation trajectory · {metric}", fontweight="bold") fig.savefig(out, dpi=140) plt.close(fig) print(f"wrote {out}") @@ -222,7 +219,7 @@ def main(): args = _parse_args() args.out_dir.mkdir(parents=True, exist_ok=True) - # Claim 1 — self-repair holds (margin), plus the gt_logit off-manifold view. + # Trajectory shape (margin), plus the gt_logit off-manifold view. for model in args.models: _pair_fig(args.in_dir, model, "margin", args.out_dir / f"v2_pair_{model}.png") _pair_fig(args.in_dir, model, "gt_logit", args.out_dir / f"v2_pair_gtlogit_{model}.png") diff --git a/scripts/draw_ce_methods.py b/scripts/draw_ce_methods.py new file mode 100644 index 0000000..eb5ad64 --- /dev/null +++ b/scripts/draw_ce_methods.py @@ -0,0 +1,190 @@ +"""Draw the two CE-computation methods in the Pearl/Hydra causal-graph style +(⊕ = residual add, red = changed by ablation). For the notes, not an experiment. + +Residual stream: x → (⊕ with branch a) → (⊕ with branch b) → y + a = the ablated layer m ; b = a downstream layer ; y = final logit (readout f). + + uv run --group viz python scripts/draw_ce_methods.py +""" + +from pathlib import Path + +import matplotlib.pyplot as plt +from matplotlib.patches import FancyArrowPatch + +RED = "#e8563f" +BLACK = "#1a1a1a" + +# node positions (shared layout): spine at x=0.9, branch column at x=0.0 +POS = { + "x": (0.90, 0.20), + "a": (0.00, 1.05), + "p1": (0.90, 1.50), + "b": (0.00, 2.55), + "p2": (0.90, 3.00), + "y": (0.90, 3.85), +} + + +def _arrow(ax, p, q, color, rad=0.0, lw=1.6): + ax.add_patch( + FancyArrowPatch( + POS[p], + POS[q], + connectionstyle=f"arc3,rad={rad}", + arrowstyle="-|>", + mutation_scale=11, + lw=lw, + color=color, + zorder=2, + ) + ) + + +# where each node's letter label sits, relative to the node (dx, dy, ha, va) +_LAB = { + "x": (0.18, 0.0, "left", "center"), + "y": (0.18, 0.0, "left", "center"), + "a": (0.0, 0.30, "center", "bottom"), # above (branch column, left side) + "b": (0.0, 0.30, "center", "bottom"), +} + + +def _node(ax, key, color, filled=True, label=None, lab_color=None): + x, y = POS[key] + if key in ("p1", "p2"): # the ⊕ residual-add node + ax.scatter([x], [y], s=230, facecolor="white", edgecolor=BLACK, lw=1.4, zorder=3) + ax.text(x, y, "+", ha="center", va="center", fontsize=13, zorder=4) + else: + ax.scatter( + [x], + [y], + s=150, + facecolor=(color if filled else "white"), + edgecolor=color, + lw=1.8, + zorder=3, + ) + if label: + dx, dy, ha, va = _LAB[key] + ax.text( + x + dx, y + dy, label, ha=ha, va=va, fontsize=12, color=lab_color or color, zorder=4 + ) + + +def _skeleton(ax, red=(), read_y=False, read_edge=False, do_b=None): + """Draw the residual graph. red = set of nodes shown changed. + read_y: highlight the output read (method A). read_edge: highlight b→y (method B).""" + + def col(n): + return RED if n in red else BLACK + + # branch edges (curved) + skip edges (straight) + _arrow(ax, "x", "a", col("a"), rad=0.45) + _arrow(ax, "a", "p1", col("a"), rad=0.45) + _arrow(ax, "x", "p1", col("x") if "x" in red else BLACK) # skip x→⊕ + _arrow(ax, "p1", "b", col("b"), rad=0.45) + _arrow( + ax, + "b", + "p2", + RED if read_edge or "b" in red else BLACK, + rad=0.45, + lw=3.0 if read_edge else 1.6, + ) # b→⊕ (method B reads here) + _arrow(ax, "p1", "p2", col("a") if "a" in red else BLACK) # skip ⊕→⊕ + _arrow(ax, "p2", "y", RED if red else BLACK) + + _node(ax, "x", BLACK, label="x") + _node(ax, "a", col("a"), label=("ã" if "a" in red else "a")) + _node(ax, "p1", BLACK) + _node(ax, "b", col("b"), label=("b'" if "b" in red else "b")) + _node(ax, "p2", BLACK) + _node(ax, "y", RED if red else BLACK, label="y") + + if do_b: + bx, by = POS["b"] + ax.annotate( + do_b, + (bx, by), + xytext=(bx - 1.55, by - 0.05), + fontsize=10, + va="center", + color=BLACK, + arrowprops=dict(arrowstyle="-|>", color=BLACK, lw=1.2), + ) + # ablation label on a + ax.annotate( + "do(m=ã)", + POS["a"], + xytext=(POS["a"][0] - 1.55, POS["a"][1] - 0.05), + fontsize=10, + color=RED, + va="center", + arrowprops=dict(arrowstyle="-|>", color=RED, lw=1.2), + ) + + if read_y: + yx, yy = POS["y"] + ax.scatter([yx], [yy], s=520, facecolor="none", edgecolor="tab:blue", lw=2.2, zorder=5) + ax.text(yx + 0.30, yy + 0.02, "read here", color="tab:blue", fontsize=10, va="center") + if read_edge: + ex, ey = 0.34, POS["b"][1] + 0.42 # on the b→⊕ curve, well below ⊕ + ax.scatter([ex], [ey], s=430, facecolor="none", edgecolor="tab:blue", lw=2.2, zorder=5) + ax.annotate( + "read b's\nwrite here", + (ex, ey), + xytext=(ex - 1.9, ey), + fontsize=9.5, + color="tab:blue", + va="center", + arrowprops=dict(arrowstyle="-|>", color="tab:blue", lw=1.2), + ) + + ax.set_xlim(-1.75, 1.55) + ax.set_ylim(-0.15, 4.30) + ax.axis("off") + ax.set_aspect("equal") + + +def main(): + fig, axes = plt.subplots(1, 4, figsize=(16, 5.2)) + + # Panel 1 — clean + _skeleton(axes[0]) + axes[0].set_title("① Clean\ny = x + a + b", fontsize=11) + + # Panel 2 — Method A, run 1: freeze b at clean (DE side) + _skeleton(axes[1], red={"a"}, read_y=True, do_b="do(b = clean)") + axes[1].set_title("② Method A · run FROZEN\nablate m, hold b at clean → read y", fontsize=11) + + # Panel 3 — Method A, run 2: b reacts (TE side) + _skeleton(axes[2], red={"a", "b"}, read_y=True) + axes[2].set_title( + "③ Method A · run FREE\nablate m, b reacts (b') → read y\nCE_A = y[FREE] − y[FROZEN]", + fontsize=11, + ) + + # Panel 4 — Method B: read b's own write, clean vs reacted + _skeleton(axes[3], red={"a", "b"}, read_edge=True) + axes[3].set_title( + "④ Method B · independent\nablate m, b reacts (b')\n" + "CE_B = (b' 's write to y) − (b's write to y)", + fontsize=11, + ) + + fig.suptitle( + "Two ways to compute CE —— A: read output y (two runs, subtract) " + "vs B: read b's own write (one run, per-layer)", + fontsize=13, + fontweight="bold", + ) + fig.tight_layout(rect=(0, 0, 1, 0.94)) + out = Path("out/ce_methods_diagram.png") + out.parent.mkdir(exist_ok=True) + fig.savefig(out, dpi=140) + print(f"wrote {out}") + + +if __name__ == "__main__": + main() diff --git a/scripts/plot_cross_dataset.py b/scripts/plot_balef_exp6_cross_dataset.py similarity index 95% rename from scripts/plot_cross_dataset.py rename to scripts/plot_balef_exp6_cross_dataset.py index abccf96..a186681 100644 --- a/scripts/plot_cross_dataset.py +++ b/scripts/plot_balef_exp6_cross_dataset.py @@ -1,4 +1,4 @@ -"""Cross-dataset / cross-model views of the exp6 self-repair sweep. +"""Cross-dataset / cross-model views of the balef2026 Exp6 ablation sweep. Drills the per-depth sweep down to the task level and plots the net output effect of skipping each layer: @@ -9,8 +9,8 @@ - **blue / TE < 0** — delete layer -> output *damaged* (layer does irreplaceable work). - **white / TE ~ 0** — output unchanged (repaired *or* redundant). -Reuses ``plot_self_repair._normalize_task`` so numbers match the trajectory plots. -Reads one sweep JSON per model (``run_self_repair_sweep.py`` output). +Reuses ``plot_balef_exp6_trajectory._normalize_task`` so numbers match the trajectory plots. +Reads one sweep JSON per model (``run_balef_exp6_sweep.py`` output). Draws three figures: - ``cross_dataset_heatmap.png`` — 2x2 task x layer TE heatmap, one panel per model, @@ -20,7 +20,7 @@ - ``cross_dataset__L.png`` — per-task fan for one (model, layer); shows when a near-zero *average* TE hides a bimodal help/hurt split. - uv run --group viz python scripts/plot_cross_dataset.py + uv run --group viz python scripts/plot_balef_exp6_cross_dataset.py """ import argparse @@ -43,9 +43,9 @@ def _load_psr(): - """Import the sibling ``plot_self_repair.py`` script for its normalizers.""" + """Import the sibling ``plot_balef_exp6_trajectory.py`` script for its normalizers.""" here = Path(__file__).resolve().parent - spec = importlib.util.spec_from_file_location("psr", here / "plot_self_repair.py") + spec = importlib.util.spec_from_file_location("psr", here / "plot_balef_exp6_trajectory.py") psr = importlib.util.module_from_spec(spec) spec.loader.exec_module(psr) return psr diff --git a/scripts/plot_self_repair.py b/scripts/plot_balef_exp6_trajectory.py similarity index 88% rename from scripts/plot_self_repair.py rename to scripts/plot_balef_exp6_trajectory.py index 8ce46f9..06ace1a 100644 --- a/scripts/plot_self_repair.py +++ b/scripts/plot_balef_exp6_trajectory.py @@ -1,6 +1,6 @@ -"""Plot the exp6 self-repair trajectories (paper Figure 8) from the sweep JSON. +"""Plot the balef2026 Exp6 ablation trajectories (paper Figure 8) from the sweep JSON. -Reads the sweep JSON (run_self_repair_sweep.py); per dataset normalize each +Reads the sweep JSON (run_balef_exp6_sweep.py); per dataset normalize each trajectory for ``--metric`` (auc: / native final AUC floored 0.5; gt_logit: per-task z-score with the clean-baseline mu/sigma), then average across datasets. Draws: @@ -9,9 +9,10 @@ - one colored line per ablated layer m, from depth m+1 onward (post-skip). - a dashed red-x connector marking the immediate drop at depth m+1. -Self-repair = a sharp drop right after the skipped layer that later layers recover. +The dip-then-recover shape the paper reads as self-repair — equally the signature of +passive redundancy; separating the two needs the frozen-downstream DE (``path_patching``). - uv run --group viz python scripts/plot_self_repair.py --metric gt_logit + uv run --group viz python scripts/plot_balef_exp6_trajectory.py --metric gt_logit """ import argparse @@ -37,8 +38,8 @@ def _parse_args(): p = argparse.ArgumentParser(description=__doc__) - p.add_argument("--in", dest="inp", type=Path, default=Path("out/self_repair.json")) - p.add_argument("--out", type=Path, default=Path("out/self_repair_fig8.png")) + p.add_argument("--in", dest="inp", type=Path, default=Path("out/balef_exp6.json")) + p.add_argument("--out", type=Path, default=Path("out/balef_exp6_fig8.png")) p.add_argument("--model", choices=list(_MODEL_LABELS), default="limix_2m", help="title label") p.add_argument("--metric", choices=list(_METRIC_YLABEL), default="auc", help="y-axis metric") return p.parse_args() @@ -132,7 +133,7 @@ def main(): if args.metric == "auc": ax.set_ylim(0.5, 1.0) # fixed chance..perfect band; auto-zoom hides ceiling effects ax.set_title( - f"{_MODEL_LABELS[args.model]} self-repair · {args.metric} " + f"{_MODEL_LABELS[args.model]} ablation trajectory · {args.metric} " f"({len(results)} TabArena binary tasks)" ) ax.legend(loc="lower right") diff --git a/scripts/plot_ce_vs_de.py b/scripts/plot_ce_vs_de.py new file mode 100644 index 0000000..4c89a70 --- /dev/null +++ b/scripts/plot_ce_vs_de.py @@ -0,0 +1,121 @@ +"""D2 — CE vs DE (Hydra's compensation-law plot), from ``run_path_patching_sweep.py``. + +Puts the two quantities on one axis: x = DE (direct effect), y = CE = DE − TE +(compensation). Self-repair = CE>0 among DE>0 layers. A Hydra-style compensation +law shows as a positive slope (fraction of DE compensated) with high R². + +- grey stripe |DE|≤tol = overshoot zone (CE>0 there is not self-repair) — excluded from the fit. +- dashed y=x = full compensation (TE=0); y=0 = no compensation. +- fit on |DE|>tol only → slope + R² printed per model (compare Hydra 0.69 / 0.92). + + uv run --group viz python scripts/plot_ce_vs_de.py \ + --models limix_2m mitra tabicl_v2 tabfm --in-dir out --coord margin +""" + +import argparse +import json +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np + +_MODEL_LABELS = { + "limix_2m": "LimiX-2M", + "tabicl_v2": "TabICLv2", + "mitra": "Mitra", + "tabfm": "TabFM", +} +_TOL = 0.1 # |DE| below this = overshoot zone, not self-repair; excluded from the fit + + +def _parse_args(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--models", nargs="+", default=["limix_2m", "mitra", "tabicl_v2", "tabfm"]) + p.add_argument("--in-dir", type=Path, default=Path("out")) + p.add_argument("--out", type=Path, default=Path("out/ce_vs_de.png")) + p.add_argument("--coord", choices=["gt_logit", "margin"], default="margin") + return p.parse_args() + + +def _points(results, coord): + de, ce, layer = [], [], [] + for r in results.values(): + e = r["effects"] + scale = max( + e["zscore"]["sigma"] if coord == "gt_logit" else abs(e["clean"]["margin"]), 1e-6 + ) + for k, d in e["de"].items(): + x = d[coord] / scale + de.append(x) + ce.append(x - e["te"][k][coord] / scale) + layer.append(int(k)) + return np.array(de), np.array(ce), np.array(layer) + + +def _fit(de, ce): + """OLS CE~DE on the DE>tol points → (slope, r2, n). None if too few.""" + keep = np.abs(de) > _TOL + if keep.sum() < 3: + return None + x, y = de[keep], ce[keep] + slope, intercept = np.polyfit(x, y, 1) + pred = slope * x + intercept + ss_res = np.sum((y - pred) ** 2) + ss_tot = np.sum((y - y.mean()) ** 2) + r2 = 1 - ss_res / ss_tot if ss_tot > 0 else float("nan") + return slope, intercept, r2, int(keep.sum()) + + +def _draw(ax, de, ce, layer, title): + lim = float(np.nanmax(np.abs(np.concatenate([de, ce])))) * 1.1 + 1e-6 + ax.axvspan(-_TOL, _TOL, color="0.88", zorder=0) # overshoot zone + ax.axhline(0, color="0.6", lw=0.7, zorder=1) + ax.plot([-lim, lim], [-lim, lim], "--", color="0.5", lw=1, zorder=1, label="CE=DE (TE=0)") + ax.scatter(de, ce, c=layer, cmap="viridis", s=14, alpha=0.7, zorder=2) + fit = _fit(de, ce) + if fit: + slope, intercept, r2, n = fit + xs = np.array([_TOL, lim]) + ax.plot(xs, slope * xs + intercept, "-", color="crimson", lw=1.8, zorder=3) + title = f"{title}\nslope={slope:+.2f} R²={r2:.2f} (n={n}, DE>{_TOL})" + ax.set_xlim(-lim, lim) + ax.set_ylim(-lim, lim) + ax.set_aspect("equal") + ax.set_xlabel("DE (direct effect)") + ax.set_title(title, fontsize=10) + + +def main(): + args = _parse_args() + loaded = {} + for m in args.models: + p = args.in_dir / f"de_{m}.json" + if p.exists(): + loaded[m] = json.loads(p.read_text()) + else: + print(f"skip {m}: {p} not found") + if not loaded: + raise SystemExit("no input files (expected in-dir/de_.json)") + + fig, axes = plt.subplots( + 1, len(loaded), figsize=(4.6 * len(loaded), 5.0), squeeze=False, constrained_layout=True + ) + print(f"\nCE~DE fit ({args.coord}, on DE>{_TOL}) — vs Hydra slope 0.69 / R² 0.92:") + for ax, (m, res) in zip(axes[0], loaded.items(), strict=True): + de, ce, layer = _points(res, args.coord) + _draw(ax, de, ce, layer, _MODEL_LABELS.get(m, m)) + fit = _fit(de, ce) + msg = ( + f"slope {fit[0]:+.2f} R² {fit[2]:.2f} n={fit[3]}" if fit else "too few DE>tol points" + ) + print(f" {m:10s} {msg}") + axes[0][0].set_ylabel("CE = DE − TE (compensation)") + axes[0][0].legend(loc="lower right", fontsize=8) + fig.suptitle(f"Compensation law: CE vs DE · {args.coord}", fontweight="bold") + args.out.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(args.out, dpi=140) + print(f"\nwrote {args.out}") + + +if __name__ == "__main__": + main() diff --git a/scripts/plot_ce_de_law.py b/scripts/plot_compensation_fit.py similarity index 97% rename from scripts/plot_ce_de_law.py rename to scripts/plot_compensation_fit.py index 1e27be7..f59bf7c 100644 --- a/scripts/plot_ce_de_law.py +++ b/scripts/plot_compensation_fit.py @@ -21,8 +21,8 @@ default by mean-over-rows → **different (non-commuting) estimators** (see layerwise_margin). Use ``--coord gt_logit`` for the exact aggregate == mean(per-row) correspondence. - uv run --group viz python scripts/plot_ce_de_law.py --coord margin - uv run --group viz python scripts/plot_ce_de_law.py --coord margin --agg --apex-scatter + uv run --group viz python scripts/plot_compensation_fit.py --coord margin + uv run --group viz python scripts/plot_compensation_fit.py --coord margin --agg --apex-scatter """ import argparse @@ -40,7 +40,7 @@ def _parse_args(): p = argparse.ArgumentParser(description=__doc__) p.add_argument("--models", nargs="+", default=["limix_2m", "mitra", "tabicl_v2", "tabfm"]) p.add_argument("--in-dir", type=Path, default=Path("out")) - p.add_argument("--out", type=Path, default=Path("out/ce_de_law.png")) + p.add_argument("--out", type=Path, default=Path("out/compensation_fit.png")) p.add_argument("--coord", choices=["gt_logit", "margin"], default="margin") p.add_argument( "--agg", diff --git a/scripts/plot_de_te_scatter.py b/scripts/plot_de_te_scatter.py index 1f0f3e0..d73a979 100644 --- a/scripts/plot_de_te_scatter.py +++ b/scripts/plot_de_te_scatter.py @@ -1,4 +1,4 @@ -"""D1 — the DE–TE scatter (Hydra Fig 2c analog), from ``run_direct_effect_sweep.py``. +"""D1 — the DE–TE scatter (Hydra Fig 2c analog), from ``run_path_patching_sweep.py``. Each point = one (layer, dataset). Axes (both through the native decoder, same ruler): diff --git a/scripts/plot_hydra_layer.py b/scripts/plot_hydra_layer.py new file mode 100644 index 0000000..4fbaeb2 --- /dev/null +++ b/scripts/plot_hydra_layer.py @@ -0,0 +1,103 @@ +"""Hydra Fig-b analog: for the **best-repair layer** of each model, scatter DE vs CE +**across the 15 tasks** + OLS fit (Hydra's Layer-23 plot: r²=0.92, slope=0.69). + +Per model: pick the layer (excluding the last, which is degenerate — no downstream so +CE≡0) with the highest mean CE = the strongest repair candidate. Each point = one task. + +Caveat: our CE = DE − TE (mechanical DE term); Hydra measures the downstream response +independently. A clean replication needs the per-layer CE decomposition (Part 2). + + uv run --group viz python scripts/plot_hydra_layer.py --coord margin +""" + +import argparse +import json +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np + +_MODEL_LABELS = { + "limix_2m": "LimiX-2M", + "tabicl_v2": "TabICLv2", + "mitra": "Mitra", + "tabfm": "TabFM", +} + + +def _parse_args(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--models", nargs="+", default=["limix_2m", "mitra", "tabicl_v2", "tabfm"]) + p.add_argument("--in-dir", type=Path, default=Path("out")) + p.add_argument("--out", type=Path, default=Path("out/hydra_layer.png")) + p.add_argument("--coord", choices=["gt_logit", "margin"], default="margin") + return p.parse_args() + + +def _layer_points(results, layer, coord): + de, ce = [], [] + for r in results.values(): + e = r["effects"] + sc = max(e["zscore"]["sigma"] if coord == "gt_logit" else abs(e["clean"]["margin"]), 1e-6) + x = e["de"][str(layer)][coord] / sc + de.append(x) + ce.append(x - e["te"][str(layer)][coord] / sc) + return np.array(de), np.array(ce) + + +def _best_repair_layer(results, coord): + """The non-final layer with the highest mean CE (strongest repair candidate).""" + n = len(next(iter(results.values()))["effects"]["de"]) + means = [(L, _layer_points(results, L, coord)[1].mean()) for L in range(n - 1)] + return max(means, key=lambda t: t[1])[0] + + +def _draw(ax, de, ce, title): + lim = float(np.nanmax(np.abs(np.concatenate([de, ce, [0.5]])))) * 1.1 + ax.axhline(0, color="0.6", lw=0.7, ls="--", zorder=1) + ax.plot([-lim, lim], [-lim, lim], "--", color="0.4", lw=1, zorder=1, label="1:1 (TE=0)") + ax.scatter(de, ce, s=28, alpha=0.75, color="tab:blue", zorder=2) + slope, intercept = np.polyfit(de, ce, 1) + pred = slope * de + intercept + r2 = 1 - np.sum((ce - pred) ** 2) / max(np.sum((ce - ce.mean()) ** 2), 1e-12) + xs = np.array([-lim, lim]) + ax.plot(xs, slope * xs + intercept, "-", color="crimson", lw=1.6, zorder=3) + ax.set_xlim(-lim, lim) + ax.set_ylim(-lim, lim) + ax.set_aspect("equal") + ax.set_xlabel("Direct effect of ablated layer (DE)") + ax.set_title(f"{title}\nslope={slope:+.2f} r²={r2:.2f}", fontsize=10) + return slope, r2 + + +def main(): + args = _parse_args() + loaded = {} + for m in args.models: + p = args.in_dir / f"de_{m}.json" + if p.exists(): + loaded[m] = json.loads(p.read_text()) + if not loaded: + raise SystemExit("no input files") + + fig, axes = plt.subplots( + 1, len(loaded), figsize=(4.4 * len(loaded), 4.8), squeeze=False, constrained_layout=True + ) + print(f"\nbest-repair layer per model ({args.coord}), DE vs CE across 15 tasks:") + for ax, (m, res) in zip(axes[0], loaded.items(), strict=True): + L = _best_repair_layer(res, args.coord) + de, ce = _layer_points(res, L, args.coord) + slope, r2 = _draw(ax, de, ce, f"{_MODEL_LABELS.get(m, m)} · L{L}") + print(f" {m:10s} best-repair L{L}: slope {slope:+.2f} r² {r2:.2f}") + axes[0][0].set_ylabel("Compensatory response CE = DE − TE") + axes[0][0].legend(loc="lower right", fontsize=8) + fig.suptitle( + f"Hydra Fig-b analog · best-repair layer, across tasks · {args.coord}", fontweight="bold" + ) + args.out.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(args.out, dpi=140) + print(f"\nwrote {args.out}") + + +if __name__ == "__main__": + main() diff --git a/scripts/plot_layer_ce.py b/scripts/plot_layer_ce.py new file mode 100644 index 0000000..6f9bfc6 --- /dev/null +++ b/scripts/plot_layer_ce.py @@ -0,0 +1,105 @@ +"""D4 — per-layer compensation CE = DE − TE, from ``run_path_patching_sweep.py``. + +Mean CE over all layers ≈ 0 can still hide a few layers with real CE>0 (self-repair). +This resolves CE **per layer**: for each layer, the per-task CE points + their mean, +z-scored per task (same normalization as the scatter). CE>0 (above the 0 line) = +downstream compensates = self-repair; CE<0 = amplification. + + uv run --group viz python scripts/plot_layer_ce.py \ + --models limix_2m mitra tabicl_v2 tabfm --in-dir out --coord margin +""" + +import argparse +import json +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np + +_MODEL_LABELS = { + "limix_2m": "LimiX-2M", + "tabicl_v2": "TabICLv2", + "mitra": "Mitra", + "tabfm": "TabFM", +} + + +def _parse_args(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--models", nargs="+", default=["limix_2m", "mitra", "tabicl_v2", "tabfm"]) + p.add_argument("--in-dir", type=Path, default=Path("out")) + p.add_argument("--out", type=Path, default=Path("out/layer_ce.png")) + p.add_argument("--coord", choices=["gt_logit", "margin"], default="margin") + p.add_argument( + "--shared-y", action="store_true", help="share y across panels (default: per-panel)" + ) + return p.parse_args() + + +def _layer_ce(results, coord): + """-> ce[layer] = list of per-task CE (normalized per task). Keyed by int layer.""" + ce = {} + for r in results.values(): + e = r["effects"] + sigma = e["zscore"]["sigma"] if coord == "gt_logit" else abs(e["clean"]["margin"]) + scale = max(sigma, 1e-6) + for m_str, d in e["de"].items(): + m = int(m_str) + val = (d[coord] - e["te"][m_str][coord]) / scale + ce.setdefault(m, []).append(val) + return ce + + +def _draw(ax, ce, title): + layers = sorted(ce) + means = [float(np.mean(ce[m])) for m in layers] + p25 = [float(np.percentile(ce[m], 25)) for m in layers] + p75 = [float(np.percentile(ce[m], 75)) for m in layers] + ax.axhspan(0, max(max(p75), 0.01) * 1.2, color="tab:green", alpha=0.06) # repair side + ax.axhline(0, color="0.5", lw=0.8) + for m in layers: + ax.plot([m] * len(ce[m]), ce[m], ".", color="0.7", ms=3, alpha=0.5, zorder=1) + ax.fill_between(layers, p25, p75, color="tab:blue", alpha=0.18, zorder=2) + ax.plot(layers, means, "-o", color="tab:blue", lw=1.5, ms=3, zorder=3, label="mean CE") + ax.set_title(title) + ax.set_xlabel("layer") + + +def main(): + args = _parse_args() + loaded = {} + for model in args.models: + path = args.in_dir / f"de_{model}.json" + if path.exists(): + loaded[model] = json.loads(path.read_text()) + else: + print(f"skip {model}: {path} not found") + if not loaded: + raise SystemExit("no input files (expected in-dir/de_.json)") + + fig, axes = plt.subplots( + 1, + len(loaded), + figsize=(4.5 * len(loaded), 4.2), + squeeze=False, + sharey=args.shared_y, + constrained_layout=True, + ) + print(f"\nper-layer CE ({args.coord}) — layers with mean CE > 0 (self-repair candidates):") + for ax, (model, results) in zip(axes[0], loaded.items(), strict=True): + ce = _layer_ce(results, args.coord) + _draw(ax, ce, _MODEL_LABELS.get(model, model)) + pos = [(m, float(np.mean(ce[m]))) for m in sorted(ce) if np.mean(ce[m]) > 0] + pos.sort(key=lambda x: -x[1]) + top = ", ".join(f"L{m}:{v:+.2f}" for m, v in pos[:5]) or "none" + print(f" {model:10s} {len(pos)}/{len(ce)} layers CE>0 | top: {top}") + axes[0][0].set_ylabel(f"CE = DE − TE ({args.coord}, z-scored)") + axes[0][0].legend(loc="upper right", fontsize=8) + fig.suptitle(f"Per-layer compensation CE · {args.coord} (>0 = self-repair)", fontweight="bold") + args.out.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(args.out, dpi=140) + print(f"\nwrote {args.out}") + + +if __name__ == "__main__": + main() diff --git a/scripts/run_self_repair_sweep.py b/scripts/run_balef_exp6_sweep.py similarity index 90% rename from scripts/run_self_repair_sweep.py rename to scripts/run_balef_exp6_sweep.py index ae8c339..8ed75cb 100644 --- a/scripts/run_self_repair_sweep.py +++ b/scripts/run_balef_exp6_sweep.py @@ -1,4 +1,4 @@ -"""Run the exp6 self-repair sweep over the TabArena binary tasks. +"""Run the balef2026 Exp6 ablation sweep over the TabArena binary tasks. For each dataset: load -> (optional) subsample -> preprocess -> ablation_sweep (the per-depth AUC trajectories) + native_final_auc (the per-dataset normalizer) @@ -15,27 +15,28 @@ without replacement, metric averaged over donors. On-distribution + norm-preserving; run it alongside the default ``zero`` for the skip-vs-resample cross-check. - uv run --group eval python scripts/run_self_repair_sweep.py --model limix_2m - uv run --group tabicl --group eval python scripts/run_self_repair_sweep.py \ + uv run --group eval python scripts/run_balef_exp6_sweep.py --model limix_2m + uv run --group tabicl --group eval python scripts/run_balef_exp6_sweep.py \ --model tabicl_v2 --subsample-train 0 --subsample-test 0 --skip-diffs \ - --out out/self_repair_tabicl.json - uv run --group eval python scripts/run_self_repair_sweep.py --model limix_2m \ - --ablation resample --out out/self_repair_limix_resample.json + --out out/balef_exp6_tabicl.json + uv run --group eval python scripts/run_balef_exp6_sweep.py --model limix_2m \ + --ablation resample --out out/balef_exp6_limix_resample.json """ import argparse import json from pathlib import Path +from tfm_lens.evaluation.balef_exp6 import ablation_diffs, ablation_sweep from tfm_lens.evaluation.datasets import TABARENA_BINARY_TASK_IDS, load_task_record from tfm_lens.evaluation.layerwise import load_decoders +from tfm_lens.evaluation.native_readout import native_final_auc from tfm_lens.evaluation.preprocess import ( limix_preprocess, mitra_preprocess, tabfm_preprocess, tabicl_preprocess, ) -from tfm_lens.evaluation.self_repair import ablation_diffs, ablation_sweep, native_final_auc from tfm_lens.finetune.__main__ import build_adapter SEED = 0 @@ -52,7 +53,7 @@ def _parse_args(): p.add_argument("--model", choices=list(MODELS), default="limix_2m") p.add_argument("--subsample-train", type=int, default=500, help="max train rows; 0 = all") p.add_argument("--subsample-test", type=int, default=200, help="max test rows; 0 = all") - p.add_argument("--out", type=Path, default=Path("out/self_repair.json")) + p.add_argument("--out", type=Path, default=Path("out/balef_exp6.json")) p.add_argument("--skip-diffs", action="store_true", help="skip ablation_diffs (scatter data)") p.add_argument("--device", default="cpu", help="cpu or cuda (Mitra's full tables need cuda)") p.add_argument("--tasks", default=None, help="comma-separated task ids; default all 15") diff --git a/scripts/run_direct_effect_sweep.py b/scripts/run_path_patching_sweep.py similarity index 91% rename from scripts/run_direct_effect_sweep.py rename to scripts/run_path_patching_sweep.py index 14bfad3..165b8df 100644 --- a/scripts/run_direct_effect_sweep.py +++ b/scripts/run_path_patching_sweep.py @@ -1,6 +1,6 @@ """Path-patching DE–TE sweep over the TabArena binary tasks (issue #44, Part 1). -For each dataset: load -> (optional) subsample -> preprocess -> ``direct_total_effect`` +For each dataset: load -> (optional) subsample -> preprocess -> ``layer_effects`` (per-layer DE and TE on the **native decoder**, same ruler) + ``native_final_auc`` (the per-dataset normalizer). Dumped to JSON for ``plot_de_te_scatter.py`` (the Hydra Fig-2c analog). @@ -10,8 +10,8 @@ are the other loaded tasks (leave-one-out). Unlike the tuned-decoder sweep, this reads only the model's own decoder — no ``weights/`` decoders needed. - uv run --group eval python scripts/run_direct_effect_sweep.py --model limix_2m - uv run --group eval python scripts/run_direct_effect_sweep.py --model mitra \ + uv run --group eval python scripts/run_path_patching_sweep.py --model limix_2m + uv run --group eval python scripts/run_path_patching_sweep.py --model mitra \ --device cuda --out out/de_mitra.json """ @@ -20,14 +20,14 @@ from pathlib import Path from tfm_lens.evaluation.datasets import TABARENA_BINARY_TASK_IDS, load_task_record -from tfm_lens.evaluation.direct_effect import direct_total_effect +from tfm_lens.evaluation.native_readout import native_final_auc +from tfm_lens.evaluation.path_patching import layer_effects from tfm_lens.evaluation.preprocess import ( limix_preprocess, mitra_preprocess, tabfm_preprocess, tabicl_preprocess, ) -from tfm_lens.evaluation.self_repair import native_final_auc from tfm_lens.finetune.__main__ import build_adapter SEED = 0 @@ -44,7 +44,7 @@ def _parse_args(): p.add_argument("--model", choices=list(MODELS), default="limix_2m") p.add_argument("--subsample-train", type=int, default=500, help="max train rows; 0 = all") p.add_argument("--subsample-test", type=int, default=200, help="max test rows; 0 = all") - p.add_argument("--out", type=Path, default=Path("out/direct_effect.json")) + p.add_argument("--out", type=Path, default=Path("out/path_patching.json")) p.add_argument("--device", default="cpu", help="cpu or cuda (Mitra's full tables need cuda)") p.add_argument("--tasks", default=None, help="comma-separated task ids; default all 15") p.add_argument( @@ -82,7 +82,7 @@ def main(): donor_tables = [ (r["Xtr"], r["ytr"], r["Xte"]) for tid, r in records.items() if tid != task_id ] - effects = direct_total_effect( + effects = layer_effects( adapter, rec["Xtr"], rec["ytr"], diff --git a/src/tfm_lens/adapters/mitra.py b/src/tfm_lens/adapters/mitra.py index 0f7220e..77ee193 100644 --- a/src/tfm_lens/adapters/mitra.py +++ b/src/tfm_lens/adapters/mitra.py @@ -1,7 +1,7 @@ """Mitra adapter — the vendored Tab2D as a ModelAdapter. Tab2D is a double-stream 4D model: each layer processes support (context) and -query (test) separately and returns ``(support, query)``. self-repair reads the +query (test) separately and returns ``(support, query)``. the ablation sweep reads the query stream. Loaded with the standard-attention path forced on (no flash), so the adapter maintains a single forward / residual layout. diff --git a/src/tfm_lens/adapters/tabfm.py b/src/tfm_lens/adapters/tabfm.py index 76353d3..25bd32a 100644 --- a/src/tfm_lens/adapters/tabfm.py +++ b/src/tfm_lens/adapters/tabfm.py @@ -1,7 +1,7 @@ """TabFM adapter — the vendored PyTorch TabFM as a ModelAdapter. TabFM is single-stream 3D, like TabICL: ``forward`` runs the col/row embedding -stages then the ICL transformer (``icl_predictor``); self-repair reads the ICL +stages then the ICL transformer (``icl_predictor``); the ablation sweep reads the ICL stage. Everything maps onto the base contract with no core changes: - ``layers`` -> ``icl_predictor.tf_icl.blocks`` diff --git a/src/tfm_lens/core/resample_ablation.py b/src/tfm_lens/core/donor_delta.py similarity index 98% rename from src/tfm_lens/core/resample_ablation.py rename to src/tfm_lens/core/donor_delta.py index 3e8d21a..8eae116 100644 --- a/src/tfm_lens/core/resample_ablation.py +++ b/src/tfm_lens/core/donor_delta.py @@ -14,7 +14,7 @@ draw from the donor's same-role positions (with replacement — donor/target row & feature counts differ). -The N=8 donor draws + metric averaging live in the sweep (``self_repair.py``); +The N=8 donor draws + metric averaging live in the sweep (``balef_exp6.py``); this module builds one ``donor_delta`` per (donor, layer). """ diff --git a/src/tfm_lens/core/interventions.py b/src/tfm_lens/core/interventions.py index 216265a..5c8f2f1 100644 --- a/src/tfm_lens/core/interventions.py +++ b/src/tfm_lens/core/interventions.py @@ -33,7 +33,7 @@ def inject_delta(adapter, idx, donor_delta): ``donor_delta`` must match the layer's input-residual shape (a Tensor, or a per-stream tuple for double-stream models) — build it with - ``core.resample_ablation.build_donor_delta``. ``skip_layer`` is the special case δ:=0. + ``core.donor_delta.build_donor_delta``. ``skip_layer`` is the special case δ:=0. """ layer = adapter.layers[idx] original = layer.forward diff --git a/src/tfm_lens/evaluation/self_repair.py b/src/tfm_lens/evaluation/balef_exp6.py similarity index 78% rename from src/tfm_lens/evaluation/self_repair.py rename to src/tfm_lens/evaluation/balef_exp6.py index ea5fba5..593416f 100644 --- a/src/tfm_lens/evaluation/self_repair.py +++ b/src/tfm_lens/evaluation/balef_exp6.py @@ -1,11 +1,17 @@ -"""Ablation sweep for the self-repair analysis. +"""balef2026 (`Is One Layer Enough?`) Exp6 — ablate a layer, decode every depth. -One frozen forward per condition (baseline, then skip each layer), decode every -depth with the fine-tuned decoders, reduce to all metrics off the same logits. +One frozen forward per condition (baseline, then ablate each layer), decoded at every +depth by **its own fine-tuned decoder**, reduced to all metrics off the same logits. +That per-depth-tuned-decoder design is what makes this the paper's experiment rather +than a generic ablation. - per-depth metrics: AUC / GT-logit / margin (one forward feeds all). - ``zscore``: clean-baseline (mu, sigma) for cross-task GT-logit normalization. -- skip is orthogonal — ``skip_layer`` wraps the same forward. +- the ablation is orthogonal — ``skip_layer`` / ``inject_delta`` wraps the same forward. + +The dip-then-recover this produces is a **decodability** trajectory, not evidence of +active compensation: measuring that needs the frozen-downstream DE in ``path_patching`` +(#45 — no self-repair in TFMs). """ import numpy as np @@ -14,8 +20,8 @@ from tfm_lens.adapters.base import ModelAdapter from tfm_lens.core.capture import capture_layers +from tfm_lens.core.donor_delta import build_donor_delta, donor_deltas from tfm_lens.core.interventions import inject_delta, skip_layer -from tfm_lens.core.resample_ablation import build_donor_delta, donor_deltas from tfm_lens.evaluation.layerwise import ( gt_logit_zscore_stats, layerwise_auc, @@ -23,44 +29,10 @@ layerwise_margin, predict_layers_logits, ) +from tfm_lens.evaluation.native_readout import native_final_auc from tfm_lens.utils import clone_residual -def native_final_logits( - adapter: ModelAdapter, - X_train: torch.Tensor, - y_train: torch.Tensor, - X_test: torch.Tensor, - n_classes: int, -) -> np.ndarray: - """The model's own decoder on the **final** layer → ``[n_test, n_classes]`` logits. - - The native readout (distinct from the fine-tuned probes) — the ruler the DE/TE - scatter uses. Wrap the call in a ``skip_layer`` / ``inject_delta`` context to read - it under ablation. - """ - # predict_layers_logits wants one decoder per depth; here every depth is the one - # shared native decoder (the same instance repeated). Fine because we read only the - # final layer ([-1]) and the decoder is stateless — depths never diverge. Do NOT - # reuse this list to read *intermediate* depths expecting independent decoders. - native = [adapter.decoder_template()] * (adapter.n_layers + 1) - return predict_layers_logits(adapter, native, X_train, y_train, X_test, n_classes)[-1] - - -def native_final_auc( - adapter: ModelAdapter, - X_train: torch.Tensor, - y_train: torch.Tensor, - X_test: torch.Tensor, - y_test: np.ndarray, - n_classes: int, -) -> float: - """AUC of the model's own decoder on the final layer — the paper's 'main' score - (used for normalization and final_diff; distinct from the fine-tuned probes).""" - logits = native_final_logits(adapter, X_train, y_train, X_test, n_classes) - return layerwise_auc([softmax(logits, axis=1)], y_test)[0] - - def ablation_diffs( adapter: ModelAdapter, decoders: list[torch.nn.Module], @@ -70,14 +42,15 @@ def ablation_diffs( y_test: np.ndarray, n_classes: int, ) -> list[tuple[int, float, float]]: - """Per ablated layer m: ``(m, immediate_diff, final_diff)`` — the two - normalized AUC drops that reveal self-repair. + """Per ablated layer m: ``(m, immediate_diff, final_diff)`` — the paper's two + normalized AUC drops. Both diffs are normalized by the baseline native-final AUC (floored at 0.5): ``immediate_diff`` is the fine-tuned decode right after the neutered layer (depth m+1) vs baseline at the same depth; ``final_diff`` is the model's - native final prediction under ablation vs baseline. Self-repair = large - immediate drop, small final drop. + native final prediction under ablation vs baseline. The paper reads a large + immediate drop with a small final drop as self-repair; it is equally the + signature of redundancy (see ``path_patching``). """ sweep = ablation_sweep(adapter, decoders, X_train, y_train, X_test, y_test, n_classes) baseline_ft = sweep["baseline"]["auc"] diff --git a/src/tfm_lens/evaluation/de_results.py b/src/tfm_lens/evaluation/de_results.py index 3c8498d..04c49c3 100644 --- a/src/tfm_lens/evaluation/de_results.py +++ b/src/tfm_lens/evaluation/de_results.py @@ -1,11 +1,11 @@ -"""Reader for the ``de_.json`` contract written by ``direct_total_effect``. +"""Reader for the ``de_.json`` contract written by ``layer_effects``. -Lives next to the producer (``direct_effect.py``) so both ends of the contract move +Lives next to the producer (``path_patching.py``) so both ends of the contract move together — the loader and the per-task z-score rule can't drift out of sync with the fields the sweep emits. Consumers: ``scripts/plot_de_te_scatter.py``, ``plot_de_te_perrow.py``, -``plot_ce_de_law.py``. No matplotlib here — this module only reads and scales. +``plot_compensation_fit.py``. No matplotlib here — this module only reads and scales. """ import json diff --git a/src/tfm_lens/evaluation/native_readout.py b/src/tfm_lens/evaluation/native_readout.py new file mode 100644 index 0000000..d750cea --- /dev/null +++ b/src/tfm_lens/evaluation/native_readout.py @@ -0,0 +1,48 @@ +"""Read the final output through the model's **own** decoder. + +Distinct from the fine-tuned per-depth probes (``layerwise.py``): this is the ruler +the model actually decides with, and the one both DE and TE are read on. + +Shared on purpose — used by ``path_patching`` (DE/TE) and by ``balef_exp6`` +(per-dataset normalizer), so it belongs to neither. +""" + +import numpy as np +import torch +from scipy.special import softmax + +from tfm_lens.adapters.base import ModelAdapter +from tfm_lens.evaluation.layerwise import layerwise_auc, predict_layers_logits + + +def native_final_logits( + adapter: ModelAdapter, + X_train: torch.Tensor, + y_train: torch.Tensor, + X_test: torch.Tensor, + n_classes: int, +) -> np.ndarray: + """The model's own decoder on the **final** layer → ``[n_test, n_classes]`` logits. + + Wrap the call in a ``skip_layer`` / ``inject_delta`` context to read it under ablation. + """ + # predict_layers_logits wants one decoder per depth; here every depth is the one + # shared native decoder (the same instance repeated). Fine because we read only the + # final layer ([-1]) and the decoder is stateless — depths never diverge. Do NOT + # reuse this list to read *intermediate* depths expecting independent decoders. + native = [adapter.decoder_template()] * (adapter.n_layers + 1) + return predict_layers_logits(adapter, native, X_train, y_train, X_test, n_classes)[-1] + + +def native_final_auc( + adapter: ModelAdapter, + X_train: torch.Tensor, + y_train: torch.Tensor, + X_test: torch.Tensor, + y_test: np.ndarray, + n_classes: int, +) -> float: + """AUC of the model's own decoder on the final layer — the paper's 'main' score + (used for normalization and final_diff; distinct from the fine-tuned probes).""" + logits = native_final_logits(adapter, X_train, y_train, X_test, n_classes) + return layerwise_auc([softmax(logits, axis=1)], y_test)[0] diff --git a/src/tfm_lens/evaluation/direct_effect.py b/src/tfm_lens/evaluation/path_patching.py similarity index 97% rename from src/tfm_lens/evaluation/direct_effect.py rename to src/tfm_lens/evaluation/path_patching.py index 5a5231f..2b08eae 100644 --- a/src/tfm_lens/evaluation/direct_effect.py +++ b/src/tfm_lens/evaluation/path_patching.py @@ -1,5 +1,5 @@ -"""Direct (DE) and total (TE) effect on the **final decoder** — path-patching -method B (issue #44, Part 1). +"""Path patching — direct (DE) and total (TE) effect on the **final decoder** +(method B; issue #44, Part 1). Two rulers for the Hydra DE–TE scatter, both read through the *same* native decoder so they're comparable by construction (not per-layer fine-tuned decoders): @@ -18,7 +18,7 @@ - **TE(m) — total effect.** Downstream must react → one real forward per m (``inject_delta``), read the final layer through the same native decoder - (reuses ``self_repair.native_final_logits``). + (reuses ``native_readout.native_final_logits``). Ablation is **resample**: ``ã_m`` = a role-matched donor δ, metric averaged over donors (reusing the #35 machinery). Both effects are the **drop** ``clean − ablated`` @@ -31,15 +31,15 @@ from tfm_lens.adapters.base import ModelAdapter from tfm_lens.core.capture import capture_layers +from tfm_lens.core.donor_delta import build_donor_delta, donor_deltas, layer_deltas from tfm_lens.core.interventions import inject_delta -from tfm_lens.core.resample_ablation import build_donor_delta, donor_deltas, layer_deltas from tfm_lens.evaluation.layerwise import ( gt_logit, gt_logit_zscore_stats, layerwise_gt_logit, layerwise_margin, ) -from tfm_lens.evaluation.self_repair import native_final_logits +from tfm_lens.evaluation.native_readout import native_final_logits from tfm_lens.utils import clone_residual # A residual stream: a Tensor, or a ``(support, query)`` pair for double-stream models. @@ -130,7 +130,7 @@ def _mean_over_donors_rows(per_donor: list[dict[str, np.ndarray]]) -> dict[str, return {k: np.mean([c[k] for c in per_donor], axis=0) for k in per_donor[0]} -def direct_total_effect( +def layer_effects( adapter: ModelAdapter, X_train: torch.Tensor, y_train: torch.Tensor, diff --git a/src/tfm_lens/evaluation/preprocess.py b/src/tfm_lens/evaluation/preprocess.py index c911fee..f9f4770 100644 --- a/src/tfm_lens/evaluation/preprocess.py +++ b/src/tfm_lens/evaluation/preprocess.py @@ -1,9 +1,12 @@ -"""Drive the vendored LimiX preprocessing on one real table. +"""Per-model preprocessing of one real table — one entry point per backbone +(``limix_preprocess`` / ``tabicl_preprocess`` / ``tabfm_preprocess`` / ``mitra_preprocess``). -Reproduces the no-retrieval single-worker pipeline that LimiX's predictor builds +Each drives that model's own vendored pipeline so real OpenML tables land in the +distribution the frozen model was trained on. + +LimiX: reproduces the no-retrieval single-worker pipeline its predictor builds (FilterValidFeatures -> quantile rebalance -> ordinal categorical encode -> -feature shuffle) so real OpenML tables land in the distribution the frozen model -was trained on. This is worker ① of ``cls_default_noretrieval.json``; the 4-way +feature shuffle) — worker ① of ``cls_default_noretrieval.json``; the 4-way ensemble is deferred. """ @@ -62,7 +65,7 @@ def tabicl_preprocess(X_train, y_train, X_test, categorical_idx, seed=0): ``power`` normalization (standard-scale -> yeo-johnson -> outlier-clip). Fits on the train (support) rows, transforms both. Returns ``(X_train_p, X_test_p)`` float32 arrays ready for ``predict_layers``. The 8-way ensemble (none/power x - feature shuffles) is deferred — self-repair needs a single clean forward. + feature shuffles) is deferred — the ablation sweep needs a single clean forward. """ import pandas as pd from tabicl._sklearn.preprocessing import PreprocessingPipeline, TransformToNumerical @@ -95,7 +98,7 @@ def tabfm_preprocess(X_train, y_train, X_test, categorical_idx, seed=0): Fits on the train (support) rows, transforms both. Returns ``(X_train_p, X_test_p)`` float32 arrays ready for ``predict_layers``. The ensemble (norm-method/shuffle/class-shift views) and the categorical embedding path - (cat_mask) are deferred — self-repair needs a single clean forward. + (cat_mask) are deferred — the ablation sweep needs a single clean forward. """ from sklearn.compose import ColumnTransformer from sklearn.impute import SimpleImputer diff --git a/tests/test_self_repair.py b/tests/test_balef_exp6.py similarity index 97% rename from tests/test_self_repair.py rename to tests/test_balef_exp6.py index 7d4f773..90f97ce 100644 --- a/tests/test_self_repair.py +++ b/tests/test_balef_exp6.py @@ -4,7 +4,8 @@ import pytest import torch -from tfm_lens.evaluation.self_repair import ablation_diffs, ablation_sweep, native_final_auc +from tfm_lens.evaluation.balef_exp6 import ablation_diffs, ablation_sweep +from tfm_lens.evaluation.native_readout import native_final_auc from toys import ToyAdapter3D H = ToyAdapter3D.HIDDEN diff --git a/tests/test_resample_ablation.py b/tests/test_donor_delta.py similarity index 98% rename from tests/test_resample_ablation.py rename to tests/test_donor_delta.py index 0afa21b..0c9d34d 100644 --- a/tests/test_resample_ablation.py +++ b/tests/test_donor_delta.py @@ -9,7 +9,7 @@ import torch from tfm_lens.core.capture import capture_layers -from tfm_lens.core.resample_ablation import build_donor_delta, donor_deltas, layer_deltas +from tfm_lens.core.donor_delta import build_donor_delta, donor_deltas, layer_deltas from toys import ToyAdapter3D, ToyAdapter4D H = ToyAdapter3D.HIDDEN diff --git a/tests/test_direct_effect.py b/tests/test_path_patching.py similarity index 95% rename from tests/test_direct_effect.py rename to tests/test_path_patching.py index bf031c3..82b6ddb 100644 --- a/tests/test_direct_effect.py +++ b/tests/test_path_patching.py @@ -1,4 +1,4 @@ -"""Path-patching DE / TE (method B) — ``evaluation.direct_effect``. +"""Path-patching DE / TE (method B) — ``evaluation.path_patching``. Headline invariant: the DE **arithmetic** ``patched_residual`` (``r_L − a_m + ã_m``, no per-layer forward) reproduces a real **frozen-downstream forward** (ablate m to @@ -17,11 +17,11 @@ import torch.nn as nn from tfm_lens.core.capture import capture_layers +from tfm_lens.core.donor_delta import layer_deltas from tfm_lens.core.interventions import inject_delta -from tfm_lens.core.resample_ablation import layer_deltas -from tfm_lens.evaluation.direct_effect import ( - direct_total_effect, +from tfm_lens.evaluation.path_patching import ( final_decoder_logits, + layer_effects, patched_residual, ) from tfm_lens.utils import clone_residual @@ -196,7 +196,7 @@ def test_resample_runs_with_shape_and_keys(): (dtr, dytr, dte) for dtr, dytr, dte, _ in (_inputs((ToyAdapter3D.HIDDEN,), seed=s) for s in (2, 3, 4)) ] - res = direct_total_effect(adapter, Xtr, ytr, Xte, y_test, 2, donor_tables, n_donors=3, seed=0) + res = layer_effects(adapter, Xtr, ytr, Xte, y_test, 2, donor_tables, n_donors=3, seed=0) for effect in ("de", "te"): assert set(res[effect]) == set(range(adapter.n_layers)) for m in range(adapter.n_layers): @@ -214,7 +214,7 @@ def test_per_row_shape_and_aggregate_consistency(): (dtr, dytr, dte) for dtr, dytr, dte, _ in (_inputs((ToyAdapter3D.HIDDEN,), seed=s) for s in (2, 3, 4)) ] - res = direct_total_effect( + res = layer_effects( adapter, Xtr, ytr, Xte, y_test, 2, donor_tables, n_donors=3, seed=0, per_row=True ) n_test = Xte.shape[0] @@ -235,4 +235,4 @@ def test_resample_needs_donor_tables(): adapter = BinaryToy3D() Xtr, ytr, Xte, y_test = _inputs((ToyAdapter3D.HIDDEN,)) with pytest.raises(ValueError, match="donor_tables"): - direct_total_effect(adapter, Xtr, ytr, Xte, y_test, 2, donor_tables=[]) + layer_effects(adapter, Xtr, ytr, Xte, y_test, 2, donor_tables=[]) diff --git a/tests/test_tabfm_adapter.py b/tests/test_tabfm_adapter.py index 3e4f17d..a6ed516 100644 --- a/tests/test_tabfm_adapter.py +++ b/tests/test_tabfm_adapter.py @@ -1,7 +1,7 @@ """TabFMAdapter integration (needs the TabFM checkpoint + the `tabfm` group). Runs the real frozen TabFM (single-stream 3D, like TabICL) through capture / -skip / logit_lens: self-repair reads the ICL stage (``icl_predictor.tf_icl``), +skip / logit_lens: the ablation sweep reads the ICL stage (``icl_predictor.tf_icl``), readout applies its RMSNorm then slices the test rows. Skipped without the `tabfm` group installed or a resolvable checkpoint. """