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,67 @@
+"""Estimate J_l = E_{prompt, t, t'>=t}[dh_final,t'/dh_l,t] over a pretraining-like corpus."""
+
+import argparse, json, sys, time
+from pathlib import Path
+
+import torch
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+from jlens.core import load_model, prompt_jacobian_pairsum
+
+
+def corpus_texts(n, min_chars=400):
+    from datasets import load_dataset
+    ds = load_dataset("HuggingFaceFW/fineweb-edu", name="sample-10BT",
+                      split="train", streaming=True)
+    got = 0
+    for ex in ds:
+        t = ex["text"].strip()
+        if len(t) >= min_chars:
+            yield t
+            got += 1
+            if got >= n:
+                return
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--n-prompts", type=int, default=256)
+    ap.add_argument("--seq-len", type=int, default=64)
+    ap.add_argument("--chunk", type=int, default=128)
+    ap.add_argument("--out", default="results/jbar.pt")
+    ap.add_argument("--dtype", default="float32")
+    args = ap.parse_args()
+
+    model, tok = load_model(dtype=getattr(torch, args.dtype))
+    L = len(model.model.language_model.layers)
+    d = model.config.get_text_config().hidden_size
+
+    Jsum = torch.zeros(L, d, d, device="cuda", dtype=torch.float32)
+    pairs_total = 0
+    t0 = time.time()
+    for i, text in enumerate(corpus_texts(args.n_prompts)):
+        ids = tok(text, return_tensors="pt", truncation=True,
+                  max_length=args.seq_len)["input_ids"].cuda()
+        if ids.shape[1] < args.seq_len:
+            continue
+        J, n_pairs = prompt_jacobian_pairsum(model, ids, chunk=args.chunk)
+        Jsum += J
+        pairs_total += n_pairs
+        if i % 5 == 0 or i == args.n_prompts - 1:
+            el = time.time() - t0
+            print(f"[{i+1}/{args.n_prompts}] {el:.0f}s ({el/(i+1):.1f}s/prompt)",
+                  flush=True)
+        if i % 50 == 49:  # checkpoint
+            torch.save({"Jbar": (Jsum / pairs_total).cpu(),
+                        "n_prompts": i + 1, "seq_len": args.seq_len},
+                       args.out + ".ckpt")
+
+    Jbar = (Jsum / pairs_total).cpu()
+    Path(args.out).parent.mkdir(parents=True, exist_ok=True)
+    torch.save({"Jbar": Jbar, "n_prompts": args.n_prompts,
+                "seq_len": args.seq_len, "pairs": pairs_total}, args.out)
+    print("saved", args.out)
+
+
+if __name__ == "__main__":
+    main()