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>

scripts/train_merge.py

@@ -0,0 +1,126 @@
+"""Train ONLY the merge adapter at the L13->L14 boundary (WORKSPACE_LOOPING.md §3).
+
+Frozen base model; band L14-30 unrolled k times through the adapter; loss =
+cross-entropy on answer tokens of the *direct* (no-CoT) prompt, supervised by
+gold answers on items the frozen model can reach (STaR filter from
+prep_star_data.py). Difficulty->depth curriculum:
+    k=1: easy items    k=2: easy+hard    k=4: hard items
+so hard items are only ever seen at depth, forcing the loop to be used.
+"""
+
+import json
+import random
+import time
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+
+from loop_common import BandLooper, MergeAdapter, build_train_batch
+import sys
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+from jlens.core import load_model  # noqa: E402
+
+OUT = Path(__file__).resolve().parent.parent / "results-loop"
+STEPS = 800
+BATCH = 8
+LR = 1e-3
+WARMUP = 20
+MAX_TOK = 400          # skip overlong items
+K_BUCKETS = [(1, ("easy",)), (2, ("easy", "hard")), (4, ("hard",))]
+SEED = 0
+
+
+def lr_at(step):
+    if step < WARMUP:
+        return LR * (step + 1) / WARMUP
+    import math
+    t = (step - WARMUP) / max(1, STEPS - WARMUP)
+    return 1e-4 + 0.5 * (LR - 1e-4) * (1 + math.cos(math.pi * t))
+
+
+@torch.no_grad()
+def val_answer_acc(looper, adapter, tok, items, k):
+    """Teacher-forced exact match: argmax at every answer position."""
+    ok = 0
+    for i in range(0, len(items), BATCH):
+        ids, msk, lab = build_train_batch(tok, items[i : i + BATCH])
+        logits = looper.loop_logits(adapter, ids, k, attention_mask=msk)
+        pred = logits[:, :-1].argmax(-1)
+        tgt = lab[:, 1:]
+        m = tgt != -100
+        ok += (((pred == tgt) | ~m).all(-1)).sum().item()
+    return ok / len(items)
+
+
+def main():
+    rng = random.Random(SEED)
+    torch.manual_seed(SEED)
+    data = json.load(open(OUT / "star_data.json"))
+
+    model, tok = load_model(dtype=torch.bfloat16)
+    for p in model.parameters():
+        p.requires_grad_(False)
+    looper = BandLooper(model)
+    adapter = MergeAdapter().cuda()
+    opt = torch.optim.AdamW(adapter.parameters(), lr=LR, weight_decay=0.01)
+
+    train = [it for it in data if it["split"] == "train" and it["label"] != "drop"]
+    # length filter
+    keep = []
+    for it in train:
+        n = len(tok(it["question"])["input_ids"])
+        if n + 40 <= MAX_TOK:
+            keep.append(it)
+    train = keep
+    pool = {"easy": [it for it in train if it["label"] == "easy"],
+            "hard": [it for it in train if it["label"] == "hard"]}
+    val = {lbl: rng.sample(lst, min(32, len(lst))) for lbl, lst in pool.items()}
+    for lbl in pool:
+        pool[lbl] = [it for it in pool[lbl] if it not in val[lbl]]
+    print(f"train pool: easy={len(pool['easy'])} hard={len(pool['hard'])}", flush=True)
+
+    log = []
+    t0 = time.time()
+    for step in range(STEPS):
+        k, labels = K_BUCKETS[step % len(K_BUCKETS)]
+        cand = [it for lbl in labels for it in pool[lbl]]
+        batch = rng.sample(cand, BATCH)
+        ids, msk, lab = build_train_batch(tok, batch)
+
+        for g in opt.param_groups:
+            g["lr"] = lr_at(step)
+        logits = looper.loop_logits(adapter, ids, k, attention_mask=msk,
+                                    use_checkpoint=True)
+        loss = F.cross_entropy(
+            logits[:, :-1].flatten(0, 1).float(), lab[:, 1:].flatten(),
+            ignore_index=-100)
+        opt.zero_grad(set_to_none=True)
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(adapter.parameters(), 1.0)
+        opt.step()
+
+        log.append({"step": step, "k": k, "loss": loss.item()})
+        if step % 10 == 0:
+            print(f"step {step:4d}  k={k}  loss={loss.item():.4f}  "
+                  f"({(time.time()-t0)/(step+1):.1f}s/step)", flush=True)
+        if step % 100 == 99 or step == STEPS - 1:
+            accs = {}
+            for kk in (0, 1, 2, 4):
+                accs[f"easy_k{kk}"] = val_answer_acc(looper, adapter, tok,
+                                                     val["easy"], kk)
+                accs[f"hard_k{kk}"] = val_answer_acc(looper, adapter, tok,
+                                                     val["hard"], kk)
+            print(f"  val@{step}: " +
+                  "  ".join(f"{n}={v:.2f}" for n, v in accs.items()), flush=True)
+            log.append({"step": step, "val": accs})
+            torch.save(adapter.state_dict(), OUT / "adapter.pt")
+            json.dump(log, open(OUT / "train_log.json", "w"), indent=1)
+
+    torch.save(adapter.state_dict(), OUT / "adapter.pt")
+    json.dump(log, open(OUT / "train_log.json", "w"), indent=1)
+    print("done; adapter ->", OUT / "adapter.pt")
+
+
+if __name__ == "__main__":
+    main()