spb/wpai0001_uqo
Public
1#!/usr/bin/env python32"""Tier 1 (open-weight models, MLX): cross-model transfer of causally localized task vectors.34Capability = performing an input→output task *without demonstrations* (zero-shot). The donor is run with5k demonstrations; the task-conditioned state at the last prompt token (difference few-shot − zero-shot,6averaged over demonstration sets) is the candidate organ. Localization in the donor = layer sweep of the7within-donor sufficiency effect (add the vector to the donor's zero-shot run) and necessity (remove the8vector's direction from the donor's few-shot run). Transfer = map the donor vector to the recipient with9maps fitted on task-free alignment text (Procrustes / ridge / CCA; sham = random orthogonal) and add it10to the recipient's zero-shot run at the last prompt position (single-position intervention; generation11then continues with the modified KV entry: within-sequence persistence). Controls: recipient zero-shot,12recipient few-shot (prompting baseline), recipient own vector (within-model ceiling), norm-matched random13vector, unrelated-task vector, sham map. Collateral: next-token NLL on held-out task-free text and14cross-task accuracy under each vector. Experiment id ``tier1_taskvector``.1516Usage: python scripts/run_tier1.py [--pairs qwen3b:qwen1.5b llama3b:qwen3b ...] [--split dev|confirmatory] [--quick]17"""18from __future__ import annotations1920import argparse21import json22import random23import sys24import time25from pathlib import Path2627import numpy as np2829ROOT = Path(__file__).resolve().parents[1]30sys.path.insert(0, str(ROOT / "src"))3132from wp_ai_0001.alignment.maps import FITTERS # noqa: E40233from wp_ai_0001.config import canonical_json # noqa: E40234from wp_ai_0001.datasets.tier1_tasks import alignment_prompts, alignment_sentences, build_tasks, few_shot_prompt, is_correct, zero_shot_prompt # noqa: E40235from wp_ai_0001.instrumentation.provenance import Manifest, make_run_id, sha256_text, utc_now # noqa: E40236from wp_ai_0001.orchestration import registry # noqa: E40237from wp_ai_0001.statistics.bootstrap import cluster_bootstrap_mean, cluster_bootstrap_paired_diff # noqa: E4023839MODELS = {40 "qwen1.5b": ("mlx-community/Qwen2.5-1.5B-Instruct-bf16", "Qwen2.5", "Apache-2.0"),41 "qwen3b": ("mlx-community/Qwen2.5-3B-Instruct-bf16", "Qwen2.5", "Qwen Research License"),42 "llama3b": ("mlx-community/Llama-3.2-3B-Instruct-bf16", "Llama-3.2", "Llama 3.2 Community License"),43 "qwen7b": ("mlx-community/Qwen2.5-7B-Instruct-bf16", "Qwen2.5", "Apache-2.0"), # Tier 2 (run on a 96 GB node)44}454647def load_runner(key: str):48 from mlx_lm import load4950 from wp_ai_0001.models.mlx_runner import Runner5152 m, t = load(MODELS[key][0])53 return Runner(m, t, key)545556class Tier1:57 def __init__(self, args):58 self.args = args59 self.k = 860 self.n_vec_sets = 4 if args.quick else 1061 self.tasks = build_tasks(seed=args.task_seed, split=args.split)62 if args.quick:63 for t in self.tasks.values():64 t.queries = t.queries[:12]65 self.align = alignment_prompts(240 if args.quick else 800, seed=11) # format-matched, task-free66 self.n_holdout = 40 if args.quick else 100 # last rows are never used for fitting (held-out R²)67 self.heldout_text = alignment_sentences(60, seed=12)68 self.results: dict = {"tasks": {k: {"n_queries": len(v.queries), "hash": v.hash} for k, v in self.tasks.items()}, "models": {}, "pairs": {}}69 self.vectors: dict = {} # (model, task) -> {layer: v}70 self.states: dict = {} # (model, layer) -> alignment states [N, d]71 self.runners: dict = {}72 self.n_boot = 300 if args.quick else 20007374 def runner(self, key):75 if key not in self.runners:76 self.runners[key] = load_runner(key)77 r = self.runners[key]78 self.results["models"][key] = {"repo": MODELS[key][0], "family": MODELS[key][1], "license": MODELS[key][2], "n_layers": r.n_layers, "hidden": r.d, "precision": "bf16"}79 return self.runners[key]8081 # ---------- accuracy ----------82 def accuracy(self, key, task_name, mode="zero", intervention=None, layer=None, demo_seed=0) -> tuple[float, np.ndarray, list[str]]:83 r, t = self.runner(key), self.tasks[task_name]84 rng = random.Random(f"demo|{demo_seed}")85 corr, outs = [], []86 for q, a in t.queries:87 prompt = few_shot_prompt(t.demos, q, self.k, rng) if mode == "few" else zero_shot_prompt(q)88 g = r.generate(prompt, intervention=intervention, layer=layer)89 outs.append(g)90 corr.append(is_correct(g, a, t.alternatives))91 c = np.array(corr, dtype=float)92 return float(c.mean()), c, outs9394 # ---------- vectors ----------95 def _cache_path(self, name: str) -> Path:96 d = ROOT / "artifacts" / "tier1_cache"97 d.mkdir(parents=True, exist_ok=True)98 return d / name99100 def task_vector(self, key, task_name, layers) -> dict[int, np.ndarray]:101 kk = (key, task_name)102 if kk in self.vectors:103 return self.vectors[kk]104 cp = self._cache_path(f"vec_{key}_{task_name}_{self.tasks[task_name].hash}_k{self.k}_n{self.n_vec_sets}.npz")105 if cp.exists():106 z = np.load(cp)107 self.vectors[kk] = {int(l): z[str(l)] for l in z.files}108 if all(l in self.vectors[kk] for l in layers):109 return self.vectors[kk]110 r, t = self.runner(key), self.tasks[task_name]111 few, zero = {l: [] for l in layers}, {l: [] for l in layers}112 rng = random.Random(f"vec|{key}|{task_name}")113 # vectors are built from demonstration items only (queries never enter the vector)114 for s in range(self.n_vec_sets):115 for q, _ in rng.sample(t.demos, 4):116 others = [d for d in t.demos if d[0] != q]117 hf = r.last_hidden(few_shot_prompt(others, q, self.k, rng), layers)118 hz = r.last_hidden(zero_shot_prompt(q), layers)119 for l in layers:120 few[l].append(hf[l])121 zero[l].append(hz[l])122 self.vectors[kk] = {**self.vectors.get(kk, {}), **{l: np.mean(few[l], 0) - np.mean(zero[l], 0) for l in layers}}123 np.savez(cp, **{str(l): v for l, v in self.vectors[kk].items()})124 return self.vectors[kk]125126 def align_states(self, key, layer) -> np.ndarray:127 kk = (key, layer)128 if kk not in self.states:129 cp = self._cache_path(f"align_{key}_L{layer}_n{len(self.align)}_s11.npy")130 if cp.exists():131 self.states[kk] = np.load(cp)132 else:133 r = self.runner(key)134 self.states[kk] = np.stack([r.last_hidden(s, [layer])[layer] for s in self.align])135 np.save(cp, self.states[kk])136 return self.states[kk]137138 # ---------- interventions ----------139 @staticmethod140 def add_vec(v, alpha=1.0):141 from wp_ai_0001.models.mlx_backend import MLXAddAtLast142143 return MLXAddAtLast(v.astype(np.float32), alpha)144145 @staticmethod146 def project_out(v):147 import mlx.core as mx148149 u = (v / np.linalg.norm(v)).astype(np.float32)150151 def fn(x, ctx):152 uu = mx.array(u).astype(x.dtype)153 last = x[:, -1, :]154 proj = (last * uu).sum(-1, keepdims=True) * uu155 out = mx.array(x)156 out[:, -1, :] = last - proj157 return out158159 return fn160161 # ---------- within-model localization ----------162 def localize(self, key) -> dict:163 r = self.runner(key)164 cp = self._cache_path(f"loc_{key}_{self.args.split}_{self.args.task_seed}_{'q' if self.args.quick else 'full'}.json")165 if self.args.split != "dev": # frozen rule (A2): layers and vectors are selected on the development split only166 dev_cp = self._cache_path(f"loc_{key}_dev_{self.args.task_seed}_{'q' if self.args.quick else 'full'}.json")167 if not dev_cp.exists():168 raise SystemExit(f"development localization missing for {key}: run --split dev first")169 cp = dev_cp170 self.results.setdefault("localization_source", {})[key] = str(dev_cp.name)171 if cp.exists():172 rep = json.loads(cp.read_text())173 rep["layers"] = [int(x) for x in rep["layers"]]174 for t in rep["tasks"].values():175 t["sufficiency_by_layer"] = {int(k): v for k, v in t["sufficiency_by_layer"].items()}176 print(f" [{key}] localization loaded from cache")177 return rep178 L = r.n_layers179 layers = sorted(set([int(round(x)) for x in np.linspace(2, L - 3, 6 if self.args.quick else 9)]))180 rep = {"layers": layers, "tasks": {}}181 for tn in self.tasks:182 z, zc, _ = self.accuracy(key, tn, "zero")183 f, fc, _ = self.accuracy(key, tn, "few")184 vec = self.task_vector(key, tn, layers)185 suff, nec = {}, {}186 for l in layers:187 a, _, _ = self.accuracy(key, tn, "zero", self.add_vec(vec[l]), l)188 suff[l] = a189 best = max(suff, key=suff.get)190 nz, _, _ = self.accuracy(key, tn, "few", self.project_out(vec[best]), best)191 rnd = np.random.default_rng(0).standard_normal(r.d)192 rnd = rnd / np.linalg.norm(rnd) * np.linalg.norm(vec[best])193 rz, _, _ = self.accuracy(key, tn, "zero", self.add_vec(rnd), best)194 rep["tasks"][tn] = {"zero_shot": z, "few_shot": f, "sufficiency_by_layer": suff, "best_layer": best, "necessity_fewshot_minus_direction": nz,195 "random_vector_same_norm_at_best": rz, "vector_norm": float(np.linalg.norm(vec[best]))}196 print(f" [{key}] {tn}: zero {z:.2f} few {f:.2f} best L{best} suff {suff[best]:.2f} nec {nz:.2f} rnd {rz:.2f}")197 cp.write_text(json.dumps(rep, default=float))198 return rep199200 # ---------- transfer ----------201 def transfer(self, dk, rk, loc_d, loc_r) -> dict:202 D, R = self.runner(dk), self.runner(rk)203 out = {"donor": dk, "recipient": rk, "tasks": {}, "maps": {}}204 rec_layers = loc_r["layers"]205 for tn, t in self.tasks.items():206 ld = loc_d["tasks"][tn]["best_layer"]207 lr_own = loc_r["tasks"][tn]["best_layer"]208 vd = self.task_vector(dk, tn, [ld])[ld]209 vr = self.task_vector(rk, tn, [lr_own])[lr_own]210 # recipient layers to try: own best layer and the proportional depth match211 lr_prop = int(round(ld / D.n_layers * R.n_layers))212 lr_prop = min(max(lr_prop, 1), R.n_layers - 2)213 cand = sorted({lr_own, lr_prop})214 base_z, cz, _ = self.accuracy(rk, tn, "zero")215 base_f, cf, _ = self.accuracy(rk, tn, "few")216 own, co, _ = self.accuracy(rk, tn, "zero", self.add_vec(vr), lr_own)217 tres = {"recipient_zero": base_z, "recipient_few": base_f, "recipient_own_vector": own, "donor_layer": ld, "conditions": {},218 "items": {"zero": cz.tolist(), "few": cf.tolist(), "own": co.tolist()}}219 for lr_ in cand:220 Xa, Ya = self.align_states(dk, ld), self.align_states(rk, lr_)221 h = self.n_holdout222 X, Y, Xh, Yh = Xa[:-h], Ya[:-h], Xa[-h:], Ya[-h:]223 fitted = {}224 for fam in ("procrustes", "ridge", "cca", "random_orth"):225 m = FITTERS[fam](X, Y) if fam != "random_orth" else FITTERS[fam](X, Y, seed=0)226 fitted[fam] = m227 out["maps"][f"{tn}|L{ld}->L{lr_}|{fam}"] = {"r2_fit": m.r2(X, Y), "r2_heldout": m.r2(Xh, Yh), "n_params": m.n_params, "n_fit": int(len(X))}228 mv = (vd @ m.W).astype(np.float32) # difference vector: no offset229 for alpha in ([1.0] if self.args.quick else [1.0, 2.0]):230 acc, c, outs = self.accuracy(rk, tn, "zero", self.add_vec(mv, alpha), lr_)231 cname = f"{fam}|L{lr_}|a{alpha}"232 tres["conditions"][cname] = {"accuracy": acc, "gain_vs_zero": cluster_bootstrap_paired_diff(c, cz, [str(i) for i in range(len(c))], self.n_boot).to_dict(),233 "mapped_norm": float(np.linalg.norm(mv)), "outputs_sample": outs[:5], "correct": c.tolist()}234 if fam in ("procrustes", "ridge"): # declared variant: mapped vector rescaled to the recipient's own vector norm235 mvs = (mv / max(np.linalg.norm(mv), 1e-6) * np.linalg.norm(vr)).astype(np.float32)236 acc, c, outs = self.accuracy(rk, tn, "zero", self.add_vec(mvs, 1.0), lr_)237 tres["conditions"][f"{fam}_rescaled|L{lr_}|a1.0"] = {"accuracy": acc, "gain_vs_zero": cluster_bootstrap_paired_diff(c, cz, [str(i) for i in range(len(c))], self.n_boot).to_dict(), "mapped_norm": float(np.linalg.norm(mvs)), "outputs_sample": outs[:5], "correct": c.tolist()}238 # donor-in-the-loop REPLACE graft: recipient state at lr_ <- map(donor few-shot state at ld) for the same query239 self._replace_graft(dk, rk, tn, ld, lr_, fitted, tres, cz)240 # controls at this layer241 rnd = np.random.default_rng(1).standard_normal(R.d)242 mv_ref = (vd @ FITTERS["procrustes"](X, Y).W)243 rnd = rnd / np.linalg.norm(rnd) * np.linalg.norm(mv_ref)244 acc, c, _ = self.accuracy(rk, tn, "zero", self.add_vec(rnd), lr_)245 tres["conditions"][f"ctl_random_norm_matched|L{lr_}"] = {"accuracy": acc, "gain_vs_zero": cluster_bootstrap_paired_diff(c, cz, [str(i) for i in range(len(c))], self.n_boot).to_dict(), "correct": c.tolist()}246 other = [o for o in self.tasks if o != tn][0]247 vo = self.task_vector(dk, other, [ld])[ld]248 mo = (vo @ FITTERS["procrustes"](X, Y).W).astype(np.float32)249 acc, c, _ = self.accuracy(rk, tn, "zero", self.add_vec(mo), lr_)250 tres["conditions"][f"ctl_unrelated_task_vector({other})|L{lr_}"] = {"accuracy": acc, "gain_vs_zero": cluster_bootstrap_paired_diff(c, cz, [str(i) for i in range(len(c))], self.n_boot).to_dict(), "correct": c.tolist()}251 out["tasks"][tn] = tres252 best = max(tres["conditions"].items(), key=lambda kv: kv[1]["accuracy"])253 print(f" [{dk}->{rk}] {tn}: zero {base_z:.2f} few {base_f:.2f} own {own:.2f} best transfer {best[0]} {best[1]['accuracy']:.2f}")254 # collateral: next-token NLL on held-out text with the best transferred vector of each task added at the last position255 coll = {}256 for tn, tres in out["tasks"].items():257 addc = [(k, v) for k, v in tres["conditions"].items() if not k.startswith(("ctl", "replace")) and k.count("|") == 2 and k.split("|")[1].startswith("L") and k.split("|")[2].startswith("a")]258 best_name, best = max(addc, key=lambda kv: kv[1]["accuracy"])259 fam, lstr, astr = best_name.split("|")260 fam = fam.replace("_rescaled", "")261 lr_ = int(lstr[1:])262 ld = tres["donor_layer"]263 m = FITTERS[fam](self.align_states(dk, ld), self.align_states(rk, lr_))264 mv = (self.task_vector(dk, tn, [ld])[ld] @ m.W).astype(np.float32)265 if "_rescaled" in best_name:266 vr_ = self.task_vector(rk, tn, [loc_r_best := self.results["models"][rk]["localization"]["tasks"][tn]["best_layer"]])[loc_r_best]267 mv = (mv / max(np.linalg.norm(mv), 1e-6) * np.linalg.norm(vr_)).astype(np.float32)268 nll0 = [R.next_token_nll(s) for s in self.heldout_text]269 nll1 = [R.next_token_nll(s, self.add_vec(mv, float(astr[1:])), lr_) for s in self.heldout_text]270 cross = {}271 for on in self.tasks:272 if on == tn:273 continue274 a0, _, _ = self.accuracy(rk, on, "zero")275 a1, _, _ = self.accuracy(rk, on, "zero", self.add_vec(mv, float(astr[1:])), lr_)276 cross[on] = {"zero": a0, "with_vector": a1}277 coll[tn] = {"condition": best_name, "heldout_nll_base": float(np.mean(nll0)), "heldout_nll_with_vector": float(np.mean(nll1)), "cross_task": cross}278 out["collateral"] = coll279 return out280281 def _replace_graft(self, dk, rk, tn, ld, lr_, fitted, tres, cz):282 from wp_ai_0001.models.mlx_backend import MLXReplaceAtLast283284 D, t = self.runner(dk), self.tasks[tn]285 rng = random.Random(f"replace|{tn}")286 few_states, zero_states = [], []287 for q, _ in t.queries:288 few_states.append(D.last_hidden(few_shot_prompt(t.demos, q, self.k, rng), [ld])[ld])289 zero_states.append(D.last_hidden(zero_shot_prompt(q), [ld])[ld])290 few_states, zero_states = np.stack(few_states), np.stack(zero_states)291 perm = np.random.default_rng(0).permutation(len(t.queries))292 for fam in ("procrustes", "ridge", "random_orth"):293 m = fitted[fam]294 for label, src in (("replace_donor_fewshot", few_states), ("ctl_replace_donor_zeroshot", zero_states), ("ctl_replace_donor_fewshot_shuffled_query", few_states[perm])):295 rows = m(src)296 corr = []297 outs = []298 for i, (q, a) in enumerate(t.queries):299 g = self.runner(rk).generate(zero_shot_prompt(q), intervention=MLXReplaceAtLast(rows[i : i + 1].astype(np.float32)), layer=lr_)300 outs.append(g)301 corr.append(is_correct(g, a, t.alternatives))302 c = np.array(corr, dtype=float)303 name = f"{label}|{fam}|L{lr_}"304 if fam == "random_orth":305 name = f"ctl_{label}|sham_map|L{lr_}" if not label.startswith("ctl_") else f"{label}|sham_map|L{lr_}"306 tres["conditions"][name] = {"accuracy": float(c.mean()), "gain_vs_zero": cluster_bootstrap_paired_diff(c, cz, [str(i) for i in range(len(c))], self.n_boot).to_dict(), "outputs_sample": outs[:5], "correct": c.tolist()}307308 def run(self):309 t0 = time.time()310 pair_keys = [p.split(":") for p in self.args.pairs]311 models = sorted({k for p in pair_keys for k in p})312 loc = {}313 for k in models:314 print(f"[{utc_now()}] localizing in {k}")315 loc[k] = self.localize(k)316 self.results["models"][k]["localization"] = loc[k]317 for dk, rk in pair_keys:318 print(f"[{utc_now()}] transfer {dk} -> {rk}")319 self.results["pairs"][f"{dk}->{rk}"] = self.transfer(dk, rk, loc[dk], loc[rk])320 if getattr(self, "partial_path", None):321 self.partial_path.write_text(json.dumps(self.results, indent=1, default=float))322 self.results["elapsed_seconds"] = time.time() - t0323324325def main() -> None:326 ap = argparse.ArgumentParser()327 ap.add_argument("--pairs", nargs="+", default=["qwen3b:qwen1.5b", "qwen1.5b:qwen3b", "llama3b:qwen3b", "qwen3b:llama3b"])328 ap.add_argument("--split", default="dev")329 ap.add_argument("--task-seed", type=int, default=0)330 ap.add_argument("--quick", action="store_true")331 args = ap.parse_args()332 cfg = {"pairs": args.pairs, "split": args.split, "task_seed": args.task_seed, "quick": args.quick, "k": 8,333 "localization": "dev-frozen (A2)" if args.split != "dev" else "in-split"}334 exp_id = "tier1_taskvector" + ("_quick" if args.quick else "") + ("" if args.split == "dev" else f"_{args.split}")335 chash = sha256_text(canonical_json(cfg))[:12]336 run_id = make_run_id(exp_id, chash)337 out_dir = ROOT / "experiments" / "summaries" / run_id338 out_dir.mkdir(parents=True, exist_ok=True)339 man = Manifest(run_id, exp_id, cfg, ROOT / "experiments" / "manifests")340 registry.record(run_id, exp_id, "planned", config_hash=chash)341 man.set_status("running")342 exp = Tier1(args)343 exp.partial_path = out_dir / "results_partial.json"344 try:345 exp.run()346 exp.results["run_id"] = run_id347 (out_dir / "results.json").write_text(json.dumps(exp.results, indent=1, default=float))348 man.data["data_hashes"] = exp.results["tasks"]349 man.data["outputs"] = {"results": str(out_dir.relative_to(ROOT) / "results.json")}350 man.set_status("completed")351 registry.record(run_id, exp_id, "completed")352 print(f"[{utc_now()}] completed {run_id} in {exp.results['elapsed_seconds']:.0f}s")353 except Exception as exc: # noqa: BLE001354 man.set_status("failed", repr(exc))355 registry.record(run_id, exp_id, "failed", error=repr(exc))356 raise357358359if __name__ == "__main__":360 main()361