NilsandClaude Fable 5
ef9c08966c
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>
104 lines
4.7 KiB
Python
104 lines
4.7 KiB
Python
"""Experiment 2: causal swap interventions gated by the J-lens.
|
|
|
|
At each (layer, position) where the J-lens reads the source concept, transfer
|
|
the activation content source -> target (embedding write basis; see core.py).
|
|
|
|
(a) Two-hop: swap spider<->ant in the workspace -> answer flips 8 -> 6.
|
|
(b) Broadcast: swap France->China under different templates.
|
|
(c) Capital-question grid over country pairs.
|
|
"""
|
|
|
|
import itertools, 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, load_model
|
|
|
|
RES = Path(os.environ.get("JLENS_RESULTS", "results"))
|
|
|
|
|
|
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")
|
|
|
|
# (a) spider -> ant
|
|
SP = [(" spider", " ant"), (" spiders", " ants"), ("spider", "ant")]
|
|
q = "The animal that spins webs has how many legs? Answer with just a number."
|
|
ids = chat_ids(tok, q)
|
|
print("\n--- two-hop swap: spider -> ant (baseline:",
|
|
repr(jl.generate(ids, 6)), ") ---")
|
|
for thr in (0.002, 0.005, 0.01):
|
|
for alpha in (1.0, 1.5, 2.0):
|
|
out, fired = jl.generate_swapped(ids, SP, thr=thr, alpha=alpha,
|
|
max_new_tokens=6)
|
|
print(f"thr={thr} a={alpha}: {out!r} (fired {len(fired)} slots)")
|
|
# paper write-basis for comparison
|
|
out, _ = jl.generate_swapped(ids, SP, thr=0.005, alpha=1.0, write="jlens",
|
|
max_new_tokens=6)
|
|
print(f"paper jlens-write thr=0.005 a=1.0: {out!r}")
|
|
# control: unrelated swap
|
|
out, _ = jl.generate_swapped(ids, [(" spider", " piano")], thr=0.005,
|
|
alpha=1.0, max_new_tokens=6)
|
|
print(f"control spider->piano: {out!r}")
|
|
|
|
# reverse: ant prompt -> spider
|
|
q2 = ("The insect that builds colonies and lifts many times its own "
|
|
"weight has how many legs? Answer with just a number.")
|
|
ids2 = chat_ids(tok, q2)
|
|
print("\nant prompt baseline:", repr(jl.generate(ids2, 6)))
|
|
for alpha in (1.0, 1.5, 2.0):
|
|
out, fired = jl.generate_swapped(
|
|
ids2, [(" ant", " spider"), (" ants", " spiders"), ("ant", "spider")],
|
|
thr=0.005, alpha=alpha, max_new_tokens=6)
|
|
print(f"swap ant->spider a={alpha}: {out!r} (fired {len(fired)})")
|
|
|
|
# (b) broadcast France -> China across templates
|
|
print("\n--- broadcast: France -> China (thr=0.01, a=1.0) ---")
|
|
FR = [(" France", " China"), ("France", "China"), (" French", " Chinese")]
|
|
templates = [
|
|
("capital", "What is the capital of France? Answer with just the city name."),
|
|
("language", "What language is spoken in France? Answer with one word."),
|
|
("continent", "Which continent is France on? Answer with one word."),
|
|
("currency", "What currency is used in France? Answer with one word."),
|
|
("river", "Name a famous river in France. Answer with one word."),
|
|
("food", "Name a famous dish from France. Answer with a short phrase."),
|
|
]
|
|
for tag, qq in templates:
|
|
ids = chat_ids(tok, qq)
|
|
base = jl.generate(ids, 8)
|
|
sw, fired = jl.generate_swapped(ids, FR, thr=0.01, alpha=1.0,
|
|
max_new_tokens=8)
|
|
print(f"[{tag:9}] base={base!r:30} swapped={sw!r} ({len(fired)} slots)")
|
|
|
|
# (c) capital grid over country pairs
|
|
print("\n--- capital-question swap grid (thr=0.01, a=1.0) ---")
|
|
countries = {"France": ("Paris", "French"), "China": ("Beijing", "Chinese"),
|
|
"Japan": ("Tokyo", "Japanese"), "Egypt": ("Cairo", "Egyptian"),
|
|
"Brazil": ("Brasília", "Brazilian"), "Canada": ("Ottawa", "Canadian")}
|
|
hits = tries = 0
|
|
rows = []
|
|
for (src, (scap, sadj)), (tgt, (tcap, tadj)) in itertools.permutations(
|
|
countries.items(), 2):
|
|
qq = f"What is the capital of {src}? Answer with just the city name."
|
|
ids = chat_ids(tok, qq)
|
|
pairs = [(" " + src, " " + tgt), (src, tgt), (" " + sadj, " " + tadj)]
|
|
sw, fired = jl.generate_swapped(ids, pairs, thr=0.01, alpha=1.0,
|
|
max_new_tokens=8)
|
|
ok = tcap.lower().replace("í", "i").split()[0][:5] in \
|
|
sw.lower().replace("í", "i")
|
|
hits += ok
|
|
tries += 1
|
|
rows.append((src, tgt, tcap, sw, ok))
|
|
print(f"{src:>7}->{tgt:<7} expect {tcap:<9} got {sw!r} {'OK' if ok else ''}")
|
|
print(f"\nswap grid success: {hits}/{tries}")
|
|
torch.save(rows, RES / "swap_grid.pt")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|