File size: 4,170 Bytes
f31f59b | 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 | """Paper figures from the per-step training logs (ckpt/<arm>/log.jsonl, one JSON line per optimiser step).
python plot.py runlogs/A.jsonl runlogs/B.jsonl runlogs/C.jsonl [runlogs/D.jsonl] --out figures
Figure 1: training loss vs tokens seen (smoothed over 50 steps, raw curve faint behind it), all arms.
Figure 2: validation loss vs tokens (every 200 steps), all arms.
Figure 3: the last 25% of training zoomed in, where the arms separate.
Also writes figures/curves.csv (step, tokens, loss per arm) for any plotting tool.
"""
import argparse, json, pathlib
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
NAMES = {"A": "A: learned table (268M trainable)", "B": "B: frozen table, float body (101M)",
"C": "C: frozen table, ternary body (101M)", "D": "D: frozen random codes (101M)"}
COLORS = {"A": "#4E5A63", "B": "#3DE0C2", "C": "#FF5C7A", "D": "#C9A227"}
def load(p):
rows = [json.loads(l) for l in open(p)]
by_step = {r["step"]: r for r in rows} # a resumed run may repeat steps: keep the last
return [by_step[s] for s in sorted(by_step)]
def smooth(y, k=50):
if len(y) < k:
return np.asarray(y)
c = np.cumsum(np.insert(y, 0, 0.0))
s = (c[k:] - c[:-k]) / k # centred moving average: no lag against the raw curve
h = k // 2
return np.concatenate([np.full(h, np.nan), s, np.full(len(y) - len(s) - h, np.nan)])
def main():
ap = argparse.ArgumentParser()
ap.add_argument("logs", nargs="+")
ap.add_argument("--out", default="figures")
a = ap.parse_args()
out = pathlib.Path(a.out); out.mkdir(exist_ok=True)
runs = {}
for p in a.logs:
arm = pathlib.Path(p).stem[0].upper()
runs[arm] = load(p)
fig, ax = plt.subplots(figsize=(8, 5))
for arm, R in runs.items():
t = np.array([r["tokens"] for r in R]) / 1e9; y = np.array([r["loss"] for r in R])
ax.plot(t, y, color=COLORS[arm], alpha=0.15, lw=0.6)
ax.plot(t, smooth(y), color=COLORS[arm], lw=1.8, label=NAMES[arm])
ax.set_xlabel("training tokens (billions)"); ax.set_ylabel("training loss (nats/token)")
ax.set_ylim(top=min(8, ax.get_ylim()[1])); ax.grid(alpha=0.3); ax.legend()
fig.tight_layout(); fig.savefig(out / "fig1_train_loss.png", dpi=200); fig.savefig(out / "fig1_train_loss.pdf")
fig, ax = plt.subplots(figsize=(8, 5))
for arm, R in runs.items():
V = [r for r in R if "val_loss" in r]
ax.plot([r["tokens"] / 1e9 for r in V], [r["val_loss"] for r in V], "o-", ms=3, color=COLORS[arm], label=NAMES[arm])
ax.set_xlabel("training tokens (billions)"); ax.set_ylabel("validation loss (nats/token)")
ax.set_ylim(top=min(8, ax.get_ylim()[1])); ax.grid(alpha=0.3); ax.legend()
fig.tight_layout(); fig.savefig(out / "fig2_val_loss.png", dpi=200); fig.savefig(out / "fig2_val_loss.pdf")
fig, ax = plt.subplots(figsize=(8, 5))
for arm, R in runs.items():
n = len(R); R2 = R[int(n * 0.75):]
t = np.array([r["tokens"] for r in R2]) / 1e9; y = smooth(np.array([r["loss"] for r in R2]))
ax.plot(t, y, color=COLORS[arm], lw=1.8, label=NAMES[arm])
ax.set_xlabel("training tokens (billions)"); ax.set_ylabel("training loss, last 25% (smoothed)")
ax.grid(alpha=0.3); ax.legend()
fig.tight_layout(); fig.savefig(out / "fig3_last_quarter.png", dpi=200); fig.savefig(out / "fig3_last_quarter.pdf")
steps = sorted(set().union(*[{r["step"] for r in R} for R in runs.values()]))
idx = {arm: {r["step"]: r for r in R} for arm, R in runs.items()}
with open(out / "curves.csv", "w") as f:
f.write("step,tokens," + ",".join(f"loss_{a},val_{a}" for a in runs) + "\n")
for s in steps:
tok = next((idx[a][s]["tokens"] for a in runs if s in idx[a]), "")
cells = []
for a_ in runs:
r = idx[a_].get(s, {})
cells += [str(r.get("loss", "")), str(r.get("val_loss", ""))]
f.write(f"{s},{tok}," + ",".join(cells) + "\n")
print("saved", sorted(p.name for p in out.iterdir()))
if __name__ == "__main__":
main()
|