File size: 5,835 Bytes
8f46582 | 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 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | #!/usr/bin/env python3
"""Does a latent encode ONE node at depth m, or BOTH frontier nodes at once?
At stage k the training target is random.choice(neighbor_k[k]); in a 2-arm star
that set has 2 members (one per arm), so the supervision is ambiguous. Two
hypotheses for what the model learns:
(H1) COMMIT: the latent picks one arm and puts nearly all mass on it.
-> p(top1) >> p(top2), and top2 is often not the sibling frontier node.
(H2) SUPERPOSITION: the latent represents BOTH depth-m nodes simultaneously;
argmax then breaks the tie ~arbitrarily.
-> top-2 tokens are exactly the two frontier nodes, with comparable mass.
We also ask whether the model has any PREFERENCE for the arm that leads to the
true target leaf (i.e. does it secretly know the answer early?):
-> compare p(target-arm node) vs p(other-arm node) at each hop.
Reported per hop m:
frontier_mass mean total prob on the 2 true depth-m nodes
top2_is_frontier fraction where the top-2 tokens are exactly those 2 nodes
p1/p2 ratio mean ratio of larger to smaller of the two frontier probs
target_arm_win fraction where the target-arm node outranks the other arm
p_target_share mean p(target arm) / (p(target)+p(other)) [0.5 = no preference]
"""
import argparse
import json
import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoConfig
from stokenizer import STokenizer
from coconut import Coconut
from scripts.probe_latents import build_prefix_tokens
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 run(ckpt, val_path, model_id, L, device, batch_size=64):
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))
acc = {m: {"mass": [], "top2": 0, "ratio": [], "twin": 0, "share": [], "n": 0}
for m in range(1, L + 1)}
n_frontier_sizes = {}
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
probs = F.softmax(logits.float(), dim=-1)
for bi, s in enumerate(batch):
root_pos = len(build_prefix_tokens(s, tok)) - 1
steps = s["steps"]
for m in range(1, L + 1):
front = s["neighbor_k"].get(str(m), [])
n_frontier_sizes[len(front)] = n_frontier_sizes.get(len(front), 0) + 1
if len(front) != 2:
continue
p = probs[bi, root_pos + (m - 1)]
pf = [p[int(n)].item() for n in front]
a = acc[m]
a["n"] += 1
a["mass"].append(sum(pf))
top2 = set(torch.topk(p, 2).indices.tolist())
if top2 == {int(front[0]), int(front[1])}:
a["top2"] += 1
hi, lo = max(pf), min(pf)
a["ratio"].append(hi / lo if lo > 0 else float("inf"))
# which of the two is on the true shortest path to the target leaf?
if m - 1 < len(steps):
tgt = int(steps[m - 1])
other = [int(n) for n in front if int(n) != tgt]
if other:
pt, po = p[tgt].item(), p[other[0]].item()
if pt > po:
a["twin"] += 1
if pt + po > 0:
a["share"].append(pt / (pt + po))
return acc, n_frontier_sizes
def mean(x):
return sum(x) / len(x) if x else float("nan")
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")
ap.add_argument("--only", default=None, help="restrict to one arm slug")
args = ap.parse_args()
for label, slug in ARMS:
if args.only and slug != args.only:
continue
ck = latest_ckpt(slug)
acc, sizes = run(ck, args.val, args.model_id, args.L, args.device)
print(f"\n===== {label} ({ck}) =====")
print(f"frontier-set sizes seen: {sizes}")
print(f"{'hop':>4} {'frontier_mass':>14} {'top2_is_frontier':>17} "
f"{'p_hi/p_lo':>10} {'target_arm_win':>15} {'p_target_share':>15}")
for m in range(1, args.L + 1):
a = acc[m]
if not a["n"]:
continue
print(f"{m:>4} {mean(a['mass']):>14.3f} {a['top2']/a['n']:>17.3f} "
f"{mean(a['ratio']):>10.2f} {a['twin']/a['n']:>15.3f} "
f"{mean(a['share']):>15.3f}")
if __name__ == "__main__":
main()
|