File size: 5,680 Bytes
f930dac | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 | #!/usr/bin/env python3
"""Separate DELAY from INTERFERENCE in the Stage-A positive result.
The preregistered mechanistic signature is: if identity-local writes protect a
binding, accuracy should depend on rewrites of the QUERIED binding, not on how
many unrelated bindings were written in between. The pooled slope conflates
the two because n_unrelated correlates with delay. This computes the 2-D
table acc(delay bucket x unrelated-write bucket).
"""
import json
import sys
from pathlib import Path
import numpy as np
import torch
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
sys.path.insert(0, str(Path(__file__).resolve().parent))
from pns.common import atomic_write_json, eval_root, shards_root # noqa
from pns.checkpoint import load_model # noqa: E402
from pns.model.modules import enum_legal_mask # noqa: E402
from pns.train.loader import iter_eval_batches # noqa: E402
from pns.world.schema import Fam # noqa: E402
from stage_a import to_dev # noqa: E402
SEM = (int(Fam.SEM_LATEST), int(Fam.SEM_2HOP))
@torch.no_grad()
def rows_for(run, dev, n_lifetimes=400):
m, _, _ = load_model(run, dev)
out = []
for b in iter_eval_batches("e3_dev", shards_root(), 24, n_lifetimes):
g = to_dev(b, dev)
B, L = g["etype"].shape
state = m.initial_state(B, dev)
wall = torch.zeros(B, L, device=dev) # cumulative ALL binding writes
wown = {} # cumulative writes per slot
cum = torch.zeros(B, L, m.cfg.n_slots, device=dev)
run_c = torch.zeros(B, m.cfg.n_slots, device=dev)
acc_all = torch.zeros(B, device=dev)
with torch.autocast("cuda", dtype=torch.bfloat16):
for t in range(L):
state, o = m.step(state, g["tok"][:, t], g["etype"][:, t],
g["dt"][:, t], g["bind_write"][:, t],
g["bind_read"][:, t], g["bind_slot_ent"],
g["bind_slot_attr"])
w = g["bind_write"][:, t]
hit = (w >= 0)
acc_all = acc_all + hit.float()
wall[:, t] = acc_all
if hit.any():
run_c[torch.arange(B, device=dev), w.clamp(min=0)] += hit.float()
cum[:, t] = run_c
sel = torch.isin(g["family"][:, t], torch.tensor(SEM, device=dev))
if sel.any():
legal = enum_legal_mask(g["enum_legal"][sel, t])
pred = o["enum"][sel].masked_fill(~legal, -1e9).argmax(-1)
ok = (pred == g["enum_gold"][sel, t]).float().cpu().numpy()
idx = torch.nonzero(sel).flatten().tolist()
for j, i in enumerate(idx):
d = int(g["delay"][i, t])
e0 = max(0, t - d)
rs = int(g["bind_read"][i, t])
total = float(wall[i, t] - wall[i, e0])
own = float(cum[i, t, rs] - cum[i, e0, rs]) if rs >= 0 else 0.0
out.append((float(ok[j]), d, total - own, own))
return np.array(out) if out else np.zeros((0, 4))
def table(a, name):
ok, d, unrel, own = a[:, 0], a[:, 1], a[:, 2], a[:, 3]
dq = [0, 16, 64, 10**9]
uq = np.quantile(unrel, [0.33, 0.66])
lines = [f"### {name}", "",
"| delay \\ unrelated writes | few | mid | many | row slope |",
"| --- | --- | --- | --- | --- |"]
res = {}
for i in range(3):
dm = (d > dq[i]) & (d <= dq[i + 1])
cells, vals = [], []
for lo, hi in ((-1, uq[0]), (uq[0], uq[1]), (uq[1], 1e18)):
m = dm & (unrel > lo) & (unrel <= hi)
if m.sum() > 25:
v = float(ok[m].mean()); vals.append(v)
cells.append(f"{v:.3f} (n={int(m.sum())})")
else:
vals.append(np.nan); cells.append("-")
sl = (vals[2] - vals[0]) if not np.isnan(vals[2] + vals[0]) else np.nan
res[f"delay_{dq[i]}_{dq[i+1]}"] = dict(few=vals[0], mid=vals[1],
many=vals[2], slope=sl)
lines.append(f"| d {dq[i]}-{dq[i+1] if i < 2 else 'inf'} | "
+ " | ".join(cells) + f" | {sl:+.3f} |")
# dependence on rewrites of the QUERIED binding, holding delay wide
lines += ["", "| own-binding rewrites since evidence | acc |", "| --- | --- |"]
for lo, hi, nm in ((-0.5, 0.5, "0"), (0.5, 1.5, "1"), (1.5, 1e9, "2+")):
m = (own > lo) & (own <= hi)
if m.sum() > 25:
lines.append(f"| {nm} | {float(ok[m].mean()):.3f} (n={int(m.sum())}) |")
res[f"own_{nm}"] = float(ok[m].mean())
return lines, res
def main():
dev = "cuda"
runs = sys.argv[1:] or ["E3A_bind_s1", "E3A_unbound_s1"]
rep, lines = {}, ["# Experiment 2 Stage-A: delay vs interference, decomposed", "",
"The pooled interference slope conflates delay with "
"unrelated writes because the two correlate. This holds "
"delay fixed within rows.", ""]
for run in runs:
a = rows_for(run, dev)
L, r = table(a, run)
lines += L + [""]
rep[run] = r
print(run, json.dumps(r, default=str)[:400], flush=True)
out = Path(__file__).resolve().parent.parent / "results" / "reproduced"
out.mkdir(parents=True, exist_ok=True)
(out / "interference.md").write_text("\n".join(lines) + "\n")
atomic_write_json(eval_root() / "E3_INTERFERENCE.json", rep)
print("->", out / "interference.md")
if __name__ == "__main__":
main()
|