#!/usr/bin/env python3 """Tier 1 (open-weight models, MLX): cross-model transfer of causally localized task vectors. Capability = performing an input→output task *without demonstrations* (zero-shot). The donor is run with k demonstrations; the task-conditioned state at the last prompt token (difference few-shot − zero-shot, averaged over demonstration sets) is the candidate organ. Localization in the donor = layer sweep of the within-donor sufficiency effect (add the vector to the donor's zero-shot run) and necessity (remove the vector's direction from the donor's few-shot run). Transfer = map the donor vector to the recipient with maps fitted on task-free alignment text (Procrustes / ridge / CCA; sham = random orthogonal) and add it to the recipient's zero-shot run at the last prompt position (single-position intervention; generation then continues with the modified KV entry: within-sequence persistence). Controls: recipient zero-shot, recipient few-shot (prompting baseline), recipient own vector (within-model ceiling), norm-matched random vector, unrelated-task vector, sham map. Collateral: next-token NLL on held-out task-free text and cross-task accuracy under each vector. Experiment id ``tier1_taskvector``. Usage: python scripts/run_tier1.py [--pairs qwen3b:qwen1.5b llama3b:qwen3b ...] [--split dev|confirmatory] [--quick] """ from __future__ import annotations import argparse import json import random import sys import time from pathlib import Path import numpy as np ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(ROOT / "src")) from wp_ai_0001.alignment.maps import FITTERS # noqa: E402 from wp_ai_0001.config import canonical_json # noqa: E402 from wp_ai_0001.datasets.tier1_tasks import alignment_prompts, alignment_sentences, build_tasks, few_shot_prompt, is_correct, zero_shot_prompt # noqa: E402 from wp_ai_0001.instrumentation.provenance import Manifest, make_run_id, sha256_text, utc_now # noqa: E402 from wp_ai_0001.orchestration import registry # noqa: E402 from wp_ai_0001.statistics.bootstrap import cluster_bootstrap_mean, cluster_bootstrap_paired_diff # noqa: E402 MODELS = { "qwen1.5b": ("mlx-community/Qwen2.5-1.5B-Instruct-bf16", "Qwen2.5", "Apache-2.0"), "qwen3b": ("mlx-community/Qwen2.5-3B-Instruct-bf16", "Qwen2.5", "Qwen Research License"), "llama3b": ("mlx-community/Llama-3.2-3B-Instruct-bf16", "Llama-3.2", "Llama 3.2 Community License"), "qwen7b": ("mlx-community/Qwen2.5-7B-Instruct-bf16", "Qwen2.5", "Apache-2.0"), # Tier 2 (run on a 96 GB node) } def load_runner(key: str): from mlx_lm import load from wp_ai_0001.models.mlx_runner import Runner m, t = load(MODELS[key][0]) return Runner(m, t, key) class Tier1: def __init__(self, args): self.args = args self.k = 8 self.n_vec_sets = 4 if args.quick else 10 self.tasks = build_tasks(seed=args.task_seed, split=args.split) if args.quick: for t in self.tasks.values(): t.queries = t.queries[:12] self.align = alignment_prompts(240 if args.quick else 800, seed=11) # format-matched, task-free self.n_holdout = 40 if args.quick else 100 # last rows are never used for fitting (held-out R²) self.heldout_text = alignment_sentences(60, seed=12) self.results: dict = {"tasks": {k: {"n_queries": len(v.queries), "hash": v.hash} for k, v in self.tasks.items()}, "models": {}, "pairs": {}} self.vectors: dict = {} # (model, task) -> {layer: v} self.states: dict = {} # (model, layer) -> alignment states [N, d] self.runners: dict = {} self.n_boot = 300 if args.quick else 2000 def runner(self, key): if key not in self.runners: self.runners[key] = load_runner(key) r = self.runners[key] 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"} return self.runners[key] # ---------- accuracy ---------- def accuracy(self, key, task_name, mode="zero", intervention=None, layer=None, demo_seed=0) -> tuple[float, np.ndarray, list[str]]: r, t = self.runner(key), self.tasks[task_name] rng = random.Random(f"demo|{demo_seed}") corr, outs = [], [] for q, a in t.queries: prompt = few_shot_prompt(t.demos, q, self.k, rng) if mode == "few" else zero_shot_prompt(q) g = r.generate(prompt, intervention=intervention, layer=layer) outs.append(g) corr.append(is_correct(g, a, t.alternatives)) c = np.array(corr, dtype=float) return float(c.mean()), c, outs # ---------- vectors ---------- def _cache_path(self, name: str) -> Path: d = ROOT / "artifacts" / "tier1_cache" d.mkdir(parents=True, exist_ok=True) return d / name def task_vector(self, key, task_name, layers) -> dict[int, np.ndarray]: kk = (key, task_name) if kk in self.vectors: return self.vectors[kk] cp = self._cache_path(f"vec_{key}_{task_name}_{self.tasks[task_name].hash}_k{self.k}_n{self.n_vec_sets}.npz") if cp.exists(): z = np.load(cp) self.vectors[kk] = {int(l): z[str(l)] for l in z.files} if all(l in self.vectors[kk] for l in layers): return self.vectors[kk] r, t = self.runner(key), self.tasks[task_name] few, zero = {l: [] for l in layers}, {l: [] for l in layers} rng = random.Random(f"vec|{key}|{task_name}") # vectors are built from demonstration items only (queries never enter the vector) for s in range(self.n_vec_sets): for q, _ in rng.sample(t.demos, 4): others = [d for d in t.demos if d[0] != q] hf = r.last_hidden(few_shot_prompt(others, q, self.k, rng), layers) hz = r.last_hidden(zero_shot_prompt(q), layers) for l in layers: few[l].append(hf[l]) zero[l].append(hz[l]) self.vectors[kk] = {**self.vectors.get(kk, {}), **{l: np.mean(few[l], 0) - np.mean(zero[l], 0) for l in layers}} np.savez(cp, **{str(l): v for l, v in self.vectors[kk].items()}) return self.vectors[kk] def align_states(self, key, layer) -> np.ndarray: kk = (key, layer) if kk not in self.states: cp = self._cache_path(f"align_{key}_L{layer}_n{len(self.align)}_s11.npy") if cp.exists(): self.states[kk] = np.load(cp) else: r = self.runner(key) self.states[kk] = np.stack([r.last_hidden(s, [layer])[layer] for s in self.align]) np.save(cp, self.states[kk]) return self.states[kk] # ---------- interventions ---------- @staticmethod def add_vec(v, alpha=1.0): from wp_ai_0001.models.mlx_backend import MLXAddAtLast return MLXAddAtLast(v.astype(np.float32), alpha) @staticmethod def project_out(v): import mlx.core as mx u = (v / np.linalg.norm(v)).astype(np.float32) def fn(x, ctx): uu = mx.array(u).astype(x.dtype) last = x[:, -1, :] proj = (last * uu).sum(-1, keepdims=True) * uu out = mx.array(x) out[:, -1, :] = last - proj return out return fn # ---------- within-model localization ---------- def localize(self, key) -> dict: r = self.runner(key) cp = self._cache_path(f"loc_{key}_{self.args.split}_{self.args.task_seed}_{'q' if self.args.quick else 'full'}.json") if self.args.split != "dev": # frozen rule (A2): layers and vectors are selected on the development split only dev_cp = self._cache_path(f"loc_{key}_dev_{self.args.task_seed}_{'q' if self.args.quick else 'full'}.json") if not dev_cp.exists(): raise SystemExit(f"development localization missing for {key}: run --split dev first") cp = dev_cp self.results.setdefault("localization_source", {})[key] = str(dev_cp.name) if cp.exists(): rep = json.loads(cp.read_text()) rep["layers"] = [int(x) for x in rep["layers"]] for t in rep["tasks"].values(): t["sufficiency_by_layer"] = {int(k): v for k, v in t["sufficiency_by_layer"].items()} print(f" [{key}] localization loaded from cache") return rep L = r.n_layers layers = sorted(set([int(round(x)) for x in np.linspace(2, L - 3, 6 if self.args.quick else 9)])) rep = {"layers": layers, "tasks": {}} for tn in self.tasks: z, zc, _ = self.accuracy(key, tn, "zero") f, fc, _ = self.accuracy(key, tn, "few") vec = self.task_vector(key, tn, layers) suff, nec = {}, {} for l in layers: a, _, _ = self.accuracy(key, tn, "zero", self.add_vec(vec[l]), l) suff[l] = a best = max(suff, key=suff.get) nz, _, _ = self.accuracy(key, tn, "few", self.project_out(vec[best]), best) rnd = np.random.default_rng(0).standard_normal(r.d) rnd = rnd / np.linalg.norm(rnd) * np.linalg.norm(vec[best]) rz, _, _ = self.accuracy(key, tn, "zero", self.add_vec(rnd), best) rep["tasks"][tn] = {"zero_shot": z, "few_shot": f, "sufficiency_by_layer": suff, "best_layer": best, "necessity_fewshot_minus_direction": nz, "random_vector_same_norm_at_best": rz, "vector_norm": float(np.linalg.norm(vec[best]))} print(f" [{key}] {tn}: zero {z:.2f} few {f:.2f} best L{best} suff {suff[best]:.2f} nec {nz:.2f} rnd {rz:.2f}") cp.write_text(json.dumps(rep, default=float)) return rep # ---------- transfer ---------- def transfer(self, dk, rk, loc_d, loc_r) -> dict: D, R = self.runner(dk), self.runner(rk) out = {"donor": dk, "recipient": rk, "tasks": {}, "maps": {}} rec_layers = loc_r["layers"] for tn, t in self.tasks.items(): ld = loc_d["tasks"][tn]["best_layer"] lr_own = loc_r["tasks"][tn]["best_layer"] vd = self.task_vector(dk, tn, [ld])[ld] vr = self.task_vector(rk, tn, [lr_own])[lr_own] # recipient layers to try: own best layer and the proportional depth match lr_prop = int(round(ld / D.n_layers * R.n_layers)) lr_prop = min(max(lr_prop, 1), R.n_layers - 2) cand = sorted({lr_own, lr_prop}) base_z, cz, _ = self.accuracy(rk, tn, "zero") base_f, cf, _ = self.accuracy(rk, tn, "few") own, co, _ = self.accuracy(rk, tn, "zero", self.add_vec(vr), lr_own) tres = {"recipient_zero": base_z, "recipient_few": base_f, "recipient_own_vector": own, "donor_layer": ld, "conditions": {}, "items": {"zero": cz.tolist(), "few": cf.tolist(), "own": co.tolist()}} for lr_ in cand: Xa, Ya = self.align_states(dk, ld), self.align_states(rk, lr_) h = self.n_holdout X, Y, Xh, Yh = Xa[:-h], Ya[:-h], Xa[-h:], Ya[-h:] fitted = {} for fam in ("procrustes", "ridge", "cca", "random_orth"): m = FITTERS[fam](X, Y) if fam != "random_orth" else FITTERS[fam](X, Y, seed=0) fitted[fam] = m 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))} mv = (vd @ m.W).astype(np.float32) # difference vector: no offset for alpha in ([1.0] if self.args.quick else [1.0, 2.0]): acc, c, outs = self.accuracy(rk, tn, "zero", self.add_vec(mv, alpha), lr_) cname = f"{fam}|L{lr_}|a{alpha}" 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(), "mapped_norm": float(np.linalg.norm(mv)), "outputs_sample": outs[:5], "correct": c.tolist()} if fam in ("procrustes", "ridge"): # declared variant: mapped vector rescaled to the recipient's own vector norm mvs = (mv / max(np.linalg.norm(mv), 1e-6) * np.linalg.norm(vr)).astype(np.float32) acc, c, outs = self.accuracy(rk, tn, "zero", self.add_vec(mvs, 1.0), lr_) 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()} # donor-in-the-loop REPLACE graft: recipient state at lr_ <- map(donor few-shot state at ld) for the same query self._replace_graft(dk, rk, tn, ld, lr_, fitted, tres, cz) # controls at this layer rnd = np.random.default_rng(1).standard_normal(R.d) mv_ref = (vd @ FITTERS["procrustes"](X, Y).W) rnd = rnd / np.linalg.norm(rnd) * np.linalg.norm(mv_ref) acc, c, _ = self.accuracy(rk, tn, "zero", self.add_vec(rnd), lr_) 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()} other = [o for o in self.tasks if o != tn][0] vo = self.task_vector(dk, other, [ld])[ld] mo = (vo @ FITTERS["procrustes"](X, Y).W).astype(np.float32) acc, c, _ = self.accuracy(rk, tn, "zero", self.add_vec(mo), lr_) 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()} out["tasks"][tn] = tres best = max(tres["conditions"].items(), key=lambda kv: kv[1]["accuracy"]) print(f" [{dk}->{rk}] {tn}: zero {base_z:.2f} few {base_f:.2f} own {own:.2f} best transfer {best[0]} {best[1]['accuracy']:.2f}") # collateral: next-token NLL on held-out text with the best transferred vector of each task added at the last position coll = {} for tn, tres in out["tasks"].items(): 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")] best_name, best = max(addc, key=lambda kv: kv[1]["accuracy"]) fam, lstr, astr = best_name.split("|") fam = fam.replace("_rescaled", "") lr_ = int(lstr[1:]) ld = tres["donor_layer"] m = FITTERS[fam](self.align_states(dk, ld), self.align_states(rk, lr_)) mv = (self.task_vector(dk, tn, [ld])[ld] @ m.W).astype(np.float32) if "_rescaled" in best_name: vr_ = self.task_vector(rk, tn, [loc_r_best := self.results["models"][rk]["localization"]["tasks"][tn]["best_layer"]])[loc_r_best] mv = (mv / max(np.linalg.norm(mv), 1e-6) * np.linalg.norm(vr_)).astype(np.float32) nll0 = [R.next_token_nll(s) for s in self.heldout_text] nll1 = [R.next_token_nll(s, self.add_vec(mv, float(astr[1:])), lr_) for s in self.heldout_text] cross = {} for on in self.tasks: if on == tn: continue a0, _, _ = self.accuracy(rk, on, "zero") a1, _, _ = self.accuracy(rk, on, "zero", self.add_vec(mv, float(astr[1:])), lr_) cross[on] = {"zero": a0, "with_vector": a1} coll[tn] = {"condition": best_name, "heldout_nll_base": float(np.mean(nll0)), "heldout_nll_with_vector": float(np.mean(nll1)), "cross_task": cross} out["collateral"] = coll return out def _replace_graft(self, dk, rk, tn, ld, lr_, fitted, tres, cz): from wp_ai_0001.models.mlx_backend import MLXReplaceAtLast D, t = self.runner(dk), self.tasks[tn] rng = random.Random(f"replace|{tn}") few_states, zero_states = [], [] for q, _ in t.queries: few_states.append(D.last_hidden(few_shot_prompt(t.demos, q, self.k, rng), [ld])[ld]) zero_states.append(D.last_hidden(zero_shot_prompt(q), [ld])[ld]) few_states, zero_states = np.stack(few_states), np.stack(zero_states) perm = np.random.default_rng(0).permutation(len(t.queries)) for fam in ("procrustes", "ridge", "random_orth"): m = fitted[fam] for label, src in (("replace_donor_fewshot", few_states), ("ctl_replace_donor_zeroshot", zero_states), ("ctl_replace_donor_fewshot_shuffled_query", few_states[perm])): rows = m(src) corr = [] outs = [] for i, (q, a) in enumerate(t.queries): g = self.runner(rk).generate(zero_shot_prompt(q), intervention=MLXReplaceAtLast(rows[i : i + 1].astype(np.float32)), layer=lr_) outs.append(g) corr.append(is_correct(g, a, t.alternatives)) c = np.array(corr, dtype=float) name = f"{label}|{fam}|L{lr_}" if fam == "random_orth": name = f"ctl_{label}|sham_map|L{lr_}" if not label.startswith("ctl_") else f"{label}|sham_map|L{lr_}" 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()} def run(self): t0 = time.time() pair_keys = [p.split(":") for p in self.args.pairs] models = sorted({k for p in pair_keys for k in p}) loc = {} for k in models: print(f"[{utc_now()}] localizing in {k}") loc[k] = self.localize(k) self.results["models"][k]["localization"] = loc[k] for dk, rk in pair_keys: print(f"[{utc_now()}] transfer {dk} -> {rk}") self.results["pairs"][f"{dk}->{rk}"] = self.transfer(dk, rk, loc[dk], loc[rk]) if getattr(self, "partial_path", None): self.partial_path.write_text(json.dumps(self.results, indent=1, default=float)) self.results["elapsed_seconds"] = time.time() - t0 def main() -> None: ap = argparse.ArgumentParser() ap.add_argument("--pairs", nargs="+", default=["qwen3b:qwen1.5b", "qwen1.5b:qwen3b", "llama3b:qwen3b", "qwen3b:llama3b"]) ap.add_argument("--split", default="dev") ap.add_argument("--task-seed", type=int, default=0) ap.add_argument("--quick", action="store_true") args = ap.parse_args() cfg = {"pairs": args.pairs, "split": args.split, "task_seed": args.task_seed, "quick": args.quick, "k": 8, "localization": "dev-frozen (A2)" if args.split != "dev" else "in-split"} exp_id = "tier1_taskvector" + ("_quick" if args.quick else "") + ("" if args.split == "dev" else f"_{args.split}") chash = sha256_text(canonical_json(cfg))[:12] run_id = make_run_id(exp_id, chash) out_dir = ROOT / "experiments" / "summaries" / run_id out_dir.mkdir(parents=True, exist_ok=True) man = Manifest(run_id, exp_id, cfg, ROOT / "experiments" / "manifests") registry.record(run_id, exp_id, "planned", config_hash=chash) man.set_status("running") exp = Tier1(args) exp.partial_path = out_dir / "results_partial.json" try: exp.run() exp.results["run_id"] = run_id (out_dir / "results.json").write_text(json.dumps(exp.results, indent=1, default=float)) man.data["data_hashes"] = exp.results["tasks"] man.data["outputs"] = {"results": str(out_dir.relative_to(ROOT) / "results.json")} man.set_status("completed") registry.record(run_id, exp_id, "completed") print(f"[{utc_now()}] completed {run_id} in {exp.results['elapsed_seconds']:.0f}s") except Exception as exc: # noqa: BLE001 man.set_status("failed", repr(exc)) registry.record(run_id, exp_id, "failed", error=repr(exc)) raise if __name__ == "__main__": main()