spb/modelmap Public License

1#!/usr/bin/env python32# =============================================================================3#  Project   : modelmap4#  File      : experiments/micro/expA_probe_reliability/implementation/benchmark.py5#  Purpose   : expA run #1 — probe noise floor on a real 4-bit model + nulls6#  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"""expA run #1 (hypothesis registered in hypothesis.md before this run).1516Per-layer linear probes with the full controls doctrine on Qwen3-0.6B-4bit:173 properties x 2 disjoint promptsets x 28 layers x 5 seeds, shuffled-label18controls built into every probe, random-init architecture twin as the19structure-from-architecture null, bootstrap p-values + BH-FDR across layer20scans, replication rate on top-k layer sets.21"""2223from __future__ import annotations2425import json26import 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_mean_pooled, install_taps, random_init_twin39from modelmap.probes.linear import probe_with_control40from modelmap.stats.replication import bh_fdr, bootstrap_ci, replication_rate4142MODEL = "mlx-community/Qwen3-0.6B-4bit"43PROPERTIES = ("lang_id", "code_prose", "arith")44SETS = ("A", "B")45SEEDS = (0, 1, 2, 3, 4)46TOP_K = 547FDR_Q = 0.0548PROMPTS = ROOT / "benchmarks" / "promptsets"495051def load_set(name: str) -> tuple[list[str], np.ndarray]:52    items = [json.loads(l) for l in (PROMPTS / f"{name}.jsonl").read_text().splitlines()]53    return [it["text"] for it in items], np.array([it["label"] for it in items])545556def probe_grid(reps: np.ndarray, labels: np.ndarray) -> list[dict]:57    """All (layer, seed) probe cells for one representation tensor."""58    out = []59    n_layers = reps.shape[1]60    for layer in range(n_layers):61        for seed in SEEDS:62            r = probe_with_control(reps[:, layer, :], labels, seed=seed)63            out.append({"layer": layer, "seed": seed,64                        "task_acc": r.task_accuracy,65                        "control_acc": r.control_accuracy,66                        "selectivity": r.selectivity})67    return out686970def summarize(cells: list[dict], n_layers: int) -> dict:71    """Per-layer aggregation: means, seed SD, bootstrap CI + p on selectivity."""72    rng = np.random.default_rng(0)73    layers = []74    pvals = []75    for layer in range(n_layers):76        sel = np.array([c["selectivity"] for c in cells if c["layer"] == layer])77        acc = np.array([c["task_acc"] for c in cells if c["layer"] == layer])78        point, lo, hi = bootstrap_ci(sel, rng=rng)79        boots = rng.choice(sel, size=(10_000, sel.size)).mean(axis=1)80        p = float(max((boots <= 0).mean(), 1 / 10_000))81        pvals.append(p)82        layers.append({"layer": layer,83                       "task_acc_mean": float(acc.mean()),84                       "task_acc_seed_sd": float(acc.std(ddof=1)),85                       "selectivity_mean": point,86                       "selectivity_ci": [lo, hi],87                       "p_boot": p})88    disc = bh_fdr(np.array(pvals), q=FDR_Q)89    for row, d in zip(layers, disc):90        row["fdr_significant"] = bool(d)91    top_sets = []92    for seed in SEEDS:93        acc_by_layer = [(c["layer"], c["task_acc"]) for c in cells if c["seed"] == seed]94        top = {l for l, _ in sorted(acc_by_layer, key=lambda t: -t[1])[:TOP_K]}95        top_sets.append(top)96    rep_point, rep_lo, rep_hi = replication_rate(top_sets)97    return {"layers": layers,98            "replication_rate_topk": {"k": TOP_K, "point": rep_point, "ci": [rep_lo, rep_hi]},99            "n_fdr_significant": int(disc.sum())}100101102def main() -> int:103    import mlx.core as mx104    from mlx_lm import load105    from mlx_lm.utils import hf_repo_to_path106107    t_start = time.time()108    mx.random.seed(0)109    model, tokenizer = load(MODEL)110    taps = install_taps(model)111    model_path = hf_repo_to_path(MODEL)112113    twin = random_init_twin(model_path)114    twin_taps = install_taps(twin)115116    results = {"real": {}, "twin": {}}117    grids = {"real": {}, "twin": {}}118119    for prop in PROPERTIES:120        for s in SETS:121            name = f"{prop}_{s}"122            texts, labels = load_set(name)123            toks = [tokenizer.encode(t) for t in texts]124            print(f"capture real  {name} ({len(texts)} prompts)…", flush=True)125            reps = capture_mean_pooled(model, taps, toks)126            grids["real"][name] = (probe_grid(reps, labels), reps.shape[1])127            if s == "A":  # architecture null on the A sets128                print(f"capture twin  {name}…", flush=True)129                reps_t = capture_mean_pooled(twin, twin_taps, toks)130                grids["twin"][name] = (probe_grid(reps_t, labels), reps_t.shape[1])131132    for kind in ("real", "twin"):133        for name, (cells, n_layers) in grids[kind].items():134            results[kind][name] = summarize(cells, n_layers)135            results[kind][name]["cells"] = cells136137    # -------- headline numbers138    summary = {}139    for prop in PROPERTIES:140        a = results["real"][f"{prop}_A"]["layers"]141        b = results["real"][f"{prop}_B"]["layers"]142        seed_sd = float(np.mean([r["task_acc_seed_sd"] for r in a + b]))143        shifts = [abs(ra["task_acc_mean"] - rb["task_acc_mean"]) for ra, rb in zip(a, b)]144        n_shift = int(sum(s > max(seed_sd, 1e-9) for s in shifts))145        twin_sel = float(max(r["selectivity_mean"] for r in results["twin"][f"{prop}_A"]["layers"]))146        summary[prop] = {147            "mean_seed_sd": seed_sd,148            "mean_dataset_shift": float(np.mean(shifts)),149            "layers_shift_gt_seed_sd": n_shift,150            "n_layers": len(a),151            "max_task_acc_A": float(max(r["task_acc_mean"] for r in a)),152            "max_task_acc_B": float(max(r["task_acc_mean"] for r in b)),153            "twin_max_selectivity": twin_sel,154            "replication_topk_A": results["real"][f"{prop}_A"]["replication_rate_topk"]["point"],155            "replication_topk_B": results["real"][f"{prop}_B"]["replication_rate_topk"]["point"],156        }157        print(f"{prop:12s} seedSD={seed_sd:.4f} shift={summary[prop]['mean_dataset_shift']:.4f} "158              f"layers(shift>sd)={n_shift}/{len(a)} maxAccA={summary[prop]['max_task_acc_A']:.3f} "159              f"twinMaxSel={twin_sel:.3f} repl={summary[prop]['replication_topk_A']:.2f}")160161    commit = subprocess.run(["git", "rev-parse", "HEAD"], cwd=ROOT,162                            capture_output=True, text=True, check=False).stdout.strip()163    ts = time.strftime("%Y%m%dT%H%M%SZ", time.gmtime())164    outdir = ROOT / "results" / "expA_probe_reliability" / ts165    outdir.mkdir(parents=True)166    doc = {167        "experiment": "expA_probe_reliability",168        "run": 1,169        "scope": "probe noise floor, 3 properties x 2 sets x 28 layers x 5 seeds + twin null",170        "commit": commit,171        "config": {"model": MODEL, "seeds": list(SEEDS), "top_k": TOP_K, "fdr_q": FDR_Q,172                   "promptsets": json.loads((PROMPTS / "manifest.json").read_text())},173        "manifest": manifest(),174        "summary": summary,175        "results": results,176        "wall_seconds": round(time.time() - t_start, 1),177    }178    (outdir / "results.json").write_text(json.dumps(doc, indent=2) + "\n")179    print(f"results -> {outdir / 'results.json'}  ({doc['wall_seconds']} s)")180    return 0181182183if __name__ == "__main__":184    sys.exit(main())185