156 lines
6.4 KiB
Python
156 lines
6.4 KiB
Python
"""Experiment 4: locate sensor / workspace / motor regimes by depth.
|
|
|
|
Per-layer diagnostics averaged over a prompt set:
|
|
sensor = frac. of positions where top J-lens token == current input token
|
|
motor = frac. where top J-lens token == NEXT input token (teacher-forced)
|
|
persist = mean Jaccard overlap of top-10 lens tokens at adjacent positions
|
|
(workspace content should persist across positions)
|
|
content = frac. of positions whose top lens token is a content word
|
|
(alphabetic, len>=3, not in a junk list)
|
|
|
|
Ignition test (paper: ambiguous inputs produce sharp binary commitment at
|
|
workspace onset): replace one token's embedding (main + per-layer) with a
|
|
w-mixture of two concept embeddings and track lens commitment
|
|
C = (P1-P2)/(P1+P2) by layer at a downstream position.
|
|
"""
|
|
|
|
import os, sys
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
|
from jlens.core import JLens, chat_ids, collect_residuals, load_model
|
|
|
|
RES = Path(os.environ.get("JLENS_RESULTS", "results"))
|
|
|
|
PROMPTS = [
|
|
"The animal that spins webs has how many legs? Answer with just a number.",
|
|
"What is the capital of France? Answer with just the city name.",
|
|
"Write one sentence about the ocean.",
|
|
"Explain photosynthesis in one sentence.",
|
|
"Name a famous river in Egypt.",
|
|
"What language is spoken in Brazil? Answer with one word.",
|
|
"Summarize the plot of Romeo and Juliet in one sentence.",
|
|
"If I have 3 apples and eat one, how many are left?",
|
|
]
|
|
|
|
JUNK = set("·.<>|/\\()[]{}:;,!?\"'`~*#-_=+ \n\t")
|
|
|
|
|
|
def regime_profile(jl, tok, model):
|
|
L = len(jl.tm.layers)
|
|
sensor = torch.zeros(L)
|
|
motor = torch.zeros(L)
|
|
persist = torch.zeros(L)
|
|
content = torch.zeros(L)
|
|
n_pos = 0
|
|
n_adj = 0
|
|
for p in PROMPTS:
|
|
ids = chat_ids(tok, p)
|
|
hs = collect_residuals(model, ids)
|
|
T = ids.shape[1]
|
|
sl = slice(4, T - 5) # skip bos/turn tokens and trailing template
|
|
cur = ids[0, sl].cuda()
|
|
nxt = ids[0, 4 + 1:T - 4].cuda()
|
|
for l in range(L):
|
|
idx, _ = jl.read(hs[l, sl], l, topk=10) # (T', 10)
|
|
top1 = idx[:, 0]
|
|
sensor[l] += (top1 == cur).sum().item()
|
|
motor[l] += (top1 == nxt).sum().item()
|
|
sets = [set(r.tolist()) for r in idx]
|
|
for a, b in zip(sets, sets[1:]):
|
|
persist[l] += len(a & b) / len(a | b)
|
|
for t in top1.tolist():
|
|
s = tok.decode([t]).strip()
|
|
content[l] += (len(s) >= 3 and s.isalpha())
|
|
n_pos += cur.numel()
|
|
n_adj += cur.numel() - 1
|
|
return sensor / n_pos, motor / n_pos, persist / n_adj, content / n_pos
|
|
|
|
|
|
def ignition(jl, tok, model, pairs, ws=(0.0, 0.25, 0.5, 0.75, 1.0)):
|
|
"""Mix two concept embeddings at one position; commitment by layer."""
|
|
tm = jl.tm
|
|
L = len(tm.layers)
|
|
template = ("My favorite thing in the world is the X ."
|
|
" I think about it every single day because")
|
|
out = {}
|
|
for w1, w2 in pairs:
|
|
id1 = tok.encode(" " + w1, add_special_tokens=False)[0]
|
|
id2 = tok.encode(" " + w2, add_special_tokens=False)[0]
|
|
ids = tok(template, return_tensors="pt")["input_ids"].cuda()
|
|
pos = (ids[0] == tok.encode(" X", add_special_tokens=False)[0]) \
|
|
.nonzero()[0].item()
|
|
has_ple = getattr(tm, "embed_tokens_per_layer", None) is not None
|
|
with torch.no_grad():
|
|
pair_ids = torch.tensor([[id1, id2]], device="cuda")
|
|
rows_main = tm.embed_tokens(pair_ids)[0].detach()
|
|
rows_ple = tm.embed_tokens_per_layer(pair_ids)[0].detach() if has_ple else None
|
|
curves = torch.zeros(len(ws), L)
|
|
for wi, w in enumerate(ws):
|
|
def make_hook(rows, w=w, pos=pos):
|
|
def mix_hook(mod, inp, out_e):
|
|
e = out_e.clone()
|
|
e[0, pos] = w * rows[0] + (1 - w) * rows[1]
|
|
return e
|
|
return mix_hook
|
|
h1 = tm.embed_tokens.register_forward_hook(make_hook(rows_main))
|
|
h2 = tm.embed_tokens_per_layer.register_forward_hook(make_hook(rows_ple)) if has_ple else None
|
|
try:
|
|
hs = collect_residuals(model, ids)
|
|
finally:
|
|
h1.remove()
|
|
if h2: h2.remove()
|
|
rpos = pos + 3 # downstream read position
|
|
for l in range(L):
|
|
logits = jl._readout(hs[l, rpos].float() @ jl.Jbar[l].T)
|
|
probs = torch.softmax(logits.float(), -1)
|
|
p1, p2 = probs[id1].item(), probs[id2].item()
|
|
curves[wi, l] = (p1 - p2) / (p1 + p2 + 1e-12)
|
|
out[(w1, w2)] = curves
|
|
return out
|
|
|
|
|
|
def main():
|
|
model, tok = load_model(dtype=torch.bfloat16)
|
|
ck = torch.load(sys.argv[1] if len(sys.argv) > 1 else "results/jbar.pt",
|
|
map_location="cuda")
|
|
jl = JLens(model, tok, ck["Jbar"].cuda())
|
|
|
|
sensor, motor, persist, content = regime_profile(jl, tok, model)
|
|
print("layer | sensor(top==cur) motor(top==next) persist(top10 Jaccard) content-word")
|
|
for l in range(len(sensor)):
|
|
bars = lambda x: "#" * int(20 * x)
|
|
print(f"L{l:>2} | {sensor[l]:.2f} {bars(sensor[l]):<20} | "
|
|
f"{motor[l]:.2f} {bars(motor[l]):<20} | "
|
|
f"{persist[l]:.2f} | {content[l]:.2f}")
|
|
|
|
pairs = [("dog", "piano"), ("ocean", "violin"), ("dragon", "bicycle")]
|
|
ign = ignition(jl, tok, model, pairs)
|
|
print("\nignition: commitment C=(P1-P2)/(P1+P2) at read pos, by layer")
|
|
print("pair | w: " + " ".join(f"{w:4.2f}" for w in (0.0, .25, .5, .75, 1.0)))
|
|
for (w1, w2), curves in ign.items():
|
|
for l in range(2, len(sensor), 4):
|
|
print(f"{w1}/{w2:<9} L{l:>2} | " +
|
|
" ".join(f"{curves[wi, l]:+.2f}" for wi in range(5)))
|
|
if len(sys.argv) > 2:
|
|
out = Path(sys.argv[2])
|
|
elif len(sys.argv) > 1:
|
|
# default: save next to the input jbar (that dir is what node
|
|
# sidecars sync; CWD-relative "results/" has crashed 3 scans)
|
|
ckp = Path(sys.argv[1])
|
|
out = ckp.parent / ckp.name.replace("jbar", "regimes")
|
|
if out == ckp:
|
|
out = RES / "regimes.pt"
|
|
else:
|
|
out = RES / "regimes.pt"
|
|
out.parent.mkdir(parents=True, exist_ok=True)
|
|
torch.save({"sensor": sensor, "motor": motor, "persist": persist,
|
|
"content": content, "ignition": ign}, out)
|
|
print("saved", out)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|