spb/modelmap Public License

1#!/usr/bin/env python32# =============================================================================3#  Project   : modelmap4#  File      : experiments/micro/expC_causal_verification/implementation/benchmark_v3.py5#  Purpose   : Run #3 — replication of the causal direction-erasure profile6#  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 #3 (hypothesis registered before this run).1516Five direction sources (set A, set B, 3 bootstrap resamples of A); shared17random-direction controls; pairwise Spearman replication of the per-layer18specific-damage profile + the 2–15 vs 20–27 band claim per source.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 = 77942BOOTS = 34344NOUN_PAIRS = [("key", "keys"), ("crate", "crates"), ("report", "reports"), ("valve", "valves"),45              ("ticket", "tickets"), ("ladder", "ladders"), ("sample", "samples"), ("cable", "cables"),46              ("permit", "permits"), ("beacon", "beacons"), ("filter", "filters"), ("stamp", "stamps")]47NEW_NEAR = ["across the yard", "above the workbench", "below the landing", "outside the depot",48            "around the corner", "beyond the fence", "beneath the awning", "atop the cabinet"]495051def spearman(a, b):52    ra = np.argsort(np.argsort(a)).astype(float)53    rb = np.argsort(np.argsort(b)).astype(float)54    ra -= ra.mean(); rb -= rb.mean()55    return float((ra * rb).sum() / np.sqrt((ra**2).sum() * (rb**2).sum()))565758def directions_from(reps, labels, idx=None):59    """Per-layer unit diff-of-means directions (+grand means)."""60    if idx is not None:61        reps, labels = reps[idx], labels[idx]62    dirs, mus = [], []63    for layer in range(reps.shape[1]):64        x = reps[:, layer, :]65        mu = x.mean(0)66        u = x[labels == "correct"].mean(0) - x[labels == "violated"].mean(0)67        u = u / (np.linalg.norm(u) + 1e-8)68        dirs.append(u); mus.append(mu)69    return dirs, mus707172def main() -> int:73    import mlx.core as mx74    from mlx_lm import load7576    t0 = time.time()77    model, tokenizer = load(MODEL)78    taps = install_taps(model)79    n_layers = len(taps)8081    rng = random.Random(SEED)82    combos = [(n, loc) for n in NOUN_PAIRS for loc in NEW_NEAR]83    rng.shuffle(combos)84    pairs = []85    for (sg, pl), loc in combos:86        pairs.append({"prefix": f"The {sg} {loc}", "singular": True})87        pairs.append({"prefix": f"The {pl} {loc}", "singular": False})88    prefix_ids = [tokenizer.encode(p["prefix"]) for p in pairs]89    id_is, id_are = tokenizer.encode(" is")[0], tokenizer.encode(" are")[0]9091    def margin() -> float:92        out = []93        for ids, p in zip(prefix_ids, pairs):94            logits = model(mx.array([ids]))[0, -1, :]95            mx.eval(logits)96            m = float(logits[id_is] - logits[id_are])97            out.append(m if p["singular"] else -m)98        return float(np.mean(out))99100    def erase_fn(u_np, mu_np):101        u = mx.array(u_np.astype(np.float32))102        mu = mx.array(mu_np.astype(np.float32))103        def fn(out):104            h = out.astype(mx.float32)105            coef = ((h - mu) * u).sum(axis=-1, keepdims=True)106            return (h - coef * u).astype(out.dtype)107        return fn108109    base_m = margin()110    print(f"baseline margin {base_m:+.4f} | pairs {len(pairs)}")111112    # ---- direction sources113    sources: dict[str, tuple] = {}114    reps_cache = {}115    for setname in ("A", "B"):116        items = [json.loads(l) for l in117                 (ROOT / "benchmarks" / "promptsets" / f"agreement_{setname}.jsonl").read_text().splitlines()]118        toks = [tokenizer.encode(it["text"]) for it in items]119        labels = np.array([it["label"] for it in items])120        print(f"capture agreement_{setname}…", flush=True)121        reps = capture_pooled(model, taps, toks)["mean"]122        reps_cache[setname] = (reps, labels)123        sources[f"set{setname}"] = directions_from(reps, labels)124    repsA, labelsA = reps_cache["A"]125    for b in range(BOOTS):126        idx = np.random.default_rng(100 + b).integers(0, len(labelsA), len(labelsA))127        sources[f"bootA{b}"] = directions_from(repsA, labelsA, idx)128129    # ---- shared random-direction control per layer130    rand_damage = []131    for layer in range(n_layers):132        ms = []133        for s in range(N_RANDOM_DIRS):134            ru = np.random.default_rng(1000 * layer + s).standard_normal(repsA.shape[-1])135            ru /= np.linalg.norm(ru)136            taps[layer].edit = erase_fn(ru.astype(np.float32), sources["setA"][1][layer])137            ms.append(margin())138            taps[layer].edit = None139        rand_damage.append(base_m - float(np.mean(ms)))140    print("random-direction controls done", flush=True)141142    # ---- per-source agreement-direction scans143    profiles = {}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        band_early = float(np.mean(prof[2:16]))153        band_late = float(np.mean(prof[20:28]))154        print(f"{name:8s} early(2-15)={band_early:+.3f} late(20-27)={band_late:+.3f} "155              f"ratio={band_early / max(band_late, 1e-6):.1f}", flush=True)156157    names = list(profiles)158    rhos = [spearman(np.array(profiles[a]), np.array(profiles[b]))159            for i, a in enumerate(names) for b in names[i + 1:]]160    mean_rho = float(np.mean(rhos))161    band_ok = all(np.mean(profiles[n][2:16]) >= 3 * max(np.mean(profiles[n][20:28]), 1e-6)162                  for n in names)163    survives = bool(mean_rho >= 0.7 and band_ok)164    print(f"replication: mean pairwise rho={mean_rho:.3f} (min {min(rhos):.3f}) "165          f"band-claim-all-sources={band_ok} -> publishable={survives}")166167    commit = subprocess.run(["git", "rev-parse", "HEAD"], cwd=ROOT,168                            capture_output=True, text=True, check=False).stdout.strip()169    ts = time.strftime("%Y%m%dT%H%M%SZ", time.gmtime())170    outdir = ROOT / "results" / "expC_causal_verification" / ts171    outdir.mkdir(parents=True)172    (outdir / "results.json").write_text(json.dumps({173        "experiment": "expC_causal_verification", "run": 3,174        "scope": "replication of the direction-erasure causal profile (5 sources)",175        "commit": commit,176        "config": {"model": MODEL, "sources": names, "n_random_dirs": N_RANDOM_DIRS,177                   "n_pairs": len(pairs), "seed": SEED, "boots": BOOTS},178        "manifest": manifest(),179        "baseline_margin": base_m,180        "random_direction_damage": rand_damage,181        "profiles": profiles,182        "replication": {"pairwise_rhos": rhos, "mean_rho": mean_rho,183                        "min_rho": float(min(rhos)), "band_claim_all": band_ok,184                        "publishable": survives},185        "wall_seconds": round(time.time() - t0, 1),186    }, indent=2) + "\n")187    print(f"results -> {outdir / 'results.json'}")188    return 0189190191if __name__ == "__main__":192    sys.exit(main())193