File size: 8,910 Bytes
84bcfb1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 | #!/usr/bin/env python
"""Draw the figures for Labs 2, 4 and 5 from the finished runs.
python scripts/summarize.py results/runs # first: writes results/summary.json
uv run --with matplotlib python scripts/plot_runs.py # any machine, a few seconds
Reads results/summary.json, the answer lengths in results/runs/**/merged.jsonl and Lab 7's
results/lab7/side_effects.json, and writes results/figures/lab2_*.png, lab4_*.png and lab5_*.png.
Runs that have not finished yet are left out.
"""
import json
import os
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt # noqa: E402
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
RES = os.path.join(ROOT, "results")
FIG = os.path.join(RES, "figures")
os.makedirs(FIG, exist_ok=True)
S = {r["name"]: r["summary"] for r in json.load(open(os.path.join(RES, "summary.json")))}
SIDE = json.load(open(os.path.join(RES, "lab7", "side_effects.json")))["models"]
# the blog's think35 numbers: accuracy %, cut at the cap %, mean tokens
PAPER = {"math500": {"base": (7.9, 93.0, 30868), "p50c": (48.1, 33.05, 22037)},
"gsm8k": {"base": (13.84, 87.5, 29289), "p50c": (67.1, 7.24, 11993)},
"ifeval": {"base": (21.0, 74.9, 25296), "p50c": (62.3, 2.2, 4241)}}
GRAY, GREEN, RED, PURPLE, BLUE = "#8c959f", "#1a7f37", "#cf222e", "#8250df", "#0969da"
METRICS = [("acc", "accuracy %", 100), ("cut_share", "answers cut at the token cap %", 100),
("mean_tokens", "mean answer length (thousand tokens)", 1e-3)]
def val(s, k, scale):
return None if s is None or s.get(k) is None else s[k] * scale
def err(s, k):
return 100 * s["acc_se"] if k == "acc" and s and s.get("acc_se") is not None else 0
def lengths(name):
p = os.path.join(RES, "runs", name, "merged.jsonl")
if not os.path.exists(p):
return None
return [s["tokens"] for line in open(p) for s in json.loads(line)["samples"]]
made = []
# Lab 2: the paper's benchmarks, ours next to the blog's
bench = [b for b in PAPER if f"bench/base/{b}" in S or f"bench/p50c/{b}" in S]
if bench:
fig, ax = plt.subplots(1, 3, figsize=(15, 4.2))
w = 0.2
for i, (k, lab, scale) in enumerate(METRICS):
for j, b in enumerate(bench):
for o, (model, src, col, hatch) in enumerate([("base", "ours", GRAY, None), ("base", "paper", GRAY, "//"),
("p50c", "ours", GREEN, None),
("p50c", "paper", GREEN, "//")]):
if src == "ours":
s = S.get(f"bench/{model}/{b}")
v, e = val(s, k, scale), err(s, k)
else:
v, e = PAPER[b][model][i] * (1e-3 if k == "mean_tokens" else 1), 0
if v is None:
continue
ax[i].bar(j + (o - 1.5) * w, v, w, yerr=e or None, color=col, hatch=hatch, edgecolor="white",
label=f"{'unedited' if model == 'base' else 'edited'}, {src}" if j == 0 else None)
ax[i].set_xticks(range(len(bench)), [b.upper().replace("MATH500", "MATH-500") for b in bench])
ax[i].set_title(lab)
ax[0].legend(fontsize=8)
fig.suptitle("The paper's three benchmarks: our runs (solid) next to the blog's numbers (hatched)")
fig.tight_layout()
fig.savefig(os.path.join(FIG, "lab2_benchmarks.png"), dpi=130)
made.append("lab2_benchmarks.png")
pairs = [(n, c) for n, c in (("dev/base", GRAY), ("dev/p50c", GREEN)) if lengths(n)]
if pairs:
fig, ax = plt.subplots(figsize=(10, 3.8))
for n, c in pairs:
ax.hist(lengths(n), bins=64, range=(0, 32768), color=c, alpha=0.6,
label=f"{'unedited' if 'base' in n else 'edited'} ({n})")
ax.set_yscale("log")
ax.set_xlabel("answer length in tokens (the cap is 32,768)")
ax.set_ylabel("answers")
ax.set_title("Where the answers end: the tall bar at the right edge is answers that never stopped")
ax.legend()
fig.tight_layout()
fig.savefig(os.path.join(FIG, "lab2_lengths.png"), dpi=130)
made.append("lab2_lengths.png")
# Lab 4: turn the knob, then the controls
DOSE = [(-4, "a-4"), (0, "base"), (1, "a1"), (2, "a2"), (4, "a4"), (6, "p50c"), (8, "a8"), (12, "a12")]
dose = [(x, S[f"dev/{n}"], SIDE.get(n)) for x, n in DOSE if f"dev/{n}" in S]
if len(dose) >= 3:
fig, ax = plt.subplots(1, 4, figsize=(17, 4))
xs = [d[0] for d in dose]
for i, (k, lab, scale) in enumerate(METRICS):
ax[i].errorbar(xs, [val(d[1], k, scale) for d in dose], yerr=[err(d[1], k) for d in dose], fmt="o-",
color=BLUE, capsize=3, label="released u and 64 neurons")
for n, dx, mk, col, lab2 in (("rebuilt", 0.35, "s", RED, "your Lab 3 rebuild"),
("capA-a6", 0.7, "^", PURPLE, "capstone, recipe A"),
("capB-a6", 1.05, "v", PURPLE, "capstone, recipe B")):
if f"dev/{n}" in S:
ax[i].errorbar([6 + dx], [val(S[f"dev/{n}"], k, scale)], yerr=[err(S[f"dev/{n}"], k)], fmt=mk,
color=col, mfc="none", capsize=3, label=f"{lab2} (alpha 6, drawn to the right)")
ax[i].set_title(lab)
sx = [d[0] for d in dose if d[2]]
ax[3].plot(sx, [d[2]["ppl_change_pct"] for d in dose if d[2]], "o-", color=RED)
if SIDE.get("rebuilt"):
ax[3].plot([6.35], [SIDE["rebuilt"]["ppl_change_pct"]], "s", color=RED, mfc="none")
ax[3].axhline(0, color="black", lw=0.5)
ax[3].set_title("side effect: perplexity change % (Lab 7)")
for a_ in ax:
a_.set_xlabel("alpha (0 = no edit, 1 = erase, 2 = mirror, 6 = released)")
a_.set_xticks([-4, 0, 1, 2, 4, 6, 8, 12])
h, lb = ax[0].get_legend_handles_labels()
fig.legend(h, lb, loc="lower center", ncol=len(lb), fontsize=9, frameon=False)
fig.suptitle("Turning the knob: same direction and neurons, different alpha (256 unseen practice problems)")
fig.tight_layout(rect=(0, 0.07, 1, 1))
fig.savefig(os.path.join(FIG, "lab4_dose.png"), dpi=130)
made.append("lab4_dose.png")
CTRL = [("base", "no edit", GRAY), ("p50c", "released edit", GREEN), ("rebuilt", "your Lab 3 rebuild", GREEN),
("capA-a6", "capstone A: own loops, just before", PURPLE), ("capB-a6", "capstone B: own loops, well before",
PURPLE),
("l15only", "layer 15 only", BLUE), ("l19only", "layer 19 only", BLUE),
("all1", "erase u from all 3,584 neurons", BLUE), ("all2", "mirror u in all 3,584 neurons", BLUE),
("randdir", "control: random direction", RED), ("randcols", "control: 64 random neurons", RED)]
ctrl = [(lab, S[f"dev/{n}"], c) for n, lab, c in CTRL if f"dev/{n}" in S]
if len(ctrl) >= 3:
fig, ax = plt.subplots(1, 3, figsize=(17, 0.45 * len(ctrl) + 1.5), sharey=True)
ys = list(range(len(ctrl)))[::-1]
for i, (k, lab, scale) in enumerate(METRICS):
ax[i].barh(ys, [val(s, k, scale) for _, s, _ in ctrl],
xerr=[err(s, k) for _, s, _ in ctrl] if k == "acc" else None,
color=[c for *_, c in ctrl], capsize=2)
ax[i].set_title(lab)
ax[0].set_yticks(ys, [lab for lab, *_ in ctrl])
fig.suptitle("Which parts of the edit matter? Same 256 unseen practice problems, 2 answers each")
fig.tight_layout()
fig.savefig(os.path.join(FIG, "lab4_controls.png"), dpi=130)
made.append("lab4_controls.png")
# Lab 5: decoding fixes on the unedited model vs the weight edit
DEC = [("base", "unedited, paper settings", GRAY), ("base+temp1.0", "temperature 1.0 (from 0.6)", BLUE),
("base+pp0", "presence penalty 0 (from 1.5)", BLUE), ("base+pp2", "presence penalty 2.0", BLUE),
("base+rep1.1", "repetition penalty 1.1", BLUE), ("base+max4k", "cap at 4,096 tokens", BLUE),
("base+nothink", "thinking off (cap 8,192)", BLUE), ("p50c", "edited, paper settings", GREEN),
("p50c+nothink", "edited, thinking off", GREEN)]
dec = [(lab, S[f"dev/{n}"], c) for n, lab, c in DEC if f"dev/{n}" in S]
if len(dec) >= 3:
fig, ax = plt.subplots(1, 3, figsize=(17, 0.45 * len(dec) + 1.5), sharey=True)
ys = list(range(len(dec)))[::-1]
for i, (k, lab, scale) in enumerate(METRICS):
ax[i].barh(ys, [val(s, k, scale) for _, s, _ in dec],
xerr=[err(s, k) for _, s, _ in dec] if k == "acc" else None,
color=[c for *_, c in dec], capsize=2)
ax[i].set_title(lab)
ax[0].set_yticks(ys, [lab for lab, *_ in dec])
fig.suptitle("Can decoding settings do what the weight edit does? Same 256 unseen practice problems")
fig.tight_layout()
fig.savefig(os.path.join(FIG, "lab5_decoding.png"), dpi=130)
made.append("lab5_decoding.png")
print("wrote", ", ".join(made) if made else "nothing yet (no finished runs)")
|