diff --git a/scripts/run_self_repair_sweep.py b/scripts/run_self_repair_sweep.py index 10a85b1..71336eb 100644 --- a/scripts/run_self_repair_sweep.py +++ b/scripts/run_self_repair_sweep.py @@ -10,10 +10,17 @@ flags; ``--skip-diffs`` drops the (expensive) scatter data so a full pass only pays for the Figure-8 trajectories. +``--ablation resample`` (#35) swaps zero/skip for **cross-table resample**: each +target's donors are the other loaded tasks (leave-one-out), ``--n-donors`` drawn +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 \ --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 """ import argparse @@ -52,6 +59,18 @@ def _parse_args(): 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") + p.add_argument( + "--ablation", choices=["zero", "resample"], default="zero", help="ablation mode (#35)" + ) + p.add_argument( + "--donor", + choices=["table", "row", "feature"], + default="table", + help="resample donor source", + ) + p.add_argument( + "--n-donors", type=int, default=8, help="resample: donors drawn per target (LOO)" + ) return p.parse_args() @@ -63,8 +82,34 @@ def _subsample(X, y, n): return X[idx], y[idx] +def _load_record(task_id, preprocess, args): + """Load -> subsample -> 0-index labels -> preprocess -> tensors for one task.""" + X_train, y_train, X_test, y_test, cat_idx = load_tabarena_task(task_id) + X_train, y_train = _subsample(X_train, y_train, args.subsample_train) + X_test, y_test = _subsample(X_test, y_test, args.subsample_test) + # 0-index the labels: TabICL's one-hot y_encoder needs contiguous ints; harmless + # for LimiX since binary labels are already 0/1. + classes = np.unique(y_train) + y_train = np.searchsorted(classes, y_train) + y_test = np.searchsorted(classes, y_test) + X_train_p, X_test_p = preprocess(X_train, y_train, X_test, cat_idx) + return { + "Xtr": torch.tensor(X_train_p), + "ytr": torch.tensor(y_train).float(), + "Xte": torch.tensor(X_test_p), + "y_test": y_test, + "n_classes": len(classes), + "n_train": int(len(X_train)), + "n_test": int(len(X_test)), + } + + def main(): args = _parse_args() + if args.ablation == "resample" and args.donor != "table": + raise NotImplementedError( + f"--donor {args.donor} not implemented; only 'table' (cross-table)" + ) cfg = MODELS[args.model] adapter = build_adapter(args.model, None, args.device) # ckpt=None -> download / cache decoders = [d.to(args.device) for d in load_decoders(cfg["weights"], adapter)] @@ -72,45 +117,63 @@ def main(): task_ids = [int(t) for t in args.tasks.split(",")] if args.tasks else TABARENA_BINARY_TASK_IDS n_tasks = len(task_ids) + + # Load every task up front — resample needs the other tables as leave-one-out donors. + records = {} + for task_id in task_ids: + try: + records[task_id] = _load_record(task_id, preprocess, args) + except Exception as exc: # keep going so one bad table can't sink the run + print(f"load task {task_id}: FAILED {type(exc).__name__}: {exc}", flush=True) + results = {} - for i, task_id in enumerate(task_ids): + for i, (task_id, rec) in enumerate(records.items()): try: - X_train, y_train, X_test, y_test, cat_idx = load_tabarena_task(task_id) - X_train, y_train = _subsample(X_train, y_train, args.subsample_train) - X_test, y_test = _subsample(X_test, y_test, args.subsample_test) - # 0-index the labels: TabICL's one-hot y_encoder needs contiguous ints; - # harmless for LimiX since binary labels are already 0/1. - classes = np.unique(y_train) - y_train = np.searchsorted(classes, y_train) - y_test = np.searchsorted(classes, y_test) - X_train_p, X_test_p = preprocess(X_train, y_train, X_test, cat_idx) - - n_classes = len(classes) - Xtr_t = torch.tensor(X_train_p) - ytr_t = torch.tensor(y_train).float() - Xte_t = torch.tensor(X_test_p) - sweep_args = (adapter, decoders, Xtr_t, ytr_t, Xte_t, y_test, n_classes) - - sweep = ablation_sweep(*sweep_args) - native_final = native_final_auc(adapter, Xtr_t, ytr_t, Xte_t, y_test, n_classes) + sweep_args = ( + adapter, + decoders, + rec["Xtr"], + rec["ytr"], + rec["Xte"], + rec["y_test"], + rec["n_classes"], + ) + sweep_kw = {} + if args.ablation == "resample": + donor_tables = [ + (r["Xtr"], r["ytr"], r["Xte"]) for tid, r in records.items() if tid != task_id + ] + sweep_kw = dict( + ablation="resample", + donor_tables=donor_tables, + n_donors=args.n_donors, + seed=SEED, + ) + + sweep = ablation_sweep(*sweep_args, **sweep_kw) + native_final = native_final_auc( + adapter, rec["Xtr"], rec["ytr"], rec["Xte"], rec["y_test"], rec["n_classes"] + ) entry = { "sweep": sweep, "native_final": native_final, - "n_train": int(len(X_train)), - "n_test": int(len(X_test)), + "n_train": rec["n_train"], + "n_test": rec["n_test"], } - if not args.skip_diffs: + # diffs is a zero-ablation scatter concept; only meaningful in zero mode. + if not args.skip_diffs and args.ablation == "zero": entry["diffs"] = ablation_diffs(*sweep_args) results[str(task_id)] = entry final_auc = sweep["baseline"]["auc"][-1] print( - f"[{i + 1}/{n_tasks}] task {task_id}: final AUC {final_auc:.3f} | " - f"native {native_final:.3f} | rows {len(X_train)}+{len(X_test)}", + f"[{i + 1}/{len(records)}] task {task_id} ({args.ablation}): " + f"final AUC {final_auc:.3f} | native {native_final:.3f} | " + f"rows {rec['n_train']}+{rec['n_test']}", flush=True, ) except Exception as exc: # keep going so one bad table can't sink the run print( - f"[{i + 1}/{n_tasks}] task {task_id}: FAILED {type(exc).__name__}: {exc}", + f"[{i + 1}/{len(records)}] task {task_id}: FAILED {type(exc).__name__}: {exc}", flush=True, ) diff --git a/scripts/validate_resample_donor_health.py b/scripts/validate_resample_donor_health.py new file mode 100644 index 0000000..35e5983 --- /dev/null +++ b/scripts/validate_resample_donor_health.py @@ -0,0 +1,244 @@ +"""Donor-δ health check for resample ablation (#35). + +Question: is resample **on-manifold + norm-preserving**? Zero ablation is the +"broken" reference. All three checks are read off the *decoded position* (label +token 4D / row vector 3D), over the test rows: + +- **③ residual-norm preservation (money plot):** ``r(m) = ‖resid_ablated‖/‖resid_clean‖``. + resample ≈ 1 (keeps the norm) · zero dips (drops it → the LayerNorm artifact). +- **① δ-norm ratio:** ``‖donor_δ‖ / ‖native_δ‖`` — want ≈ 1 (else rescale donor δ). +- **② direction cosine:** ``cos(donor_δ, native_δ)`` vs within-table native pairs — + cross-table should look like "another normal δ", not orthogonal. + +Cheap: the post-ablation residual *at* layer m needs only layer m's clean input + +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 \ + --models limix_2m mitra tabicl_v2 --tasks 363621,363671,363696 +""" + +import argparse +import json +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np +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.evaluation.datasets import TABARENA_BINARY_TASK_IDS, load_tabarena_task +from tfm_lens.evaluation.preprocess import ( + limix_preprocess, + mitra_preprocess, + tabfm_preprocess, + tabicl_preprocess, +) +from tfm_lens.finetune.__main__ import build_adapter + +SEED = 0 +PREPROCESS = { + "limix_2m": limix_preprocess, + "tabicl_v2": tabicl_preprocess, + "mitra": mitra_preprocess, + "tabfm": tabfm_preprocess, +} + + +def _parse_args(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--models", nargs="+", default=["limix_2m", "mitra", "tabicl_v2"]) + p.add_argument("--tasks", default=None, help="comma-separated task ids; default all 15") + p.add_argument("--n-donors", type=int, default=8) + p.add_argument("--subsample-train", type=int, default=500) + p.add_argument("--subsample-test", type=int, default=200) + p.add_argument("--device", default="cpu") + p.add_argument("--out-dir", type=Path, default=Path("out")) + return p.parse_args() + + +def _subsample(X, y, n): + if n <= 0 or n >= len(X): + return X, y + idx = np.random.RandomState(SEED).choice(len(X), n, replace=False) + return X[idx], y[idx] + + +def _load(task_id, preprocess, args): + X_train, y_train, X_test, y_test, cat_idx = load_tabarena_task(task_id) + X_train, y_train = _subsample(X_train, y_train, args.subsample_train) + X_test, y_test = _subsample(X_test, y_test, args.subsample_test) + classes = np.unique(y_train) + y_train = np.searchsorted(classes, y_train) + Xtr_p, Xte_p = preprocess(X_train, y_train, X_test, cat_idx) + return torch.tensor(Xtr_p), torch.tensor(y_train).float(), torch.tensor(Xte_p) + + +def _clone(c): + if isinstance(c, tuple): + return tuple(t.clone() if torch.is_tensor(t) else t for t in c) + return c.clone() + + +def _forward_cache(adapter, Xtr, ytr, Xte, device): + """Clean per-layer capture cache + eval_pos for one table.""" + eval_pos = Xtr.shape[0] + X = torch.cat([Xtr, Xte], dim=0).unsqueeze(0).to(device) + y = ytr.unsqueeze(0).to(device) + with torch.no_grad(), capture_layers(adapter) as cache: + adapter.forward_frozen(X, y, eval_pos) + return [_clone(c) for c in cache], eval_pos + + +def _add(a, b): + if isinstance(a, tuple): + return tuple(x + y for x, y in zip(a, b, strict=True)) + return a + b + + +def _sub(a, b): + if isinstance(a, tuple): + return tuple(x - y for x, y in zip(a, b, strict=True)) + return a - b + + +def _decoded(residual, eval_pos, label_index): + """The decoded position over test rows -> ``[N, hidden]``: query stream / test + rows, then the label token (4D). This is where the norm/cosine are measured.""" + h = residual[1] if isinstance(residual, tuple) else residual[:, eval_pos:] + if label_index is not None: + h = h[..., label_index, :] + return h.reshape(-1, h.shape[-1]) + + +def _norm(residual, eval_pos, label_index): + return torch.linalg.norm(_decoded(residual, eval_pos, label_index), dim=-1) # [N] + + +def _median(t): + return float(torch.median(t).item()) + + +def _model_stats(adapter, tables, target_idxs, args): + """Per-layer ③ (residual-norm) + ② (cosine) for one model. + + ③ is aggregated **per table** (the 8 donors are averaged within a table, per the + N=8 design), and every table's per-layer value is kept — so the plot can draw a + median line + IQR/min-max band showing cross-dataset spread. zero is + donor-independent → one value per table. ② stays a mean over table x donor. + """ + device = args.device + li = adapter.label_token_index + n_layers = adapter.n_layers + zero_tbl = [[] for _ in range(n_layers)] # zero_tbl[m] = one ratio per table + res_tbl = [[] for _ in range(n_layers)] # res_tbl[m] = donor-mean ratio per table + nr, cosA, cosB = ([[] for _ in range(n_layers)] for _ in range(3)) + + for ti in target_idxs: + Xtr, ytr, Xte = tables[ti] + cache, ep_t = _forward_cache(adapter, Xtr, ytr, Xte, device) + r_in = [adapter.residual_of(cache[m]) for m in range(n_layers)] + r_out = [adapter.residual_of(cache[m + 1]) for m in range(n_layers)] + native = [_sub(r_out[m], r_in[m]) for m in range(n_layers)] + clean_n = [_norm(r_out[m], ep_t, li) for m in range(n_layers)] + + donors = [j for j in range(len(tables)) if j != ti] + picks = np.random.default_rng(SEED).permutation(donors)[: args.n_donors] + res_per_donor = [] # [n_donors, n_layers] for this table + for di, dj in enumerate(picks): + Xd_tr, yd_tr, Xd_te = tables[dj] + d_deltas = donor_deltas( + adapter, + torch.cat([Xd_tr, Xd_te]).unsqueeze(0).to(device), + yd_tr.unsqueeze(0).to(device), + Xd_tr.shape[0], + ) + draw = np.random.default_rng(SEED + 1 + di) + res_m = [] + for m in range(n_layers): + dd = build_donor_delta(r_in[m], d_deltas[m], ep_t, Xd_tr.shape[0], li, draw) + res_m.append(_median(_norm(_add(r_in[m], dd), ep_t, li) / clean_n[m])) + nat_dec, don_dec = _decoded(native[m], ep_t, li), _decoded(dd, ep_t, li) + nr[m].append( + _median(torch.linalg.norm(don_dec, dim=-1) / torch.linalg.norm(nat_dec, dim=-1)) + ) + cosB[m].append(_median(torch.cosine_similarity(don_dec, nat_dec, dim=-1))) + perm = torch.randperm(nat_dec.shape[0]) + cosA[m].append(_median(torch.cosine_similarity(nat_dec, nat_dec[perm], dim=-1))) + res_per_donor.append(res_m) + res_mean = np.mean(res_per_donor, axis=0) # donor-mean per layer, this table + for m in range(n_layers): + zero_tbl[m].append(_median(_norm(r_in[m], ep_t, li) / clean_n[m])) # donor-independent + res_tbl[m].append(float(res_mean[m])) + + agg = lambda rows: [float(np.mean(r)) for r in rows] # noqa: E731 + return { + "ratio_zero_tbl": zero_tbl, # [n_layers][n_tables] — for median + band + "ratio_resample_tbl": res_tbl, + "ratio_zero": [float(np.median(r)) for r in zero_tbl], # summary median line + "ratio_resample": [float(np.median(r)) for r in res_tbl], + "cos_within": agg(cosA), + "cos_cross": agg(cosB), + } + + +def _plot(stats, out): + models = list(stats) + fig, axes = plt.subplots(1, len(models), figsize=(5.2 * len(models), 4.4), squeeze=False) + for ax, name in zip(axes.ravel(), models, strict=True): + s = stats[name] + n = len(s["ratio_zero"]) + x = np.arange(n) / max(n - 1, 1) + ax.axhline(1.0, ls=":", color="0.5", lw=1) + ax.plot(x, s["ratio_zero"], "--o", color="tab:red", ms=3, label="zero (skip)") + ax.plot(x, s["ratio_resample"], "-o", color="tab:blue", ms=3, label="resample") + ax.set_title(f"{name} · ‖resid_ablated‖ / ‖resid_clean‖") + ax.set_xlabel("relative depth (0=first · 1=last)") + ax.set_ylabel("residual-norm ratio @ decoded pos") + ax.legend(loc="lower left", fontsize=8) + fig.suptitle( + "V1 · residual-norm preservation — resample keeps the norm, zero drops it", + fontweight="bold", + ) + fig.tight_layout() + fig.savefig(out, dpi=140) + print(f"wrote {out}") + + +def main(): + args = _parse_args() + args.out_dir.mkdir(parents=True, exist_ok=True) + task_ids = [int(t) for t in args.tasks.split(",")] if args.tasks else TABARENA_BINARY_TASK_IDS + + stats = {} + for model in args.models: + adapter = build_adapter(model, None, args.device) + tables, ok = [], [] + for tid in task_ids: + try: + tables.append(_load(tid, PREPROCESS[model], args)) + ok.append(tid) + except Exception as exc: + print(f"{model} load {tid}: FAILED {type(exc).__name__}: {exc}", flush=True) + s = _model_stats(adapter, tables, list(range(len(tables))), args) + stats[model] = s + # ②/③ one-line summary (median over layers / tables) + cc, cw = np.median(s["cos_cross"]), np.median(s["cos_within"]) + zmin = min(s["ratio_zero"]) + print( + f"{model}: ② cos cross={cc:.2f} vs within={cw:.2f} | " + f"③ resample median-line min={min(s['ratio_resample']):.2f} " + f"zero median-line min={zmin:.2f}", + flush=True, + ) + del adapter + if args.device == "cuda": + torch.cuda.empty_cache() + + (args.out_dir / "v1_stats.json").write_text(json.dumps(stats, indent=2)) + _plot(stats, args.out_dir / "v1_residual_norm.png") + + +if __name__ == "__main__": + main() diff --git a/scripts/validate_resample_self_repair.py b/scripts/validate_resample_self_repair.py new file mode 100644 index 0000000..a5e6a17 --- /dev/null +++ b/scripts/validate_resample_self_repair.py @@ -0,0 +1,245 @@ +"""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 + 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. + +Inputs: ``out/v2_{model}_{zero,resample}.json`` (``run_self_repair_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): + +- **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. + +Figures: + +- ``v2_pair_{model}.png`` margin, zero | resample side-by-side (shared y). +- ``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 \ + --models limix_2m mitra tabicl_v2 tabfm --in-dir out +""" + +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", +} + +_METRIC_YLABEL = { + "auc": "Normalized performance (ROC-AUC)", + "gt_logit": "GT-logit (z-scored, σ)", + "margin": "Margin z_true - z_other (fraction of final)", +} + +# overshoot = ablation makes the margin *better* than baseline (imm < 0); a small +# tolerance keeps averaging noise from counting as an overshoot. +_OVERSHOOT_TOL = 0.02 + + +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-dir", type=Path, default=Path("out")) + return p.parse_args() + + +def _load(in_dir, model, mode): + """One (model, ablation) sweep: ``out/v2_{model}_{mode}.json`` (mode = zero|resample).""" + return json.loads((in_dir / f"v2_{model}_{mode}.json").read_text()) + + +def _normalize_task(sweep, native_final, metric): + """One task's ``(baseline, skips)`` trajectories, normalized for ``metric``. + + - auc: divide by native final AUC (floored 0.5). + - gt_logit: per-task z-score with the stored clean-baseline (mu, sigma). + - margin: divide by the baseline final-layer margin (0 = boundary, final -> 1). + """ + margin_final = max(abs(sweep["baseline"]["margin"][-1]), 1e-6) + + def norm(arr): + arr = np.array(arr, dtype=float) + if metric == "auc": + return arr / max(native_final, 0.5) + if metric == "margin": + return arr / margin_final + z = sweep["zscore"] + return (arr - z["mu"]) / z["sigma"] + + skip = sweep["skip"] + baseline = norm(sweep["baseline"][metric]) + skips = np.array([norm(skip[str(m)][metric]) for m in range(len(skip))]) + return baseline, skips + + +def _normalized(results, metric): + """-> (baseline_mean [n_depths], skip_mean [n_layers, n_depths]) averaged over tasks.""" + baselines, skips = [], [] + for r in results.values(): + b, s = _normalize_task(r["sweep"], r["native_final"], metric) + baselines.append(b) + skips.append(s) + return np.mean(baselines, axis=0), np.mean(skips, axis=0) + + +def _imm(baseline, skip): + """imm(m) = baseline[m+1] - ablated_m[m+1], the immediate drop right after layer m.""" + return np.array([baseline[m + 1] - skip[m][m + 1] for m in range(skip.shape[0])]) + + +def _draw_panel(ax, baseline, skip, title, metric, ylim=None): + """Baseline (black) + one trajectory per ablated layer, each with its red-x immediate drop.""" + n_depths = len(baseline) + n_layers = skip.shape[0] + depths = np.arange(n_depths) + colors = plt.get_cmap("viridis")(np.linspace(0, 1, n_layers)) + ax.plot(depths, baseline, "-o", color="black", lw=2, ms=3, zorder=5, label="baseline") + for m in range(n_layers): + ax.plot(depths[m + 1 :], skip[m][m + 1 :], "-o", color=colors[m], lw=1, ms=2, alpha=0.7) + ax.plot( + [m + 1, m + 1], + [baseline[m + 1], skip[m][m + 1]], + "--", + color=colors[m], + lw=1.2, + alpha=0.9, + ) + ax.plot(m + 1, skip[m][m + 1], "x", color="red", ms=4, mew=1) + ax.set_title(title) + ax.set_xlabel("Layer (forward-pass order)") + if metric == "auc": + ax.set_ylim(0.5, 1.0) + elif ylim is not None: + ax.set_ylim(*ylim) + ax.legend(loc="lower right", fontsize=8) + + +def _pair_fig(in_dir, model, metric, out): + """zero | resample side-by-side for one model, shared y.""" + zb, zs = _normalized(_load(in_dir, model, "zero"), metric) + rb, rs = _normalized(_load(in_dir, model, "resample"), metric) + lo = min(zb.min(), zs.min(), rb.min(), rs.min()) + hi = max(zb.max(), zs.max(), rb.max(), rs.max()) + pad = 0.05 * (hi - lo + 1e-9) + ylim = None if metric == "auc" else (lo - pad, hi + pad) + fig, axes = plt.subplots(1, 2, figsize=(12, 5), sharey=True, constrained_layout=True) + _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.savefig(out, dpi=140) + plt.close(fig) + print(f"wrote {out}") + + +def _stability(in_dir, model): + """Per-model stability of the immediate effect, on the margin (the decision metric).""" + out = {} + for mode in ("zero", "resample"): + b, s = _normalized(_load(in_dir, model, mode), "margin") + imm = _imm(b, s) + out[mode] = dict( + std=float(np.std(imm)), + n_overshoot=int(np.sum(imm < -_OVERSHOOT_TOL)), # ablation improved the margin + max_overshoot=float(max(0.0, (-imm).max())), + ) + return out + + +def _stability_fig(stats, out): + """Grouped bars: per-model std(imm), zero vs resample; overshoot count annotated.""" + models = list(stats) + x = np.arange(len(models)) + w = 0.38 + fig, ax = plt.subplots(figsize=(1.8 * len(models) + 2, 5), constrained_layout=True) + zb = ax.bar( + x - w / 2, [stats[m]["zero"]["std"] for m in models], w, color="tab:red", label="zero" + ) + rb = ax.bar( + x + w / 2, + [stats[m]["resample"]["std"] for m in models], + w, + color="tab:blue", + label="resample", + ) + for m, bz, br in zip(models, zb, rb, strict=False): + ax.annotate( + f"{stats[m]['zero']['n_overshoot']} ovr", + (bz.get_x() + bz.get_width() / 2, bz.get_height()), + ha="center", + va="bottom", + fontsize=7, + color="tab:red", + ) + ax.annotate( + f"{stats[m]['resample']['n_overshoot']} ovr", + (br.get_x() + br.get_width() / 2, br.get_height()), + ha="center", + va="bottom", + fontsize=7, + color="tab:blue", + ) + ax.set_xticks(x, [_MODEL_LABELS[m] for m in models]) + ax.set_ylabel("std(imm) over layers · margin (fraction of final)") + ax.set_title( + "Stability: resample perturbs consistently; zero swings + overshoots", fontweight="bold" + ) + ax.legend() + fig.savefig(out, dpi=140) + plt.close(fig) + print(f"wrote {out}") + + +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. + 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") + + # Claim 2 — stability. + stats = {m: _stability(args.in_dir, m) for m in args.models} + _stability_fig(stats, args.out_dir / "v2_stability.png") + + print("\nstability (margin) — std(imm) over layers, overshoot count / max:") + for m in args.models: + z, r = stats[m]["zero"], stats[m]["resample"] + print( + f" {m:10s} std: zero {z['std']:.3f} resample {r['std']:.3f} | " + f"overshoot layers: zero {z['n_overshoot']:2d} resample {r['n_overshoot']:2d} | " + f"max overshoot: zero {z['max_overshoot']:.2f} resample {r['max_overshoot']:.2f}" + ) + + +if __name__ == "__main__": + main() diff --git a/src/tfm_lens/adapters/base.py b/src/tfm_lens/adapters/base.py index ebc68e7..024fd3f 100644 --- a/src/tfm_lens/adapters/base.py +++ b/src/tfm_lens/adapters/base.py @@ -18,6 +18,9 @@ class ModelAdapter(ABC): # Capability declaration — static, overridden by subclasses as needed. needs_transpose: bool = False # whether the decoder wants (seq, batch, hidden) + # Token-axis index of the label token (the only position the decoder reads), for + # resample ablation's label←label bucketing. None = no token axis (3D residual). + label_token_index: int | None = None # ---- must be implemented ---- @property @@ -63,6 +66,34 @@ def identity_forward(self, *args, **kwargs): """ return args[0] if args else next(iter(kwargs.values())) + # ---- δ-injection hooks (overridable; default suits single-stream layers) ---- + def residual_of(self, layer_output): + """The residual stream(s) a layer carries forward, in *raw layer-output* + coordinates — stream picked, but **no token-slice and no norm** (unlike + ``readout``, which prepares the decoder input). This is the coordinate the + contribution ``δ = residual_of(out) − residual_of(in)`` is captured and a + donor ``δ`` is applied in. + + Default: single stream — unwrap a tuple (LimiX's ``(res, …)``, or the + capture cache's input tuple ``(x,)``), else the tensor itself. Override for + double-stream models (Mitra returns both ``(support, query)``). + """ + return layer_output[0] if isinstance(layer_output, tuple) else layer_output + + def inject_delta_forward(self, donor_delta, *args, **kwargs): + """What a δ-injected layer returns: the input residual **plus** ``donor_delta``, + in the layer's output shape. Mirrors ``identity_forward`` (same tuple layout) + but adds the donor contribution instead of passing the input through. + + Default: single stream — ``input + donor_delta``. Override for tuple-returning + layers: + + - LimiX: ``return args[0] + donor_delta, None, None`` + - Mitra: ``return args[0] + donor_delta[0], args[1] + donor_delta[1]`` + """ + x = args[0] if args else next(iter(kwargs.values())) + return x + donor_delta + @property def n_layers(self) -> int: return len(self.layers) diff --git a/src/tfm_lens/adapters/limix.py b/src/tfm_lens/adapters/limix.py index 15ea773..ea455f9 100644 --- a/src/tfm_lens/adapters/limix.py +++ b/src/tfm_lens/adapters/limix.py @@ -14,6 +14,7 @@ class LimixAdapter(ModelAdapter): needs_transpose = True # cls_y_decoder consumes (seq, batch, hidden) + label_token_index = -1 # label token is last on the token axis (see transformer.py) def __init__(self, model): self.model = model @@ -53,3 +54,7 @@ def readout(self, out, eval_pos: int): def identity_forward(self, *args, **kwargs): return args[0], None, None # LimiX layers return (residual, feat_attn, sample_attn) + + # residual_of: base default (out[0]) already picks the residual from the 3-tuple. + def inject_delta_forward(self, donor_delta, *args, **kwargs): + return args[0] + donor_delta, None, None # match the layer's 3-tuple return diff --git a/src/tfm_lens/adapters/mitra.py b/src/tfm_lens/adapters/mitra.py index 78d3fa4..0f7220e 100644 --- a/src/tfm_lens/adapters/mitra.py +++ b/src/tfm_lens/adapters/mitra.py @@ -19,6 +19,8 @@ class MitraAdapter(ModelAdapter): + label_token_index = 0 # label token is first on the token axis (einops.pack((y, x))) + def __init__(self, model): self.model = model # Force the standard-attention path (see module doc): the flash path has a @@ -67,3 +69,13 @@ def readout(self, out, eval_pos: int): def identity_forward(self, *args, **kwargs): # A skipped layer returns its two streams unchanged: (support, query). return args[0], args[1] + + def residual_of(self, layer_output): + # Double-stream: both (support, query) carry forward -> δ is captured/applied + # on both. layer_output is the 2-tuple (from a layer output or the capture + # cache's input tuple). + return layer_output[0], layer_output[1] + + def inject_delta_forward(self, donor_delta, *args, **kwargs): + # donor_delta is a (support_δ, query_δ) pair, one per stream. + return args[0] + donor_delta[0], args[1] + donor_delta[1] diff --git a/src/tfm_lens/core/interventions.py b/src/tfm_lens/core/interventions.py index 983002b..216265a 100644 --- a/src/tfm_lens/core/interventions.py +++ b/src/tfm_lens/core/interventions.py @@ -24,5 +24,27 @@ def skip_layer(adapter, idx): layer.forward = original +@contextmanager +def inject_delta(adapter, idx, donor_delta): + """Replace layer ``idx``'s contribution with ``donor_delta`` for the context: + the layer returns ``input_residual + donor_delta`` instead of its own output. + Policy-agnostic: resample ablation (#35) supplies a real donor δ; ``skip_layer`` + is the special case δ:=0. Upstream is untouched; downstream reacts freely. + + ``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. + """ + layer = adapter.layers[idx] + original = layer.forward + layer.forward = lambda *args, **kwargs: adapter.inject_delta_forward( + donor_delta, *args, **kwargs + ) + try: + yield + finally: + layer.forward = original + + # TODO(exp5): loop_layer(adapter, idx, n) — call a layer's forward n times in a # row, for the looping / repeated-layer experiment. diff --git a/src/tfm_lens/core/resample_ablation.py b/src/tfm_lens/core/resample_ablation.py new file mode 100644 index 0000000..3e8d21a --- /dev/null +++ b/src/tfm_lens/core/resample_ablation.py @@ -0,0 +1,113 @@ +"""Donor δ for resample ablation (#35). + +Resample ablation replaces a layer's contribution δ with a **real donor's** δ, +role-matched (issue #35 "Design settled"): + +- **row axis (all models):** context←context, query←query. +- **token axis (4D only):** label←label (decoder reads only the label token), + feature←feature random. 3D models have no token axis (``label_token_index=None``). + +Two steps: + +- ``donor_deltas`` — one donor table's per-layer δ, ``δ = residual_of(out) − residual_of(in)``. +- ``build_donor_delta`` — a **target-shaped** δ: each position filled by a random + 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``); +this module builds one ``donor_delta`` per (donor, layer). +""" + +import numpy as np +import torch + +from tfm_lens.core.capture import capture_layers + + +def _sub(out, inp): + """δ = out − in, preserving stream structure (Tensor, or ``(support, query)``).""" + if isinstance(out, tuple): + return tuple(o - i for o, i in zip(out, inp, strict=True)) + return out - inp + + +def layer_deltas(adapter, cache) -> list: + """Per-layer contribution δ from a capture cache: ``residual_of(out) − residual_of(in)``. + + ``len == n_layers``; each entry a Tensor (single stream) or ``(support, query)`` + tuple (double stream), in raw residual coordinates (all rows / tokens, no norm). + """ + return [ + _sub(adapter.residual_of(cache[m + 1]), adapter.residual_of(cache[m])) + for m in range(adapter.n_layers) + ] + + +def donor_deltas(adapter, X, y_train, eval_pos) -> list: + """Run one donor table through the frozen model → its per-layer δ (``layer_deltas``).""" + with torch.no_grad(), capture_layers(adapter) as cache: + adapter.forward_frozen(X, y_train, eval_pos) + return layer_deltas(adapter, cache) + + +def _blocks(residual, eval_pos) -> list: + """Split a residual into ``[context, query]`` blocks by row-role. + + - single stream: split the seq axis (dim 1) at ``eval_pos``. + - double stream ``(support, query)``: the streams *are* the split. + """ + if isinstance(residual, tuple): + return [residual[0], residual[1]] + return [residual[:, :eval_pos], residual[:, eval_pos:]] + + +def _fill_like(target, pool, rng) -> torch.Tensor: + """Tensor shaped like ``target`` (``[..., hidden]``), each cell = a random row of + ``pool`` (``[N, hidden]``), drawn uniformly with replacement.""" + lead = target.shape[:-1] + idx = torch.from_numpy(rng.integers(0, pool.shape[0], size=int(np.prod(lead)))).to(pool.device) + return pool[idx].reshape(*lead, pool.shape[-1]) + + +def _fill_block(t_block, d_block, label_index, rng) -> torch.Tensor: + """Fill one role block: each target cell ← a random donor cell of matching token-role. + + - ``label_index is None`` (3D): one bucket over rows. + - else (4D): label column vs feature columns are separate buckets. + """ + hidden = t_block.shape[-1] + if label_index is None: + return _fill_like(t_block, d_block.reshape(-1, hidden), rng) + + # Target and donor token counts differ across tables (different #features), so + # resolve the label position and feature indices *per block* — never index the + # donor with the target's token indices (that overruns when the donor is narrower). + t_label = label_index % t_block.shape[-2] + d_label = label_index % d_block.shape[-2] + t_feats = [t for t in range(t_block.shape[-2]) if t != t_label] + d_feats = [t for t in range(d_block.shape[-2]) if t != d_label] + out = torch.empty_like(t_block) + out[..., t_label, :] = _fill_like( + t_block[..., t_label, :], d_block[..., d_label, :].reshape(-1, hidden), rng + ) + out[..., t_feats, :] = _fill_like( + t_block[..., t_feats, :], d_block[..., d_feats, :].reshape(-1, hidden), rng + ) + return out + + +def build_donor_delta(target_residual, donor_delta, eval_pos_t, eval_pos_d, label_index, rng): + """A target-shaped donor δ: every position filled by a random draw from the + donor's **same-role** positions (row-role always; token-role for 4D). + + - ``target_residual`` — ``residual_of(target layer input)``: gives shape + the row split. + - ``donor_delta`` — one donor's δ at this layer (from :func:`donor_deltas`). + - ``eval_pos_t`` / ``eval_pos_d`` — target / donor train/test split (row-role boundary). + - ``label_index`` — ``adapter.label_token_index`` (None for 3D). + """ + t_blocks = _blocks(target_residual, eval_pos_t) + d_blocks = _blocks(donor_delta, eval_pos_d) + filled = [_fill_block(t, d, label_index, rng) for t, d in zip(t_blocks, d_blocks, strict=True)] + if isinstance(target_residual, tuple): + return tuple(filled) + return torch.cat(filled, dim=1) diff --git a/src/tfm_lens/evaluation/layerwise.py b/src/tfm_lens/evaluation/layerwise.py index cbf663f..615f442 100644 --- a/src/tfm_lens/evaluation/layerwise.py +++ b/src/tfm_lens/evaluation/layerwise.py @@ -129,7 +129,7 @@ def gt_logit_zscore_stats( - raises if sigma == 0 (degenerate task). """ y_test = np.asarray(y_test) - gt = _gt_logit(clean_logits[-1], y_test) + gt = _gt_logit(clean_logits[-1], y_test) # [-1] = final layer (the task's output scale) mu, sigma = float(gt.mean()), float(gt.std()) if sigma == 0.0: raise ValueError("degenerate task: final-layer GT-logit has zero spread across rows") diff --git a/src/tfm_lens/evaluation/self_repair.py b/src/tfm_lens/evaluation/self_repair.py index da4b45f..992a7dc 100644 --- a/src/tfm_lens/evaluation/self_repair.py +++ b/src/tfm_lens/evaluation/self_repair.py @@ -13,7 +13,9 @@ from scipy.special import softmax from tfm_lens.adapters.base import ModelAdapter -from tfm_lens.core.interventions import skip_layer +from tfm_lens.core.capture import capture_layers +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, @@ -89,8 +91,12 @@ def ablation_sweep( X_test: torch.Tensor, y_test: np.ndarray, n_classes: int, + ablation: str = "zero", + donor_tables: list[tuple] | None = None, + n_donors: int = 8, + seed: int = 0, ) -> dict: - """Baseline + skip-each-layer per-depth metric trajectories for one table. + """Baseline + ablate-each-layer per-depth metric trajectories for one table. One forward per condition feeds every metric. Returns:: @@ -99,7 +105,16 @@ def ablation_sweep( "zscore": {"mu": float, "sigma": float}} ``metric`` in {auc, gt_logit, margin}; ``zscore`` = clean-baseline stats for - cross-task GT-logit normalization. + cross-task GT-logit normalization. (Key stays ``"skip"`` for both modes so the + plots read either unchanged.) + + ``ablation`` (#35): + + - ``"zero"`` — skip the layer (δ:=0). Off-distribution; kept as the cross-check. + - ``"resample"`` — δ:= a role-matched real-donor δ (on-manifold, norm-preserving). + Needs ``donor_tables`` (leave-one-out real tables, each ``(X_train, y_train, + X_test)``); draws ``n_donors`` without replacement, **averages the metric over + donors** (not the δ), and adds ``{metric}_p25/_p75`` per layer for the IQR band. """ def _logits(): @@ -108,8 +123,95 @@ def _logits(): base_logits = _logits() baseline = _all_metrics(base_logits, y_test) mu, sigma = gt_logit_zscore_stats(base_logits, y_test) - skip = {} - for m in range(adapter.n_layers): - with skip_layer(adapter, m): - skip[m] = _all_metrics(_logits(), y_test) + + if ablation == "zero": + skip = {} + for m in range(adapter.n_layers): + with skip_layer(adapter, m): + skip[m] = _all_metrics(_logits(), y_test) + elif ablation == "resample": + if not donor_tables: + raise ValueError("resample ablation needs donor_tables (leave-one-out real tables)") + skip = _resample_skip( + adapter, + decoders, + X_train, + y_train, + X_test, + y_test, + n_classes, + donor_tables, + n_donors, + seed, + ) + else: + raise ValueError(f"unknown ablation {ablation!r}; expected 'zero' or 'resample'") + return {"baseline": baseline, "skip": skip, "zscore": {"mu": mu, "sigma": sigma}} + + +def _clone_residual(r): + return tuple(t.clone() for t in r) if isinstance(r, tuple) else r.clone() + + +def _target_residuals(adapter, decoders, X_train, y_train, X_test): + """Clean per-layer *input* residuals of the target forward — the shape + row-split + the donor draw fills against (values unused). One extra forward per task.""" + device = next(decoders[0].parameters()).device + eval_pos = X_train.shape[0] + X = torch.cat([X_train, X_test], dim=0).unsqueeze(0).to(device) + y = y_train.unsqueeze(0).to(device) + with torch.no_grad(), capture_layers(adapter) as cache: + adapter.forward_frozen(X, y, eval_pos) + return [_clone_residual(adapter.residual_of(cache[m])) for m in range(adapter.n_layers)] + + +def _average_over_donors(per_donor: list[dict], n_layers: int) -> dict: + """Per layer, mean each per-depth metric across donors (+ p25/p75 for the IQR band).""" + names = list(per_donor[0][0].keys()) + out = {} + for m in range(n_layers): + entry = {} + for name in names: + stack = np.array([d[m][name] for d in per_donor]) # [n_donors, n_depths] + entry[name] = stack.mean(0).tolist() + entry[f"{name}_p25"] = np.percentile(stack, 25, axis=0).tolist() + entry[f"{name}_p75"] = np.percentile(stack, 75, axis=0).tolist() + out[m] = entry + return out + + +def _resample_skip( + adapter, decoders, X_train, y_train, X_test, y_test, n_classes, donor_tables, n_donors, seed +) -> dict: + """Resample-ablate each layer: for ``n_donors`` leave-one-out donor tables, swap + the layer's δ for a role-matched donor δ, then average the metric over donors.""" + + def _logits(): + return predict_layers_logits(adapter, decoders, X_train, y_train, X_test, n_classes) + + device = next(decoders[0].parameters()).device + eval_pos_t = X_train.shape[0] + target_res = _target_residuals(adapter, decoders, X_train, y_train, X_test) + + k = min(n_donors, len(donor_tables)) + picks = np.random.default_rng(seed).permutation(len(donor_tables))[:k] # without replacement + + per_donor = [] + for di, i in enumerate(picks): + Xd_train, yd_train, Xd_test = donor_tables[i] + eval_pos_d = Xd_train.shape[0] + Xd = torch.cat([Xd_train, Xd_test], dim=0).unsqueeze(0).to(device) + yd = yd_train.unsqueeze(0).to(device) + d_deltas = donor_deltas(adapter, Xd, yd, eval_pos_d) + draw = np.random.default_rng(seed + 1 + di) # per-donor position draw + by_layer = {} + for m in range(adapter.n_layers): + donor_delta = build_donor_delta( + target_res[m], d_deltas[m], eval_pos_t, eval_pos_d, adapter.label_token_index, draw + ) + with inject_delta(adapter, m, donor_delta): + by_layer[m] = _all_metrics(_logits(), y_test) + per_donor.append(by_layer) + + return _average_over_donors(per_donor, adapter.n_layers) diff --git a/tests/conftest.py b/tests/conftest.py index 8308485..3eb78d6 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -101,6 +101,12 @@ def readout(self, out, eval_pos): def identity_forward(self, *args, **kwargs): return args[0], args[1] + def residual_of(self, layer_output): + return layer_output[0], layer_output[1] # both streams (support, query) + + def inject_delta_forward(self, donor_delta, *args, **kwargs): + return args[0] + donor_delta[0], args[1] + donor_delta[1] + @pytest.fixture(scope="session") def limix_ckpt() -> str: diff --git a/tests/test_adapter_base.py b/tests/test_adapter_base.py index aaaa434..c7f0ccc 100644 --- a/tests/test_adapter_base.py +++ b/tests/test_adapter_base.py @@ -9,7 +9,7 @@ import torch from tfm_lens.adapters.base import ModelAdapter -from toys import ToyAdapter3D +from toys import ToyAdapter3D, ToyAdapter4D class TestModelAdapter: @@ -56,3 +56,66 @@ class Incomplete(ModelAdapter): with pytest.raises(TypeError): Incomplete() + + +class TestResampleHooks: + """residual_of / inject_delta_forward — the coordinate resample ablation (#35) + captures a contribution δ in and applies a donor δ in. + + Key invariant: ``δ = residual_of(out) − residual_of(in)`` and + ``inject_delta_forward(δ, in) == out`` — re-applying a layer's *own* δ reproduces + its output, so a *donor* δ is a clean drop-in (only the contribution swaps). + """ + + # ---- residual_of: pick the stream, keep every row/token (no slice, no norm) ---- + def test_residual_of_single_stream(self, toy_adapter): + x = torch.randn(2, 5, ToyAdapter3D.HIDDEN) + # bare tensor (a layer output) -> itself; capture's input tuple (x,) -> x. + torch.testing.assert_close(toy_adapter.residual_of(x), x) + torch.testing.assert_close(toy_adapter.residual_of((x,)), x) + + def test_residual_of_4d_keeps_token_axis(self, toy_adapter_4d): + # bucketing must be able to index the label token, so the token axis stays. + res = torch.randn(2, 5, ToyAdapter4D.TOKENS, ToyAdapter4D.HIDDEN) + r = toy_adapter_4d.residual_of((res, None, None)) # LimiX-style 3-tuple + torch.testing.assert_close(r, res) + assert r.shape == (2, 5, ToyAdapter4D.TOKENS, ToyAdapter4D.HIDDEN) + + def test_residual_of_double_stream_returns_both(self, toy_adapter_double_stream): + s = torch.randn(2, 3, ToyAdapter3D.HIDDEN) + q = torch.randn(2, 2, ToyAdapter3D.HIDDEN) + rs, rq = toy_adapter_double_stream.residual_of((s, q)) + torch.testing.assert_close(rs, s) + torch.testing.assert_close(rq, q) + + # ---- inject_delta_forward: input + donor δ, in the layer's own output shape ---- + def test_resample_roundtrip_single_stream(self, toy_adapter): + x = torch.randn(2, 5, ToyAdapter3D.HIDDEN) + out = toy_adapter.layers[1](x) + delta = toy_adapter.residual_of(out) - toy_adapter.residual_of((x,)) + torch.testing.assert_close(toy_adapter.inject_delta_forward(delta, x), out) + + def test_resample_roundtrip_4d_keeps_tuple_shape(self, toy_adapter_4d): + x = torch.randn(2, 5, ToyAdapter4D.TOKENS, ToyAdapter4D.HIDDEN) + out = toy_adapter_4d.layers[1](x) # (residual, None, None) + delta = toy_adapter_4d.residual_of(out) - toy_adapter_4d.residual_of((x,)) + res = toy_adapter_4d.inject_delta_forward(delta, x) + assert isinstance(res, tuple) and len(res) == 3 and res[1] is None and res[2] is None + torch.testing.assert_close(res[0], out[0]) + + def test_resample_roundtrip_double_stream(self, toy_adapter_double_stream): + a = toy_adapter_double_stream + s = torch.randn(2, 3, ToyAdapter3D.HIDDEN) + q = torch.randn(2, 2, ToyAdapter3D.HIDDEN) + out = a.layers[1](s, q) # (lin(s), lin(q)) + r_out, r_in = a.residual_of(out), a.residual_of((s, q)) + delta = (r_out[0] - r_in[0], r_out[1] - r_in[1]) + res = a.inject_delta_forward(delta, s, q) + torch.testing.assert_close(res[0], out[0]) + torch.testing.assert_close(res[1], out[1]) + + def test_resample_applies_donor_delta(self, toy_adapter): + # a *foreign* δ swaps the contribution: output = input + donor δ, exactly. + x = torch.randn(2, 5, ToyAdapter3D.HIDDEN) + donor = torch.randn(2, 5, ToyAdapter3D.HIDDEN) + torch.testing.assert_close(toy_adapter.inject_delta_forward(donor, x), x + donor) diff --git a/tests/test_interventions.py b/tests/test_interventions.py index fb7379b..2764f11 100644 --- a/tests/test_interventions.py +++ b/tests/test_interventions.py @@ -8,7 +8,7 @@ import torch from tfm_lens.core.capture import capture_layers -from tfm_lens.core.interventions import skip_layer +from tfm_lens.core.interventions import inject_delta, skip_layer class TestSkipLayer: @@ -59,3 +59,50 @@ def test_restores_on_exception(self, toy_adapter, toy_input): after = self._capture(toy_adapter, toy_input) for a, b in zip(baseline, after, strict=True): torch.testing.assert_close(a, b) + + +class TestResampleLayer: + """inject_delta makes a layer output ``input + donor_delta``; δ:=0 == skip.""" + + def _capture(self, adapter, x, eval_pos=5): + with capture_layers(adapter) as cache: + adapter.forward_frozen(x, None, eval_pos) + return cache + + def test_zero_delta_is_skip(self, toy_adapter, toy_input): + # donor_delta == 0 → layer output equals its input (same as skip_layer). + delta = torch.zeros_like(toy_input) + with inject_delta(toy_adapter, 1, delta): + cache = self._capture(toy_adapter, toy_input) + torch.testing.assert_close(cache[2], cache[1]) # out(layer 1) == in(layer 1) + + def test_output_is_input_plus_delta(self, toy_adapter, toy_input): + delta = torch.randn_like(toy_input) + with inject_delta(toy_adapter, 1, delta): + cache = self._capture(toy_adapter, toy_input) + torch.testing.assert_close(cache[2], cache[1] + delta) # residual + donor δ + + def test_4d_keeps_tuple_shape(self, toy_adapter_4d, toy_input_4d): + delta = torch.randn_like(toy_input_4d) + with inject_delta(toy_adapter_4d, 1, delta): + cache = self._capture(toy_adapter_4d, toy_input_4d) + r_in = toy_adapter_4d.residual_of(cache[1]) + r_out = toy_adapter_4d.residual_of(cache[2]) + torch.testing.assert_close(r_out, r_in + delta) + + def test_double_stream_delta_per_stream(self, toy_adapter_double_stream, toy_input): + adapter = toy_adapter_double_stream + # a (support_δ, query_δ) pair; support has 3 rows, query 2 (eval_pos=3). + delta = (torch.randn(2, 3, adapter.HIDDEN), torch.randn(2, 2, adapter.HIDDEN)) + with inject_delta(adapter, 1, delta), capture_layers(adapter) as cache: + adapter.forward_frozen(toy_input, None, eval_pos=3) + torch.testing.assert_close(cache[2][0], cache[1][0] + delta[0]) # support + torch.testing.assert_close(cache[2][1], cache[1][1] + delta[1]) # query + + def test_restores_forward_on_exit(self, toy_adapter, toy_input): + baseline = self._capture(toy_adapter, toy_input) + with inject_delta(toy_adapter, 1, torch.randn_like(toy_input)): + pass + after = self._capture(toy_adapter, toy_input) + for a, b in zip(baseline, after, strict=True): + torch.testing.assert_close(a, b) diff --git a/tests/test_resample_ablation.py b/tests/test_resample_ablation.py new file mode 100644 index 0000000..0afa21b --- /dev/null +++ b/tests/test_resample_ablation.py @@ -0,0 +1,162 @@ +"""Donor δ capture + role bucketing for resample ablation (#35). + +Role integrity is checked by *tagging*: give every donor cell in a bucket a +constant value, then assert each target cell lands on its bucket's tag — proving +context←context / query←query and (4D) label←label / feature←feature. +""" + +import numpy as np +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 toys import ToyAdapter3D, ToyAdapter4D + +H = ToyAdapter3D.HIDDEN + + +def _rng(): + return np.random.default_rng(0) + + +class TestLayerDeltas: + def test_one_delta_per_layer_3d(self, toy_adapter, toy_input): + with capture_layers(toy_adapter) as cache: + toy_adapter.forward_frozen(toy_input, None, 5) + deltas = layer_deltas(toy_adapter, cache) + assert len(deltas) == toy_adapter.n_layers + assert all(d.shape == toy_input.shape for d in deltas) # [b, seq, hidden] + + def test_delta_reconstructs_output_4d(self, toy_adapter_4d, toy_input_4d): + # δ = residual_of(out) − residual_of(in) → residual_of(in) + δ == residual_of(out). + with capture_layers(toy_adapter_4d) as cache: + toy_adapter_4d.forward_frozen(toy_input_4d, None, 5) + deltas = layer_deltas(toy_adapter_4d, cache) + for m, d in enumerate(deltas): + r_in = toy_adapter_4d.residual_of(cache[m]) + r_out = toy_adapter_4d.residual_of(cache[m + 1]) + torch.testing.assert_close(r_in + d, r_out) + + def test_double_stream_delta_is_a_pair(self, toy_adapter_double_stream, toy_input): + with capture_layers(toy_adapter_double_stream) as cache: + toy_adapter_double_stream.forward_frozen(toy_input, None, eval_pos=3) + deltas = layer_deltas(toy_adapter_double_stream, cache) + assert len(deltas) == toy_adapter_double_stream.n_layers + assert all(isinstance(d, tuple) and len(d) == 2 for d in deltas) + + +class TestBuildDonorDelta3D: + def test_row_role_match(self): + # target: 3 context + 2 query; donor: context δ tagged +1, query δ tagged −1. + target = torch.zeros(1, 5, H) + donor = torch.empty(1, 4, H) + donor[:, :2] = 1.0 # 2 context rows + donor[:, 2:] = -1.0 # 2 query rows + out = build_donor_delta( + target, donor, eval_pos_t=3, eval_pos_d=2, label_index=None, rng=_rng() + ) + assert out.shape == (1, 5, H) + torch.testing.assert_close(out[:, :3], torch.ones(1, 3, H)) # context ← +1 + torch.testing.assert_close(out[:, 3:], -torch.ones(1, 2, H)) # query ← −1 + + def test_reproducible_with_seed(self): + target, donor = torch.zeros(1, 5, H), torch.randn(1, 4, H) + kw = dict(eval_pos_t=3, eval_pos_d=2, label_index=None) + a = build_donor_delta(target, donor, rng=np.random.default_rng(7), **kw) + b = build_donor_delta(target, donor, rng=np.random.default_rng(7), **kw) + torch.testing.assert_close(a, b) + + +class TestBuildDonorDelta4D: + def test_four_bucket_match(self): + # tokens: label at −1 (index T-1), features = the rest. Tag all four buckets. + T = ToyAdapter4D.TOKENS + target = torch.zeros(1, 4, T, H) # 2 context + 2 query rows (eval_pos_t=2) + donor = torch.empty(1, 3, T, H) # 1 context + 2 query rows (eval_pos_d=1) + donor[:, :1, -1, :] = 10.0 # context · label + donor[:, 1:, -1, :] = 20.0 # query · label + donor[:, :1, :-1, :] = 1.0 # context · feature + donor[:, 1:, :-1, :] = 2.0 # query · feature + out = build_donor_delta( + target, donor, eval_pos_t=2, eval_pos_d=1, label_index=-1, rng=_rng() + ) + assert out.shape == (1, 4, T, H) + assert torch.all(out[:, :2, -1, :] == 10.0) # context label ← context label + assert torch.all(out[:, 2:, -1, :] == 20.0) # query label ← query label + assert torch.all(out[:, :2, :-1, :] == 1.0) # context feats ← context feats + assert torch.all(out[:, 2:, :-1, :] == 2.0) # query feats ← query feats + + def test_label_never_drawn_from_features(self): + # the hard constraint: a feature δ must never land on the label token. + T = ToyAdapter4D.TOKENS + target = torch.zeros(1, 4, T, H) + donor = torch.zeros(1, 4, T, H) + donor[..., -1, :] = 99.0 # label column only + out = build_donor_delta( + target, donor, eval_pos_t=2, eval_pos_d=2, label_index=-1, rng=_rng() + ) + assert torch.all(out[..., -1, :] == 99.0) # label ← label (99) + assert torch.all(out[..., :-1, :] == 0.0) # features never see the 99 label tag + + def test_differing_token_counts(self): + # regression (smoke): donor narrower than target on the token axis (fewer + # features). Resolve label/feature indices per block — never index the donor + # with the target's token indices. + target = torch.zeros(1, 4, 8, H) # target: 7 feature tokens + 1 label (idx -1) + donor = torch.zeros(1, 3, 4, H) # donor: only 3 feature tokens + 1 label + donor[..., -1, :] = 5.0 # label column + donor[..., :-1, :] = 1.0 # feature columns + out = build_donor_delta( + target, donor, eval_pos_t=2, eval_pos_d=1, label_index=-1, rng=_rng() + ) + assert out.shape == (1, 4, 8, H) # keeps the *target* token count + assert torch.all(out[..., -1, :] == 5.0) # label ← donor label + assert torch.all(out[..., :-1, :] == 1.0) # all 7 target feats ← donor's 3 feats + + def test_donor_wider_than_target(self): + # the other direction: donor has more feature tokens than the target. + target = torch.zeros(1, 2, 3, H) # 2 feature tokens + 1 label + donor = torch.full((1, 2, 9, H), 1.0) # 8 feature tokens + 1 label + donor[..., -1, :] = 7.0 + out = build_donor_delta( + target, donor, eval_pos_t=1, eval_pos_d=1, label_index=-1, rng=_rng() + ) + assert out.shape == (1, 2, 3, H) + assert torch.all(out[..., -1, :] == 7.0) and torch.all(out[..., :-1, :] == 1.0) + + +class TestBuildDonorDeltaDoubleStream: + def test_3d_two_streams(self): + target = (torch.zeros(1, 3, H), torch.zeros(1, 2, H)) + donor = (torch.ones(1, 5, H), -torch.ones(1, 4, H)) # support=+1, query=−1 + out = build_donor_delta( + target, donor, eval_pos_t=0, eval_pos_d=0, label_index=None, rng=_rng() + ) + assert isinstance(out, tuple) and len(out) == 2 + torch.testing.assert_close(out[0], torch.ones(1, 3, H)) # support (context) + torch.testing.assert_close(out[1], -torch.ones(1, 2, H)) # query + + def test_4d_two_streams_with_label_token(self): + # Mitra-like: double stream AND a token axis, label at index 0. + T = 3 + target = (torch.zeros(1, 2, T, H), torch.zeros(1, 2, T, H)) + sup, qry = torch.empty(1, 4, T, H), torch.empty(1, 4, T, H) + sup[:, :, 0, :], sup[:, :, 1:, :] = 10.0, 1.0 # support: label / feature + qry[:, :, 0, :], qry[:, :, 1:, :] = 20.0, 2.0 # query: label / feature + out = build_donor_delta( + (target[0], target[1]), + (sup, qry), + eval_pos_t=0, + eval_pos_d=0, + label_index=0, + rng=_rng(), + ) + assert torch.all(out[0][:, :, 0, :] == 10.0) and torch.all(out[0][:, :, 1:, :] == 1.0) + assert torch.all(out[1][:, :, 0, :] == 20.0) and torch.all(out[1][:, :, 1:, :] == 2.0) + + +class TestDonorDeltasIntegration: + def test_donor_deltas_runs_forward_and_captures(self, toy_adapter_4d, toy_input_4d): + deltas = donor_deltas(toy_adapter_4d, toy_input_4d, None, eval_pos=5) + assert len(deltas) == toy_adapter_4d.n_layers + assert all(d.shape == toy_input_4d.shape for d in deltas) diff --git a/tests/test_self_repair.py b/tests/test_self_repair.py index a76511e..7d4f773 100644 --- a/tests/test_self_repair.py +++ b/tests/test_self_repair.py @@ -7,6 +7,16 @@ from tfm_lens.evaluation.self_repair import ablation_diffs, ablation_sweep, native_final_auc from toys import ToyAdapter3D +H = ToyAdapter3D.HIDDEN + + +def _donor_tables(): + """Two leave-one-out donor tables of differing row counts (binary).""" + return [ + (torch.randn(5, H), torch.tensor([0.0, 1.0, 0.0, 1.0, 0.0]), torch.randn(3, H)), + (torch.randn(7, H), torch.tensor([1.0, 0.0, 1.0, 0.0, 1.0, 0.0, 1.0]), torch.randn(2, H)), + ] + @pytest.fixture def toy_table(): @@ -49,3 +59,65 @@ def test_ablation_diffs_one_per_ablated_layer(toy_adapter, toy_decoders, toy_tab assert isinstance(m, int) assert isinstance(immediate, float) assert isinstance(final, float) + + +class TestResampleAblation: + def test_resample_sweep_shape_and_iqr(self, toy_adapter, toy_decoders, toy_table): + X_train, y_train, X_test, y_test = toy_table + res = ablation_sweep( + toy_adapter, + toy_decoders, + X_train, + y_train, + X_test, + y_test, + n_classes=2, + ablation="resample", + donor_tables=_donor_tables(), + n_donors=2, + seed=0, + ) + n_depths = toy_adapter.n_layers + 1 + assert set(res["skip"]) == set(range(toy_adapter.n_layers)) + for traj in res["skip"].values(): + for name in ("auc", "gt_logit", "margin"): + assert len(traj[name]) == n_depths # donor-averaged trajectory + assert len(traj[f"{name}_p25"]) == n_depths # IQR band kept + assert len(traj[f"{name}_p75"]) == n_depths + assert set(res["zscore"]) == {"mu", "sigma"} + + def test_resample_is_reproducible(self, toy_adapter, toy_decoders, toy_table): + X_train, y_train, X_test, y_test = toy_table + donors = _donor_tables() + kw = dict(n_classes=2, ablation="resample", donor_tables=donors, n_donors=2, seed=3) + a = ablation_sweep(toy_adapter, toy_decoders, X_train, y_train, X_test, y_test, **kw) + b = ablation_sweep(toy_adapter, toy_decoders, X_train, y_train, X_test, y_test, **kw) + assert a["skip"][0]["auc"] == b["skip"][0]["auc"] # fixed seed → identical + + def test_resample_requires_donor_tables(self, toy_adapter, toy_decoders, toy_table): + X_train, y_train, X_test, y_test = toy_table + with pytest.raises(ValueError, match="donor_tables"): + ablation_sweep( + toy_adapter, + toy_decoders, + X_train, + y_train, + X_test, + y_test, + n_classes=2, + ablation="resample", + ) + + def test_unknown_ablation_raises(self, toy_adapter, toy_decoders, toy_table): + X_train, y_train, X_test, y_test = toy_table + with pytest.raises(ValueError, match="unknown ablation"): + ablation_sweep( + toy_adapter, + toy_decoders, + X_train, + y_train, + X_test, + y_test, + n_classes=2, + ablation="mean", + ) diff --git a/tests/toys.py b/tests/toys.py index 30805c9..0bc1315 100644 --- a/tests/toys.py +++ b/tests/toys.py @@ -85,6 +85,7 @@ class ToyAdapter4D(ModelAdapter): layers return a 3-tuple. Exercises the 4D decode path with no checkpoint.""" needs_transpose = False + label_token_index = -1 # label token last on the token axis (LimiX convention) HIDDEN = 8 N_CLASSES = 4 @@ -116,3 +117,7 @@ def readout(self, out, eval_pos): def identity_forward(self, *args, **kwargs): return args[0], None, None # match the layer's 3-tuple return + + # residual_of: base default (out[0]) already picks the residual from the 3-tuple. + def inject_delta_forward(self, donor_delta, *args, **kwargs): + return args[0] + donor_delta, None, None # match the layer's 3-tuple return