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>
58 lines
2.4 KiB
Python
58 lines
2.4 KiB
Python
"""Two-panel figure: accuracy vs loop depth k, and J-lens concept sharpening."""
|
||
import json
|
||
from pathlib import Path
|
||
|
||
import matplotlib
|
||
matplotlib.use("Agg")
|
||
import matplotlib.pyplot as plt
|
||
|
||
OUT = Path(__file__).resolve().parent.parent / "results-loop"
|
||
tr = json.load(open(OUT / "eval_trained.json"))
|
||
un = json.load(open(OUT / "eval_untrained.json"))
|
||
|
||
ks = sorted(int(k) for k in tr["ks"])
|
||
BLUE, GRAY = "#2b6cb0", "#8a8f98"
|
||
|
||
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11, 4.2))
|
||
|
||
def curve(res, sel):
|
||
return [res["ks"][str(k)]["acc"] if sel == "all"
|
||
else res["ks"][str(k)]["by_label"].get(sel, 0.0) for k in ks]
|
||
|
||
for ax in (ax1, ax2):
|
||
ax.grid(True, color="#e5e5e5", lw=0.7, zorder=0)
|
||
ax.set_axisbelow(True)
|
||
for s in ("top", "right"):
|
||
ax.spines[s].set_visible(False)
|
||
ax.set_xticks(ks)
|
||
ax.set_xlabel("loop depth k")
|
||
|
||
# --- panel 1: GSM8K accuracy vs k ---
|
||
ax1.plot(ks, curve(tr, "all"), "-o", color=BLUE, lw=2, ms=5, label="trained · all")
|
||
ax1.plot(ks, curve(un, "all"), "-o", color=GRAY, lw=2, ms=5, label="untrained · all")
|
||
ax1.plot(ks, curve(tr, "hard"), "--s", color=BLUE, lw=2, ms=5, label="trained · hard (CoT-only)")
|
||
ax1.plot(ks, curve(un, "hard"), "--s", color=GRAY, lw=2, ms=5, label="untrained · hard")
|
||
ax1.set_ylabel("accuracy (greedy, 256 held-out items)")
|
||
ax1.set_title("GSM8K accuracy vs loop depth", fontsize=11)
|
||
ax1.legend(fontsize=8, frameon=False)
|
||
ax1.annotate("hard items: 0.8% → 6.3%", xy=(2, 0.063), xytext=(3.2, 0.115),
|
||
fontsize=8, color=BLUE,
|
||
arrowprops=dict(arrowstyle="-", color=BLUE, lw=0.8))
|
||
|
||
# --- panel 2: J-lens sharpening ---
|
||
sp_tr = {r["k"]: r["P_spider_lens"] for r in tr["spider"]}
|
||
sp_un = {r["k"]: r["P_spider_lens"] for r in un["spider"]}
|
||
ax2.plot(ks, [sp_tr[k] for k in ks], "-o", color=BLUE, lw=2, ms=5, label="trained")
|
||
ax2.plot(ks, [sp_un[k] for k in ks], "-o", color=GRAY, lw=2, ms=5, label="untrained")
|
||
ax2.set_ylabel("P('spider') under J-lens at L30")
|
||
ax2.set_title("Latent concept sharpening across loops", fontsize=11)
|
||
ax2.legend(fontsize=8, frameon=False)
|
||
ax2.text(ks[-1], sp_tr[ks[-1]] + 0.006, "8× over control", fontsize=8,
|
||
color=BLUE, ha="right")
|
||
|
||
fig.suptitle("Trained merge adapter: loop depth buys latent sharpening and some "
|
||
"CoT-only answers, not overall accuracy", fontsize=12, y=1.02)
|
||
fig.tight_layout()
|
||
fig.savefig(OUT / "loop_eval.png", dpi=140, bbox_inches="tight", facecolor="white")
|
||
print("wrote", OUT / "loop_eval.png")
|