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>
2026-07-14 00:54:12 +02:00
commit ef9c08966c
42 changed files with 5068 additions and 0 deletions
@@ -0,0 +1,75 @@
+"""J-lens sharpening battery on MBPP prompts: does the trained loop
+concentrate the workspace readout while reading the problem?
+
+For N test prompts, at each k: J-lens distribution at L30, last prompt
+position -> entropy + top-1 probability. Trained vs untrained adapter.
+"""
+
+import json
+import sys
+from pathlib import Path
+
+import torch
+
+from loop_common import BandLooper, MergeAdapter
+from prep_mbpp import DIRECT_SUFFIX, mbpp_prompt
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+from jlens.core import JLens, load_model  # noqa: E402
+
+OUT = Path(__file__).resolve().parent.parent / "results-loop"
+ROOT = OUT.parent
+N_ITEMS = 20
+KS = (0, 1, 2, 4)
+
+
+@torch.no_grad()
+def battery(looper, adapter, lens, tok, items):
+    rows = []
+    for it in items:
+        ids = tok(mbpp_prompt(tok, it, DIRECT_SUFFIX), return_tensors="pt",
+                  add_special_tokens=False)["input_ids"].cuda()
+        calls, _ = looper.capture(ids, logits_to_keep=1)
+        e = looper._hin[looper.l0]
+        s = looper.band(e, calls)
+        state = {0: s}
+        for k in range(1, max(KS) + 1):
+            s = looper.band(adapter(e, s), calls)
+            state[k] = s
+        for k in KS:
+            probs = torch.softmax(lens._readout(
+                state[k][0, -1].float() @ lens.Jbar[looper.l1].T).float(), -1)
+            ent = -(probs * (probs + 1e-12).log()).sum().item()
+            rows.append({"task_id": it["task_id"], "k": k, "entropy": ent,
+                         "top1": probs.max().item()})
+    return rows
+
+
+def main():
+    model, tok = load_model(dtype=torch.bfloat16)
+    looper = BandLooper(model)
+    jbar = torch.load(ROOT / "results" / "jbar.pt", map_location="cuda")
+    lens = JLens(model, tok, jbar["Jbar"].float())
+    items = [it for it in json.load(open(OUT / "mbpp_data.json"))
+             if it["split"] == "test"][:N_ITEMS]
+
+    res = {}
+    for tag, path in (("untrained", None),
+                      ("trained", OUT / "adapter_code_e399.pt")):
+        adapter = MergeAdapter().cuda()
+        if path:
+            adapter.load_state_dict(torch.load(path, map_location="cuda"))
+        rows = battery(looper, adapter, lens, tok, items)
+        res[tag] = rows
+        for k in KS:
+            sel = [r for r in rows if r["k"] == k]
+            ent = sum(r["entropy"] for r in sel) / len(sel)
+            top = sum(r["top1"] for r in sel) / len(sel)
+            print(f"[{tag}] k={k}: mean_entropy={ent:.3f}  mean_top1={top:.3f}",
+                  flush=True)
+    json.dump(res, open(OUT / "lens_battery.json", "w"), indent=1)
+    print("wrote", OUT / "lens_battery.json")
+
+
+if __name__ == "__main__":
+    main()