spb/modelmap Public License

1#!/usr/bin/env python32# =============================================================================3#  Project   : modelmap4#  File      : experiments/micro/expC_causal_verification/implementation/benchmark_v4.py5#  Purpose   : Run #4 — the narrowed early-band claim, everything fresh6#  Author    : Simon-Pierre Boucher7#  Contact   : contact@spboucher.ai8#  Website   : https://modelmap.io9#  Created   : 2026-08-1210#  Modified  : 2026-08-1211#  Platform  : macOS / Apple Silicon (arm64) — MLX / Metal12#  License   : All rights reserved (research code)13# =============================================================================14"""expC run #4 (hypothesis registered before this run).1516Band claim: early-band (layers 2-15) mean specific damage >= 2.5 in EVERY17of six fresh direction sources (disjoint halves of A and B + two fresh18bootstraps), on a fresh behavioral bank. Late band reported, no claim.19"""2021from __future__ import annotations2223import json24import random25import subprocess26import sys27import time28from pathlib import Path2930import numpy as np3132ROOT = Path(__file__).resolve().parents[4]33sys.path.insert(0, str(ROOT / "src"))34sys.path.insert(0, str(ROOT / "benchmarks"))35from hardware_manifest import manifest3637from modelmap.capture.mlx_capture import capture_pooled, install_taps3839MODEL = "mlx-community/Qwen3-0.6B-4bit"40N_RANDOM_DIRS = 341SEED = 78042EARLY = slice(2, 16)43LATE = slice(20, 28)44BAR = 2.54546NOUN_PAIRS = [("key", "keys"), ("crate", "crates"), ("report", "reports"), ("valve", "valves"),47              ("ticket", "tickets"), ("ladder", "ladders"), ("sample", "samples"), ("cable", "cables"),48              ("permit", "permits"), ("beacon", "beacons"), ("filter", "filters"), ("stamp", "stamps")]49FRESH_NEAR = ["past the boiler room", "near the loading dock", "behind the ticket booth",50              "under the mezzanine", "beside the flagpole", "opposite the greenhouse",51              "inside the stairwell", "along the towpath"]525354def directions_from(reps, labels, idx):55    dirs, mus = [], []56    r, l = reps[idx], labels[idx]57    for layer in range(reps.shape[1]):58        x = r[:, layer, :]59        mu = x.mean(0)60        u = x[l == "correct"].mean(0) - x[l == "violated"].mean(0)61        u = u / (np.linalg.norm(u) + 1e-8)62        dirs.append(u); mus.append(mu)63    return dirs, mus646566def main() -> int:67    import mlx.core as mx68    from mlx_lm import load6970    t0 = time.time()71    model, tokenizer = load(MODEL)72    taps = install_taps(model)73    n_layers = len(taps)7475    rng = random.Random(SEED)76    combos = [(n, loc) for n in NOUN_PAIRS for loc in FRESH_NEAR]77    rng.shuffle(combos)78    pairs = []79    for (sg, pl), loc in combos:80        pairs.append({"prefix": f"The {sg} {loc}", "singular": True})81        pairs.append({"prefix": f"The {pl} {loc}", "singular": False})82    prefix_ids = [tokenizer.encode(p["prefix"]) for p in pairs]83    id_is, id_are = tokenizer.encode(" is")[0], tokenizer.encode(" are")[0]8485    def margin() -> float:86        out = []87        for ids, p in zip(prefix_ids, pairs):88            logits = model(mx.array([ids]))[0, -1, :]89            mx.eval(logits)90            m = float(logits[id_is] - logits[id_are])91            out.append(m if p["singular"] else -m)92        return float(np.mean(out))9394    def erase_fn(u_np, mu_np):95        u = mx.array(u_np.astype(np.float32))96        mu = mx.array(mu_np.astype(np.float32))97        def fn(out):98            h = out.astype(mx.float32)99            coef = ((h - mu) * u).sum(axis=-1, keepdims=True)100            return (h - coef * u).astype(out.dtype)101        return fn102103    base_m = margin()104    print(f"fresh bank: {len(pairs)} pairs | baseline margin {base_m:+.4f}", flush=True)105106    # fresh direction sources: disjoint halves + fresh bootstraps107    reps_cache = {}108    for s in ("A", "B"):109        items = [json.loads(l) for l in110                 (ROOT / "benchmarks" / "promptsets" / f"agreement_{s}.jsonl").read_text().splitlines()]111        toks = [tokenizer.encode(it["text"]) for it in items]112        labels = np.array([it["label"] for it in items])113        print(f"capture agreement_{s}…", flush=True)114        reps_cache[s] = (capture_pooled(model, taps, toks)["mean"], labels)115116    sources = {}117    for s in ("A", "B"):118        reps, labels = reps_cache[s]119        n = len(labels)120        perm = np.random.default_rng(SEED + ord(s)).permutation(n)121        sources[f"{s}half1"] = directions_from(reps, labels, perm[: n // 2])122        sources[f"{s}half2"] = directions_from(reps, labels, perm[n // 2:])123        boot = np.random.default_rng(200 + ord(s)).integers(0, n, n)124        sources[f"boot{s}"] = directions_from(reps, labels, boot)125    del sources["bootB"]  # keep 6 by hypothesis? A/B halves (4) + bootA + bootB = 6 — keep both126    sources["bootB"] = directions_from(reps_cache["B"][0], reps_cache["B"][1],127                                       np.random.default_rng(202).integers(0, len(reps_cache["B"][1]),128                                                                           len(reps_cache["B"][1])))129130    rand_damage = []131    d_model = reps_cache["A"][0].shape[-1]132    for layer in range(n_layers):133        ms = []134        for s in range(N_RANDOM_DIRS):135            ru = np.random.default_rng(2000 * layer + s).standard_normal(d_model)136            ru /= np.linalg.norm(ru)137            taps[layer].edit = erase_fn(ru.astype(np.float32), sources["Ahalf1"][1][layer])138            ms.append(margin())139            taps[layer].edit = None140        rand_damage.append(base_m - float(np.mean(ms)))141    print("random-direction controls done", flush=True)142143    profiles, bands = {}, {}144    for name, (dirs, mus) in sources.items():145        prof = []146        for layer in range(n_layers):147            taps[layer].edit = erase_fn(dirs[layer], mus[layer])148            m = margin()149            taps[layer].edit = None150            prof.append((base_m - m) - rand_damage[layer])151        profiles[name] = prof152        bands[name] = {"early_mean": float(np.mean(prof[EARLY])),153                       "late_mean": float(np.mean(prof[LATE]))}154        print(f"{name:8s} early={bands[name]['early_mean']:+.3f} late={bands[name]['late_mean']:+.3f}",155              flush=True)156157    early_means = [bands[n]["early_mean"] for n in sources]158    passes = bool(all(e >= BAR for e in early_means))159    print(f"BAND CLAIM (early mean >= {BAR} in every source): "160          f"min={min(early_means):+.3f} -> passes={passes}")161162    arr = np.array(list(profiles.values()))163    commit = subprocess.run(["git", "rev-parse", "HEAD"], cwd=ROOT,164                            capture_output=True, text=True, check=False).stdout.strip()165    ts = time.strftime("%Y%m%dT%H%M%SZ", time.gmtime())166    outdir = ROOT / "results" / "expC_causal_verification" / ts167    outdir.mkdir(parents=True)168    (outdir / "results.json").write_text(json.dumps({169        "experiment": "expC_causal_verification", "run": 4,170        "scope": "narrowed early-band claim, fresh sources + fresh behavioral bank",171        "commit": commit,172        "config": {"model": MODEL, "sources": list(sources), "bar": BAR,173                   "early_layers": [2, 15], "late_layers": [20, 27],174                   "n_random_dirs": N_RANDOM_DIRS, "n_pairs": len(pairs), "seed": SEED},175        "manifest": manifest(),176        "baseline_margin": base_m,177        "random_direction_damage": rand_damage,178        "profiles": profiles,179        "bands": bands,180        "profile_mean": arr.mean(axis=0).tolist(),181        "profile_min": arr.min(axis=0).tolist(),182        "band_claim": {"bar": BAR, "early_means": early_means,183                       "min_early_mean": float(min(early_means)), "passes": passes},184        "wall_seconds": round(time.time() - t0, 1),185    }, indent=2) + "\n")186    print(f"results -> {outdir / 'results.json'}")187    return 0188189190if __name__ == "__main__":191    sys.exit(main())192