spb/modelmap Public License

1#!/usr/bin/env python32# =============================================================================3#  Project   : modelmap4#  File      : experiments/micro/expC_causal_verification/implementation/benchmark_v2.py5#  Purpose   : Run #2 — specificity-normalized skip + direction-level erasure6#  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 #2 (hypothesis registered before this run).1516P1: layer-skip conditions rerun with general-damage (NLL) normalization.17P2: surgical test — erase the layer-l agreement direction (diff-of-means)18from every position of that layer's output; compare per-layer specific19damage (vs random-direction control) with the differential probe profile.20"""2122from __future__ import annotations2324import json25import random26import subprocess27import sys28import time29from pathlib import Path3031import numpy as np3233ROOT = Path(__file__).resolve().parents[4]34sys.path.insert(0, str(ROOT / "src"))35sys.path.insert(0, str(ROOT / "benchmarks"))36from hardware_manifest import manifest3738from modelmap.capture.mlx_capture import capture_pooled, install_taps3940MODEL = "mlx-community/Qwen3-0.6B-4bit"41MAP_JSON = ROOT / "atlas" / "qwen3-0.6b-4bit" / "probes" / "v2" / "map.json"42K = 543N_RANDOM_DRAWS = 2044N_RANDOM_DIRS = 345SEED = 7784647NOUN_PAIRS = [("key", "keys"), ("crate", "crates"), ("report", "reports"), ("valve", "valves"),48              ("ticket", "tickets"), ("ladder", "ladders"), ("sample", "samples"), ("cable", "cables"),49              ("permit", "permits"), ("beacon", "beacons"), ("filter", "filters"), ("stamp", "stamps")]50NEW_NEAR = ["across the yard", "above the workbench", "below the landing", "outside the depot",51            "around the corner", "beyond the fence", "beneath the awning", "atop the cabinet"]52NEUTRAL = ["The committee reviewed the plans before the meeting.",53           "A local historian kept detailed notes for years.",54           "The lead engineer questioned the original estimate.",55           "Her assistant preferred the older method.",56           "The night watchman described the process in a letter.",57           "An early visitor returned before the first frost.",58           "The town council approved the request after some debate.",59           "The apprentice carried the tools across the yard.",60           "According to the survey, the harbor required constant maintenance.",61           "The archive near the market attracted visitors from the region.",62           "The restored lighthouse stood at the edge of town.",63           "The workshop changed hands twice last century.",64           "The observatory remained open despite the storm.",65           "The vineyard was documented in the annual report.",66           "The glacier trail closed early in the season.",67           "The orchard supplied the market for decades.",68           "The library extended its hours during the recess.",69           "The clerk signed the manifest without a word.",70           "The surveyor traced the boundary along the quay.",71           "The curator shelved the samples behind the annex."]727374def spearman(a, b):75    ra = np.argsort(np.argsort(a)).astype(float)76    rb = np.argsort(np.argsort(b)).astype(float)77    ra -= ra.mean(); rb -= rb.mean()78    return float((ra * rb).sum() / np.sqrt((ra**2).sum() * (rb**2).sum()))798081def main() -> int:82    import mlx.core as mx83    from mlx_lm import load8485    t0 = time.time()86    model, tokenizer = load(MODEL)87    taps = install_taps(model)88    n_layers = len(taps)8990    mdoc = json.loads(MAP_JSON.read_text())91    agree = mdoc["properties"]["agreement"]92    diff_profile = np.array([a["selectivity_mean"] - t["selectivity_mean"]93                             for a, t in zip(agree["per_layer"]["A"],94                                             agree["twin_null_per_layer_A"])])95    ranked = list(np.argsort(-diff_profile))96    top_k, bottom_k = [int(x) for x in ranked[:K]], [int(x) for x in ranked[-K:]]9798    # ---- enlarged held-out bank (new locations -> no dedup collisions)99    rng = random.Random(SEED)100    combos = [(n, loc) for n in NOUN_PAIRS for loc in NEW_NEAR]101    rng.shuffle(combos)102    pairs = []103    for (sg, pl), loc in combos:104        pairs.append({"prefix": f"The {sg} {loc}", "singular": True})105        pairs.append({"prefix": f"The {pl} {loc}", "singular": False})106    print(f"held-out pairs: {len(pairs)}")107    prefix_ids = [tokenizer.encode(p["prefix"]) for p in pairs]108    id_is, id_are = tokenizer.encode(" is")[0], tokenizer.encode(" are")[0]109    neutral_ids = [tokenizer.encode(s) for s in NEUTRAL]110111    def margin() -> float:112        out = []113        for ids, p in zip(prefix_ids, pairs):114            logits = model(mx.array([ids]))[0, -1, :]115            mx.eval(logits)116            m = float(logits[id_is] - logits[id_are])117            out.append(m if p["singular"] else -m)118        return float(np.mean(out))119120    def nll() -> float:121        tot, cnt = 0.0, 0122        for ids in neutral_ids:123            x = mx.array([ids])124            logits = model(x)[0]125            logp = logits - mx.logsumexp(logits, axis=-1, keepdims=True)126            tgt = mx.array(ids[1:])127            picked = mx.take_along_axis(logp[:-1], tgt[:, None], axis=-1)128            mx.eval(picked)129            tot += float(-picked.sum()); cnt += len(ids) - 1130        return tot / cnt131132    def with_skip(layers: set[int], fn):133        for i, t in enumerate(taps):134            t.skip = i in layers135        try:136            return fn()137        finally:138            for t in taps:139                t.skip = False140141    base_m, base_nll = margin(), nll()142    print(f"baseline margin={base_m:+.4f} nll={base_nll:.4f}")143144    # ---------------- P1: skip conditions with NLL normalization145    def skip_cell(layers):146        m = with_skip(set(layers), margin)147        n = with_skip(set(layers), nll)148        return {"layers": sorted(int(x) for x in layers),149                "margin_damage": base_m - m,150                "nll_damage": n - base_nll,151                "specificity": (base_m - m) / max(n - base_nll, 1e-3)}152153    p1 = {"top": skip_cell(top_k), "bottom": skip_cell(bottom_k), "random": []}154    cand = [i for i in range(n_layers) if i not in set(top_k)]155    for d in range(N_RANDOM_DRAWS):156        p1["random"].append(skip_cell(random.Random(SEED + 1 + d).sample(cand, K)))157    rspec = np.array([c["specificity"] for c in p1["random"]])158    print(f"P1 specificity: top={p1['top']['specificity']:.3f} bottom={p1['bottom']['specificity']:.3f} "159          f"random mean={rspec.mean():.3f} p95={np.percentile(rspec, 95):.3f}")160161    # ---------------- P2: direction-level erasure, all layers162    # directions from agreement_A mean-pooled reps at each layer163    items = [json.loads(l) for l in164             (ROOT / "benchmarks" / "promptsets" / "agreement_A.jsonl").read_text().splitlines()]165    toks = [tokenizer.encode(it["text"]) for it in items]166    labels = np.array([it["label"] for it in items])167    reps = capture_pooled(model, taps, toks)["mean"]  # (n, L, d)168    dirs, mus = [], []169    for layer in range(n_layers):170        x = reps[:, layer, :]171        mu = x.mean(0)172        u = x[labels == "correct"].mean(0) - x[labels == "violated"].mean(0)173        u = u / (np.linalg.norm(u) + 1e-8)174        mus.append(mu); dirs.append(u)175176    def erase_fn(u_np, mu_np):177        u = mx.array(u_np.astype(np.float32))178        mu = mx.array(mu_np.astype(np.float32))179        def fn(out):180            h = out.astype(mx.float32)181            coef = ((h - mu) * u).sum(axis=-1, keepdims=True)182            return (h - coef * u).astype(out.dtype)183        return fn184185    per_layer = []186    for layer in range(n_layers):187        taps[layer].edit = erase_fn(dirs[layer], mus[layer])188        m_agree = margin()189        taps[layer].edit = None190        rms = []191        for s in range(N_RANDOM_DIRS):192            ru = np.random.default_rng(1000 * layer + s).standard_normal(dirs[layer].shape)193            ru /= np.linalg.norm(ru)194            taps[layer].edit = erase_fn(ru.astype(np.float32), mus[layer])195            rms.append(margin())196            taps[layer].edit = None197        specific = (base_m - m_agree) - (base_m - float(np.mean(rms)))198        per_layer.append({"layer": layer, "agree_dir_damage": base_m - m_agree,199                          "random_dir_damage_mean": base_m - float(np.mean(rms)),200                          "specific_damage": specific})201        print(f"  L{layer:02d} agreeDir={base_m - m_agree:+.3f} randDir={base_m - float(np.mean(rms)):+.3f} "202              f"specific={specific:+.3f}", flush=True)203204    spec_profile = np.array([r["specific_damage"] for r in per_layer])205    rho = spearman(spec_profile, diff_profile)206    perm_rng = np.random.default_rng(0)207    perms = np.array([spearman(perm_rng.permutation(spec_profile), diff_profile)208                      for _ in range(10_000)])209    p_perm = float((perms >= rho).mean())210    survives = bool(rho >= 0.4 and p_perm < 0.05)211    print(f"P2: Spearman rho={rho:+.3f} perm-p={p_perm:.4f} -> survives={survives}")212213    commit = subprocess.run(["git", "rev-parse", "HEAD"], cwd=ROOT,214                            capture_output=True, text=True, check=False).stdout.strip()215    ts = time.strftime("%Y%m%dT%H%M%SZ", time.gmtime())216    outdir = ROOT / "results" / "expC_causal_verification" / ts217    outdir.mkdir(parents=True)218    (outdir / "results.json").write_text(json.dumps({219        "experiment": "expC_causal_verification", "run": 2,220        "scope": "NLL-normalized skip specificity + direction-level erasure scan",221        "commit": commit,222        "config": {"model": MODEL, "k": K, "n_random_draws": N_RANDOM_DRAWS,223                   "n_random_dirs": N_RANDOM_DIRS, "n_pairs": len(pairs), "seed": SEED,224                   "top_layers": sorted(top_k), "bottom_layers": sorted(bottom_k)},225        "manifest": manifest(),226        "baseline": {"margin": base_m, "nll": base_nll},227        "p1_skip_specificity": p1,228        "p2_direction_erasure": {"per_layer": per_layer, "spearman_rho": rho,229                                 "perm_p": p_perm, "survives": survives},230        "wall_seconds": round(time.time() - t0, 1),231    }, indent=2) + "\n")232    print(f"results -> {outdir / 'results.json'}")233    return 0234235236if __name__ == "__main__":237    sys.exit(main())238