#!/usr/bin/env python3 """Probe whether a memory source vector is directly decodable. This diagnostic freezes the trained CEDL checkpoint and trains a temporary linear vocab readout from one selected source vector: h_d, h_e, q_attractor, or q_mem It bypasses the external bank bottleneck: source -> direct_head -> vocab instead of: source -> bank_q_proj -> 256-slot mem_keys/mem_vals -> mem_head_bank -> vocab If direct_head generalizes on B2 while the bank path fails, the bottleneck is the learned 256-slot bank readout. If direct_head also fails, the selected source does not robustly encode the current-token target under this synthetic held-out probe. """ import argparse import json import os import sys from typing import Dict, List, Tuple import numpy as np import torch import torch.nn as nn import torch.nn.functional as F sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) sys.path.insert(0, "/content") import CEDL import data_v4c_pairs as v4c SOURCES = ("h_d", "h_e", "q_attractor", "q_mem") MEMORY_GATE_ATTR = "v" + "6_lambda_head" STORE_SOURCES_ATTR = "store_v" + "6_bank_sources" NEEDS_SOURCES_METHOD = "_v" + "6_needs_dstage_bank_sources" SOURCE_INPUT_METHOD = "_v" + "6_bank_query_input" SOURCE_ALIASES = { "contextual_memory_state": "q_mem", "decoder_state": "h_d", "expanded_state": "h_e", "attractor_state": "q_attractor", "q_mem": "q_mem", "h_d": "h_d", "h_e": "h_e", "q_attractor": "q_attractor", } def load_sidecar_constructor_kwargs( sidecar_path: str, source: str | None, ) -> Dict[str, object]: 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(source or mem.get("source", "contextual_memory_state")) 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_ALIASES.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 load_model(args, device: torch.device): source = None if args.source == "sidecar" else args.source model_kwargs = load_sidecar_constructor_kwargs(args.sidecar, source) model = CEDL.build_model("CEDL", vocab=50257, max_seq=1024, **model_kwargs) model = model.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 = model.load_state_dict(msd, strict=True) print(f"[load] strict OK missing={len(res.missing_keys)} " f"unexpected={len(res.unexpected_keys)}") model.feedback_alpha.fill_(1.0) if hasattr(model, "sl_alpha"): model.sl_alpha.fill_(1.0) for p in model.parameters(): p.requires_grad_(False) return model def make_batch(tokenizer, batch_size: int, max_seq: int, seed: int, split: str, hard_collision_frac: float): items = v4c.generate( tokenizer, n=batch_size, seed=seed, split=split, hard_collision_frac=hard_collision_frac, family_weights={ "but_update": 0.50, "however_revision": 0.20, "temporal_update": 0.15, "paraphrased_equiv": 0.05, "neutral_control": 0.10, }, ) return CEDL._pad_v4c_items(items, max_seq), items @torch.no_grad() def answer_source(model, ids: torch.Tensor, items) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: h_C = model.c_stage(ids, feedback=None) h_E, h_E_sparse = model.e_stage(h_C, feedback=None) v_vec, _v_scalar = model.salience(h_E) prev = getattr(model.d_stage, STORE_SOURCES_ATTR, False) setattr(model.d_stage, STORE_SOURCES_ATTR, getattr(model, NEEDS_SOURCES_METHOD)()) try: h_D, _ = model.d_stage(h_E, h_E_sparse, v_vec=v_vec) finally: setattr(model.d_stage, STORE_SOURCES_ATTR, prev) b_idx, p_idx, cur_tok, stale_tok = CEDL._v4c_collect_answer_quads(ids, items) if b_idx.numel() == 0: empty = torch.empty(0, model.d_model, device=ids.device) return empty, cur_tok, stale_tok src = getattr(model, SOURCE_INPUT_METHOD)(h_D, h_E, b_idx, p_idx) return src.detach(), cur_tok, stale_tok def init_direct_head(model, init_from: str) -> nn.Linear: head = nn.Linear(model.d_model, 50257, bias=True).to(next(model.parameters()).device) with torch.no_grad(): if init_from == "mem_head_bank" and hasattr(model, "mem_head_bank"): head.weight.copy_(model.mem_head_bank.weight) head.bias.copy_(model.mem_head_bank.bias) elif init_from == "tok_emb": head.weight.copy_(model.c_stage.tok_emb.weight) if getattr(model.l_stage.mem_head, "bias", None) is not None: head.bias.copy_(model.l_stage.mem_head.bias) else: head.bias.zero_() elif init_from == "random": nn.init.normal_(head.weight, std=0.02) head.bias.zero_() else: raise ValueError(f"unknown --init-from {init_from!r}") return head def row_metrics(logits: torch.Tensor, cur_tok: torch.Tensor, stale_tok: torch.Tensor): logp = F.log_softmax(logits, dim=-1) top1 = logp.argmax(dim=-1) entropy = -(logp.exp() * logp).sum(dim=-1) cur_rank = (logp > logp.gather(1, cur_tok[:, None])).sum(dim=1) + 1 stale_rank = (logp > logp.gather(1, stale_tok[:, None])).sum(dim=1) + 1 margin = ( logp.gather(1, cur_tok[:, None]).squeeze(1) - logp.gather(1, stale_tok[:, None]).squeeze(1) ) return logp, top1, entropy, cur_rank, stale_rank, margin @torch.no_grad() def evaluate_replay(model, head, tokenizer, *, n_items: int, batch_size: int, max_seq: int, seed: int, split: str) -> Dict[str, float]: rows: List[Dict[str, float]] = [] produced = 0 batch_seed = seed while produced < n_items: need = min(batch_size, n_items - produced) ids, items = make_batch( tokenizer, need, max_seq, batch_seed, split, hard_collision_frac=0.0, ) ids = ids.to(next(model.parameters()).device) src, cur_tok, stale_tok = answer_source(model, ids, items) batch_seed += 1009 if src.numel() == 0: continue logits = head(src) _logp, top1, entropy, cur_rank, stale_rank, margin = row_metrics( logits, cur_tok, stale_tok) for i in range(cur_tok.numel()): rows.append({ "top1_cur": float(top1[i].item() == cur_tok[i].item()), "top1_stale": float(top1[i].item() == stale_tok[i].item()), "entropy": float(entropy[i].item()), "cur_rank": float(cur_rank[i].item()), "stale_rank": float(stale_rank[i].item()), "margin": float(margin[i].item()), }) produced += need if not rows: raise RuntimeError("evaluation generated no valid V4c answer rows") return { "n": float(len(rows)), "top1_cur": float(np.mean([r["top1_cur"] for r in rows])), "top1_stale": float(np.mean([r["top1_stale"] for r in rows])), "entropy": float(np.mean([r["entropy"] for r in rows])), "cur_rank_mean": float(np.mean([r["cur_rank"] for r in rows])), "cur_rank_median": float(np.median([r["cur_rank"] for r in rows])), "stale_rank_mean": float(np.mean([r["stale_rank"] for r in rows])), "stale_rank_median": float(np.median([r["stale_rank"] for r in rows])), "margin": float(np.mean([r["margin"] for r in rows])), } def print_metrics(label: str, metrics: Dict[str, float]): print(f"[{label}] n={int(metrics['n'])} " f"top1_cur={metrics['top1_cur'] * 100:5.1f}% " f"top1_stale={metrics['top1_stale'] * 100:5.1f}% " f"entropy={metrics['entropy']:.3f} " f"cur_rank_mean={metrics['cur_rank_mean']:.0f} " f"cur_rank_med={metrics['cur_rank_median']:.0f} " f"stale_rank_med={metrics['stale_rank_median']:.0f} " f"margin={metrics['margin']:+.3f}") def token_decode(tok, token_id): s = tok.decode([int(token_id)]) return repr(s) if len(s) < 16 else repr(s[:13]) + "..." class ConstantLambdaHead(nn.Module): def __init__(self, logit_value): super().__init__() self.logit_value = float(logit_value) def forward(self, h): return torch.full( h.shape[:-1] + (1,), self.logit_value, device=h.device, dtype=h.dtype, ) def model_forward_logits(m, ids_b): out = m(ids_b) return out[0] if isinstance(out, tuple) else out @torch.no_grad() def evaluate_b2_direct(model, head, tokenizer, *, n_items: int, seed: int): device = next(model.parameters()).device items = v4c.generate(tokenizer, n=n_items, seed=seed) records = [] saved_head = getattr(model, MEMORY_GATE_ATTR) setattr(model, MEMORY_GATE_ATTR, ConstantLambdaHead(-100.0).to(device)) try: for it_idx, it in enumerate(items): if it.family == "neutral_control": continue if not it.current or not it.stale: continue ol = getattr(it, "original_length", 0) if ol <= 1 or ol > 1024: continue cur_t = int(it.ids[it.current[0][0]]) stale_t = int(it.ids[it.stale[0][0]]) if cur_t == stale_t: continue ans_p = ol - 1 ids_b = torch.tensor([it.ids[:ol]], device=device, dtype=torch.long) trunk_logits = model_forward_logits(model, ids_b)[0, ans_p] src, cur_tok, stale_tok = answer_source(model, ids_b, [it]) if src.numel() == 0: continue direct_logits = head(src)[0] results = {"item": it_idx, "cur_t": cur_t, "stale_t": stale_t} for path, row in (("trunk", trunk_logits), ("direct", direct_logits)): logp = F.log_softmax(row, dim=-1) top5_p, top5_t = torch.topk(logp, k=5) results[path] = { "top5_t": top5_t.tolist(), "top5_lp": top5_p.tolist(), "top1_t": int(top5_t[0].item()), "ent": float(-(logp.exp() * logp).sum().item()), "cur_rank": int((logp > logp[cur_t]).sum().item()) + 1, "stl_rank": int((logp > logp[stale_t]).sum().item()) + 1, "cur_lp": float(logp[cur_t].item()), "stl_lp": float(logp[stale_t].item()), } records.append(results) finally: setattr(model, MEMORY_GATE_ATTR, saved_head) n = len(records) print(f"\n[b2-direct] generated={len(items)} used={n}") if n == 0: return {} for path in ("trunk", "direct"): top1_cur = sum(1 for r in records if r[path]["top1_t"] == r["cur_t"]) top1_stl = sum(1 for r in records if r[path]["top1_t"] == r["stale_t"]) top1_other = n - top1_cur - top1_stl ents = np.array([r[path]["ent"] for r in records]) cur_ranks = np.array([r[path]["cur_rank"] for r in records]) stl_ranks = np.array([r[path]["stl_rank"] for r in records]) cur_lps = np.array([r[path]["cur_lp"] for r in records]) stl_lps = np.array([r[path]["stl_lp"] for r in records]) print(f"[{path:>6}] top1: cur={top1_cur:>2}/{n} ({top1_cur/n*100:5.1f}%) " f"stl={top1_stl:>2}/{n} ({top1_stl/n*100:5.1f}%) " f"other={top1_other:>2}/{n} ({top1_other/n*100:5.1f}%)") print(f" entropy: mean={ents.mean():.3f} std={ents.std():.3f}") print(f" cur_rank: mean={cur_ranks.mean():.0f} median={np.median(cur_ranks):.0f} " f"min={cur_ranks.min()} max={cur_ranks.max()}") print(f" stl_rank: mean={stl_ranks.mean():.0f} median={np.median(stl_ranks):.0f} " f"min={stl_ranks.min()} max={stl_ranks.max()}") print(f" logp(cur): mean={cur_lps.mean():.3f} logp(stl): mean={stl_lps.mean():.3f}\n") direct_eq_trunk = sum( 1 for r in records if r["direct"]["top1_t"] == r["trunk"]["top1_t"]) direct_in_trunk_top5 = sum( 1 for r in records if r["direct"]["top1_t"] in r["trunk"]["top5_t"]) trunk_in_direct_top5 = sum( 1 for r in records if r["trunk"]["top1_t"] in r["direct"]["top5_t"]) print("[cross-path]") print(f" direct top1 == trunk top1: {direct_eq_trunk:>2}/{n} ({direct_eq_trunk/n*100:5.1f}%)") print(f" direct top1 in trunk top5: {direct_in_trunk_top5:>2}/{n} ({direct_in_trunk_top5/n*100:5.1f}%)") print(f" trunk top1 in direct top5: {trunk_in_direct_top5:>2}/{n} ({trunk_in_direct_top5/n*100:5.1f}%)") print("\n[first 10 - qualitative direct vs trunk top-3 at answer position]") for r in records[:10]: cur_s = token_decode(tokenizer, r["cur_t"]) stl_s = token_decode(tokenizer, r["stale_t"]) print(f"\n --- item {r['item']} cur={cur_s} stale={stl_s} ---") for path in ("trunk", "direct"): top3_t = r[path]["top5_t"][:3] top3_lp = r[path]["top5_lp"][:3] top3_decoded = [ (token_decode(tokenizer, t), f"{lp:+.2f}") for t, lp in zip(top3_t, top3_lp) ] print(f" {path:>6}: top3={top3_decoded} " f"cur_rank={r[path]['cur_rank']} stl_rank={r[path]['stl_rank']}") direct_top1_cur = sum( 1 for r in records if r["direct"]["top1_t"] == r["cur_t"]) / n trunk_top1_cur = sum( 1 for r in records if r["trunk"]["top1_t"] == r["cur_t"]) / n print(f"\n{'='*64}") print("Direct-source readout verdict") print(f"{'='*64}") if direct_top1_cur >= trunk_top1_cur: print(f" DIRECT >= TRUNK ({direct_top1_cur*100:.1f}% vs {trunk_top1_cur*100:.1f}%).") print(" -> Source is decodable; 256-slot bank readout is the bottleneck.") elif direct_top1_cur >= 0.5 * trunk_top1_cur: print(f" DIRECT PARTIAL ({direct_top1_cur*100:.1f}% vs trunk {trunk_top1_cur*100:.1f}%).") print(" -> Source has signal, but direct readout is still not robust enough.") else: print(f" DIRECT FAIL ({direct_top1_cur*100:.1f}% vs trunk {trunk_top1_cur*100:.1f}%).") print(" -> Selected source is not robustly decodable under B2.") print(f"{'='*64}") return { "n": n, "direct_top1_cur": direct_top1_cur, "trunk_top1_cur": trunk_top1_cur, } def save_outputs(head, args, before, after, b2): os.makedirs(args.out_dir, exist_ok=True) stem = f"CEDL_direct_source_{args.source}" head_path = os.path.join(args.out_dir, f"{stem}_head.pt") metrics_path = os.path.join(args.out_dir, f"{stem}_metrics.json") torch.save({ "direct_head": head.state_dict(), "checkpoint": args.checkpoint, "sidecar": args.sidecar, "source": args.source, "init_from": args.init_from, "steps": args.steps, "lr": args.lr, "batch_size": args.batch_size, "seed": args.seed, }, head_path) with open(metrics_path, "w") as f: json.dump({"before": before, "after": after, "b2": b2}, f, indent=2) print(f"[save] direct_head={head_path}") print(f"[save] metrics={metrics_path}") def main(): p = argparse.ArgumentParser() p.add_argument("--checkpoint", required=True) p.add_argument("--sidecar", required=True) p.add_argument("--out-dir", default="outputs/memory_source_readout") p.add_argument("--source", choices=("sidecar",) + SOURCES, default="sidecar") p.add_argument("--init-from", choices=("mem_head_bank", "tok_emb", "random"), default="mem_head_bank") p.add_argument("--steps", type=int, default=2000) p.add_argument("--batch-size", type=int, default=64) p.add_argument("--max-seq", type=int, default=128) p.add_argument("--lr", type=float, default=5e-4) p.add_argument("--weight-decay", type=float, default=0.0) p.add_argument("--seed", type=int, default=7234) p.add_argument("--split", choices=("all", "train", "test"), default="all") p.add_argument("--hard-collision-frac", type=float, default=0.2) p.add_argument("--eval-items", type=int, default=256) p.add_argument("--b2-items", type=int, default=50) p.add_argument("--b2-seed", type=int, default=0) p.add_argument("--log-every", type=int, default=100) args = p.parse_args() torch.manual_seed(args.seed) 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] out_dir={args.out_dir}") print(f"[setup] source={args.source}") print(f"[setup] init_from={args.init_from}") from transformers import GPT2TokenizerFast tokenizer = GPT2TokenizerFast.from_pretrained("gpt2") model = load_model(args, device) print(f"[setup] resolved_source={model_kwargs['memory_query_source']}") head = init_direct_head(model, args.init_from) trainable = sum(p.numel() for p in head.parameters()) print(f"[freeze] trainable params={trainable:,} (direct_head only)") before = evaluate_replay( model, head, tokenizer, n_items=args.eval_items, batch_size=args.batch_size, max_seq=args.max_seq, seed=args.seed + 17, split=args.split, ) print_metrics("before", before) opt = torch.optim.AdamW( head.parameters(), lr=args.lr, weight_decay=args.weight_decay, betas=(0.9, 0.95), ) seen_rows = 0 ema_loss = None for step in range(1, args.steps + 1): ids, items = make_batch( tokenizer, args.batch_size, args.max_seq, args.seed + step * 7919, args.split, args.hard_collision_frac, ) ids = ids.to(device) src, cur_tok, _stale_tok = answer_source(model, ids, items) if src.numel() == 0: continue logits = head(src) loss = F.cross_entropy(logits, cur_tok) opt.zero_grad(set_to_none=True) loss.backward() torch.nn.utils.clip_grad_norm_(head.parameters(), max_norm=1.0) opt.step() seen_rows += int(cur_tok.numel()) loss_val = float(loss.detach().item()) ema_loss = loss_val if ema_loss is None else 0.98 * ema_loss + 0.02 * loss_val if step == 1 or step % args.log_every == 0: with torch.no_grad(): logp = F.log_softmax(logits, dim=-1) top1_cur = (logp.argmax(dim=-1) == cur_tok).float().mean().item() entropy = (-(logp.exp() * logp).sum(dim=-1)).mean().item() print(f"[train] step={step}/{args.steps} rows={seen_rows} " f"loss={loss_val:.4f} ema={ema_loss:.4f} " f"batch_top1_cur={top1_cur * 100:5.1f}% " f"entropy={entropy:.3f}", flush=True) after = evaluate_replay( model, head, tokenizer, n_items=args.eval_items, batch_size=args.batch_size, max_seq=args.max_seq, seed=args.seed + 17, split=args.split, ) print_metrics("after", after) b2 = evaluate_b2_direct( model, head, tokenizer, n_items=args.b2_items, seed=args.b2_seed) save_outputs(head, args, before, after, b2) if __name__ == "__main__": main()