251 lines
10 KiB
Python
251 lines
10 KiB
Python
"""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)
|
||
# untrained E2B control (250-item era, n_hard=28)
|
||
d = json.load(open(OUT / "eval_code_untrained.json"))
|
||
hu = max(v["by_label"]["hard"] for k, v in d["ks"].items() if k != "0")
|
||
add("untrained loop (best k)", hu, 28, GRAY)
|
||
d = json.load(open(OUT / "eval_code_ff.json"))
|
||
add("trained FF (no recurrence)", d["ks"]["1"]["by_label"]["hard"], 28,
|
||
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)),
|
||
]),
|
||
]
|
||
fig, axes = plt.subplots(1, 2, figsize=(7.6, 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()
|