#!/usr/bin/env python3 """ make_figures.py -- Generate all ACML 2026 paper figures. Output: paper/figures/fig_*.png (300 dpi, black-and-white, JMLR column widths) Run from the paper/ directory: python make_figures.py """ import csv, json, math, sys import numpy as np import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import matplotlib.ticker as mticker from pathlib import Path # --------------------------------------------------------------------------- # Global ACML / JMLR style (black-and-white, 10pt serif) # --------------------------------------------------------------------------- plt.rcParams.update({ "font.family": "serif", "font.size": 9, "axes.titlesize": 9, "axes.labelsize": 9, "xtick.labelsize": 8, "ytick.labelsize": 8, "legend.fontsize": 7.5, "legend.framealpha": 0.85, "figure.dpi": 300, "savefig.dpi": 300, "savefig.bbox": "tight", "axes.linewidth": 0.8, "grid.linewidth": 0.4, "lines.linewidth": 1.4, "lines.markersize": 4, "text.usetex": False, "axes.spines.top": False, "axes.spines.right": False, }) COL1 = 3.25 # single column (inches) COL2 = 6.75 # full width DATA = Path(__file__).parent / "data" OUT = Path(__file__).parent / "figures" OUT.mkdir(exist_ok=True) LS = ["-", "--", "-.", ":"] MK = ["o", "s", "^", "D", "v", "P"] GR = ["#000000", "#444444", "#888888", "#AAAAAA"] def save(name): p = OUT / f"fig_{name}.png" plt.savefig(p) plt.close() print(f" saved {p.name}") # --------------------------------------------------------------------------- # Fig 1 -- VRAM vs sequence length # --------------------------------------------------------------------------- def fig_vram(): rows = list(csv.DictReader(open(DATA / "results/roofline/vram_comparison.csv"))) fno_x, fno_y = [], [] sdpa_real_x, sdpa_real_y = [], [] sdpa_th_x, sdpa_th_y = [], [] for r in rows: if not r["vram_mb"]: continue T = int(r["seq_len"]) v = int(r["vram_mb"]) / 1024 if "FNO+GLA" in r["model"]: fno_x.append(T); fno_y.append(v) elif "theoretical" in r["model"]: sdpa_th_x.append(T); sdpa_th_y.append(v) elif "SDPA" in r["model"]: sdpa_real_x.append(T); sdpa_real_y.append(v) fig, ax = plt.subplots(figsize=(COL1, 2.6)) ax.plot(fno_x, fno_y, LS[0]+MK[0], color="k", label="FELA 1.13B") if sdpa_real_x: ax.plot(sdpa_real_x, sdpa_real_y, MK[1], color="k", markerfacecolor="white", ms=6, zorder=5, label="GPT-2 XL (OOM at 2K)") if sdpa_th_x: ax.plot(sdpa_th_x, sdpa_th_y, LS[1], color="k", alpha=0.7, label="GPT-2 XL (extrap.)") ax.axhline(24, color="k", lw=0.7, ls=":", alpha=0.5) ax.text(512, 25.5, "A10G 24 GB", fontsize=7, color="gray") ax.set_xscale("log", base=2) ax.set_yscale("log") ax.set_xlabel("Sequence length (tokens)") ax.set_ylabel("Peak VRAM (GB)") ax.set_title("VRAM vs. Sequence Length") ax.xaxis.set_major_formatter(mticker.FuncFormatter(lambda x, _: f"{int(x):,}")) ax.legend(loc="upper left") ax.grid(True, which="both", alpha=0.2) save("vram_measured") # --------------------------------------------------------------------------- # Fig 2 -- Long-context BPB # --------------------------------------------------------------------------- def fig_longctx(): rows = list(csv.DictReader( open(DATA / "results/paper_eval_22b/longctx_bpb.tsv"), delimiter="\t")) xs = [int(r["ctx_len"]) for r in rows] ys = [float(r["bpb"]) / 4.0 for r in rows] # bits/token -> bits/byte fig, ax = plt.subplots(figsize=(COL1, 2.5)) ax.plot(xs, ys, LS[0]+MK[0], color="k") dy = ys[0] - ys[-1] ax.annotate("", xy=(xs[-1], ys[-1]), xytext=(xs[-1], ys[0]), arrowprops=dict(arrowstyle="<->", color="k", lw=0.8)) ax.text(xs[-1]*1.08, (ys[0]+ys[-1])/2, f"-{dy:.2f}\n({dy/ys[0]*100:.0f}%)", fontsize=7, va="center") ax.set_xscale("log", base=2) ax.set_xlabel("Context length (tokens)") ax.set_ylabel("WikiText-103 BPB (bits/byte)") ax.set_title("Long-Context Utilisation") ax.xaxis.set_major_formatter(mticker.FuncFormatter(lambda x, _: f"{int(x):,}")) ax.grid(True, alpha=0.2) save("longctx_bpb") # --------------------------------------------------------------------------- # Fig 3 -- Throughput # --------------------------------------------------------------------------- def fig_throughput(): rows = list(csv.DictReader( open(DATA / "results/paper_eval_22b/throughput.tsv"), delimiter="\t")) xs = [int(r["seq_len"]) for r in rows] ys = [int(r["tok_per_sec"]) / 1000 for r in rows] base = ys[0] sdpa_xs = [512, 1024, 2048, 4096, 8192, 16384, 32768] sdpa_ys = [base * (512 / T) for T in sdpa_xs] fig, ax = plt.subplots(figsize=(COL1, 2.5)) ax.plot(xs, ys, LS[0]+MK[0], color="k", label="FELA 1.13B") ax.plot(sdpa_xs, sdpa_ys, LS[1]+MK[1], color="k", markerfacecolor="white", label="SDPA (theoretical)") ax.set_xscale("log", base=2) ax.set_xlabel("Sequence length (tokens)") ax.set_ylabel("Throughput (K tok/s)") ax.set_title("Prefill Throughput") ax.xaxis.set_major_formatter(mticker.FuncFormatter(lambda x, _: f"{int(x):,}")) ax.legend() ax.grid(True, alpha=0.2) save("throughput") # --------------------------------------------------------------------------- # Fig 4 -- Needle-in-haystack heatmap # --------------------------------------------------------------------------- def fig_needle(): data = json.load(open(DATA / "results/paper_eval_22b/needle_heatmap.json")) rows = data if isinstance(data, list) else data.get("results", []) ctx_lens = sorted(set(int(r["ctx_len"]) for r in rows)) depths = sorted(set(float(r["depth"]) for r in rows)) mat = np.full((len(depths), len(ctx_lens)), np.nan) for r in rows: ci = ctx_lens.index(int(r["ctx_len"])) di = depths.index(float(r["depth"])) acc = r.get("accuracy") if acc is not None and not (isinstance(acc, float) and math.isnan(acc)): mat[di, ci] = float(acc) fig, ax = plt.subplots(figsize=(COL1, 2.8)) im = ax.imshow(mat, aspect="auto", cmap="Greys", vmin=0, vmax=1, origin="lower", interpolation="nearest") plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04, label="Accuracy") ax.set_xticks(range(len(ctx_lens))) ax.set_xticklabels( [f"{c//1024}K" if c >= 1024 else str(c) for c in ctx_lens], rotation=45, ha="right") ax.set_yticks(range(len(depths))) ax.set_yticklabels([f"{int(d*100)}%" for d in depths]) ax.set_xlabel("Context length") ax.set_ylabel("Needle depth") ax.set_title("Passkey Retrieval (FELA 1.13B)") save("needle_heatmap") # --------------------------------------------------------------------------- # Fig 5 -- Chunk-size ablation: 3-panel # --------------------------------------------------------------------------- def fig_chunk(): chunks = [64, 256, 512, 1024] # BPB from per-variant wikitext.json bpb_vals = [] for c in chunks: wt = json.load(open(DATA / f"results/paper_eval_22b/chunk_ablation/variant_{c}/wikitext.json")) # bpb field is bits/token; divide by 4 (avg bytes/BPE token) -> bits/byte raw = wt.get("bpb") or wt.get("val_bpb") or wt.get("wikitext_bpb") bpb_vals.append(float(raw) / 4.0 if raw is not None else float("nan")) # Retrieval from chunk_ablation_FIXED needle_fixed.json (4K ctx, depth 90%) retr_vals = [] for c in chunks: needle = json.load(open( DATA / f"results/chunk_ablation_fixed/variant_{c}/needle_fixed.json")) hit = [float(r["accuracy"]) for r in needle if int(r["ctx_len"]) == 4096 and abs(float(r["depth"]) - 0.9) < 0.01 and r.get("accuracy") is not None and not math.isnan(float(r["accuracy"]))] retr_vals.append(hit[0] if hit else float("nan")) # Throughput: max tok/s from per-variant throughput.json tput_vals = [] for c in chunks: tp = json.load(open(DATA / f"results/paper_eval_22b/chunk_ablation/variant_{c}/throughput.json")) vals = [] def _collect(obj): if isinstance(obj, dict): for v in obj.values(): _collect(v) elif isinstance(obj, list): for v in obj: _collect(v) elif isinstance(obj, (int, float)) and 1000 < float(obj) < 1e7: vals.append(float(obj)) _collect(tp) tput_vals.append(max(vals) / 1000 if vals else float("nan")) print(f" chunk BPB: {bpb_vals}") print(f" chunk retr: {retr_vals}") print(f" chunk tput: {tput_vals}") fig, axes = plt.subplots(1, 3, figsize=(COL2, 2.4)) xs = list(range(len(chunks))) xlabels = [str(c) for c in chunks] hatches = ["", "//", "xx", ".."] specs = [ (bpb_vals, "BPB (lower is better)", "WikiText-103 BPB"), (retr_vals, "Retrieval accuracy", "Retrieval (4K, 90% depth)"), (tput_vals, "Throughput (K tok/s)", "Max Throughput"), ] for ax, (vals, ylabel, title) in zip(axes, specs): # Only plot bars where value is not nan for xi, (v, h) in enumerate(zip(vals, hatches)): if not math.isnan(v): bar = ax.bar(xi, v, color="white", edgecolor="black", linewidth=0.8, hatch=h) ax.text(xi, v * 1.015, f"{v:.3f}" if v < 10 else f"{v:.0f}", ha="center", va="bottom", fontsize=7) ax.set_xticks(xs) ax.set_xticklabels(xlabels) ax.set_xlabel("GLA chunk size") ax.set_ylabel(ylabel) ax.set_title(title) ax.grid(True, axis="y", alpha=0.25) if not all(math.isnan(v) for v in vals): valid = [v for v in vals if not math.isnan(v)] pad = (max(valid) - min(valid)) * 0.15 or max(valid) * 0.1 ax.set_ylim(min(valid) * 0.92, max(valid) + pad * 3) fig.suptitle("GLA Chunk Size Ablation (109M, 1B tokens)", y=1.02) plt.tight_layout() save("chunk_ablation") # --------------------------------------------------------------------------- # Fig 6 -- FNO filter spectrum # --------------------------------------------------------------------------- def fig_fno_spectrum(): data = json.load(open(DATA / "results/interpretability/fno_filter_spectrum.json")) groups = [ ("Early (L0-2)", ["layer_0", "layer_1", "layer_2"]), ("Mid (L4-6)", ["layer_4", "layer_5", "layer_6"]), ("Late (L8-10)", ["layer_8", "layer_9", "layer_10"]), ] fig, axes = plt.subplots(1, 3, figsize=(COL2, 2.5), sharey=False) for ax, (title, lnames) in zip(axes, groups): for i, lname in enumerate(lnames): if lname not in data: continue mags = np.array(data[lname]["magnitudes"]) freqs = np.arange(len(mags)) / len(mags) lnum = lname.split("_")[1] ax.plot(freqs, mags, ls=LS[i], color=GR[i], label=f"L{lnum}", alpha=0.9) ax.set_yscale("log") ax.set_xlabel("Normalised frequency") if ax is axes[0]: ax.set_ylabel("|H(f)| (mean over channels)") ax.set_title(title) ax.legend(fontsize=7) ax.grid(True, alpha=0.2) fig.suptitle("FNO Filter Spectra: Coarse-to-Fine Hierarchy", y=1.02) plt.tight_layout() save("fno_spectrum") # --------------------------------------------------------------------------- # Fig 7 -- GLA gate hierarchy # --------------------------------------------------------------------------- def fig_gla_gates(): data = json.load(open(DATA / "results/interpretability/gla_gate_heatmap.json")) texts = list(data.keys()) layer_names = list(data[texts[0]].keys()) layer_nums = [int(ln.split("_")[1]) for ln in layer_names] fig, axes = plt.subplots(1, len(texts), figsize=(COL2, 2.4), sharey=True) if len(texts) == 1: axes = [axes] for ax, text in zip(axes, texts): for li, (lname, lnum) in enumerate(zip(layer_names, layer_nums)): gates = np.array(data[text][lname]) # [T, H] decay = np.exp(gates) mean_decay = decay.mean(axis=0) ax.plot(np.arange(len(mean_decay)), mean_decay, ls=LS[li], marker=MK[li], color=GR[li], label=f"L{lnum}", markersize=3) ax.set_xlabel("Head") ax.set_ylim(0, 1.05) short = text[:24] + "..." if len(text) > 24 else text ax.set_title(short, fontsize=7) ax.grid(True, alpha=0.2) ax.axhline(1.0, color="k", lw=0.4, ls=":") axes[0].set_ylabel("Mean gate decay (1=remember)") axes[0].legend(fontsize=7, title="Layer") fig.suptitle("GLA Gate Hierarchy: Deeper Layers Retain More Context", y=1.02) plt.tight_layout() save("gla_gates") # --------------------------------------------------------------------------- # Fig 8 -- Logit lens # --------------------------------------------------------------------------- def fig_logit_lens(): data = json.load(open(DATA / "results/explain/logit_lens.json")) texts = list(data.keys()) fig, ax = plt.subplots(figsize=(COL1, 2.8)) all_layer_ids = None for ti, text in enumerate(texts[:3]): layers_data = data[text]["layers"] # final token = top-1 at last layer final_token = layers_data[-1]["top_k"][0]["token"] layer_ids = [l["layer"] for l in layers_data] # 1 if top-1 token already matches final prediction, 0 otherwise matches = [1 if l["top_k"][0]["token"] == final_token else 0 for l in layers_data] short = text[:20] + "..." ax.plot(layer_ids, matches, LS[ti]+MK[ti], color=GR[ti], label=short, lw=1.5) all_layer_ids = layer_ids ax.set_xlabel("Layer") ax.set_ylabel("Top 1 token matches final prediction") ax.set_title("Logit Lens: When Does Prediction Commit?") ax.set_xticks(all_layer_ids) ax.set_yticks([0, 1]) ax.set_yticklabels(["No", "Yes"]) ax.set_ylim(-0.15, 1.25) ax.axvline(6, color="k", lw=0.6, ls=":", alpha=0.5) ax.text(6.15, 0.08, "all commit\nby layer 6", fontsize=7, color="gray") ax.legend(fontsize=7) ax.grid(True, alpha=0.2) save("logit_lens") # --------------------------------------------------------------------------- # Fig 9 -- Training curve # --------------------------------------------------------------------------- def fig_training(): rows = list(csv.DictReader(open(DATA / "results/gpu_v1/training_data.csv"))) steps = [int(r["step"]) for r in rows if r["step"]] loss = [float(r["loss"]) for r in rows if r["loss"]] val_steps = [int(r["step"]) for r in rows if r.get("val_bpb")] val_bpb = [float(r["val_bpb"]) for r in rows if r.get("val_bpb")] fig, ax1 = plt.subplots(figsize=(COL1, 2.5)) ax1.plot(steps, loss, color="k", lw=0.6, alpha=0.4) if len(loss) > 50: smooth = np.convolve(loss, np.ones(50)/50, mode="valid") ax1.plot(steps[49:], smooth, color="k", lw=1.5, label="Train loss") ax1.set_xlabel("Step") ax1.set_ylabel("Loss (nats)") ax1.set_title("Training Curve -- FELA 1.13B") ax1.grid(True, alpha=0.2) if val_bpb: ax2 = ax1.twinx() ax2.plot(val_steps, val_bpb, LS[1]+MK[1], color="k", markerfacecolor="white", ms=5, label="Val BPB") ax2.set_ylabel("Val BPB") lines1, labs1 = ax1.get_legend_handles_labels() lines2, labs2 = ax2.get_legend_handles_labels() ax1.legend(lines1+lines2, labs1+labs2, fontsize=7, loc="upper right") else: ax1.legend(fontsize=7) save("training_curve") # --------------------------------------------------------------------------- # Fig 10 -- Wall-clock throughput: FELA vs SDPA (measured head-to-head) # --------------------------------------------------------------------------- def fig_wallclock(): rows = list(csv.DictReader(open(DATA / "results/gpu_v1/gpu_crossover_v1.csv"))) fela_x, fela_y = [], [] sdpa_x, sdpa_y = [], [] for r in rows: T = int(r["seq_len"]) if r["fno_oom"] == "True" or not r["fno_gla_tok_s"]: pass else: fela_x.append(T); fela_y.append(float(r["fno_gla_tok_s"]) / 1000) if r["sdpa_oom"] == "True" or not r["sdpa_tok_s"]: pass else: sdpa_x.append(T); sdpa_y.append(float(r["sdpa_tok_s"]) / 1000) fig, ax = plt.subplots(figsize=(COL1, 2.6)) ax.plot(fela_x, fela_y, LS[0]+MK[0], color="k", label="FELA 1.13B") ax.plot(sdpa_x, sdpa_y, LS[1]+MK[1], color="k", markerfacecolor="white", label="GPT-2 XL SDPA") # Mark SDPA OOM point if sdpa_x: ax.axvline(sdpa_x[-1], color="k", lw=0.7, ls=":", alpha=0.5) ax.text(sdpa_x[-1]*1.05, min(sdpa_y)*1.1, "SDPA\nOOM", fontsize=7, color="gray") ax.set_xscale("log", base=2) ax.set_xlabel("Sequence length (tokens)") ax.set_ylabel("Throughput (K tok/s)") ax.set_title("Measured Throughput: FELA vs. SDPA") ax.xaxis.set_major_formatter(mticker.FuncFormatter(lambda x, _: f"{int(x):,}")) ax.legend() ax.grid(True, alpha=0.2) save("wallclock") # --------------------------------------------------------------------------- # Fig 11 -- Inference FLOPs # --------------------------------------------------------------------------- def fig_flops(): rows = list(csv.DictReader( open(DATA / "results/paper_eval_22b/gmacs.tsv"), delimiter="\t")) fela = [r for r in rows if "FELA" in r.get("model", "")] xs_f = [int(r["seq_len"]) for r in fela] ys_f = [float(r["gmacs"]) / 1000 for r in fela] # GMAC -> TFLOP approx sdpa_xs = sorted(set(xs_f + [512, 2048, 8192, 32768, 65536, 131072])) # GPT-2 XL: 48 layers, 25 heads, 64 head-dim; 2*N*L*H*D per-token inference sdpa_ys = [2 * T * 48 * 25 * 64 / 1e9 for T in sdpa_xs] fig, ax = plt.subplots(figsize=(COL1, 2.5)) ax.plot(xs_f, ys_f, LS[0]+MK[0], color="k", label="FELA (prefill)") ax.plot(sdpa_xs, sdpa_ys, LS[1]+MK[1], color="k", markerfacecolor="white", label="SDPA per-token") ax.set_xscale("log", base=2) ax.set_yscale("log") ax.set_xlabel("Sequence length (tokens)") ax.set_ylabel("GFLOPs") ax.set_title("Compute vs. Sequence Length") ax.xaxis.set_major_formatter(mticker.FuncFormatter(lambda x, _: f"{int(x):,}")) ax.legend() ax.grid(True, which="both", alpha=0.2) save("flops") # --------------------------------------------------------------------------- if __name__ == "__main__": target = sys.argv[1] if len(sys.argv) > 1 else None FIGS = [ ("vram_measured", fig_vram), ("longctx_bpb", fig_longctx), ("throughput", fig_throughput), ("wallclock", fig_wallclock), ("needle_heatmap", fig_needle), ("chunk_ablation", fig_chunk), ("fno_spectrum", fig_fno_spectrum), ("gla_gates", fig_gla_gates), ("logit_lens", fig_logit_lens), ("training_curve", fig_training), ("flops", fig_flops), ] for name, fn in FIGS: if target and target != name: continue print(f"[{name}]") try: fn() except Exception as e: import traceback print(f" ERROR: {e}") traceback.print_exc() print("done")