| |
| """Characterize the +-1 off-diagonal entries in the latent confusion matrices. |
| |
| Question: when latent slot m decodes a node at TRUE depth m-1 or m+1, is that |
| (a) a PHASE DRIFT -- the BFS chain slipped a step and stays shifted, or |
| (b) an ISOLATED SLIP -- slot m errs but slot m+1 is back on the diagonal, or |
| (c) a STALL -- slot m decodes the *same node* slot m-1 already decoded |
| (the chain failed to advance a hop)? |
| |
| For each arm we report: |
| * histogram of offset = (true depth of decoded node) - m |
| * repeat rate: decoded(m) == decoded(m-1) [stall signature] |
| * drift persistence: P(slot m+1 also off by the same offset | slot m off by d) |
| vs the base rate. High persistence => real drift; low => isolated noise. |
| """ |
| import argparse |
| import json |
| from collections import Counter |
|
|
| import torch |
| from transformers import AutoModelForCausalLM, AutoConfig |
|
|
| from stokenizer import STokenizer |
| from coconut import Coconut |
| from scripts.probe_latents import build_prefix_tokens, node_depth_maps |
|
|
| ARMS = [ |
| ("Backtracking", "backtrack"), |
| ("Current-stage-only", "curstage"), |
| ("Retention-gated (no repair)", "accstage-nobt"), |
| ] |
|
|
|
|
| def latest_ckpt(slug): |
| import os |
| d = f"ckpts/star-coconut-L10-bfs-{slug}" |
| cks = sorted((f for f in os.listdir(d) if f.startswith("checkpoint_")), |
| key=lambda x: int(x.split("_")[1])) |
| return os.path.join(d, cks[-1]) |
|
|
|
|
| @torch.no_grad() |
| def decode_chains(ckpt, val_path, model_id, L, device, batch_size=64): |
| """Return list of per-sample dicts: decoded node id + its true depth per slot.""" |
| tok = STokenizer() |
| latent_id = tok.convert_tokens_to_ids("<|latent|>") |
| base = AutoModelForCausalLM.from_config(AutoConfig.from_pretrained(model_id)) |
| model = Coconut(base, latent_id, |
| tok.convert_tokens_to_ids("<|start-latent|>"), |
| tok.convert_tokens_to_ids("<|end-latent|>"), |
| tok.eos_token_id) |
| model.load_state_dict(torch.load(ckpt, map_location="cpu"), strict=False) |
| model.to(device).eval() |
|
|
| data = json.load(open(val_path)) |
| chains = [] |
| for i in range(0, len(data), batch_size): |
| batch = data[i:i + batch_size] |
| input_ids = torch.tensor([build_prefix_tokens(s, tok) + [latent_id] * L |
| for s in batch], device=device) |
| attn = torch.ones_like(input_ids) |
| pos = torch.arange(input_ids.shape[1], device=device).unsqueeze(0).expand(len(batch), -1) |
| logits = model.forward(input_ids, attn, input_ids.clone(), pos).logits |
| for bi, s in enumerate(batch): |
| role = node_depth_maps(s, L) |
| root_pos = len(build_prefix_tokens(s, tok)) - 1 |
| nodes, depths = [], [] |
| for m in range(1, L + 1): |
| pred = int(torch.argmax(logits[bi, root_pos + (m - 1)]).item()) |
| r = role.get(pred) |
| nodes.append(pred) |
| |
| depths.append(r[1] if (r is not None and r[0] == "pos") else None) |
| chains.append({"nodes": nodes, "depths": depths}) |
| return chains |
|
|
|
|
| def analyze(chains, L): |
| offsets = Counter() |
| repeat = 0 |
| slot_total = 0 |
| |
| off_then_off_same = 0 |
| off_then_diag = 0 |
| off_events = 0 |
| for c in chains: |
| for m in range(1, L + 1): |
| d = c["depths"][m - 1] |
| slot_total += 1 |
| if d is None: |
| offsets["decoy/off"] += 1 |
| continue |
| offsets[d - m] += 1 |
| if m >= 2 and c["nodes"][m - 1] == c["nodes"][m - 2]: |
| repeat += 1 |
| |
| for m in range(1, L): |
| d, d2 = c["depths"][m - 1], c["depths"][m] |
| if d is None or d2 is None: |
| continue |
| o = d - m |
| if o == 0: |
| continue |
| off_events += 1 |
| if d2 - (m + 1) == o: |
| off_then_off_same += 1 |
| elif d2 - (m + 1) == 0: |
| off_then_diag += 1 |
| return offsets, repeat, slot_total, off_events, off_then_off_same, off_then_diag |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--val", default="data/star_2arm_L10_valid_fo_bfs.json") |
| ap.add_argument("--model_id", default="configs/symbol-2layer-8head-768dim-L20.json") |
| ap.add_argument("--L", type=int, default=10) |
| ap.add_argument("--device", default="cuda:0") |
| args = ap.parse_args() |
|
|
| for label, slug in ARMS: |
| ck = latest_ckpt(slug) |
| chains = decode_chains(ck, args.val, args.model_id, args.L, args.device) |
| offs, repeat, tot, oe, same, diag = analyze(chains, args.L) |
| print(f"\n===== {label} ({ck}) =====") |
| print(f"total latent slots decoded: {tot}") |
| print("offset (true_depth - slot) histogram:") |
| for k in sorted([k for k in offs if isinstance(k, int)]): |
| print(f" {k:+d}: {offs[k]:>6} ({offs[k]/tot:.4%})") |
| print(f" decoy/off-graph: {offs['decoy/off']:>6} ({offs['decoy/off']/tot:.4%})") |
| print(f"stall (decoded node == previous slot's node): {repeat} " |
| f"({repeat/tot:.4%})") |
| if oe: |
| print(f"drift persistence: of {oe} off-diagonal in-component events, " |
| f"{same} ({same/oe:.1%}) kept the SAME offset at slot m+1, " |
| f"{diag} ({diag/oe:.1%}) snapped back to the diagonal") |
| else: |
| print("drift persistence: no off-diagonal in-component events") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
