Files
jspace/scripts/figures_final.py

254 lines
10 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""Final paper figures (CPU only), regenerated from consolidated results.
fig_placement.png : band-entrance cliff + structural nulls
fig_ladder.png : MBPP-E2B attribution ladder, hard bucket, Wilson CIs
fig_scale.png : cross-scale / cross-task attribution grid
fig_transfer.png : substrate transfer panel (HumanEval / Rust / BW)
"""
import json
import math
from pathlib import Path
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
ROOT = Path(__file__).resolve().parent.parent
OUT = ROOT / "results-loop"
BLUE, GREEN, ORANGE, GRAY, RED = ("#2b6cb0", "#2f855a", "#dd6b20",
"#8a8f98", "#c53030")
def wilson(c, n, z=1.96):
p = c / n
d = 1 + z * z / n
ctr = (p + z * z / (2 * n)) / d
hw = z * math.sqrt(p * (1 - p) / n + z * z / (4 * n * n)) / d
return ctr - hw, ctr + hw
def ev(path, k):
d = json.load(open(path))
v = d["ks"][str(k)]
return v["acc"], v["by_label"].get("hard", float("nan"))
def best_hard(path, exclude0=True):
d = json.load(open(path))
items = [(int(k), v) for k, v in d["ks"].items()
if not (exclude0 and k == "0")]
k, v = max(items, key=lambda kv: kv[1]["by_label"].get("hard", 0))
return k, v["acc"], v["by_label"].get("hard", 0)
def style(ax):
ax.grid(True, color="#e8e8e8", lw=0.7)
ax.set_axisbelow(True)
for s in ("top", "right"):
ax.spines[s].set_visible(False)
# ---------------- fig 1: placement ----------------
def fig_placement():
entr = [] # (entrance_layer, overall@bestk, hard@bestk)
for lo in (9, 11, 12, 13, 14):
f = ROOT / f"results-band-{lo}_30" / f"eval_code_band_{lo}_30.json"
if lo == 14:
f = OUT / "eval_code_code_s0_full.json"
k, a, h = best_hard(f)
entr.append((lo, a, h))
base_a, base_h = ev(OUT / "eval_code_code_s0_full.json", 0)
fig, ax = plt.subplots(figsize=(6.4, 4.2))
style(ax)
xs = [e[0] for e in entr]
ax.plot(xs, [e[2] for e in entr], "-s", color=BLUE, lw=2, ms=6,
label="hard bucket (best k)")
ax.plot(xs, [e[1] for e in entr], "-o", color=GRAY, lw=2, ms=5,
label="overall (same k)")
ax.axhline(base_h, color=BLUE, lw=1, ls=":", alpha=0.6)
ax.axhline(base_a, color=GRAY, lw=1, ls=":", alpha=0.6)
ax.annotate("base hard", (9.1, base_h + 0.012), fontsize=8, color=BLUE)
ax.annotate("base overall", (9.1, base_a + 0.012), fontsize=8, color=GRAY)
ax.axvspan(13.5, 14.5, color="#ebf4ff", zorder=0)
ax.annotate("lens-identified\nworkspace entrance", (13.55, 0.60),
fontsize=8, color=BLUE)
ax.set_xticks(xs)
ax.set_xlabel("loop entrance layer (exit fixed at L30)")
ax.set_ylabel("MBPP pass@1")
ax.set_title("Placement cliff: the retrofit works only at the "
"lens boundary (L14)", fontsize=11)
ax.legend(fontsize=8, frameon=False, loc="center left")
fig.text(0.13, 0.005,
"Entrances 17/24 (not shown): structurally null — KV sharing "
"makes k>0 bit-identical to k=0.", fontsize=7.5, color="#666")
fig.tight_layout(rect=(0, 0.03, 1, 1))
fig.savefig(OUT / "fig_placement.png", dpi=140, facecolor="white",
bbox_inches="tight")
print("wrote fig_placement.png")
# ---------------- fig 2: attribution ladder ----------------
def fig_ladder():
NH = 55
rows = [] # (label, hard_acc, n, color, note)
def add(label, h, n, color, note=""):
rows.append((label, h, n, color, note))
_, h0 = ev(OUT / "eval_code_code_s0_full.json", 0)
add("base (k=0, exact)", h0, NH, GRAY)
_, hu = ev(OUT / "eval_code_untrained_full.json", 4)
add("untrained loop (perturbation floor)", hu, NH, GRAY)
_, hf = ev(OUT / "eval_code_code_ff_full.json", 1)
add("trained FF (no recurrence)", hf, NH, ORANGE)
_, hp = ev(OUT / "eval_code_pause16.json", 1)
add("pause-16 registers (width)", hp, NH, ORANGE)
_, hl = ev(OUT / "eval_code_code_s0_full.json", 4)
add("loop k=4 (depth)", hl, NH, BLUE, "seed mean 0.375 ± 0.055")
d = json.load(open(OUT / "eval_rung2_s1.json"))
add("rung-2: + band LoRA (k=4)", d["4"]["by_label"]["hard"], 28, BLUE)
_, hd = ev(OUT / "eval_code_distill_s1.json", 1)
add("plan-distilled FF", hd, NH, GREEN, "8-run mean 0.457 ± 0.046")
d = json.load(open(OUT / "eval_budgetcot.json"))
add("budget-CoT (50 visible tok)", d["by_label"]["hard"], NH, "#805ad5")
d = json.load(open(OUT / "eval_bestof3.json"))
add("best-of-3 sampling (~matched FLOPs)", d["by_label"]["hard"], NH,
"#805ad5")
d = json.load(open(OUT / "eval_plan_baseline.json"))
add("explicit plan in context (ceiling)", d["by_label"]["hard"], NH,
"#1a202c")
fig, ax = plt.subplots(figsize=(7.4, 4.8))
style(ax)
ys = range(len(rows))[::-1]
for y, (label, h, n, color, note) in zip(ys, rows):
lo, hi = wilson(round(h * n), n)
ax.barh(y, h, color=color, height=0.62, alpha=0.88)
ax.plot([lo, hi], [y, y], color="#333", lw=1.2)
txt = f"{h:.2f}"
if note:
txt += f" ({note})"
ax.text(hi + 0.015, y, txt, va="center", fontsize=8)
ax.set_yticks(list(ys))
ax.set_yticklabels([r[0] for r in rows], fontsize=9)
ax.set_xlim(0, 1.02)
ax.set_xlabel("pass@1, MBPP hard bucket (plan-dependent items)")
ax.set_title("Attribution ladder: what closes the plan gap "
"(bars: point estimate, whiskers: Wilson 95%)", fontsize=11)
fig.tight_layout()
fig.savefig(OUT / "fig_ladder.png", dpi=140, facecolor="white",
bbox_inches="tight")
print("wrote fig_ladder.png")
# ---------------- fig 3: cross-scale grid ----------------
def fig_scale():
N2 = ROOT / "results-node2-final/results-12b"
N1 = ROOT / "results-node-final/results-12b"
panels = {
("MBPP", "E2B"): [
("base", *ev(OUT / "eval_code_code_s0_full.json", 0)),
("loop k=4", *ev(OUT / "eval_code_code_s0_full.json", 4)),
("adaptive k=4", *ev(OUT / "eval_code_e2b_adaptive.json", 4)),
("distill", *ev(OUT / "eval_code_distill_s1.json", 1)),
],
("MBPP", "12B"): [
("base", *ev(N1 / "eval_code_12b_trained.json", 0)),
("loop k=4 (α=.3)", *ev(N1 / "eval_code_12b_trained.json", 4)),
("adaptive k=4", *ev(N2 / "eval_code_12b_adaptive.json", 4)),
("distill", *ev(OUT / "eval_code_12b_distill.json", 1)),
],
("GSM8K", "E2B"): [
("base", *ev(OUT / "eval_uni.json", 0)),
("loop k=2", *ev(OUT / "eval_uni.json", 2)),
("adaptive k=2", *ev(OUT / "eval_gsm_e2b_adaptive.json", 2)),
("distill", *ev(OUT / "eval_gsm_distill_retry.json", 1)),
],
("GSM8K", "12B"): [
("base", *ev(N1 / "eval_12b_gsm_loop.json", 0)),
("loop k=2 (α=.3)", *ev(N1 / "eval_12b_gsm_loop.json", 2)),
("adaptive k=2", *ev(OUT / "eval_12b_gsm_adaptive.json", 2)),
("distill*", *ev(OUT / "eval_12b_gsm_distill.json", 1)),
],
}
fig, axes = plt.subplots(2, 2, figsize=(9.6, 6.6))
for ax, ((task, scale), arms) in zip(axes.flat, panels.items()):
if (task, scale) == ("GSM8K", "12B"):
ax.annotate("*training collapse (0.00)", (3, 0.03), fontsize=7.5,
ha="center", color=RED)
style(ax)
x = range(len(arms))
ax.bar([i - 0.19 for i in x], [a[1] for a in arms], width=0.36,
color=GRAY, alpha=0.85, label="overall")
ax.bar([i + 0.19 for i in x], [a[2] for a in arms], width=0.36,
color=BLUE, alpha=0.85, label="hard")
ax.axhline(arms[0][1], color=GRAY, lw=1, ls=":")
ax.set_xticks(list(x))
ax.set_xticklabels([a[0] for a in arms], fontsize=8)
ax.set_title(f"{task} · {scale}", fontsize=10)
ax.set_ylim(0, 1.0)
if ax is axes.flat[0]:
ax.legend(fontsize=8, frameon=False)
fig.suptitle("Cross-scale attribution: constant-α destroys the 12B "
"substrate; state-dependent α restores it (MBPP) but not "
"everywhere", fontsize=11.5)
fig.tight_layout(rect=(0, 0, 1, 0.96))
fig.savefig(OUT / "fig_scale.png", dpi=140, facecolor="white",
bbox_inches="tight")
print("wrote fig_scale.png")
# ---------------- fig 4: transfer panel ----------------
def fig_transfer():
def he_hard(path, k):
d = json.load(open(OUT / path))
return d["ks"][str(k)]["acc"], d["ks"][str(k)]["hard_acc"]
groups = [
("HumanEval\n(MBPP-trained)", [
("base", *he_hard("eval_humaneval_trained.json", 0)),
("loop k=4", *he_hard("eval_humaneval_trained.json", 4)),
("distill", *he_hard("eval_humaneval_distill_transfer.json", 1)),
]),
("Rust / MultiPL-E\n(Python-trained)", [
("base", *ev(OUT / "eval_rust_py_transfer.json", 0)),
("loop k=4", *ev(OUT / "eval_rust_py_transfer.json", 4)),
]),
("LiveCodeBench (far)\nuntrained wins", [
("base", *he_hard("eval_lcb_untrained.json", 0)),
("untrained k=4", *he_hard("eval_lcb_untrained.json", 4)),
("loop k=4", *he_hard("eval_lcb_loop.json", 4)),
("distill", *he_hard("eval_lcb_distill.json", 1)),
]),
]
fig, axes = plt.subplots(1, 3, figsize=(10.8, 3.8))
for ax, (title, arms) in zip(axes, groups):
style(ax)
x = range(len(arms))
ax.bar([i - 0.19 for i in x], [a[1] for a in arms], width=0.36,
color=GRAY, alpha=0.85, label="overall")
ax.bar([i + 0.19 for i in x],
[a[2] if a[2] is not None else 0 for a in arms],
width=0.36, color=BLUE, alpha=0.85, label="hard")
ax.set_xticks(list(x))
ax.set_xticklabels([a[0] for a in arms], fontsize=8.5)
ax.set_title(title, fontsize=9.5)
ax.set_ylim(0, 1.0)
axes[0].legend(fontsize=8, frameon=False)
fig.suptitle("Transfer: the implant moves with the substrate, "
"not the task", fontsize=11.5)
fig.tight_layout(rect=(0, 0, 1, 0.94))
fig.savefig(OUT / "fig_transfer.png", dpi=140, facecolor="white",
bbox_inches="tight")
print("wrote fig_transfer.png")
if __name__ == "__main__":
fig_placement()
fig_ladder()
fig_scale()
fig_transfer()