J-lens workspace reproduction + loop retrofit: lens, band looping, adapters, controls, multi-task evals
Reproduction of the 2026 workspace/J-lens paper on gemma-4 (E2B/12B/26B), plus the workspace-loop retrofit line: merge adapter, prompt-only latent planning (MBPP), carry variant, attribution controls (FF/pause/untrained), band-location ablation, Blocksworld harness, 12B replication scripts. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Nils
Claude Fable 5
@@ -0,0 +1,130 @@
|
||||
"""Experiment 1: J-lens readouts.
|
||||
|
||||
(a) Two-hop reasoning: unspoken intermediate 'spider' visible mid-network.
|
||||
(b) Multilingual: English intermediates during a Chinese task.
|
||||
(c) Directed modulation: 'hold citrus in mind' while copying unrelated text.
|
||||
(d) Layer profile: where in depth the lens carries abstract content.
|
||||
"""
|
||||
|
||||
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,
|
||||
generate_with_residuals, load_model)
|
||||
|
||||
RES = Path(os.environ.get("JLENS_RESULTS", "results"))
|
||||
|
||||
|
||||
def show_table(jl, tok, ids, layers, positions, topk=6, header=""):
|
||||
"""Print top J-lens tokens at (layer, position) grid."""
|
||||
hs = collect_residuals(jl.model, ids)
|
||||
toks = [tok.decode([t]) for t in ids[0].tolist()]
|
||||
print(f"\n=== {header} ===")
|
||||
print("positions:", {p: repr(toks[p]) for p in positions})
|
||||
for l in layers:
|
||||
row = []
|
||||
for p in positions:
|
||||
idx, prob = jl.read(hs[l, p], l, topk=topk)
|
||||
row.append(" ".join(
|
||||
(tok.decode([i]).strip() or "·") for i in idx.tolist()[:topk]))
|
||||
print(f"L{l:>2} | " + " || ".join(row))
|
||||
|
||||
|
||||
def concept_heatmap(jl, tok, ids, words, tag, hs=None):
|
||||
"""P(concept tokens) by (layer, position); save tensor + print peak."""
|
||||
if hs is None:
|
||||
hs = collect_residuals(jl.model, ids)
|
||||
tids = []
|
||||
for w in words:
|
||||
for v in (w, " " + w, w.capitalize(), " " + w.capitalize()):
|
||||
e = tok.encode(v, add_special_tokens=False)
|
||||
if len(e) == 1:
|
||||
tids.append(e[0])
|
||||
tids = sorted(set(tids))
|
||||
L, T, _ = hs.shape
|
||||
out = torch.zeros(L, T)
|
||||
for l in range(L):
|
||||
logits = jl._readout(hs[l].float() @ jl.Jbar[l].T)
|
||||
probs = torch.softmax(logits.float(), dim=-1)
|
||||
out[l] = probs[:, tids].sum(-1).cpu()
|
||||
peak = out.max().item()
|
||||
lmax, pmax = divmod(out.argmax().item(), T)
|
||||
toks = [tok.decode([t]) for t in ids[0].tolist()] if ids is not None else None
|
||||
print(f"[{tag}] words={words} peak P={peak:.3f} at layer {lmax}, "
|
||||
f"pos {pmax}" + (f" ({toks[pmax]!r})" if toks and pmax < len(toks) else ""))
|
||||
RES.mkdir(exist_ok=True)
|
||||
torch.save({"map": out, "words": words,
|
||||
"tokens": [tok.decode([t]) for t in ids[0].tolist()] if ids is not None else None},
|
||||
RES / f"heat_{tag}.pt")
|
||||
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())
|
||||
print(f"Jbar from {ck.get('n_prompts')} prompts")
|
||||
L = len(model.model.language_model.layers)
|
||||
mid = range(max(2, L // 6), L - 4, max(2, L // 16))
|
||||
|
||||
# (a) two-hop: spider
|
||||
ids = chat_ids(tok, "The animal that spins webs has how many legs? "
|
||||
"Answer with just a number.")
|
||||
print("answer:", jl.generate(ids, 5))
|
||||
show_table(jl, tok, ids, mid, list(range(ids.shape[1] - 10, ids.shape[1])),
|
||||
header="two-hop spider (last 10 positions)")
|
||||
concept_heatmap(jl, tok, ids, ["spider", "spiders"], "spider")
|
||||
concept_heatmap(jl, tok, ids, ["eight", "8"], "eight")
|
||||
|
||||
# (b) multilingual: Chinese antonym of small -> big
|
||||
ids = chat_ids(tok, "小的反义词是什么？只用一个字回答。")
|
||||
print("\nanswer:", jl.generate(ids, 5))
|
||||
show_table(jl, tok, ids, mid, list(range(ids.shape[1] - 8, ids.shape[1])),
|
||||
header="Chinese antonym of 小")
|
||||
concept_heatmap(jl, tok, ids, ["big", "large", "bigger"], "big_en")
|
||||
|
||||
# (c) directed modulation: think of citrus while copying text
|
||||
copy_text = "The committee will meet on Thursday to review the budget."
|
||||
for cond, instr in [
|
||||
("citrus", "While you copy the text, silently think about citrus "
|
||||
"fruits the entire time. Copy this text exactly, output "
|
||||
f"nothing else: \"{copy_text}\""),
|
||||
("control", f"Copy this text exactly, output nothing else: \"{copy_text}\""),
|
||||
]:
|
||||
ids = chat_ids(tok, instr)
|
||||
text, all_ids, hs = generate_with_residuals(model, tok, ids, 30)
|
||||
print(f"\n[{cond}] output: {text!r}")
|
||||
# only look at generated positions
|
||||
gen0 = ids.shape[1]
|
||||
heat = concept_heatmap(jl, tok, all_ids.unsqueeze(0),
|
||||
["citrus", "lemon", "orange", "lime"],
|
||||
f"citrus_{cond}", hs=hs)
|
||||
print(f" citrus P over generated positions: mean "
|
||||
f"{heat[:, gen0:].mean():.4f} max {heat[:, gen0:].max():.4f}")
|
||||
|
||||
# (d) layer profile: lens entropy + top-token agreement with input/output
|
||||
ids = chat_ids(tok, "Write one sentence about the ocean.")
|
||||
text, all_ids, hs = generate_with_residuals(model, tok, ids, 20)
|
||||
L, T, _ = hs.shape
|
||||
ent, agree_in, agree_next = [], [], []
|
||||
for l in range(L):
|
||||
logits = jl._readout(hs[l].float() @ jl.Jbar[l].T)
|
||||
probs = torch.softmax(logits.float(), dim=-1)
|
||||
e = -(probs * probs.clamp_min(1e-12).log()).sum(-1).mean()
|
||||
top = probs.argmax(-1)
|
||||
ent.append(e.item())
|
||||
agree_in.append((top[:-1] == all_ids[:len(top) - 1].cuda()).float().mean().item())
|
||||
agree_next.append((top[:-1] == all_ids[1:len(top)].cuda()).float().mean().item())
|
||||
print("\nlayer | lens entropy | top==current tok | top==next tok")
|
||||
for l in range(0, L, 2):
|
||||
print(f"L{l:>2} | {ent[l]:8.2f} | {agree_in[l]:.2f} | {agree_next[l]:.2f}")
|
||||
torch.save({"entropy": ent, "agree_in": agree_in, "agree_next": agree_next},
|
||||
RES / "layer_profile.pt")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user