"""Three figures from tables the card shipped but never drew. 20 where the parameters actually live, and why the vocabulary was cut after the layers 21 the distillation ladder's own losses: the difficulty is not monotone 22 the threshold the protocol selects moves between seeds, and every arm over-predicts Run from benchmark/figures/ . Reads ../compression_ladder.csv and ../seed_metrics.csv, both of which ship in this repository. """ import csv import json import os import statistics as st import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt from matplotlib.lines import Line2D from matplotlib.patches import Patch HERE = os.path.dirname(os.path.abspath(__file__)) BENCH = os.path.join(HERE, "..") + os.sep # benchmark/ holds the CSVs this reads OUT = HERE + os.sep INK, MUTED, FAINT, GRID = "#1a1a1a", "#6a6a72", "#9a9aa2", "#e6e6ea" ROOTC, GRIDC, QWENC, ENCC, DOWN = "#1a5496", "#0d7a52", "#b05512", "#9a9aa2", "#c0392b" H = 1024 # hidden size, unchanged down the whole ladder def style(ax, xgrid=True, ygrid=True): ax.grid(axis="both" if (xgrid and ygrid) else ("x" if xgrid else "y"), color=GRID, lw=.8) ax.set_axisbelow(True) for s in ("top", "right"): ax.spines[s].set_visible(False) for s in ("bottom", "left"): ax.spines[s].set_color("#c9c9cf") ax.tick_params(length=0, colors="#4a4a52") def footer(fig, text): fig.text(.006, .012, text, fontsize=8.6, color=MUTED, ha="left") def fig_parameters(): """Two cuts in sequence: depth first, then vocabulary. Embedding share is the story.""" ladder = list(csv.DictReader(open(BENCH + "compression_ladder.csv"))) emb248 = 248320 * H depth = [(f"{r['layers']}L", int(r["parameters"]), emb248) for r in ladder] layer_at_4l = depth[-1][1] - emb248 vocab = [("248,320", 248320), ("128,000", 128000), ("72,455", 72455), ("39,866", 39866), ("23,551", 23551), ("15,380", 15380)] vocab = [(lab, v * H + layer_at_4l + 1024, v * H) for lab, v in vocab] fig, (ax, ax2) = plt.subplots(1, 2, figsize=(14.6, 6.4)) fig.patch.set_facecolor("white") for axis, data, xlab, title, sub in ( (ax, depth, "layers kept (vocabulary held at 248,320)", "Cut 1 — depth: 24 layers to 4", "the layer stack loses 84% of its parameters and the embedding does not move"), (ax2, vocab, "vocabulary entries kept (depth held at 4 layers)", "Cut 2 — vocabulary: 248,320 entries to 39,866", "by then the embedding was three quarters of the model \u2014 which is why it came second")): labs = [d[0] for d in data] tot = [d[1] / 1e6 for d in data] embm = [d[2] / 1e6 for d in data] lays = [t - e for t, e in zip(tot, embm)] axis.bar(labs, embm, width=.62, color=ROOTC, zorder=3, label="embedding table") axis.bar(labs, lays, bottom=embm, width=.62, color=ENCC, zorder=3, label="layer stack") for i, (t, e) in enumerate(zip(tot, embm)): axis.text(i, t + max(tot) * .035, f"{t:,.0f}M", ha="center", fontsize=9.2, color=INK, fontweight="bold") axis.text(i, e / 2, f"{e/t:.0%}", ha="center", va="center", fontsize=9, color="white", fontweight="bold") axis.set_xlabel(xlab, fontsize=10) axis.set_ylabel("parameters (millions)", fontsize=10) axis.set_ylim(0, max(tot) * 1.24) style(axis, xgrid=False) axis.set_title(f"{title}\n{sub}", fontsize=11.5, loc="left", color=INK, pad=12) # mark the two endpoints that actually shipped ax.get_xticklabels()[-1].set_color(ROOTC) ax.get_xticklabels()[-1].set_fontweight("bold") ax2.get_xticklabels()[3].set_color(ROOTC) ax2.get_xticklabels()[3].set_fontweight("bold") ax2.annotate("this root", (3, vocab[3][1] / 1e6), textcoords="offset points", xytext=(0, 52), ha="center", fontsize=9.6, color=ROOTC, fontweight="bold", arrowprops=dict(arrowstyle="-|>", color=ROOTC, lw=1.3, shrinkA=2, shrinkB=22)) ax.legend(handles=[Patch(color=ROOTC, label="embedding table"), Patch(color=ENCC, label="layer stack")], fontsize=9, frameon=False, loc="upper right") footer(fig, "Percentages are the embedding's share of the model. Cutting depth alone raised that " "share from 33.8% to 76.1%; cutting the vocabulary brought it back to 33.8%, the " "same balance the 24-layer teacher had. Hidden size is 1,024 throughout.") fig.tight_layout(rect=(0, .055, 1, 1)) fig.savefig(OUT + "20_where_the_parameters_live.png", dpi=170, facecolor="white") plt.close(fig) def fig_ladder_loss(): rows = [r for r in csv.DictReader(open(BENCH + "compression_ladder.csv")) if r["kd_stage"]] stages = [r["kd_stage"].replace("to", " → ") for r in rows] tot = [float(r["mean_total_loss"]) for r in rows] itf = [float(r["mean_interface_loss"]) for r in rows] fin = [float(r["mean_final_loss"]) for r in rows] dropped = [16, 2, 2] fig, (ax, ax2) = plt.subplots(1, 2, figsize=(13.8, 5.8), gridspec_kw={"width_ratios": [1.12, 1]}) fig.patch.set_facecolor("white") x = range(len(stages)) w = .26 for off, vals, c, lab in ((-w, itf, ROOTC, "interface loss"), (0.0, tot, GRIDC, "total loss"), (w, fin, QWENC, "final-layer loss")): ax.bar([i + off for i in x], vals, width=w * .92, color=c, zorder=3, label=lab) for i, v in zip(x, vals): ax.text(i + off, v + .004, f"{v:.3f}", ha="center", fontsize=8.4, color=INK, fontweight="bold") ax.set_xticks(list(x)) ax.set_xticklabels([f"{s}\n{d} layers removed" for s, d in zip(stages, dropped)], fontsize=9.4) ax.set_ylabel("mean KD loss over the distillation run", fontsize=10) ax.set_ylim(0, max(itf) * 1.22) style(ax, xgrid=False) ax.legend(fontsize=9, frameon=False, loc="upper right", ncol=1) ax.set_title("The ladder's own losses are not monotone\n" "24\u21928 drops 16 layers for 0.176; 8\u21926 drops 2 for 0.034; then 6\u21924 " "drops 2 for 0.047", fontsize=11.5, loc="left", color=INK, pad=12) # per-layer-removed difficulty per = [t / d for t, d in zip(tot, dropped)] ax2.plot(list(x), per, "-o", color=ROOTC, lw=2.4, ms=11, mec="white", mew=1.7, zorder=4) for i, v in zip(x, per): ax2.annotate(f"{v:.4f}", (i, v), textcoords="offset points", xytext=(0, 14), ha="center", fontsize=9.2, color=INK, fontweight="bold") ax2.set_xticks(list(x)) ax2.set_xticklabels(stages, fontsize=9.6) ax2.set_xlim(-.45, len(stages) - .55) ax2.set_ylabel("total loss per layer removed", fontsize=10) ax2.set_ylim(0, max(per) * 1.30) style(ax2, xgrid=False) ax2.set_title("Per layer removed, each step is harder than the last\n" "the final two layers cost 2.1x per layer what the first sixteen did", fontsize=11.5, loc="left", color=INK, pad=12) footer(fig, "Interface loss matches hidden states at the kept-layer boundaries; final loss matches " "the last hidden state. No logits and no labels enter either. One run per stage, so " "there is no spread to report and no significance is claimed.") fig.tight_layout(rect=(0, .055, 1, 1)) fig.savefig(OUT + "21_ladder_loss.png", dpi=170, facecolor="white") plt.close(fig) def fig_thresholds(): rows = list(csv.DictReader(open(BENCH + "seed_metrics.csv"))) NAMES = {"task_agnostic_base": "task-free 4L root", "structural_copy_control": "structural copy control", "existing_specialized_separate_lineage": "specialised 4L, separate lineage"} arms = list(dict.fromkeys(r["arm"] for r in rows)) actual = [(int(r["tp"]) + int(r["fn"])) / (55 * 14) for r in rows] assert max(actual) - min(actual) < 1e-9 actual = actual[0] fig, (ax, ax2) = plt.subplots(1, 2, figsize=(14.2, 5.8)) fig.patch.set_facecolor("white") colours = {a: c for a, c in zip(arms, (ROOTC, QWENC, GRIDC))} for a in arms: rs = [r for r in rows if r["arm"] == a] th = [float(r["threshold"]) for r in rs] pp = [float(r["predicted_positive_rate"]) for r in rs] ax.plot(th, pp, "o", ms=11, color=colours[a], mec="white", mew=1.6, zorder=4) for r, t, p in zip(rs, th, pp): ax.annotate(f"seed {r['seed']}", (t, p), textcoords="offset points", xytext=(11, -3), fontsize=8.2, color=MUTED) ax.axhline(actual, color=DOWN, ls=(0, (5, 3)), lw=1.6, zorder=3) ax.text(.328, actual + .006, f"actual positive rate {actual:.3f}", fontsize=8.8, color=DOWN, ha="right", va="bottom") ax.set_xlabel("threshold the frozen grid selected for that run", fontsize=10) ax.set_ylabel("fraction of (document, label) pairs predicted positive", fontsize=10) ax.set_xlim(.02, .335) ax.set_ylim(.335, .775) style(ax) ax.legend(handles=[Line2D([], [], marker="o", ls="", color=colours[a], ms=9, label=NAMES[a]) for a in arms], fontsize=8.8, frameon=False, loc="upper right") ax.set_title("Every run predicts more positives than the data has\n" "and the threshold the grid selects moves up to 4x between seeds of one arm", fontsize=11.5, loc="left", color=INK, pad=12) labels, tps, fps, fns, cols = [], [], [], [], [] for a in arms: for r in [r for r in rows if r["arm"] == a]: labels.append(f"{NAMES[a].split(',')[0]}\nseed {r['seed']} thr {r['threshold']}") tps.append(int(r["tp"])) fps.append(int(r["fp"])) fns.append(int(r["fn"])) cols.append(colours[a]) y = range(len(labels)) ax2.barh(list(y), tps, height=.6, color=ROOTC, zorder=3, label="true positives") ax2.barh(list(y), fps, left=tps, height=.6, color=QWENC, zorder=3, label="false positives") ax2.barh(list(y), fns, left=[t + f for t, f in zip(tps, fps)], height=.6, color=ENCC, zorder=3, label="false negatives") for i, (t, f) in enumerate(zip(tps, fps)): ax2.text(t + f + 12, i, f"{f/t:.2f} fp per tp", va="center", fontsize=8.4, color=DOWN if f > t else MUTED, fontweight="bold" if f > t else "normal") ax2.set_yticks(list(y)) ax2.set_yticklabels(labels, fontsize=7.8) ax2.set_xlabel("(document, label) pairs, out of 770", fontsize=10) ax2.set_xlim(0, 720) ax2.invert_yaxis() style(ax2, xgrid=True, ygrid=False) ax2.spines["left"].set_visible(False) ax2.legend(fontsize=8.6, frameon=False, loc="lower right", ncol=1) ax2.set_title("What that costs, run by run\n" "one control run returns more false positives than true ones", fontsize=11.5, loc="left", color=INK, pad=12) footer(fig, "Opened 56-document calibration split, 55 scored articles x 14 labels. If you train a " "head on this backbone, expect to pick the threshold yourself: the protocol's own grid " "lands on a different one for every seed.") fig.tight_layout(rect=(0, .055, 1, 1)) fig.savefig(OUT + "22_threshold_and_overprediction.png", dpi=170, facecolor="white") plt.close(fig) if __name__ == "__main__": os.makedirs(OUT, exist_ok=True) for fn in (fig_parameters, fig_ladder_loss, fig_thresholds): fn() print("ok", fn.__name__)