SPIRA v2: entropy-patch tokenization + hierarchical dopamine-gated spiking cortex + hippocampal one-shot memory with replay consolidation
032a961 verified | #!/usr/bin/env python3 | |
| """Generate all paper figures from real experiment artifacts.""" | |
| import json | |
| import sys | |
| import numpy as np | |
| import os as _os | |
| ROOT = _os.path.dirname(_os.path.dirname(_os.path.abspath(__file__))) | |
| ART = f"{ROOT}/artifacts" | |
| sys.path.insert(0, ROOT) | |
| import matplotlib | |
| matplotlib.use("Agg") | |
| import matplotlib.font_manager as fm | |
| fm.fontManager.addfont('/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf') | |
| import matplotlib.pyplot as plt | |
| plt.rcParams.update({ | |
| "font.family": "DejaVu Sans", | |
| "font.size": 9, | |
| "axes.spines.top": False, | |
| "axes.spines.right": False, | |
| "axes.grid": True, | |
| "grid.alpha": 0.25, | |
| "figure.dpi": 200, | |
| }) | |
| OUT = f"{ART}/figures" | |
| import os | |
| os.makedirs(OUT, exist_ok=True) | |
| metrics = json.load(open(f"{ART}/metrics.json")) | |
| ACCENT = "#2e7d5b" | |
| ACCENT2 = "#b8860b" | |
| GRAY = "#666666" | |
| # ---------------------------------------------------------------- fig: EPT | |
| from spira import SPIRA | |
| model = SPIRA.load(f"{ART}/spira_model.npz") | |
| sample = (b"It was the best of times, it was the worst of times, it was the age " | |
| b"of wisdom, it was the age of foolishness, it was the epoch of belief, " | |
| b"it was the epoch of incredulity, it was the season of Light, " | |
| b"it was the season of Darkness") | |
| sur = model.ept.surprise(sample) | |
| patches = model.ept.segment(sample) | |
| fig, axes = plt.subplots(2, 1, figsize=(7.0, 3.2), height_ratios=[1.2, 1], | |
| constrained_layout=True, sharex=True) | |
| axes[0].fill_between(range(len(sur)), sur, color=ACCENT, alpha=0.55, step="mid") | |
| axes[0].plot(range(len(sur)), sur, color=ACCENT, lw=0.8) | |
| axes[0].set_ylabel("surprisal (bits)") | |
| axes[0].set_title("EPT: per-byte surprisal and constant-information patch boundaries", fontsize=9) | |
| b = model.ept.boundaries(sample) | |
| for bx in b[:-1]: | |
| axes[0].axvline(bx, color=ACCENT2, lw=0.7, alpha=0.85) | |
| # patch length bars | |
| axes[1].bar(b[:-1], [b[i+1]-b[i] for i in range(len(b)-1)], width=1.0, | |
| color=ACCENT, alpha=0.7, align="edge") | |
| axes[1].axhline(float(np.mean([b[i+1]-b[i] for i in range(len(b)-1)])), | |
| color=ACCENT2, lw=1, ls="--", label=f"mean {np.mean([b[i+1]-b[i] for i in range(len(b)-1)]):.1f} B") | |
| axes[1].set_xlabel("byte position") | |
| axes[1].set_ylabel("patch length (B)") | |
| axes[1].legend(frameon=False, fontsize=8) | |
| fig.savefig(f"{OUT}/fig_ept.pdf") | |
| plt.close(fig) | |
| print("fig_ept done:", len(patches), "patches, avg", np.mean([len(p) for p in patches])) | |
| # ------------------------------------------------------- fig: compression | |
| comp = metrics["compression"] | |
| methods = ["gzip -9", "BPE-4k\n(unigram)", "BPE-4k\n(bigram)", "BPE-16k\n(bigram)", "EPT\n(bigram)", "EPT\n(unigram)"] | |
| vals = [comp["gzip"]["bits_per_byte"], | |
| comp["bpe_4096"]["bits_per_byte_unigram"], | |
| comp["bpe_4096"]["bits_per_byte_bigram"], | |
| comp["bpe_16384"]["bits_per_byte_bigram"], | |
| comp["ept"]["bits_per_byte_bigram"], | |
| comp["ept"]["bits_per_byte_unigram"]] | |
| colors = [GRAY, "#9aa5b1", "#9aa5b1", "#9aa5b1", ACCENT, ACCENT] | |
| fig, ax = plt.subplots(figsize=(4.6, 2.9), constrained_layout=True) | |
| bars = ax.bar(range(len(vals)), vals, color=colors, width=0.65) | |
| for i, v in enumerate(vals): | |
| ax.text(i, v + 0.08, f"{v:.2f}", ha="center", fontsize=8) | |
| ax.set_xticks(range(len(methods))) | |
| ax.set_xticklabels(methods, fontsize=7.5) | |
| ax.set_ylabel("bits per byte (with dictionary)") | |
| ax.set_title("Coded size on the held-out test stream", fontsize=9) | |
| fig.savefig(f"{OUT}/fig_compression.pdf") | |
| plt.close(fig) | |
| print("fig_compression done") | |
| # -------------------------------------------------- fig: one-shot+corrupt | |
| one = metrics["one_shot"] | |
| corr = metrics["correction"] | |
| corr_sp = metrics.get("correction_sparse_am", {}) | |
| ks = [0, 1, 2, 3] | |
| am_curve = [corr[f"corrupt_{k}_of_4"]["am_acc"] for k in ks] | |
| fig, ax = plt.subplots(figsize=(4.6, 2.9), constrained_layout=True) | |
| ax.plot(ks, [v * 100 for v in am_curve], "o-", color=ACCENT, lw=1.6, | |
| label="full hippocampus (46% load)") | |
| if corr_sp: | |
| am_sp = [corr_sp[f"corrupt_{k}_of_4"]["am_acc"] for k in ks] | |
| ax.plot(ks, [v * 100 for v in am_sp], "s--", color=ACCENT2, lw=1.4, | |
| label="sparse hippocampus (book only)") | |
| for k, v in zip(ks, am_sp): | |
| ax.annotate(f"{v*100:.1f}%", (k, v * 100), textcoords="offset points", | |
| xytext=(0, 7), ha="center", fontsize=8, color=ACCENT2) | |
| for k, v in zip(ks, am_curve): | |
| ax.annotate(f"{v*100:.1f}%", (k, v * 100), textcoords="offset points", | |
| xytext=(0, -13), ha="center", fontsize=8) | |
| ax.set_xlabel("corrupted context slots (of 4)") | |
| ax.set_ylabel("exact recall (%)") | |
| ax.set_ylim(-14, 110) | |
| ax.set_xticks(ks) | |
| ax.set_title("Pattern completion under context corruption\n(completion trades against hippocampal load)", fontsize=9) | |
| ax.legend(frameon=False, fontsize=8, loc="center right") | |
| fig.savefig(f"{OUT}/fig_correction.pdf") | |
| plt.close(fig) | |
| print("fig_correction done:", am_curve) | |
| # ------------------------------------------------- fig: vocab coverage | |
| vc = metrics["vocab_coverage_curve"] | |
| vs = [c["vocab"] for c in vc] | |
| ct = [c["cov_test"] * 100 for c in vc] | |
| co = [c["cov_oneshot"] * 100 for c in vc] | |
| fig, ax = plt.subplots(figsize=(4.6, 2.9), constrained_layout=True) | |
| ax.plot(vs, ct, "o-", color=ACCENT, lw=1.6, ms=3.5, label="held-out test") | |
| ax.plot(vs, co, "s-", color=ACCENT2, lw=1.4, ms=3.5, label="unseen book (one-shot)") | |
| ax.axvline(4096, color=GRAY, lw=0.9, ls=":") | |
| ax.text(4096 * 1.06, 88, "v1: 4096", fontsize=7.5, color=GRAY) | |
| ax.axvline(16384, color=GRAY, lw=0.9, ls=":") | |
| ax.text(16384 * 1.06, 88, "v2: 16384", fontsize=7.5, color=GRAY) | |
| ax.annotate(f"{ct[-1]:.1f}%", (vs[-1], ct[-1]), textcoords="offset points", | |
| xytext=(2, 6), fontsize=8) | |
| ax.set_xscale("log", base=2) | |
| ax.set_xticks(vs) | |
| ax.set_xticklabels([str(v) for v in vs], fontsize=7.5) | |
| ax.set_xlabel("readout vocabulary size (top-V by train frequency)") | |
| ax.set_ylabel("held-out target coverage (%)") | |
| ax.set_title("Zipf knee of the readout vocabulary\n(constant-information patches flatten the tail)", fontsize=9) | |
| ax.legend(frameon=False, fontsize=8, loc="lower right") | |
| fig.savefig(f"{OUT}/fig_vocab.pdf") | |
| plt.close(fig) | |
| print("fig_vocab done") | |
| # ------------------------------------------------- fig: replay consolidation | |
| rep = metrics["replay"] | |
| xs = [0] + [c["replayed_tokens"] for c in rep["curve"]] | |
| ys = [rep["before"]["hier_acc"] * 100] + [c["hier_acc"] * 100 for c in rep["curve"]] | |
| ysi = [rep["before"]["hier_acc_invocab"] * 100] + [c["hier_acc_invocab"] * 100 for c in rep["curve"]] | |
| fig, ax = plt.subplots(figsize=(4.9, 3.0), constrained_layout=True) | |
| ax.plot(xs, ys, "o-", color=ACCENT, lw=1.6, ms=3.5, label="cortex-only recall") | |
| ax.plot(xs, ysi, "s--", color=ACCENT2, lw=1.4, ms=3.2, label="cortex-only, in-vocab only") | |
| ax.axhline(rep["am_recall_after_replay"]["am_acc"] * 100, color=GRAY, lw=1.1, ls=":") | |
| ax.text(xs[-1] * 0.06, rep["am_recall_after_replay"]["am_acc"] * 100 - 4.5, | |
| f"hippocampus (unchanged): {rep['am_recall_after_replay']['am_acc']*100:.1f}%", | |
| fontsize=7.5, color=GRAY) | |
| ax.axhline(rep["after"]["hier_acc"] * 100, color=ACCENT, lw=0.8, ls="--", alpha=0.5) | |
| ax.annotate(f"full eval after replay: {rep['after']['hier_acc']*100:.1f}%\n" | |
| f"(in-vocab {rep['after']['hier_acc_invocab']*100:.1f}%)", | |
| (xs[-1], rep["after"]["hier_acc"] * 100), textcoords="offset points", | |
| xytext=(-96, 6), fontsize=7.5, color=ACCENT) | |
| ax.set_xlabel("hippocampal tokens replayed through the cortex") | |
| ax.set_ylabel("exact next-patch recall (%)") | |
| ax.set_title("Systems consolidation: the cortex learns the book\nonly through hippocampal replay (never seen awake)", fontsize=9) | |
| ax.legend(frameon=False, fontsize=8, loc="center right") | |
| fig.savefig(f"{OUT}/fig_replay.pdf") | |
| plt.close(fig) | |
| print("fig_replay done") | |
| # ------------------------------------------------------------ fig: raster | |
| z = np.load(f"{ART}/spira_model.npz.raster.npz") | |
| ticks, neurons = z["ticks"], z["neurons"] | |
| fig, ax = plt.subplots(figsize=(6.8, 2.7), constrained_layout=True) | |
| ax.scatter(ticks, neurons, s=0.55, c=ACCENT, linewidths=0) | |
| ax.axhline(512 * 0.8, color=ACCENT2, lw=0.8, ls="--") | |
| ax.text(5, 512 * 0.8 + 6, "inhibitory population (Dale's law)", fontsize=7, color=ACCENT2) | |
| ax.set_xlabel("tick (token)") | |
| ax.set_ylabel("reservoir neuron") | |
| ax.set_title(f"Spike raster — {len(ticks)} spikes over 1500 ticks " | |
| f"({len(ticks)/1500:.0f} spikes/tick)", fontsize=9) | |
| fig.savefig(f"{OUT}/fig_raster.pdf") | |
| plt.close(fig) | |
| print("fig_raster done") | |
| # ------------------------------------------------- fig: learning curve | |
| import glob | |
| logs = [] | |
| for line in open(f"{ART}/chunks.jsonl"): | |
| try: | |
| logs.append(json.loads(line)) | |
| except Exception: | |
| pass | |
| # training acc EMAs were logged inside msg lines; rebuild a coarse curve from | |
| # chunk results recorded in the metrics/history + chunks.jsonl messages | |
| accs = [] | |
| for e in logs: | |
| if "acc=" in e.get("msg", ""): | |
| try: | |
| accs.append(float(e["msg"].split("acc=")[1].split()[0])) | |
| except Exception: | |
| pass | |
| fig, ax = plt.subplots(figsize=(4.6, 2.9), constrained_layout=True) | |
| if accs: | |
| xs = np.arange(len(accs)) * 25000 | |
| ax.plot(xs, [a * 100 for a in accs], "-", color=ACCENT, lw=1.4) | |
| ax.set_xlabel("tick (token)") | |
| ax.set_ylabel("online next-patch accuracy (%)") | |
| ax.set_title("Online learning during training (vocab-masked, incl. UNK targets)", fontsize=9) | |
| fig.savefig(f"{OUT}/fig_learning.pdf") | |
| plt.close(fig) | |
| print("fig_learning done:", len(accs), "points") | |