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,59 @@
+"""Render heatmaps (concept P by layer x position) and the layer profile."""
+
+import os, sys
+from pathlib import Path
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import torch
+
+RES = Path(__file__).resolve().parent.parent / os.environ.get("JLENS_RESULTS", "results")
+
+
+def plot_heat(path):
+    d = torch.load(path)
+    m, words, toks = d["map"].detach(), d["words"], d["tokens"]
+    fig, ax = plt.subplots(figsize=(min(16, 0.28 * m.shape[1] + 2), 6))
+    im = ax.imshow(m.numpy(), aspect="auto", origin="lower", cmap="magma",
+                   vmin=0)
+    ax.set_ylabel("layer")
+    ax.set_xlabel("position")
+    if toks and len(toks) <= 60:
+        ax.set_xticks(range(len(toks)))
+        ax.set_xticklabels([t.replace("\n", "\\n") for t in toks],
+                           rotation=90, fontsize=6)
+    ax.set_title(f"J-lens P({'/'.join(words)})")
+    fig.colorbar(im)
+    out = path.with_suffix(".png")
+    fig.tight_layout()
+    fig.savefig(out, dpi=130)
+    plt.close(fig)
+    print("wrote", out)
+
+
+def plot_profile(path):
+    d = torch.load(path)
+    fig, ax1 = plt.subplots(figsize=(8, 4.5))
+    L = len(d["entropy"])
+    ax1.plot(range(L), d["entropy"], "k-", label="lens entropy")
+    ax1.set_xlabel("layer")
+    ax1.set_ylabel("entropy (nats)")
+    ax2 = ax1.twinx()
+    ax2.plot(range(L), d["agree_in"], "b--", label="top tok == current input")
+    ax2.plot(range(L), d["agree_next"], "r--", label="top tok == next output")
+    ax2.set_ylabel("agreement")
+    fig.legend(loc="upper center", ncol=3, fontsize=8)
+    fig.tight_layout()
+    out = path.with_suffix(".png")
+    fig.savefig(out, dpi=130)
+    plt.close(fig)
+    print("wrote", out)
+
+
+if __name__ == "__main__":
+    for p in sorted(RES.glob("heat_*.pt")):
+        plot_heat(p)
+    lp = RES / "layer_profile.pt"
+    if lp.exists():
+        plot_profile(lp)