#!/usr/bin/env python3 # ============================================================================= # Project : modelmap # File : experiments/micro/expC_causal_verification/implementation/benchmark.py # Purpose : Run #1 — layer-skip ablation test of the agreement probe map # Author : Simon-Pierre Boucher # Contact : contact@spboucher.ai # Website : https://modelmap.io # Created : 2026-08-12 # Modified : 2026-08-12 # Platform : macOS / Apple Silicon (arm64) — MLX / Metal # License : All rights reserved (research code) # ============================================================================= """expC run #1 (hypothesis registered before this run). Does the agreement DIFFERENTIAL probe map (expA run #2) survive causal testing? Skip the map's top-5 layers vs 20 random-5 draws vs bottom-5, measure the drop in grammatical-agreement logit margin on held-out minimal pairs. """ from __future__ import annotations import json import random import subprocess import sys import time from pathlib import Path import numpy as np ROOT = Path(__file__).resolve().parents[4] sys.path.insert(0, str(ROOT / "src")) sys.path.insert(0, str(ROOT / "benchmarks")) from hardware_manifest import manifest from modelmap.capture.mlx_capture import install_taps MODEL = "mlx-community/Qwen3-0.6B-4bit" MAP_JSON = ROOT / "atlas" / "qwen3-0.6b-4bit" / "probes" / "v2" / "map.json" K = 5 N_RANDOM_DRAWS = 20 N_PAIRS = 200 SEED = 777 NOUN_PAIRS = [("key", "keys"), ("crate", "crates"), ("report", "reports"), ("valve", "valves"), ("ticket", "tickets"), ("ladder", "ladders"), ("sample", "samples"), ("cable", "cables"), ("permit", "permits"), ("beacon", "beacons"), ("filter", "filters"), ("stamp", "stamps")] NEAR = ["near the entrance", "beside the counter", "under the shelf", "behind the gate", "next to the archive", "along the corridor", "opposite the office", "inside the vault", "across the yard", "above the workbench"] def held_out_pairs(rng: random.Random, forbidden: set[str]): """Minimal pairs (prefix, singular?) deduped against every v2 probe text.""" pairs, seen = [], set() combos = [(n, loc) for n in NOUN_PAIRS for loc in NEAR] rng.shuffle(combos) for (sg, pl), loc in combos * 4: for noun, singular in ((sg, True), (pl, False)): prefix = f"The {noun} {loc}" probe_like = f"{prefix} is" if singular else f"{prefix} are" if prefix in seen or any(probe_like in f for f in forbidden): continue seen.add(prefix) pairs.append({"prefix": prefix, "singular": singular}) if len(pairs) >= N_PAIRS: return pairs return pairs def main() -> int: import mlx.core as mx from mlx_lm import load t0 = time.time() model, tokenizer = load(MODEL) taps = install_taps(model) n_layers = len(taps) # top/bottom differential layers from the published map (mean pooling) mdoc = json.loads(MAP_JSON.read_text()) agree = mdoc["properties"]["agreement"] diff = [(a["layer"], a["selectivity_mean"] - t["selectivity_mean"]) for a, t in zip(agree["per_layer"]["A"], agree["twin_null_per_layer_A"])] ranked = [l for l, _ in sorted(diff, key=lambda x: -x[1])] top_k, bottom_k = ranked[:K], ranked[-K:] print(f"top-{K} differential layers: {sorted(top_k)} | bottom-{K}: {sorted(bottom_k)}") # held-out minimal pairs forbidden = set() for f in (ROOT / "benchmarks" / "promptsets").glob("agreement_*.jsonl"): forbidden.update(json.loads(l)["text"] for l in f.read_text().splitlines()) rng = random.Random(SEED) pairs = held_out_pairs(rng, forbidden) print(f"held-out minimal pairs: {len(pairs)}") tok_is = tokenizer.encode(" is") tok_are = tokenizer.encode(" are") assert len(tok_is) == 1 and len(tok_are) == 1, "verb forms must be single tokens" id_is, id_are = tok_is[0], tok_are[0] prefix_ids = [tokenizer.encode(p["prefix"]) for p in pairs] def margin(skip_layers: set[int]) -> float: for i, t in enumerate(taps): t.skip = i in skip_layers margins = [] for ids, p in zip(prefix_ids, pairs): logits = model(mx.array([ids]))[0, -1, :] mx.eval(logits) m = float(logits[id_is] - logits[id_are]) margins.append(m if p["singular"] else -m) for t in taps: t.skip = False return float(np.mean(margins)) baseline = margin(set()) print(f"baseline margin: {baseline:+.4f}") top_m = margin(set(top_k)) bottom_m = margin(set(bottom_k)) candidates = [i for i in range(n_layers) if i not in set(top_k)] random_ms = [] for d in range(N_RANDOM_DRAWS): draw = set(random.Random(SEED + 1 + d).sample(candidates, K)) random_ms.append(margin(draw)) dmg = lambda m: baseline - m rd = np.array([dmg(m) for m in random_ms]) p95 = float(np.percentile(rd, 95)) verdict = bool(dmg(top_m) >= p95 and dmg(top_m) >= 2 * rd.mean()) print(f"damage: top-{K}={dmg(top_m):+.4f} bottom-{K}={dmg(bottom_m):+.4f} " f"random mean={rd.mean():+.4f} p95={p95:+.4f} -> survives={verdict}") commit = subprocess.run(["git", "rev-parse", "HEAD"], cwd=ROOT, capture_output=True, text=True, check=False).stdout.strip() ts = time.strftime("%Y%m%dT%H%M%SZ", time.gmtime()) outdir = ROOT / "results" / "expC_causal_verification" / ts outdir.mkdir(parents=True) doc = { "experiment": "expC_causal_verification", "run": 1, "scope": "layer-skip ablation of the agreement differential probe map", "commit": commit, "config": {"model": MODEL, "k": K, "n_random_draws": N_RANDOM_DRAWS, "n_pairs": len(pairs), "seed": SEED, "map_source": str(MAP_JSON.relative_to(ROOT)), "top_layers": sorted(top_k), "bottom_layers": sorted(bottom_k)}, "manifest": manifest(), "results": { "baseline_margin": baseline, "top_k_margin": top_m, "top_k_damage": dmg(top_m), "bottom_k_margin": bottom_m, "bottom_k_damage": dmg(bottom_m), "random_margins": random_ms, "random_damage_mean": float(rd.mean()), "random_damage_p95": p95, "survives_causal_test": verdict, }, "wall_seconds": round(time.time() - t0, 1), } (outdir / "results.json").write_text(json.dumps(doc, indent=2) + "\n") print(f"results -> {outdir / 'results.json'} ({doc['wall_seconds']} s)") return 0 if __name__ == "__main__": sys.exit(main())