spb/wpai0001_uqo
Public
1#!/usr/bin/env python32"""Sprints 3-5 on Tier 0: transplant families, full control battery, sweeps, mechanism tests.34Experiment id ``graft_tier0``. Requires the model matrix from scripts/train_models.py.5Usage: python scripts/run_transplant.py [--quick] [--device auto] [--seed-base 30000] [--tag main|confirmatory]6Every condition is evaluated on the same held-out splits; per-item outputs are saved; CIs are cluster7bootstraps over programs; control-adjusted effects are paired on identical items.8"""9from __future__ import annotations1011import argparse12import gzip13import json14import sys15import time16from pathlib import Path1718import numpy as np19import torch2021ROOT = Path(__file__).resolve().parents[1]22sys.path.insert(0, str(ROOT / "src"))23sys.path.insert(0, str(ROOT / "scripts"))2425from train_models import MODELS_DIR, ROLES, model_dir # noqa: E40226from wp_ai_0001.alignment.maps import FITTERS, AffineMap, fit_ridge # noqa: E40227from wp_ai_0001.config import OpSeqTaskConfig, canonical_json # noqa: E40228from wp_ai_0001.datasets import opseq # noqa: E40229from wp_ai_0001.datasets.common import Split, assert_disjoint # noqa: E40230from wp_ai_0001.evaluation.metrics import chance_level # noqa: E40231from wp_ai_0001.instrumentation.provenance import Manifest, make_run_id, sha256_text, utc_now # noqa: E40232from wp_ai_0001.models.hooks import Zero # noqa: E40233from wp_ai_0001.models.toy_transformer import resolve_device # noqa: E40234from wp_ai_0001.orchestration import registry # noqa: E40235from wp_ai_0001.statistics.bootstrap import cluster_bootstrap_mean, cluster_bootstrap_paired_diff, seeds_summary # noqa: E40236from wp_ai_0001.training.train_toy import load as load_model # noqa: E40237from wp_ai_0001.transplantation.experiment import ( # noqa: E40238 MAX_SEQ, OMEGA, build_eval_splits, collect_rows, distill, evaluate_condition, fit_adapter_from_donor, omega_step_acc, train_supervised_adapter,39)40from wp_ai_0001.transplantation.grafts import ( # noqa: E40241 ActivationGraft, AdapterGraft, GraftSpec, NormMatchedRandomGraft, SteeringGraft, SubspaceGraft, copy_embedding_row, weight_graft,42)4344EVAL = ["target_iid", "control_iid", "comp_heldout"]454647class Exp:48 def __init__(self, args):49 self.args = args50 self.device = resolve_device(args.device)51 self.task = ROLES["donor"][0]52 self.vocab = opseq.vocab_for(self.task)53 self.n_boot = 300 if args.quick else 200054 self.models: dict[str, object] = {}55 self.results: dict = {"conditions": {}, "pairs": {}, "notes": {}, "unsupported": {}}56 self.items: dict = {}5758 # ---------- models ----------59 def model(self, key: str):60 if key not in self.models:61 m, v = load_model(self._dir(key), self.device)62 assert v.itos == self.vocab.itos, f"vocab mismatch for {key}"63 m.eval()64 self.models[key] = m65 return self.models[key]6667 @staticmethod68 def _dir(key: str) -> Path:69 role, seed = key.rsplit("_s", 1)70 return model_dir(role, int(seed))7172 # ---------- evaluation bookkeeping ----------73 def record(self, cond: str, pair: str, res: dict, extra: dict | None = None) -> None:74 rec = {"pair": pair, **(extra or {})}75 for n, r in res.items():76 om = omega_step_acc(self.splits[n], r)77 valid = ~np.isnan(om)78 rec[n] = {79 "accuracy": r.accuracy, "trace_accuracy": r.trace_accuracy,80 "ci95": cluster_bootstrap_mean(r.correct.astype(float), r.clusters, self.n_boot, seed=0).to_dict(),81 "omega_step_acc": float(om[valid].mean()) if valid.any() else None,82 }83 self.items[f"{cond}|{pair}|{n}"] = {"correct": r.correct, "trace": r.trace_correct, "omega": om, "clusters": r.clusters, "pred": r.predictions}84 self.results["conditions"].setdefault(cond, {})[pair] = rec8586 def paired(self, cond_a: str, cond_b: str, pair: str, n: str = "target_iid", metric: str = "correct") -> dict | None:87 ka, kb = f"{cond_a}|{pair}|{n}", f"{cond_b}|{pair}|{n}"88 if ka not in self.items or kb not in self.items:89 return None90 a, b = self.items[ka], self.items[kb]91 if metric == "omega":92 va, vb = a["omega"], b["omega"]93 ok = ~np.isnan(va) & ~np.isnan(vb)94 cl = [c for c, o in zip(a["clusters"], ok) if o]95 return cluster_bootstrap_paired_diff(va[ok], vb[ok], cl, self.n_boot, seed=1).to_dict()96 return cluster_bootstrap_paired_diff(a[metric].astype(float), b[metric].astype(float), a["clusters"], self.n_boot, seed=1).to_dict()9798 # ---------- main ----------99 def run(self):100 a = self.args101 n_eval, n_align = (300, 600) if a.quick else (1000, 3000)102 index = json.loads((MODELS_DIR / "index.json").read_text())103 train_hashes: set[str] = set()104 for k, v in index.items():105 sp = Split.load(self._dir(k) / "train_split.json")106 train_hashes |= sp.item_hashes()107 self.splits = build_eval_splits(self.task, train_hashes, a.seed_base, n_eval, n_align)108 assert_disjoint(*self.splits.values())109 self.results["splits"] = {k: {"n": len(v), "hash": v.hash, "seed": v.seed} for k, v in self.splits.items()}110 self.results["chance"] = {k: chance_level(self.splits[k]) for k in EVAL}111 self.results["leakage_audit"] = {"alignment_vs_eval_overlap": 0, "train_vs_eval_overlap": 0, "n_train_hashes": len(train_hashes)}112 seeds = [0] if a.quick else [0, 1, 2]113 pairs = [(f"donor_s{i}", f"recipient_s{j}") for i in seeds for j in seeds]114 self.results["pair_meta"] = {f"{d}->{r}": {"same_init_seed": d.rsplit("_s", 1)[1] == r.rsplit("_s", 1)[1]} for d, r in pairs}115 matched = [(f"donor_s{i}", f"recipient_s{i}") for i in seeds]116 hook = "resid_post.0"117 t0 = time.time()118119 # ---------- baselines ----------120 for key in sorted({p for pr in pairs for p in pr} | ({f"unrelated_s{i}" for i in seeds})):121 self.record("base", key, evaluate_condition(self.model(key), None, None, self.splits, self.vocab, self.device, EVAL + ["len_ood"]))122 print(f"[{utc_now()}] baselines done ({time.time()-t0:.0f}s)")123124 # ---------- maps on alignment data ----------125 self.rows: dict[tuple[str, str], object] = {}126127 def rows(key, h=hook, split="align_target"):128 k = (key, h, split)129 if k not in self.rows:130 self.rows[k] = collect_rows(self.model(key), self.splits[split], self.vocab, h, self.device)131 return self.rows[k]132133 self.results["maps"] = {}134 maps: dict[tuple, AffineMap] = {}135 for d, r in pairs:136 X, Y = rows(d).X, rows(r).X137 Xu, Yu = rows(d, split="align_unrelated").X, rows(r, split="align_unrelated").X138 for fam in ("identity", "procrustes", "ridge", "cca", "random_orth"):139 m = FITTERS[fam](X, Y) if fam != "random_orth" else FITTERS[fam](X, Y, seed=0)140 maps[(d, r, fam)] = m141 self.results["maps"][f"{d}->{r}|{fam}"] = {"r2_fit": m.r2(X, Y), "r2_unrelated_heldout": m.r2(Xu, Yu), "n_params": m.n_params, **{k: v for k, v in m.info.items() if k != "fit_seconds"}}142 mu = fit_ridge(Xu, Yu, family="unrelated_data")143 maps[(d, r, "unrelated_data")] = mu144 mup = FITTERS["procrustes"](Xu, Yu)145 maps[(d, r, "unrelated_procrustes")] = mup146 self.results["maps"][f"{d}->{r}|unrelated_procrustes"] = {"r2_fit_unrelated": mup.r2(Xu, Yu), "r2_on_target_rows": mup.r2(X, Y), "n_params": mup.n_params}147 self.results["maps"][f"{d}->{r}|unrelated_data"] = {"r2_fit_unrelated": mu.r2(Xu, Yu), "r2_on_target_rows": mu.r2(X, Y), "n_params": mu.n_params}148 print(f"[{utc_now()}] maps fitted ({time.time()-t0:.0f}s)")149150 # ---------- activation grafts: map family × targeting (all pairs) ----------151 for d, r in pairs:152 pair = f"{d}->{r}"153 R = self.model(r)154 for fam in ("identity", "procrustes", "ridge", "cca"):155 for tg in ("all", "omega"):156 spec = GraftSpec("activation", hook, hook, 1.0, tg)157 g = ActivationGraft(spec, self.model(d), maps[(d, r, fam)])158 self.record(f"act_{fam}_{tg}", pair, evaluate_condition(R, hook, g, self.splits, self.vocab, self.device, EVAL), {"map": fam, "targeting": tg, "applied": g.applied})159 # controls160 spec = GraftSpec("activation", hook, hook, 1.0, "all")161 self.record("ctl_sham_map", pair, evaluate_condition(R, hook, ActivationGraft(spec, self.model(d), maps[(d, r, "random_orth")]), self.splits, self.vocab, self.device, EVAL))162 self.record("ctl_unrelated_data_map", pair, evaluate_condition(R, hook, ActivationGraft(spec, self.model(d), maps[(d, r, "unrelated_data")]), self.splits, self.vocab, self.device, EVAL))163 self.record("act_unrelated_procrustes_all", pair, evaluate_condition(R, hook, ActivationGraft(spec, self.model(d), maps[(d, r, "unrelated_procrustes")]), self.splits, self.vocab, self.device, EVAL), {"map": "procrustes fitted on non-Ω data"})164 self.record("ctl_norm_matched_random", pair, evaluate_condition(R, hook, NormMatchedRandomGraft(spec, self.model(d), maps[(d, r, "ridge")], seed=0), self.splits, self.vocab, self.device, EVAL))165 self.record("ctl_shuffled_token", pair, evaluate_condition(R, hook, ActivationGraft(spec, self.model(d), maps[(d, r, "ridge")], shuffle_tokens=torch.Generator().manual_seed(0)), self.splits, self.vocab, self.device, EVAL))166 # shuffled layer: donor layer 3 -> recipient layer 0 with its own ridge map167 X3 = rows(d, h="resid_post.3").X168 m3 = fit_ridge(X3, rows(r).X)169 self.record("ctl_shuffled_layer", pair, evaluate_condition(R, hook, ActivationGraft(GraftSpec("activation", hook, "resid_post.3", 1.0, "all"), self.model(d), m3), self.splits, self.vocab, self.device, EVAL))170 # unrelated-skill donor (competent on vex, not on Ω)171 u = f"unrelated_s{r.rsplit('_s',1)[1]}"172 mu = fit_ridge(rows(u).X, rows(r).X)173 self.record("ctl_unrelated_donor", pair, evaluate_condition(R, hook, ActivationGraft(spec, self.model(u), mu), self.splits, self.vocab, self.device, EVAL))174 print(f"[{utc_now()}] activation grafts + controls {pair} ({time.time()-t0:.0f}s)")175176 # ---------- steering, subspace, adapters, distillation, weight grafts (matched seeds) ----------177 for d, r in matched:178 pair = f"{d}->{r}"179 R, D = self.model(r), self.model(d)180 rd, rr = rows(d), rows(r)181 REF = "unrelated_procrustes" # reference map for secondary families (D7): orthogonal, fitted on task-free rows182 mapped = maps[(d, r, REF)](rd.X)183 # steering vector: mapped mean(Ω rows) - mean(non-Ω rows) of the donor184 v = mapped[rd.is_omega].mean(0) - mapped[~rd.is_omega].mean(0)185 for alpha in ([1.0] if a.quick else [0.5, 1.0, 2.0, 4.0]):186 self.record(f"steer_a{alpha}", pair, evaluate_condition(R, hook, SteeringGraft(GraftSpec("steering", hook, None, alpha, "all"), v), self.splits, self.vocab, self.device, EVAL), {"alpha": alpha})187 # subspace grafts: PCA of (mapped donor - recipient) on Ω rows; norm-matched random subspace of the same k188 diff = (mapped - rr.X)[rd.is_omega]189 U, S, Vt = np.linalg.svd(diff - diff.mean(0), full_matrices=False)190 rng = np.random.default_rng(0)191 for k in ([4, 32] if a.quick else [1, 2, 4, 8, 16, 32, 64, 128]):192 basis = Vt[:k].T193 self.record(f"subspace_k{k}", pair, evaluate_condition(R, hook, SubspaceGraft(GraftSpec("subspace", hook, hook, 1.0, "all"), D, maps[(d, r, REF)], basis), self.splits, self.vocab, self.device, EVAL), {"k": k, "explained_var": float((S[:k] ** 2).sum() / (S ** 2).sum()), "map": REF})194 rb = rng.standard_normal((basis.shape[0], k))195 self.record(f"subspace_random_k{k}", pair, evaluate_condition(R, hook, SubspaceGraft(GraftSpec("subspace", hook, hook, 1.0, "all"), D, maps[(d, r, REF)], rb), self.splits, self.vocab, self.device, EVAL), {"k": k, "map": REF})196 # adapter graft (donor-free at inference): ridge from recipient rows to mapped donor rows197 ad = fit_adapter_from_donor(rr.X, mapped)198 self.record("adapter_graft", pair, evaluate_condition(R, hook, AdapterGraft(GraftSpec("adapter", hook, None, 1.0, "all"), ad), self.splits, self.vocab, self.device, EVAL), {"n_params": ad.n_params, "n_fit_items": len(self.splits["align_target"]), "targets": REF})199 self.adapters = getattr(self, "adapters", {})200 self.adapters[pair] = ad201 # matched-size supervised adapter (labels, no donor) on the same alignment items202 steps = 100 if a.quick else 300203 sup = train_supervised_adapter(R, hook, self.splits["align_target"], self.vocab, self.device, steps=steps)204 self.record("baseline_supervised_adapter", pair, evaluate_condition(R, hook, AdapterGraft(GraftSpec("adapter", hook, None, 1.0, "all"), sup), self.splits, self.vocab, self.device, EVAL), {"n_params": sup.n_params, "steps": steps})205 # distillation baseline (full recipient fine-tuned on donor distributions, same items)206 student = distill(R, D, self.splits["align_target"], self.vocab, self.device, steps=steps)207 self.record("baseline_distillation", pair, evaluate_condition(student, None, None, self.splits, self.vocab, self.device, EVAL), {"steps": steps, "n_params_trained": student.n_params})208 del student209 # weight grafts (same architecture): copies into a fresh copy of the recipient210 import copy as _copy211212 omega_id = self.vocab.stoi[OMEGA]213 for name, parts, src_key, src_layer in [214 ("weight_L0_attn", ["attn"], d, 0), ("weight_L0_mlp", ["mlp"], d, 0), ("weight_L0_attn_mlp", ["attn", "mlp"], d, 0),215 ("weight_L0_block_ln", ["attn", "mlp", "ln"], d, 0), ("weight_embed_omega", [], d, 0), ("weight_L0_block_embed", ["attn", "mlp", "ln"], d, 0),216 ("weight_L012_blocks_embed", ["attn", "mlp", "ln"], d, 0),217 ("ctl_weight_shuffled_layer", ["attn", "mlp", "ln"], d, 3), ("ctl_weight_unrelated_donor", ["attn", "mlp", "ln"], f"unrelated_s{r.rsplit('_s',1)[1]}", 0),218 ]:219 Rc = _copy.deepcopy(R)220 rep = {"n_params_copied": 0}221 S = self.model(src_key)222 if parts:223 rep = weight_graft(Rc, S, parts, src_layer, 0)224 if name == "weight_L012_blocks_embed":225 for l in (1, 2):226 rep["n_params_copied"] += weight_graft(Rc, S, parts, l, l)["n_params_copied"]227 if name in ("weight_embed_omega", "weight_L0_block_embed", "weight_L012_blocks_embed"):228 rep["n_params_copied"] += copy_embedding_row(Rc, S, omega_id)229 self.record(name, pair, evaluate_condition(Rc, None, None, self.splits, self.vocab, self.device, EVAL), rep)230 del Rc231 print(f"[{utc_now()}] steering/subspace/adapter/weights {pair} ({time.time()-t0:.0f}s)")232233 # ---------- sweeps (matched seed 0): layer grid, dose ----------234 d, r = matched[0]235 pair = f"{d}->{r}"236 R, D = self.model(r), self.model(d)237 L = R.cfg.n_layers238 grid = {}239 for ld in range(L):240 for lr_ in range(L):241 hd, hr = f"resid_post.{ld}", f"resid_post.{lr_}"242 m = FITTERS["procrustes"](rows(d, h=hd, split="align_unrelated").X, rows(r, h=hr, split="align_unrelated").X)243 res = evaluate_condition(R, hr, ActivationGraft(GraftSpec("activation", hr, hd, 1.0, "all"), D, m), self.splits, self.vocab, self.device, ["target_iid", "control_iid"])244 grid[f"{ld}->{lr_}"] = {"target": res["target_iid"].accuracy, "control": res["control_iid"].accuracy, "omega_step": float(np.nanmean(omega_step_acc(self.splits["target_iid"], res["target_iid"])))}245 self.results["layer_grid"] = {"pair": pair, "map": "unrelated_procrustes", "cells": grid}246 dose = {}247 for alpha in ([0.0, 0.5, 1.0] if a.quick else [0.0, 0.1, 0.25, 0.5, 0.75, 1.0, 1.25, 1.5, 2.0]):248 res = evaluate_condition(R, hook, ActivationGraft(GraftSpec("activation", hook, hook, alpha, "all"), D, maps[(d, r, "unrelated_procrustes")]), self.splits, self.vocab, self.device, ["target_iid", "control_iid"])249 dose[str(alpha)] = {"target": res["target_iid"].accuracy, "control": res["control_iid"].accuracy}250 self.results["dose"] = {"pair": pair, "map": "unrelated_procrustes", "cells": dose}251 print(f"[{utc_now()}] sweeps done ({time.time()-t0:.0f}s)")252253 # ---------- mismatched pairs ----------254 mis = {}255 for rk, layers in (("recipient_w96_s0", [(0, 0), (1, 1)]), ("recipient_L6_s0", [(0, 0), (0, 1), (1, 1), (1, 2), (2, 2)])):256 if not (self._dir(rk) / "model.pt").exists():257 self.results["unsupported"][rk] = "model not trained"258 continue259 Rm = self.model(rk)260 base = evaluate_condition(Rm, None, None, self.splits, self.vocab, self.device, EVAL)261 self.record("base", rk, base)262 for ld, lr_ in layers:263 hd, hr = f"resid_post.{ld}", f"resid_post.{lr_}"264 for fam in ("procrustes", "ridge", "cca", "unrelated_procrustes"):265 if fam == "unrelated_procrustes":266 X, Y = rows(d, h=hd, split="align_unrelated").X, rows(rk, h=hr, split="align_unrelated").X267 m = FITTERS["procrustes"](X, Y)268 else:269 X, Y = rows(d, h=hd).X, rows(rk, h=hr).X270 m = FITTERS[fam](X, Y)271 self.record(f"mismatch_{fam}_L{ld}to{lr_}", f"{d}->{rk}", evaluate_condition(Rm, hr, ActivationGraft(GraftSpec("activation", hr, hd, 1.0, "all"), D, m), self.splits, self.vocab, self.device, EVAL), {"r2_fit": m.r2(X, Y)})272 sham = FITTERS["random_orth"](rows(d, h=hd).X, rows(rk, h=hr).X, seed=0)273 self.record(f"mismatch_sham_L{ld}to{lr_}", f"{d}->{rk}", evaluate_condition(Rm, hr, ActivationGraft(GraftSpec("activation", hr, hd, 1.0, "all"), D, sham), self.splits, self.vocab, self.device, EVAL))274 mis[rk] = "done"275 self.results["mismatch"] = mis276 # stride-2 pair (multi-step composition inside one prediction)277 if (self._dir("donor_s2_s0") / "model.pt").exists() and (self._dir("recipient_s2_s0") / "model.pt").exists():278 task2 = ROLES["donor_s2"][0]279 ex2 = set()280 for k in ("donor_s2_s0", "recipient_s2_s0"):281 ex2 |= Split.load(self._dir(k) / "train_split.json").item_hashes()282 sp2 = build_eval_splits(task2, ex2, a.seed_base + 500, n_eval, n_align)283 saved = self.splits284 self.splits = sp2285 D2, R2 = self.model("donor_s2_s0"), self.model("recipient_s2_s0")286 self.rows_s2 = {}287 self.record("base", "donor_s2_s0", evaluate_condition(D2, None, None, sp2, self.vocab, self.device, EVAL))288 self.record("base", "recipient_s2_s0", evaluate_condition(R2, None, None, sp2, self.vocab, self.device, EVAL))289 for l in range(L):290 h = f"resid_post.{l}"291 X = collect_rows(D2, sp2["align_target"], self.vocab, h, self.device).X292 Y = collect_rows(R2, sp2["align_target"], self.vocab, h, self.device).X293 m = fit_ridge(X, Y)294 self.record(f"stride2_act_ridge_L{l}", "donor_s2_s0->recipient_s2_s0", evaluate_condition(R2, h, ActivationGraft(GraftSpec("activation", h, h, 1.0, "all"), D2, m), sp2, self.vocab, self.device, EVAL), {"r2_fit": m.r2(X, Y)})295 Xu = collect_rows(D2, sp2["align_unrelated"], self.vocab, h, self.device).X296 Yu = collect_rows(R2, sp2["align_unrelated"], self.vocab, h, self.device).X297 mp = FITTERS["procrustes"](Xu, Yu)298 self.record(f"stride2_act_unrelated_procrustes_L{l}", "donor_s2_s0->recipient_s2_s0", evaluate_condition(R2, h, ActivationGraft(GraftSpec("activation", h, h, 1.0, "all"), D2, mp), sp2, self.vocab, self.device, EVAL), {"r2_fit": mp.r2(Xu, Yu)})299 sham = FITTERS["random_orth"](X, Y, seed=0)300 self.record(f"stride2_sham_L{l}", "donor_s2_s0->recipient_s2_s0", evaluate_condition(R2, h, ActivationGraft(GraftSpec("activation", h, h, 1.0, "all"), D2, sham), sp2, self.vocab, self.device, EVAL))301 self.results["splits_stride2"] = {k: {"n": len(v), "hash": v.hash} for k, v in sp2.items()}302 self.splits = saved303 else:304 self.results["unsupported"]["stride2"] = "models not trained"305 print(f"[{utc_now()}] mismatched + stride-2 done ({time.time()-t0:.0f}s)")306307 # ---------- mechanism: donor ablation -> loss -> rescue; bidirectional ----------308 mech = {}309 d0, d1, r0 = matched[0][0], "donor_s1", matched[0][1]310 if not (self._dir(d1) / "model.pt").exists():311 d1 = d0312 self.results["unsupported"]["rescue"] = "intact donor_s1 missing: rescue source is the ablated model itself (invalid)"313 Dab = self.model(d0)314 zero = Zero(positions=None) # ablate attention output of layer 0 at ALL positions -> but only read positions matter315 # ablation implemented as hook on attn_out.0 restricted to read positions via positions=ctx.positions is not supported by Zero;316 # use the graft machinery: replace with zeros at read positions through a norm-0 "steering" of -x -> SteeringGraft cannot; so patch:317 from wp_ai_0001.models.hooks import HookContext318319 class _ZeroAtRead:320 def __call__(self, x, ctx: HookContext):321 if ctx.positions is None:322 return x323 out = x.clone()324 out[torch.arange(x.shape[0], device=x.device), ctx.positions] = 0325 return out326327 zr = _ZeroAtRead()328 Dab.hooks.add("attn_out.0", zr)329 abl = evaluate_condition(Dab, None, None, self.splits, self.vocab, self.device, EVAL)330 self.record("mech_donor_ablated_attn0", d0, abl)331 # rescue: graft intact donor_s1's resid_post.0 (ridge-mapped) into the ablated donor_s0 at read positions332 Xr, Yr = rows(d1).X, rows(d0).X # NOTE: Yr collected on the intact d0 (hook added after) -> re-collect under ablation333 Yab = collect_rows(Dab, self.splits["align_target"], self.vocab, hook, self.device).X334 mres = FITTERS["procrustes"](rows(d1, split="align_unrelated").X, rows(d0, split="align_unrelated").X) # task-free Procrustes into the intact donor's coordinates335 res = evaluate_condition(Dab, hook, ActivationGraft(GraftSpec("activation", hook, hook, 1.0, "all"), self.model(d1), mres), self.splits, self.vocab, self.device, EVAL)336 self.record("mech_rescue_from_intact_donor", f"{d1}->{d0}(ablated)", res)337 shamm = FITTERS["random_orth"](rows(d1, split="align_unrelated").X, rows(d0, split="align_unrelated").X, seed=0)338 res = evaluate_condition(Dab, hook, ActivationGraft(GraftSpec("activation", hook, hook, 1.0, "all"), self.model(d1), shamm), self.splits, self.vocab, self.device, EVAL)339 self.record("mech_rescue_sham", f"{d1}->{d0}(ablated)", res)340 Dab.hooks.remove("attn_out.0", zr)341 # bidirectional: recipient -> donor graft (should not harm donor on control, should remove Ω on target if replacing)342 mrd = fit_ridge(rows(r0).X, rows(d0).X)343 res = evaluate_condition(self.model(d0), hook, ActivationGraft(GraftSpec("activation", hook, hook, 1.0, "all"), self.model(r0), mrd), self.splits, self.vocab, self.device, EVAL)344 self.record("mech_reverse_recipient_to_donor", f"{r0}->{d0}", res)345 self.results["mechanism"] = {"ablated_component": "attn_out.0 at read positions (donor_s0)", "rescue_source": d1, "yab_norm_mean": float(np.linalg.norm(Yab, axis=1).mean())}346 print(f"[{utc_now()}] mechanism done ({time.time()-t0:.0f}s)")347348 # ---------- rule vs cache (H5): maps/adapters fitted on a subset of Ω input values ----------349 d, r = matched[0]350 pair = f"{d}->{r}"351 rd, rr = rows(d), rows(r)352 seen_vals = {0, 1, 2, 3, 4, 5}353 keep = (~rd.is_omega) | np.isin(rd.omega_input, list(seen_vals))354 m_seen = fit_ridge(rd.X[keep], rr.X[keep])355 mp_seen = FITTERS["procrustes"](rd.X[keep], rr.X[keep])356 ad_seen = fit_adapter_from_donor(rr.X[keep], maps[(d, r, "unrelated_procrustes")](rd.X[keep]))357 # partition target_iid by whether ALL Ω inputs are in the seen set358 from wp_ai_0001.transplantation.experiment import omega_input_values359360 oin = omega_input_values(self.splits["target_iid"])361 seen_mask = np.array([all(v in seen_vals for v in vals) for vals in oin])362 unseen_mask = np.array([all(v not in seen_vals for v in vals) for vals in oin])363 self.results["rule_vs_cache"] = {"seen_values": sorted(seen_vals), "n_items_seen": int(seen_mask.sum()), "n_items_unseen": int(unseen_mask.sum()), "n_fit_rows": int(keep.sum())}364 for cname, g in (("h5_act_ridge_seenfit", ActivationGraft(GraftSpec("activation", hook, hook, 1.0, "all"), self.model(d), m_seen)),365 ("h5_act_procrustes_seenfit", ActivationGraft(GraftSpec("activation", hook, hook, 1.0, "all"), self.model(d), mp_seen)),366 ("h5_adapter_seenfit", AdapterGraft(GraftSpec("adapter", hook, None, 1.0, "all"), ad_seen))):367 res = evaluate_condition(self.model(r), hook, g, self.splits, self.vocab, self.device, ["target_iid"])368 self.record(cname, pair, res)369 om = omega_step_acc(self.splits["target_iid"], res["target_iid"])370 cl = res["target_iid"].clusters371 self.results["rule_vs_cache"][cname] = {372 "seen": cluster_bootstrap_mean(om[seen_mask], [c for c, k in zip(cl, seen_mask) if k], self.n_boot).to_dict(),373 "unseen": cluster_bootstrap_mean(om[unseen_mask], [c for c, k in zip(cl, unseen_mask) if k], self.n_boot).to_dict(),374 }375 print(f"[{utc_now()}] rule-vs-cache done ({time.time()-t0:.0f}s)")376377 # ---------- persistent graft: save adapter for an independent process ----------378 pdir = self.out_dir / "persistent"379 pdir.mkdir(parents=True, exist_ok=True)380 ad = self.adapters[f"{matched[0][0]}->{matched[0][1]}"]381 np.savez(pdir / "adapter_graft.npz", W=ad.W, b=ad.b, hook=np.array(hook), recipient=np.array(matched[0][1]))382 self.results["persistent"] = {"adapter_file": str((pdir / "adapter_graft.npz").relative_to(ROOT)), "evaluated_in_fresh_process": False}383384 # ---------- derived contrasts ----------385 self.results["contrasts"] = {}386 for d, r in pairs:387 pair = f"{d}->{r}"388 c = {}389 for cond in [k for k in self.results["conditions"] if pair in self.results["conditions"][k] and k != "base"]:390 gain = self.paired(cond, "base", pair) if f"base|{r}|target_iid" in self.items else None391 # paired against the recipient baseline: items identical; recipient base stored under key r392 a_ = self.items.get(f"{cond}|{pair}|target_iid")393 b_ = self.items.get(f"base|{r}|target_iid")394 s_ = self.items.get(f"ctl_sham_map|{pair}|target_iid")395 if a_ is None or b_ is None:396 continue397 gain = cluster_bootstrap_paired_diff(a_["correct"].astype(float), b_["correct"].astype(float), a_["clusters"], self.n_boot, seed=2).to_dict()398 entry = {"transfer_gain": gain}399 if s_ is not None:400 adj = a_["correct"].astype(float) - b_["correct"].astype(float) - (s_["correct"].astype(float) - b_["correct"].astype(float))401 entry["control_adjusted"] = cluster_bootstrap_mean(adj, a_["clusters"], self.n_boot, seed=3).to_dict()402 cb = self.items.get(f"base|{r}|control_iid")403 cc = self.items.get(f"{cond}|{pair}|control_iid")404 if cb is not None and cc is not None:405 entry["retention_control"] = cluster_bootstrap_paired_diff(cc["correct"].astype(float), cb["correct"].astype(float), cc["clusters"], self.n_boot, seed=4).to_dict()406 c[cond] = entry407 self.results["contrasts"][pair] = c408 # seed summaries for the main conditions409 self.results["across_pairs"] = {}410 for cond, d_ in self.results["conditions"].items():411 vals = [v["target_iid"]["accuracy"] for p, v in d_.items() if "target_iid" in v and "->" in p]412 if len(vals) >= 2:413 self.results["across_pairs"][cond] = {"target": seeds_summary(vals), "control": seeds_summary([v["control_iid"]["accuracy"] for p, v in d_.items() if "control_iid" in v and "->" in p])}414 self.results["elapsed_seconds"] = time.time() - t0415416417def main() -> None:418 ap = argparse.ArgumentParser()419 ap.add_argument("--quick", action="store_true")420 ap.add_argument("--device", default="auto")421 ap.add_argument("--seed-base", type=int, default=30_000)422 ap.add_argument("--tag", default="main")423 args = ap.parse_args()424 exp_id = "graft_tier0" + ("_quick" if args.quick else "") + ("" if args.tag == "main" else f"_{args.tag}")425 cfg = {"quick": args.quick, "seed_base": args.seed_base, "tag": args.tag, "models": "matrix_v1", "protocol": "v2-unrelated-procrustes-reference"}426 chash = sha256_text(canonical_json(cfg))[:12]427 run_id = make_run_id(exp_id, chash)428 out_dir = ROOT / "experiments" / "summaries" / run_id429 art_dir = ROOT / "artifacts" / "runs" / run_id430 out_dir.mkdir(parents=True, exist_ok=True)431 art_dir.mkdir(parents=True, exist_ok=True)432 man = Manifest(run_id, exp_id, cfg, ROOT / "experiments" / "manifests")433 registry.record(run_id, exp_id, "planned", config_hash=chash)434 exp = Exp(args)435 exp.out_dir = out_dir436 man.set_status("running")437 registry.record(run_id, exp_id, "running")438 try:439 exp.run()440 exp.results["run_id"] = run_id441 exp.results["device"] = str(exp.device)442 (out_dir / "results.json").write_text(json.dumps(exp.results, indent=1, default=float))443 with gzip.open(art_dir / "items.json.gz", "wt") as f:444 json.dump({k: {"correct": v["correct"].tolist(), "trace": v["trace"].tolist(), "omega": [None if np.isnan(x) else float(x) for x in v["omega"]], "pred": v["pred"]} for k, v in exp.items.items()}, f)445 man.data["data_hashes"] = exp.results["splits"]446 man.data["outputs"] = {"results": str(out_dir.relative_to(ROOT) / "results.json"), "items": str(art_dir.relative_to(ROOT))}447 man.set_status("completed")448 registry.record(run_id, exp_id, "completed", results=str(out_dir.relative_to(ROOT) / "results.json"))449 print(f"[{utc_now()}] completed {run_id} in {exp.results['elapsed_seconds']:.0f}s")450 except Exception as exc: # noqa: BLE001451 man.set_status("failed", repr(exc))452 registry.record(run_id, exp_id, "failed", error=repr(exc))453 raise454455456if __name__ == "__main__":457 main()458