diff options
Diffstat (limited to 'experiments')
| -rw-r--r-- | experiments/plot_main_figures.py | 338 |
1 files changed, 338 insertions, 0 deletions
diff --git a/experiments/plot_main_figures.py b/experiments/plot_main_figures.py new file mode 100644 index 0000000..9306d90 --- /dev/null +++ b/experiments/plot_main_figures.py @@ -0,0 +1,338 @@ +#!/usr/bin/env python3 +"""Generate the two audited, publication-facing SDIL figures. + +Strict mode refuses to render a final figure unless every required cell has +five clean-provenance seeds. Use ``--allow-incomplete`` only for layout +previews while long experiments are still running. +""" +import argparse +import glob +import hashlib +import json +import math +import os +import statistics + +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +from matplotlib.ticker import FuncFormatter + + +COLORS = { + "BP": "#4D4D4D", + "FA": "#009E73", + "DFA": "#D55E00", + "SDIL": "#0072B2", + "EP": "#CC79A7", +} +MARKERS = {"BP": "*", "FA": "s", "DFA": "^", "SDIL": "o", "EP": "D"} +DEPTHS = [5, 10, 20, 30, 60] + + +def read_many(patterns): + paths = sorted({p for pattern in patterns for p in glob.glob(pattern)}) + rows = [] + for path in paths: + with open(path) as f: + row = json.load(f) + row["_path"] = path + rows.append(row) + return rows + + +def method(row): + args = row["args"] + return (args.get("mode") or args.get("method")).upper() + + +def mean_sd(values): + return (statistics.mean(values), + statistics.stdev(values) if len(values) > 1 else 0.0) + + +def require_clean(rows): + bad = [r["_path"] for r in rows + if r.get("provenance", {}).get("git_dirty") is not False] + if bad: + raise RuntimeError("dirty or unknown provenance:\n" + "\n".join(bad)) + + +def require_cells(groups, expected, n, label, strict): + missing = [] + for key in expected: + count = len(groups.get(key, [])) + if count != n: + missing.append(f"{key}: expected {n}, found {count}") + if missing and strict: + raise RuntimeError(f"incomplete {label}:\n" + "\n".join(missing)) + return missing + + +def scaling_records(root): + return read_many([ + os.path.join(root, "scale_v3_cifar10_*.json"), + os.path.join(root, "fa_scale_v2_cifar10_*.json"), + ]) + + +def scaling_groups(rows): + groups = {} + for row in rows: + groups.setdefault((method(row), row["args"]["depth"]), []).append(row) + return groups + + +def ep_groups(root): + rows = read_many([ + os.path.join(root, "ep_original_v2_*.json"), + os.path.join(root, "ep_archmatch_v1_*.json"), + os.path.join(root, "ep_depth2_v1_*.json"), + ]) + groups = {} + for row in rows: + args = row["args"] + # The d3 near-parameter runs are intentionally excluded: this panel is + # exact forward-architecture comparison only. + if args["depth"] in (1, 2) and args["width"] == 500: + groups.setdefault((method(row), args["depth"]), []).append(row) + return rows, groups + + +def nondominated(points): + """Return increasing-time points that improve accuracy.""" + front, best = [], -math.inf + for point in sorted(points, key=lambda p: (p[0], -p[1])): + if point[1] > best: + front.append(point) + best = point[1] + return front + + +def setup_style(): + plt.rcParams.update({ + "font.family": "DejaVu Sans", + "font.size": 9, + "axes.titlesize": 10.5, + "axes.labelsize": 9.5, + "axes.linewidth": 0.8, + "axes.spines.top": False, + "axes.spines.right": False, + "xtick.direction": "out", + "ytick.direction": "out", + "legend.frameon": False, + "pdf.fonttype": 42, + "ps.fonttype": 42, + "savefig.bbox": "tight", + }) + + +def error_point(ax, x, y, xerr, yerr, name, marker=None, alpha=1.0, + label=None, zorder=3): + ax.errorbar( + x, y, xerr=xerr, yerr=yerr, fmt=marker or MARKERS[name], + color=COLORS[name], markerfacecolor=COLORS[name], + markeredgecolor="white", markeredgewidth=0.65, markersize=7, + elinewidth=1.0, capsize=2.2, alpha=alpha, + label=label, zorder=zorder) + + +def plot_tradeoff(root, outdir, strict): + scale = scaling_records(root) + ep_rows, ep = ep_groups(root) + require_clean(scale + ep_rows) + sg = scaling_groups(scale) + miss_scale = require_cells( + sg, [(m, d) for m in ("BP", "FA", "DFA", "SDIL") for d in DEPTHS], + 5, "CIFAR scaling", strict) + miss_ep = require_cells( + ep, [(m, d) for m in ("BP", "DFA", "SDIL", "EP") for d in (1, 2)], + 5, "exact EP comparison", strict) + + fig, axes = plt.subplots(1, 2, figsize=(7.35, 3.15), + gridspec_kw={"wspace": 0.30}) + ax = axes[0] + frontier_candidates = [] + for name in ("BP", "FA", "DFA", "SDIL"): + first = True + for depth in DEPTHS: + rows = sg.get((name, depth), []) + if not rows: + continue + acc, acc_sd = mean_sd([100 * r["final"]["test_acc"] for r in rows]) + wall, wall_sd = mean_sd([r["final"]["wall_s"] for r in rows]) + error_point(ax, wall, acc, wall_sd, acc_sd, name, + label=name if first else None, + alpha=0.55 if name == "BP" else 0.92, + zorder=2 if name == "BP" else 3) + ax.annotate(f"d{depth}", (wall, acc), xytext=(3, 3), + textcoords="offset points", fontsize=6.7, + color=COLORS[name], alpha=0.85) + if name != "BP": + frontier_candidates.append((wall, acc, name, depth)) + first = False + front = nondominated(frontier_candidates) + if len(front) > 1: + ax.plot([p[0] for p in front], [p[1] for p in front], + color="#222222", linewidth=1.15, alpha=0.75, zorder=1) + ax.fill_between([p[0] for p in front], [p[1] for p in front], + min(p[1] for p in front) - 0.5, + color="#BDBDBD", alpha=0.10, zorder=0) + ax.set_xscale("log") + ax.xaxis.set_major_formatter(FuncFormatter(lambda value, _: f"{value:g}")) + ax.set_xlabel("Wall time (s, log scale)") + ax.set_ylabel("CIFAR-10 test accuracy (%)") + ax.set_title("a Local-learning Pareto frontier", loc="left", fontweight="bold") + ax.grid(True, which="major", color="#D9D9D9", linewidth=0.55, alpha=0.65) + ax.legend(ncol=2, loc="lower right", handletextpad=0.4, columnspacing=0.8) + + ax = axes[1] + for name in ("EP", "SDIL"): + xs, ys = [], [] + for depth in (1, 2): + rows = ep.get((name, depth), []) + if not rows: + continue + acc, acc_sd = mean_sd([100 * r["final"]["test_acc"] for r in rows]) + wall, wall_sd = mean_sd([r["final"]["wall_s"] for r in rows]) + error_point(ax, wall, acc, wall_sd, acc_sd, name, + marker="D" if depth == 1 else "o", + label=f"{name}, d{depth}") + xs.append(wall); ys.append(acc) + if len(xs) == 2: + ax.plot(xs, ys, color=COLORS[name], linewidth=1.0, alpha=0.55) + ax.set_xscale("log") + ax.xaxis.set_major_formatter(FuncFormatter(lambda value, _: f"{value:g}")) + ax.set_xlabel("Wall time (s, log scale)") + ax.set_ylabel("MNIST test accuracy (%)") + ax.set_title("b Against canonical EP (exact architecture)", + loc="left", fontweight="bold") + ax.grid(True, which="major", color="#D9D9D9", linewidth=0.55, alpha=0.65) + ax.legend(loc="best", fontsize=8) + fig.savefig(os.path.join(outdir, "figure1_pareto.pdf")) + fig.savefig(os.path.join(outdir, "figure1_pareto.png"), dpi=320) + plt.close(fig) + return scale + ep_rows, miss_scale + miss_ep + + +def early_alignment(row): + values = (row.get("final", {}).get("cos_r_negg") + or row.get("final", {}).get("cos_fa_negg")) + if not values: + return None + n = max(1, len(values) // 3) + finite = [v for v in values[:n] if math.isfinite(v)] + return statistics.mean(finite) if finite else None + + +def plot_scaling(root, outdir, strict): + rows = scaling_records(root) + require_clean(rows) + groups = scaling_groups(rows) + missing = require_cells( + groups, [(m, d) for m in ("BP", "FA", "DFA", "SDIL") for d in DEPTHS], + 5, "CIFAR scaling", strict) + fig, axes = plt.subplots(1, 3, figsize=(7.35, 2.75), + gridspec_kw={"wspace": 0.38}) + + ax = axes[0] + for name in ("BP", "FA", "DFA", "SDIL"): + xs, ys, es = [], [], [] + for depth in DEPTHS: + rs = groups.get((name, depth), []) + if not rs: + continue + mu, sd = mean_sd([100 * r["final"]["test_acc"] for r in rs]) + xs.append(depth); ys.append(mu); es.append(sd) + ax.errorbar(xs, ys, yerr=es, color=COLORS[name], marker=MARKERS[name], + markersize=5.5, linewidth=1.4, capsize=2.2, label=name, + markeredgecolor="white", markeredgewidth=0.5) + ax.set_xlabel("Hidden depth") + ax.set_ylabel("Test accuracy (%)") + ax.set_title("a Accuracy", loc="left", fontweight="bold") + ax.grid(True, color="#D9D9D9", linewidth=0.55, alpha=0.65) + ax.legend(fontsize=7.5, ncol=2, handlelength=1.5, + columnspacing=0.7, handletextpad=0.35) + + ax = axes[1] + bp_by_seed = {(r["args"]["depth"], r["args"]["seed"]): r + for r in rows if method(r) == "BP"} + for name in ("FA", "DFA", "SDIL"): + xs, ys, es = [], [], [] + for depth in DEPTHS: + gaps = [] + for r in groups.get((name, depth), []): + bp = bp_by_seed.get((depth, r["args"]["seed"])) + if bp: + gaps.append(100 * (bp["final"]["test_acc"] + - r["final"]["test_acc"])) + if gaps: + mu, sd = mean_sd(gaps) + xs.append(depth); ys.append(mu); es.append(sd) + ax.errorbar(xs, ys, yerr=es, color=COLORS[name], marker=MARKERS[name], + markersize=5.5, linewidth=1.4, capsize=2.2, label=name, + markeredgecolor="white", markeredgewidth=0.5) + ax.axhline(0, color="#777777", linewidth=0.8, linestyle="--") + ax.set_xlabel("Hidden depth") + ax.set_ylabel("Gap to paired BP (points)") + ax.set_title("b Optimization gap", loc="left", fontweight="bold") + ax.grid(True, color="#D9D9D9", linewidth=0.55, alpha=0.65) + + ax = axes[2] + for name in ("FA", "DFA", "SDIL"): + xs, ys, es = [], [], [] + for depth in DEPTHS: + vals = [v for r in groups.get((name, depth), []) + if (v := early_alignment(r)) is not None] + if vals: + mu, sd = mean_sd(vals) + xs.append(depth); ys.append(mu); es.append(sd) + ax.errorbar(xs, ys, yerr=es, color=COLORS[name], marker=MARKERS[name], + markersize=5.5, linewidth=1.4, capsize=2.2, label=name, + markeredgecolor="white", markeredgewidth=0.5) + ax.axhline(0, color="#777777", linewidth=0.8, linestyle="--") + ax.set_xlabel("Hidden depth") + ax.set_ylabel("Early-third cos(signal, −gradient)") + ax.set_title("c Credit alignment", loc="left", fontweight="bold") + ax.grid(True, color="#D9D9D9", linewidth=0.55, alpha=0.65) + fig.savefig(os.path.join(outdir, "figure2_scaling.pdf")) + fig.savefig(os.path.join(outdir, "figure2_scaling.png"), dpi=320) + plt.close(fig) + return rows, missing + + +def file_hash(path): + h = hashlib.sha256() + with open(path, "rb") as f: + for chunk in iter(lambda: f.read(1024 * 1024), b""): + h.update(chunk) + return h.hexdigest() + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--results", default="results") + parser.add_argument("--outdir", default="results/figs") + parser.add_argument("--allow-incomplete", action="store_true") + args = parser.parse_args() + strict = not args.allow_incomplete + os.makedirs(args.outdir, exist_ok=True) + setup_style() + rows1, missing1 = plot_tradeoff(args.results, args.outdir, strict) + rows2, missing2 = plot_scaling(args.results, args.outdir, strict) + source_paths = sorted({r["_path"] for r in rows1 + rows2}) + manifest = { + "strict": strict, + "missing_cells": missing1 + missing2, + "sources": [{"path": p, "sha256": file_hash(p)} for p in source_paths], + "outputs": ["figure1_pareto.pdf", "figure1_pareto.png", + "figure2_scaling.pdf", "figure2_scaling.png"], + } + with open(os.path.join(args.outdir, "main_figure_manifest.json"), "w") as f: + json.dump(manifest, f, indent=2) + print(f"rendered figures from {len(source_paths)} audited result files; " + f"missing cells={len(manifest['missing_cells'])}") + + +if __name__ == "__main__": + main() |