"""Memory-readout causality probe. Runs the same current-vs-stale items through the live CEDL memory path and a memory-off ablation. The paired comparison estimates whether contextual memory readout causally improves current-token selection. Usage: python probes/probe_memory_causality.py \ --checkpoint pytorch_model.bin \ --sidecar cedl_config.json \ --n-items 50 """ import argparse import json import os import sys import numpy as np import torch import torch.nn as nn sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) sys.path.insert(0, "/content") import CEDL import data_v4c_pairs as v4c MEMORY_GATE_ATTR = "v" + "6_lambda_head" class ZeroLambdaHead(nn.Module): """Forces the memory mixture gate near zero.""" def forward(self, h): return torch.full( h.shape[:-1] + (1,), -100.0, device=h.device, dtype=h.dtype, ) def load_sidecar_constructor_kwargs(sidecar_path): """Read public CEDL config and return constructor keyword arguments.""" with open(sidecar_path) as f: sc = json.load(f) mem = sc.get("memory_readout") if not isinstance(mem, dict): raise ValueError("Expected cedl_config.json with a memory_readout block.") source_name = str(mem.get("source", "contextual_memory_state")) source_map = { "contextual_memory_state": "q_mem", "decoder_state": "h_d", "expanded_state": "h_e", "attractor_state": "q_attractor", "q_mem": "q_mem", } return dict( lambda_head=bool(mem.get("lambda_head", True)), lambda_head_hidden=int(mem.get("lambda_head_hidden", 160)), lambda_head_bias_init=float(mem.get("lambda_head_bias_init", -7.0)), lambda_head_w_init_std=float( mem.get("lambda_head_w_init_std", 0.05)), bce_objective=( mem.get("selection_objective") == "binary_answer_background"), sel_weight=1.0, bg_weight=1.0, bg_target=float(mem.get("background_target", 0.01)), wt_sparsity_weight=float(mem.get("sparsity_weight", 0.05)), wt_sparsity_target=float(mem.get("sparsity_target", 0.05)), memory_head_enabled=bool(mem.get("enabled", True)), memory_ce_weight=float(mem.get("memory_ce_weight", 1.0)), memory_pair_ce_weight=float(mem.get("pair_ce_weight", 5.0)), memory_query_source=source_map.get(source_name, source_name), memory_readout_mode="direct", source_adapter=bool(mem.get("source_adapter", True)), context_adapter=bool(mem.get("context_adapter", True)), specialist_noinject=bool(mem.get("no_injection", True)), ) def model_forward_logits(m, ids_b): """Call m(ids_b) and return the logits tensor [B, T, V] regardless of whether forward returns logits-only or (logits, aux_loss).""" out = m(ids_b) if isinstance(out, tuple): logits = out[0] else: logits = out return logits def main(): p = argparse.ArgumentParser() p.add_argument("--checkpoint", type=str, default="pytorch_model.bin") p.add_argument("--sidecar", type=str, default="cedl_config.json") p.add_argument("--n-items", type=int, default=50) p.add_argument("--seed", type=int, default=0) args = p.parse_args() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"[setup] device={device}") print(f"[setup] checkpoint={args.checkpoint}") print(f"[setup] sidecar={args.sidecar}") print(f"[setup] n_items={args.n_items}") model_kwargs = load_sidecar_constructor_kwargs(args.sidecar) print(f"[setup] memory_head={model_kwargs['lambda_head']} " f"hidden={model_kwargs['lambda_head_hidden']} " f"bias_init={model_kwargs['lambda_head_bias_init']} " f"sparsity_weight={model_kwargs['wt_sparsity_weight']} " f"sparsity_target={model_kwargs['wt_sparsity_target']} " f"memory_source={model_kwargs['memory_query_source']} " f"readout_mode={model_kwargs['memory_readout_mode']}") m = CEDL.build_model("CEDL", vocab=50257, max_seq=1024, **model_kwargs) m = m.to(device).eval() state = torch.load(args.checkpoint, map_location="cpu", weights_only=True) msd = state["model"] if isinstance(state, dict) and "model" in state else state if any(k.startswith("_orig_mod.") for k in msd): msd = {k.replace("_orig_mod.", ""): v for k, v in msd.items()} res = m.load_state_dict(msd, strict=True) print(f"[load] strict OK (missing={len(res.missing_keys)} unexpected={len(res.unexpected_keys)})") m.feedback_alpha.fill_(1.0) if hasattr(m, "sl_alpha"): m.sl_alpha.fill_(1.0) from transformers import GPT2TokenizerFast tok = GPT2TokenizerFast.from_pretrained("gpt2") print(f"\n[v4c] generating {args.n_items} items (seed={args.seed})...") items = v4c.generate(tok, n=args.n_items, seed=args.seed) print(f"[v4c] generated {len(items)} raw items") saved_head = getattr(m, MEMORY_GATE_ATTR) zero_head = ZeroLambdaHead().to(device) margin_active = [] margin_off = [] pred_correct_active = [] pred_correct_off = [] lam_at_ans_active = [] n_skipped = 0 skip_reasons = {"neutral_control": 0, "no_cur_stale": 0, "bad_length": 0, "same_token": 0} for it_idx, it in enumerate(items): if it.family == "neutral_control": n_skipped += 1; skip_reasons["neutral_control"] += 1; continue if not it.current or not it.stale: n_skipped += 1; skip_reasons["no_cur_stale"] += 1; continue ol = getattr(it, "original_length", 0) if ol <= 1 or ol > 1024: n_skipped += 1; skip_reasons["bad_length"] += 1; continue cur_t = int(it.ids[it.current[0][0]]) stale_t = int(it.ids[it.stale[0][0]]) if cur_t == stale_t: n_skipped += 1; skip_reasons["same_token"] += 1; continue ans_p = ol - 1 ids_b = torch.tensor([it.ids[:ol]], device=device, dtype=torch.long) setattr(m, MEMORY_GATE_ATTR, saved_head) with torch.no_grad(): logits_a = model_forward_logits(m, ids_b) cur_lp_a = float(logits_a[0, ans_p, cur_t].item()) stl_lp_a = float(logits_a[0, ans_p, stale_t].item()) pred_a = int(logits_a[0, ans_p].argmax().item()) margin_active.append(cur_lp_a - stl_lp_a) pred_correct_active.append(pred_a == cur_t) with torch.no_grad(): h_C = m.c_stage(ids_b, feedback=None) h_E, h_E_sparse = m.e_stage(h_C, feedback=None) v_vec, _ = m.salience(h_E) if m.use_salience else (None, None) h_D, _ = m.d_stage(h_E, h_E_sparse, v_vec=v_vec) lam_logit_a = saved_head(h_D[:, ans_p:ans_p + 1, :]) lam_a = torch.sigmoid(lam_logit_a).clamp(1e-4, 1.0 - 1e-4) lam_at_ans_active.append(float(lam_a.item())) setattr(m, MEMORY_GATE_ATTR, zero_head) with torch.no_grad(): logits_o = model_forward_logits(m, ids_b) cur_lp_o = float(logits_o[0, ans_p, cur_t].item()) stl_lp_o = float(logits_o[0, ans_p, stale_t].item()) pred_o = int(logits_o[0, ans_p].argmax().item()) margin_off.append(cur_lp_o - stl_lp_o) pred_correct_off.append(pred_o == cur_t) setattr(m, MEMORY_GATE_ATTR, saved_head) n = len(margin_active) print(f"\n[items] used={n} skipped={n_skipped} reasons={skip_reasons}") if n == 0: print("No valid items — aborting.") return ma = np.array(margin_active) mo = np.array(margin_off) delta = ma - mo print(f"\n[λ at answer row, ACTIVE path] " f"mean={np.mean(lam_at_ans_active):.4f} " f"std={np.std(lam_at_ans_active):.4f} " f"min={min(lam_at_ans_active):.4f} " f"max={max(lam_at_ans_active):.4f}") print(f"\n[current − stale logit margin at answer position]") print(f" ACTIVE: mean={ma.mean():+.3f} std={ma.std():.3f} " f"median={np.median(ma):+.3f}") print(f" BANK-OFF: mean={mo.mean():+.3f} std={mo.std():.3f} " f"median={np.median(mo):+.3f}") print(f" Δ (act−off): mean={delta.mean():+.3f} std={delta.std():.3f} " f"median={np.median(delta):+.3f}") from scipy import stats as sst t_stat, p_val = sst.ttest_rel(ma, mo) print(f" paired t = {t_stat:+.3f} p = {p_val:.4f} N={n}") d_paired = float(delta.mean() / max(delta.std(), 1e-9)) print(f" Cohen's d (paired) = {d_paired:+.3f}") acc_a = float(np.mean(pred_correct_active)) acc_o = float(np.mean(pred_correct_off)) print(f"\n[top-1 accuracy: argmax(logits[ans_p]) == cur_t]") print(f" ACTIVE: {acc_a*100:5.1f}% ({sum(pred_correct_active)}/{n})") print(f" BANK-OFF: {acc_o*100:5.1f}% ({sum(pred_correct_off)}/{n})") print(f" Δ: {(acc_a-acc_o)*100:+5.1f} pp") print(f"\n[first 10 items — paired breakdown]") print(f" {'#':>3} {'m_A':>7} {'m_O':>7} {'Δ':>7} " f"{'pred_A':>6} {'pred_O':>6}") for i in range(min(10, n)): print(f" {i:>3} {ma[i]:>+7.3f} {mo[i]:>+7.3f} {delta[i]:>+7.3f} " f"{'OK' if pred_correct_active[i] else 'XX':>6} " f"{'OK' if pred_correct_off[i] else 'XX':>6}") print(f"\n{'='*64}") print(f"Memory-readout causality — verdict") print(f"{'='*64}") if t_stat > 2.0 and (acc_a - acc_o) > 0.10: verdict = ( f"ACTIVE > BANK-OFF (t={t_stat:+.2f}, Δacc=+{(acc_a-acc_o)*100:.1f}pp, " f"d={d_paired:+.2f}).\n" f" → Bank is causally useful for V4c current-vs-stale.\n" f" → PROCEED to full probe suite + manuscript update." ) elif abs(t_stat) < 1.0 and abs(acc_a - acc_o) < 0.05: verdict = ( f"ACTIVE ≈ BANK-OFF (t={t_stat:+.2f}, Δacc={(acc_a-acc_o)*100:+.1f}pp, " f"d={d_paired:+.2f}).\n" f" → λ head routes correctly but BANK READOUT IS NOT CAUSALLY USEFUL.\n" f" → If mem_head_bank is already untied and alternate sources still\n" f" fail, test direct-source readout before more bank replay." ) elif t_stat < -2.0: verdict = ( f"BANK-OFF > ACTIVE (t={t_stat:+.2f}, Δacc={(acc_a-acc_o)*100:+.1f}pp, " f"d={d_paired:+.2f}).\n" f" → BANK ACTIVELY HURTS V4c prediction. Diagnose the bank readout\n" f" bottleneck before any manuscript work." ) else: verdict = ( f"INTERMEDIATE (t={t_stat:+.2f}, Δacc={(acc_a-acc_o)*100:+.1f}pp, " f"d={d_paired:+.2f}).\n" f" → Mixed signal. Run full probe suite to see broader pattern; if\n" f" contrastive / stale-vs-current probes show clear lift, the bank\n" f" is helpful on average even if the V4c-margin metric is noisy." ) print(f" {verdict}") print(f"{'='*64}") if __name__ == "__main__": main()