File size: 3,162 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 | #!/usr/bin/env python3
"""Direct test of the mechanism behind the null interventions: does Sigma
carry lifetime-specific information at all?
If Sigma_t were a useful state it would differ between lifetimes at the same
event index. If it has collapsed to a learned constant, the between-lifetime
spread will be negligible against the state's own scale.
"""
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.train.loader import iter_eval_batches # noqa
from evaluate import load_model, to_gpu # noqa
dev = "cuda"
model, _ = load_model(sys.argv[1] if len(sys.argv) > 1 else "PNSR_K4_fin", dev)
d = model.cfg.d
PROBE_T = [31, 63, 127, 191, 255]
snaps = {t: [] for t in PROBE_T}
with torch.no_grad():
for b in iter_eval_batches("val", shards_root(), 32, 32):
g = to_gpu(b, dev)
B, L = g["etype"].shape
state = model.initial_state(B, dev)
with torch.autocast("cuda", dtype=torch.bfloat16):
h = model.recenc.static(g["rec_val_toks"], g["rec_key_toks"], g["rec_store"],
g["rec_kind"], g["rec_key"], g["rec_ent"])
for t in range(L):
live = g["live"][:, t]; mask = live >= 0; rows = live.clamp(min=0)
bank = torch.gather(h, 1, rows.unsqueeze(-1).expand(-1, -1, d))
bank = model.recenc.finalize(
bank, (t - torch.gather(g["rec_birth"], 1, rows)).clamp(min=0))
state, _ = model.step(state, g["tok"][:, t], g["etype"][:, t],
g["dt"][:, t], bank * mask.unsqueeze(-1), mask)
if t in snaps:
snaps[t].append(state.float().cpu())
break
rep = {}
for t, lst in snaps.items():
S = torch.cat(lst) # [B, 120, d]
mean = S.mean(0, keepdim=True)
between = (S - mean).norm(dim=-1).mean().item() # spread across lifetimes
scale = S.norm(dim=-1).mean().item() # typical slot norm
flat = S.flatten(1)
fn = torch.nn.functional.normalize(flat, dim=-1)
cos = (fn @ fn.T)
off = cos[~torch.eye(len(fn), dtype=torch.bool)].mean().item()
rep[str(t)] = dict(slot_norm=round(scale, 4),
between_lifetime_spread=round(between, 4),
spread_over_norm=round(between / max(scale, 1e-9), 4),
mean_pairwise_cosine=round(off, 4))
print(t, rep[str(t)], flush=True)
# drift of the mean state over time: does it even change with t?
ms = {t: torch.cat(lst).mean(0).flatten() for t, lst in snaps.items()}
ks = sorted(ms)
rep["mean_state_cosine_t31_vs_t255"] = round(float(
torch.nn.functional.cosine_similarity(ms[ks[0]], ms[ks[-1]], dim=0)), 4)
run_name = sys.argv[1] if len(sys.argv) > 1 else "PNSR_K4_fin"
atomic_write_json(eval_root() / f"SIGMA_COLLAPSE_{run_name}.json", rep)
print("mean-state cosine t=31 vs t=255:", rep["mean_state_cosine_t31_vs_t255"])
|