SPIRA v3: cortex representation repair — direct input pathway, multi-timescale trace banks, structured trigram inputs; controlled A/B (SNN alone 0.40% -> 2.07% in-vocab held-out vs bigram 4.94% at matched vocab 4096); paper, README and docs updated to third iteration
6923c51 verified | """SPIRA — Spiking Plastic Recurrent Architecture. | |
| Full system pipeline: | |
| raw bytes ──EPT──▶ content-addressed patches ──hashing trick──▶ input spikes | |
| │ │ | |
| │ LIF reservoir | |
| │ (STDP, Dale) | |
| │ │ | |
| │ three-factor | |
| │ readout (R-STDP) | |
| ▼ ▼ | |
| hippocampal associative memory ──fused──▶ prediction | |
| generation runs the same loop closed: predicted patch is fed back as | |
| input spikes, one patch (multi-byte) at a time. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import time | |
| from dataclasses import dataclass, field, asdict | |
| import numpy as np | |
| from .ept import EPT | |
| from .hashing import fnv1a64, splitmix64 | |
| from .hippocampus import Hippocampus | |
| from .spiking import SpikingCore | |
| UNK_COL = -1 # sentinel before adding the +1 offset | |
| class SPIRAConfig: | |
| # EPT | |
| target_avg_len: float = 4.5 | |
| min_len: int = 2 | |
| max_len: int = 32 | |
| alpha: float = 0.5 | |
| # spiking core | |
| d_in: int = 1024 | |
| k_in: int = 12 | |
| n_res: int = 512 | |
| vocab_size: int = 16384 | |
| seed: int = 20261022 | |
| # hierarchical multi-scale cortex (v2) | |
| hierarchy: bool = True | |
| n_g: int = 256 | |
| win: int = 8 | |
| eta_pos: float = 0.01 | |
| eta_r: float = 0.01 | |
| eta_stdp: float = 0.0015 | |
| # hippocampus | |
| am_size: int = 1 << 23 | |
| am_m: int = 4 | |
| am_min_lanes: int = 3 | |
| # training | |
| snn_tokens: int = 600_000 | |
| am_tokens: int = 1_200_000 | |
| epochs: int = 2 | |
| # hippocampal replay -> cortical consolidation | |
| replay_tokens: int = 180_000 | |
| replay_slice: int = 512 | |
| # v3 cortex fixes (Muse) | |
| direct_input: bool = True | |
| trace_taus: tuple = (0.4, 0.8, 0.95) | |
| direct_tau: float = 0.5 | |
| structured_inputs: bool = True | |
| k_tri: int = 6 | |
| # generation / fusion | |
| beta_am: float = 0.5 | |
| am_override: float = 6.0 | |
| temperature: float = 0.8 | |
| top_k: int = 8 | |
| class PatchRegistry: | |
| """Compact content-addressed patch storage. | |
| Sorted uint64 id array + offsets into a single concatenated byte buffer. | |
| Lookup is O(log n) via searchsorted (vectorizable). Also stores counts. | |
| """ | |
| def __init__(self) -> None: | |
| self.ids = np.zeros(0, dtype=np.uint64) # sorted unique | |
| self.counts = np.zeros(0, dtype=np.uint32) | |
| self.offsets = np.zeros(0, dtype=np.uint64) # len = n+1 | |
| self.buffer = b"" | |
| def build(cls, ids: np.ndarray, patches: list[bytes], keep: set[int] | None = None) -> "PatchRegistry": | |
| reg = cls() | |
| id2bytes: dict[int, bytes] = {} | |
| for i, b in zip(ids.tolist(), patches): | |
| if keep is not None and i not in keep: | |
| continue | |
| id2bytes[i] = b # last occurrence wins; identical anyway | |
| if not id2bytes: | |
| return reg | |
| reg.ids = np.fromiter(id2bytes.keys(), dtype=np.uint64, count=len(id2bytes)) | |
| reg.ids.sort() | |
| buf = bytearray() | |
| offsets = np.zeros(len(reg.ids) + 1, dtype=np.uint64) | |
| for j, pid in enumerate(reg.ids.tolist()): | |
| b = id2bytes[pid] | |
| offsets[j] = len(buf) | |
| buf.extend(b) | |
| offsets[-1] = len(buf) | |
| reg.offsets = offsets | |
| reg.buffer = bytes(buf) | |
| # counts from the stream | |
| uniq, cnt = np.unique(ids, return_counts=True) | |
| pos = np.searchsorted(reg.ids, uniq) | |
| pos_ok = (pos < len(reg.ids)) & (reg.ids[np.minimum(pos, len(reg.ids) - 1)] == uniq) | |
| reg.counts = np.zeros(len(reg.ids), dtype=np.uint32) | |
| reg.counts[pos[pos_ok]] = cnt[pos_ok].astype(np.uint32) | |
| return reg | |
| def __len__(self) -> int: | |
| return len(self.ids) | |
| def has(self, ids: np.ndarray | int) -> np.ndarray | bool: | |
| if len(self.ids) == 0: | |
| return np.zeros(len(ids), dtype=bool) if isinstance(ids, np.ndarray) else False | |
| if isinstance(ids, (int, np.integer)): | |
| pos = np.searchsorted(self.ids, np.uint64(ids)) | |
| return bool(pos < len(self.ids) and self.ids[min(pos, len(self.ids) - 1)] == ids) | |
| pos = np.searchsorted(self.ids, ids) | |
| pos_c = np.minimum(pos, len(self.ids) - 1) | |
| return (pos < len(self.ids)) & (self.ids[pos_c] == ids) | |
| def bytes_of(self, pid: int) -> bytes | None: | |
| pos = int(np.searchsorted(self.ids, np.uint64(pid))) | |
| if pos >= len(self.ids) or self.ids[pos] != np.uint64(pid): | |
| return None | |
| a, b = int(self.offsets[pos]), int(self.offsets[pos + 1]) | |
| return self.buffer[a:b] | |
| def index_of(self, pid: int) -> int: | |
| """Registry index of a patch id, or -1.""" | |
| pos = int(np.searchsorted(self.ids, np.uint64(pid))) | |
| if pos >= len(self.ids) or self.ids[pos] != np.uint64(pid): | |
| return -1 | |
| return pos | |
| def extend(self, id2bytes: dict[int, bytes]) -> int: | |
| """Append new patches (one-shot episodic learning). | |
| Returns the number of patches added. Keeps arrays sorted. | |
| """ | |
| if not id2bytes: | |
| return 0 | |
| cand = np.array(sorted(id2bytes.keys()), dtype=np.uint64) | |
| cand = cand[~self.has(cand)] | |
| if len(cand) == 0: | |
| return 0 | |
| n_old = len(self.ids) | |
| merged = np.concatenate([self.ids, cand]) | |
| order = np.argsort(merged, kind="stable") | |
| merged = merged[order] | |
| old_lens = np.diff(self.offsets).astype(np.uint64) | |
| new_lens = np.array([len(id2bytes[int(i)]) for i in cand], dtype=np.uint64) | |
| old_off = self.offsets[:-1] | |
| out = bytearray() | |
| offsets = np.zeros(len(merged) + 1, dtype=np.uint64) | |
| for k, idx in enumerate(order.tolist()): | |
| if idx < n_old: | |
| a, b = int(old_off[idx]), int(old_off[idx] + old_lens[idx]) | |
| out.extend(self.buffer[a:b]) | |
| else: | |
| out.extend(id2bytes[int(cand[idx - n_old])]) | |
| offsets[k] = len(out) | |
| offsets[-1] = len(out) | |
| self.ids = merged | |
| self.offsets = offsets | |
| self.buffer = bytes(out) | |
| self.counts = np.concatenate([self.counts, np.zeros(len(cand), dtype=np.uint32)])[order] | |
| return len(cand) | |
| class SPIRA: | |
| EPI_CAP_TOKENS = 2_000_000 # hippocampal episodic-trace ring cap | |
| def __init__(self, cfg: SPIRAConfig) -> None: | |
| self.cfg = cfg | |
| self.ept = EPT(budget_bits=16.0, min_len=cfg.min_len, max_len=cfg.max_len, alpha=cfg.alpha) | |
| self.core = SpikingCore( | |
| d_in=cfg.d_in, k_in=cfg.k_in, n_res=cfg.n_res, | |
| n_out=cfg.vocab_size + 1, seed=cfg.seed, | |
| hier=cfg.hierarchy, n_g=cfg.n_g, win=cfg.win, | |
| trace_taus=tuple(cfg.trace_taus), direct_input=cfg.direct_input, | |
| direct_tau=cfg.direct_tau, | |
| ) | |
| self._row_cache: dict[int, np.ndarray] = {} | |
| self.am = Hippocampus(m=cfg.am_m, size=cfg.am_size, min_lanes=cfg.am_min_lanes) | |
| self.registry = PatchRegistry() | |
| self.vocab_ids: np.ndarray = np.zeros(0, dtype=np.uint64) # col -> id | |
| self.n_out = cfg.vocab_size + 1 | |
| self.meta: dict = {} | |
| # inference-time fusion calibration (variance matching): the raw slow | |
| # readout operates at a different score scale than the fast one, so the | |
| # top-down prior is rescaled by sigma_fast/sigma_slow — measured once | |
| # on a reference stream (calibrate_fusion), never trained. | |
| self.fusion: dict = {"scale": 1.0, "sigma_fast": None, "sigma_slow": None} | |
| # episodic traces retained for hippocampal replay (CA3-style sequence | |
| # store; the AM itself only keeps hashed associations) | |
| self.episodes: list[np.ndarray] = [] | |
| # ================================================================ build | |
| def build(cls, train_bytes: bytes, cfg: SPIRAConfig, log=print) -> "SPIRA": | |
| m = cls(cfg) | |
| fit_slice = train_bytes[: 2_000_000] | |
| calib_slice = train_bytes[2_000_000: 2_400_000] | |
| log("[spira] fitting EPT byte-level Markov model ...") | |
| m.ept.fit(fit_slice) | |
| if calib_slice: | |
| budget = m.ept.calibrate(calib_slice, target_len=cfg.target_avg_len) | |
| log(f"[spira] calibrated surprisal budget = {budget:.2f} bits " | |
| f"(target avg patch len {cfg.target_avg_len})") | |
| log("[spira] segmenting full train stream ...") | |
| t0 = time.perf_counter() | |
| patches = m.ept.segment(train_bytes) | |
| ids = np.fromiter((fnv1a64(p) for p in patches), dtype=np.uint64, count=len(patches)) | |
| seg_s = time.perf_counter() - t0 | |
| st = m.ept.stats(patches) | |
| log(f"[spira] {st['n_patches']} patches, avg len {st['avg_len']:.2f} B " | |
| f"({st['bytes']/1e6:.1f} MB segmented in {seg_s:.1f}s, " | |
| f"{st['bytes']/1e6/max(seg_s,1e-9):.2f} MB/s)") | |
| uniq, cnt = np.unique(ids, return_counts=True) | |
| log(f"[spira] distinct patches in train: {len(uniq)}") | |
| # vocabulary: top-V by count (deterministic tie-break: larger id first) | |
| order = np.lexsort((uniq, -cnt))[: cfg.vocab_size] | |
| m.vocab_ids = np.sort(uniq[order]).astype(np.uint64) | |
| # registry: vocab + anything with count >= 2 (repeatable content) | |
| keep = set(m.vocab_ids.tolist()) | set(uniq[cnt >= 2].tolist()) | |
| log(f"[spira] building patch registry ({len(keep)} ids) ...") | |
| m.registry = PatchRegistry.build(ids, patches, keep=keep) | |
| m.meta = { | |
| "train_bytes": len(train_bytes), | |
| "n_patches": st["n_patches"], | |
| "avg_patch_len": st["avg_len"], | |
| "distinct_patches": int(len(uniq)), | |
| "registry_size": len(m.registry), | |
| "budget_bits": m.ept.budget_bits, | |
| "segmentation_MB_per_s": st["bytes"] / 1e6 / max(seg_s, 1e-9), | |
| } | |
| m._ids_train = ids | |
| m._patches_train = None # registry carries everything; free ~100MB | |
| return m | |
| # =========================================================== array prep | |
| def input_neurons_stream(ids: np.ndarray, d_in: int, k_in: int) -> np.ndarray: | |
| """Vectorized input-neuron projection: (T, k_in) int32.""" | |
| T = len(ids) | |
| out = np.empty((T, k_in), dtype=np.int32) | |
| u64 = ids.astype(np.uint64) | |
| for j in range(k_in): | |
| salt = (0x51ED270B733D4C31 + j * 0x9E3779B97F4A7C15) & 0xFFFFFFFFFFFFFFFF | |
| z = u64 ^ np.uint64(salt) | |
| z = z + np.uint64(0x9E3779B97F4A7C15) | |
| z = (z ^ (z >> np.uint64(30))) * np.uint64(0xBF58476D1CE4E5B9) | |
| z = (z ^ (z >> np.uint64(27))) * np.uint64(0x94D049BB133111EB) | |
| z = z ^ (z >> np.uint64(31)) | |
| out[:, j] = ((z >> np.uint64(32)) % np.uint64(d_in)).astype(np.int32) | |
| return out | |
| def cols_of(self, ids: np.ndarray) -> np.ndarray: | |
| """Vocabulary column per patch id (UNK -> last column).""" | |
| V = self.cfg.vocab_size | |
| pos = np.searchsorted(self.vocab_ids, ids) | |
| pos_c = np.minimum(pos, V - 1) | |
| ok = (pos < V) & (self.vocab_ids[pos_c] == ids) | |
| cols = np.where(ok, pos_c, V).astype(np.int64) # V = UNK column | |
| return cols | |
| # ------------------------------------------------------- v3 input rows | |
| def _structured_row(self, pid: int, pbytes: bytes | None) -> np.ndarray: | |
| """Input-neuron row for one patch: (k_in - k_tri) whole-patch hash | |
| neurons + k_tri byte-trigram hash neurons. Patches sharing trigrams | |
| share input neurons, so similar contexts produce similar reservoir | |
| states (v1/v2 rows were pure whole-patch hashes: fully orthogonal). | |
| Falls back to the v2 row when the patch bytes are unknown.""" | |
| from .hashing import fnv1a64, input_neurons, splitmix64 | |
| cfg = self.cfg | |
| if not cfg.structured_inputs or pbytes is None: | |
| return np.asarray(input_neurons(pid, cfg.d_in, cfg.k_in), dtype=np.int32) | |
| k_whole = cfg.k_in - cfg.k_tri | |
| row = list(input_neurons(pid, cfg.d_in, k_whole)) | |
| grams: set[bytes] = set() | |
| if len(pbytes) >= 3: | |
| grams = {pbytes[i:i + 3] for i in range(len(pbytes) - 2)} | |
| elif len(pbytes) == 2: | |
| grams = {pbytes} | |
| cand = sorted({splitmix64(fnv1a64(g)) % cfg.d_in for g in grams}) | |
| tri = cand[: cfg.k_tri] | |
| while len(tri) < cfg.k_tri: # short patches: fill from whole-hash salts | |
| tri.append(int(input_neurons(pid, cfg.d_in, cfg.k_tri + len(tri))[-1])) | |
| return np.asarray(row + tri, dtype=np.int32) | |
| def neuron_rows(self, ids: np.ndarray, patches: list[bytes] | None = None) -> np.ndarray: | |
| """(T, k_in) input rows; structured when enabled, v2-identical else.""" | |
| cfg = self.cfg | |
| if not cfg.structured_inputs: | |
| return self.input_neurons_stream(ids, cfg.d_in, cfg.k_in) | |
| by_id: dict[int, bytes] = {} | |
| if patches is not None: | |
| for pid, pb in zip(ids.tolist(), patches): | |
| by_id.setdefault(pid, pb) | |
| rows = np.empty((len(ids), cfg.k_in), dtype=np.int32) | |
| cache = self._row_cache | |
| for t, pid in enumerate(ids.tolist()): | |
| r = cache.get(pid) | |
| if r is None: | |
| pb = by_id.get(pid) | |
| if pb is None: | |
| pb = self.registry.bytes_of(pid) | |
| r = self._structured_row(pid, pb) | |
| cache[pid] = r | |
| rows[t] = r | |
| return rows | |
| def prepare(self) -> None: | |
| """Precompute training arrays for the SNN and AM subsets only.""" | |
| cfg = self.cfg | |
| ids = self._ids_train | |
| n_snn = min(cfg.snn_tokens, len(ids) - 1) | |
| n_am = min(cfg.am_tokens, len(ids) - 1) | |
| log = print | |
| log(f"[spira] precomputing projections for {n_snn} SNN + {n_am} AM tokens ...") | |
| t0 = time.perf_counter() | |
| self._neurons = self.neuron_rows(ids[:n_snn]) | |
| self._cols = self.cols_of(ids[: n_snn + 1]) | |
| # AM: val = registry index of the NEXT token (skip unregistered) | |
| am_ids = ids[:n_am] | |
| reg_idx = np.full(n_am, -1, dtype=np.int64) | |
| pos = np.searchsorted(self.registry.ids, am_ids) | |
| pos_c = np.minimum(pos, len(self.registry.ids) - 1) | |
| ok = (pos < len(self.registry.ids)) & (self.registry.ids[pos_c] == am_ids) | |
| reg_idx[ok] = pos_c[ok] | |
| self._reg_idx = reg_idx | |
| ls = self.am.lane_keys_stream(am_ids) | |
| # AM contexts (diagonal keys): ctx[r, l] = lane key of token r+l at | |
| # slot l; the context ending at t = r+m-1 continues with token r+m. | |
| m = self.am.m | |
| n_ctx = len(ls) - m | |
| ctx = np.empty((n_ctx, m), dtype=np.uint32) | |
| for lane in range(m): | |
| ctx[:, lane] = ls[lane : lane + n_ctx, lane] | |
| self._am_ctx = ctx | |
| log(f" done in {time.perf_counter()-t0:.1f}s") | |
| # ================================================================ train | |
| def train(self, log_fn=print, ckpt_path: str | None = None) -> dict: | |
| cfg = self.cfg | |
| n_snn = len(self._neurons) | |
| n_ctx_am = len(self._am_ctx) | |
| history = {"epochs": []} | |
| for ep in range(cfg.epochs): | |
| log_fn(f"[spira] === SNN epoch {ep+1}/{cfg.epochs} ({n_snn} ticks) ===") | |
| res = self.core.train_pass( | |
| self._neurons[:n_snn], self._cols[1 : n_snn + 1], | |
| log_every=10_000, | |
| log_fn=lambda d, ep=ep: log_fn( | |
| f" tick {d['tick']}/{d['total']} acc={d['acc_ema']:.3f} " | |
| f"dop={d['dopamine_ema']:+.3f} {d['ticks_per_sec']:.0f} t/s"), | |
| record_raster_from=(n_snn - 1500 if ep == cfg.epochs - 1 else None), | |
| ) | |
| raster = res.pop("raster", None) | |
| if raster and ckpt_path: | |
| ticks = np.concatenate([np.full(len(sp), t, dtype=np.int32) for t, sp in raster]) | |
| neurons_r = np.concatenate([sp for _, sp in raster]).astype(np.int32) | |
| np.savez_compressed(ckpt_path + ".raster.npz", ticks=ticks, neurons=neurons_r) | |
| history["epochs"].append({k: v for k, v in res.items()}) | |
| if ckpt_path: | |
| self.save(ckpt_path) | |
| # hippocampal one-shot pass over the (larger) AM stream | |
| log_fn(f"[spira] === hippocampus one-shot write pass ({n_ctx_am} contexts) ===") | |
| t0 = time.perf_counter() | |
| m = self.am.m | |
| ctx = self._am_ctx[:n_ctx_am] | |
| vals = self._ids_train[m : m + n_ctx_am].astype(np.uint64) | |
| self.am.write_stream(ctx, vals) | |
| log_fn(f" wrote {n_ctx_am} associations in " | |
| f"{time.perf_counter()-t0:.1f}s; {self.am.stats()}") | |
| history["am"] = self.am.stats() | |
| if ckpt_path: | |
| self.save(ckpt_path) | |
| return history | |
| # =============================================================== predict | |
| def _scores(self, x: np.ndarray) -> np.ndarray: | |
| """Behavioural scores: fast readout + variance-matched top-down slow | |
| prior (if the hierarchical cortex is active). No hippocampal vote.""" | |
| u = (x - self.core.x_mean) @ self.core.W_out | |
| if self.core.hier: | |
| u = u + np.float32(self.core.beta * self.fusion["scale"]) * self.core.u_slow | |
| return u | |
| def calibrate_fusion(self, data: bytes, n_tokens: int = 3000) -> dict: | |
| """Measure fast/slow readout score scales on a reference stream and set | |
| the fusion normalization sigma_fast/sigma_slow. | |
| The slow readout's raw scores live at a different (much larger) scale | |
| than the fast one; fusing them raw lets even a tiny beta swamp the fast | |
| argmax. Variance matching bounds the top-down prior's contribution to | |
| beta * sigma_fast — the same units the fast readout itself uses. Pure | |
| measurement, no plasticity (same spirit as the EPT budget calibration). | |
| """ | |
| patches = self.ept.segment(data) | |
| ids = np.fromiter((fnv1a64(p) for p in patches), dtype=np.uint64, | |
| count=len(patches))[: n_tokens + 1] | |
| neurons = self.neuron_rows(ids[:-1], patches[:-1]) | |
| snap = self.core.snapshot() | |
| w_rec_snap = self.core.W_rec.copy() | |
| self.core.reset_dynamics() | |
| self._win_count = 0 | |
| sig_f: list[float] = [] | |
| sig_s: list[float] = [] | |
| for t in range(len(ids) - 1): | |
| x = self.core.step(neurons[t]) | |
| if self.core.hier: | |
| self._win_count += 1 | |
| if self._win_count % self.core.win == 0: | |
| self.core.step_G() | |
| if len(sig_s) < 400 or len(sig_s) % 8 == 0: | |
| sig_s.append(float(np.std(self.core.u_slow))) | |
| u_fast = (x - self.core.x_mean) @ self.core.W_out | |
| if len(sig_f) < 400 or len(sig_f) % 8 == 0: | |
| sig_f.append(float(np.std(u_fast))) | |
| self.core.restore(snap) | |
| self.core.W_rec = w_rec_snap | |
| sf = float(np.mean(sig_f)) if sig_f else 1.0 | |
| ss = float(np.mean(sig_s)) if sig_s else 1.0 | |
| self.fusion = {"scale": sf / max(ss, 1e-9), "sigma_fast": sf, "sigma_slow": ss} | |
| return dict(self.fusion) | |
| def _feed(self, pid: int) -> np.ndarray: | |
| """Feed one token: fast step (+ slow step at window boundaries).""" | |
| x = self.core.step(self._neuron_row(pid)) | |
| if self.core.hier: | |
| self._win_count += 1 | |
| if self._win_count % self.core.win == 0: | |
| self.core.step_G() | |
| return x | |
| def _col_of_id(self, pid: int) -> int | None: | |
| """Vocab column of a patch id, or None if out of vocabulary.""" | |
| V = self.cfg.vocab_size | |
| col = int(np.searchsorted(self.vocab_ids, np.uint64(pid))) | |
| if col < V and self.vocab_ids[col] == np.uint64(pid): | |
| return col | |
| return None | |
| def fuse_scores(self, x: np.ndarray, last_ids: list[int]): | |
| """Combine behavioural scores (fast + top-down slow prior) with a | |
| hippocampal vote. | |
| Returns (fused_scores, am_val, am_weight, bypass_id). | |
| """ | |
| u = self._scores(x) | |
| am_val, am_w = None, 0.0 | |
| bypass_id = None | |
| if len(last_ids) >= self.am.m: | |
| lanes = self.am.lane_keys(last_ids[-self.am.m :]) | |
| am_val, am_w = self.am.query(lanes) | |
| if am_val is None: | |
| return u, am_val, am_w, None | |
| pid = int(am_val) | |
| col = self._col_of_id(pid) | |
| if col is not None: | |
| fused = u.copy() | |
| fused[col] += self.cfg.beta_am * am_w | |
| return fused, am_val, am_w, None | |
| if am_w >= self.cfg.am_override and self.registry.bytes_of(pid) is not None: | |
| bypass_id = pid # strong episodic recall outside vocabulary | |
| return u, am_val, am_w, bypass_id | |
| def generate(self, prompt: bytes, n_tokens: int = 120, temperature: float | None = None, | |
| top_k: int | None = None, use_am: bool = True, seed: int | None = None, | |
| on_token=None) -> bytes: | |
| """Closed-loop generation: predicted patches are fed back as spikes.""" | |
| cfg = self.cfg | |
| temperature = cfg.temperature if temperature is None else temperature | |
| top_k = cfg.top_k if top_k is None else top_k | |
| rng = np.random.default_rng(seed if seed is not None else int(time.time_ns() & 0x7FFFFFFF)) | |
| self.core.reset_dynamics() | |
| self._win_count = 0 | |
| patches = self.ept.segment(prompt) | |
| if not patches: | |
| return prompt | |
| ids = [fnv1a64(p) for p in patches] | |
| out = bytearray(prompt) | |
| V = cfg.vocab_size | |
| UNK = V # UNK column index | |
| # warm-up: feed the entire prompt (state now reflects the last patch) | |
| x = np.zeros(self.core.n_res, dtype=np.float32) | |
| for pid in ids: | |
| x = self._feed(pid) | |
| for step in range(n_tokens): | |
| if use_am: | |
| fused, am_val, am_w, bypass_id = self.fuse_scores(x, ids) | |
| else: | |
| fused, bypass_id = self._scores(x), None | |
| if bypass_id is not None: | |
| chosen_id = bypass_id | |
| else: | |
| scores = np.asarray(fused, dtype=np.float64).copy() | |
| scores[UNK] = -np.inf # cannot emit UNK bytes | |
| k = max(1, min(top_k, V)) | |
| top = np.argpartition(scores, -k)[-k:] | |
| z = (scores[top] - scores[top].max()) / max(temperature, 1e-6) | |
| w = np.exp(z) | |
| w /= w.sum() | |
| pick = int(top[rng.choice(len(top), p=w)]) | |
| chosen_id = int(self.vocab_ids[pick]) | |
| b = self.registry.bytes_of(chosen_id) | |
| if b is None: | |
| break # unknown emission: stop rather than corrupt | |
| out.extend(b) | |
| ids.append(chosen_id) | |
| x = self._feed(chosen_id) | |
| if on_token is not None: | |
| on_token(step, b) | |
| return bytes(out) | |
| def _neuron_row(self, pid: int) -> np.ndarray: | |
| r = self._row_cache.get(pid) | |
| if r is None: | |
| r = self._structured_row(pid, self.registry.bytes_of(pid)) | |
| self._row_cache[pid] = r | |
| return r | |
| # ============================================================= one-shot | |
| def teach_bytes(self, data: bytes) -> dict: | |
| """ONE-SHOT episodic teaching: segment, register unseen patches, and | |
| write every context -> continuation into the hippocampus — a single | |
| pass, no gradients, no epochs, no repetition.""" | |
| patches = self.ept.segment(data) | |
| ids = np.fromiter((fnv1a64(p) for p in patches), dtype=np.uint64, count=len(patches)) | |
| return self.teach_ids(ids, patches) | |
| def teach_ids(self, ids: np.ndarray, patches: list[bytes] | None = None, | |
| episodic: bool = True) -> dict: | |
| n = len(ids) | |
| # 0) retain the episodic trace for hippocampal replay (the sequence | |
| # the cortex will re-run during "sleep") — only for genuinely | |
| # episodic teaching, not for the corpus baseline fill | |
| if episodic: | |
| self._remember_episode(ids) | |
| # 1) register unseen patches so continuations are always storable | |
| known = self.registry.has(ids) | |
| if patches is not None: | |
| new_map: dict[int, bytes] = {} | |
| for i, b in zip(ids[~known].tolist(), (p for p, k in zip(patches, known) if not k)): | |
| if i not in new_map: | |
| new_map[i] = b | |
| n_new = self.registry.extend(new_map) | |
| else: | |
| n_new = 0 | |
| # 2) write associations (vals are stable patch ids) | |
| lane = self.am.lane_keys_stream(ids) | |
| m = self.am.m | |
| n_ctx = n - m | |
| ctx = np.empty((n_ctx, m), dtype=np.uint32) | |
| for li in range(m): | |
| ctx[:, li] = lane[li : li + n_ctx, li] | |
| vals = ids[m : m + n_ctx].astype(np.uint64) | |
| self.am.write_stream(ctx, vals) | |
| return { | |
| "tokens": int(n), | |
| "new_patches": int(n_new), | |
| "associations": int(n_ctx), | |
| "episode_tokens": int(n), | |
| } | |
| # ================================================================ replay | |
| def _remember_episode(self, ids: np.ndarray) -> None: | |
| """Append to the hippocampal episodic trace (ring-capped).""" | |
| if len(ids) < 2: | |
| return | |
| self.episodes.append(np.asarray(ids, dtype=np.uint64)) | |
| total = sum(len(e) for e in self.episodes) | |
| while total > self.EPI_CAP_TOKENS and len(self.episodes) > 1: | |
| total -= len(self.episodes[0]) | |
| self.episodes.pop(0) | |
| def replay(self, n_tokens: int | None = None, slice_len: int | None = None, | |
| seed: int = 0, log=print) -> dict: | |
| """Hippocampal replay -> cortical consolidation (closing the loop). | |
| Random slices of the episodic traces retained by the hippocampus are | |
| re-run through the cortex in full training mode: the same local rules | |
| (three-factor readout delta, modulated recurrent STDP, dopaminergic | |
| RPE, slow-area learning) now act on material the cortex never saw | |
| externally. One-shot memories are gradually consolidated into slow | |
| cortical weights — the complementary-systems loop, closed. | |
| Biologically this is Sharpewave-ripple-like offline replay: the cortex | |
| does not re-run isolated (context -> continuation) pairs but short | |
| SEQUENCES, so reservoir state statistics match the awake stream. | |
| """ | |
| cfg = self.cfg | |
| n_tokens = cfg.replay_tokens if n_tokens is None else n_tokens | |
| slice_len = cfg.replay_slice if slice_len is None else slice_len | |
| if not self.episodes: | |
| return {"replayed_tokens": 0, "note": "no episodic traces to replay"} | |
| rng = np.random.default_rng(seed) | |
| weights = np.array([len(e) for e in self.episodes], dtype=np.float64) | |
| weights /= weights.sum() | |
| chunks: list[np.ndarray] = [] | |
| got = 0 | |
| while got < n_tokens: | |
| e = self.episodes[int(rng.choice(len(self.episodes), p=weights))] | |
| start = int(rng.integers(0, max(1, len(e) - slice_len))) | |
| take = min(slice_len, len(e) - start) | |
| chunks.append(e[start : start + take]) | |
| got += take | |
| ids = np.concatenate(chunks) | |
| neurons = self.neuron_rows(ids[:-1]) | |
| cols = self.cols_of(ids)[1:] | |
| log(f"[replay] re-running {len(ids)} cortical tokens sampled from " | |
| f"{len(self.episodes)} episodic trace(s) (slice {slice_len})") | |
| res = self.core.train_pass( | |
| neurons, cols, log_every=25_000, | |
| log_fn=lambda d: log(f" tick {d['tick']}/{d['total']} acc={d['acc_ema']:.3f} " | |
| f"beta={d.get('beta')} {d['ticks_per_sec']:.0f} t/s"), | |
| ) | |
| res.pop("raster", None) | |
| return {"replayed_tokens": int(len(ids)), **res} | |
| # ================================================================ eval | |
| def eval_stream(self, data: bytes, use_am: bool = True, teach_am_first: bool = False, | |
| corrupt_lanes: int = 0, seed: int = 0, max_positions: int | None = None) -> dict: | |
| """Evaluate next-patch prediction on a byte stream. | |
| Metrics for the SNN readout alone, the hippocampus alone, and the | |
| fused system — both over all positions and over in-vocab targets only | |
| (targets outside the readout vocabulary cannot be predicted by the | |
| readout, but the hippocampus can still retrieve them). | |
| """ | |
| from .hashing import lane_key32 | |
| patches = self.ept.segment(data) | |
| ids = np.fromiter((fnv1a64(p) for p in patches), dtype=np.uint64, count=len(patches)) | |
| n = len(ids) | |
| cols = self.cols_of(ids) | |
| V = self.cfg.vocab_size | |
| if teach_am_first: | |
| # ONE-SHOT teaching pass (registers new patches too) | |
| taught = self.teach_ids(ids, patches) | |
| # SNN inference pass in the TRAINING regime: the reservoir keeps its | |
| # modulated-STDP plasticity (the operating point is part of the | |
| # learned system) while the readout stays frozen for measurement. | |
| snap = self.core.snapshot() | |
| w_rec_snap = self.core.W_rec.copy() | |
| neurons = self.neuron_rows(ids, patches) | |
| lane = self.am.lane_keys_stream(ids) | |
| # diagonal context-key stream: ctxq[r] = key of context tokens r..r+m-1 | |
| m_ctx = self.am.m | |
| n_ctxq = n - m_ctx + 1 | |
| ctxq = np.empty((n_ctxq, m_ctx), dtype=np.uint32) | |
| for li in range(m_ctx): | |
| ctxq[:, li] = lane[li : li + n_ctxq, li] | |
| rng = np.random.default_rng(seed) | |
| train_ids = self._ids_train if hasattr(self, "_ids_train") else ids | |
| tot = inv = 0 | |
| snn_ok = snn_ok_inv = am_ok = am_ok_inv = fused_ok = fused_ok_inv = 0 | |
| hier_ok = hier_ok_inv = 0 | |
| stop = n - 1 if max_positions is None else min(n - 1, max_positions) | |
| x = np.zeros(self.core.n_res, dtype=np.float32) | |
| for t in range(stop): | |
| x = self.core.step(neurons[t]) | |
| target_col = int(cols[t + 1]) | |
| target_id = int(ids[t + 1]) | |
| u_fast = (x - self.core.x_mean) @ self.core.W_out | |
| if self.core.hier: | |
| u = u_fast + np.float32(self.core.beta * self.fusion["scale"]) * self.core.u_slow | |
| hier_pred = int(np.argmax(u[:V])) # cortex (fast + top-down) | |
| else: | |
| u = u_fast | |
| hier_pred = None | |
| pred = int(np.argmax(u_fast[:V])) # fast-area-only ablation | |
| # keep training-regime reservoir dynamics (readout frozen) | |
| self.core.learn_reservoir(hier_pred if hier_pred is not None else pred, | |
| target_col, x, u=u_fast) | |
| if self.core.hier and (t + 1) % self.core.win == 0: | |
| self.core.step_G() | |
| # hippocampal query with optional context corruption | |
| am_val, am_w = None, 0.0 | |
| if use_am and t >= m_ctx - 1: | |
| ctx_lane = ctxq[t - m_ctx + 1].copy() | |
| if corrupt_lanes > 0: | |
| pidxs = rng.choice(m_ctx, size=min(corrupt_lanes, m_ctx), replace=False) | |
| for pidx in pidxs: | |
| rand_id = int(train_ids[rng.integers(0, len(train_ids))]) | |
| ctx_lane[pidx] = lane_key32(rand_id, int(pidx)) | |
| am_val, am_w = self.am.query(ctx_lane) | |
| am_id = int(am_val) if am_val is not None else None | |
| am_hit = am_id == target_id | |
| # fused prediction (hierarchical cortex + hippocampal vote) | |
| fused_pred, bypass = (hier_pred if hier_pred is not None else pred), False | |
| if am_id is not None: | |
| col = self._col_of_id(am_id) | |
| if col is not None: | |
| fused = u.copy() | |
| fused[col] += self.cfg.beta_am * am_w | |
| fused_pred = int(np.argmax(fused)) | |
| elif am_w >= self.cfg.am_override: | |
| bypass = True | |
| fused_hit = am_hit if bypass else (fused_pred == target_col) | |
| tot += 1 | |
| if target_col != V: | |
| inv += 1 | |
| if pred == target_col: | |
| snn_ok_inv += 1 | |
| if hier_pred is not None and hier_pred == target_col: | |
| hier_ok_inv += 1 | |
| if am_hit: | |
| am_ok_inv += 1 | |
| if fused_hit: | |
| fused_ok_inv += 1 | |
| if pred == target_col: | |
| snn_ok += 1 | |
| if hier_pred is not None and hier_pred == target_col: | |
| hier_ok += 1 | |
| if am_hit: | |
| am_ok += 1 | |
| if fused_hit: | |
| fused_ok += 1 | |
| self.core.restore(snap) | |
| self.core.W_rec = w_rec_snap | |
| return { | |
| "n_positions": tot, | |
| "vocab_coverage": inv / max(tot, 1), | |
| "snn_acc": snn_ok / max(tot, 1), | |
| "hier_acc": hier_ok / max(tot, 1), | |
| "am_acc": am_ok / max(tot, 1), | |
| "fused_acc": fused_ok / max(tot, 1), | |
| "snn_acc_invocab": snn_ok_inv / max(inv, 1), | |
| "hier_acc_invocab": hier_ok_inv / max(inv, 1), | |
| "am_acc_invocab": am_ok_inv / max(inv, 1), | |
| "fused_acc_invocab": fused_ok_inv / max(inv, 1), | |
| "corrupt_lanes": corrupt_lanes, | |
| } | |
| # =========================================================== save/load | |
| def save(self, path: str, am_cap: int = 1_500_000) -> None: | |
| """Checkpoint everything. ALL hippocampal entries are packed (a | |
| one-shot episodic memory must persist single-occurrence traces!) in a | |
| single memory-frugal pass. | |
| AM packing dedups EXACTLY: every slot an entry occupies is derivable | |
| from its key (primary full-key hash + leave-one-out hashes), so each | |
| logical entry is packed once with its surviving copies marked. | |
| """ | |
| from array import array | |
| cfg = self.cfg | |
| am = self.am | |
| visited = np.zeros(am.size, dtype=np.bool_) | |
| keys_buf = array("I") | |
| vals_buf = array("Q") | |
| counts_buf = array("B") | |
| for addr in np.nonzero(am.used)[0]: | |
| if visited[addr]: | |
| continue | |
| k = am.keys[addr] | |
| c = int(am.counts[addr]) | |
| for a in am._slots_of(k): | |
| if am.used[a] and np.array_equal(am.keys[a], k): | |
| visited[a] = True | |
| keys_buf.extend(k.tolist()) | |
| vals_buf.append(int(am.vals[addr])) | |
| counts_buf.append(min(c, 255)) | |
| if len(vals_buf) >= am_cap: | |
| break | |
| n_e = len(vals_buf) | |
| if n_e == 0: | |
| am_keys = np.zeros((0, am.m), dtype=np.uint32) | |
| am_vals = np.zeros(0, dtype=np.uint64) | |
| am_counts = np.zeros(0, dtype=np.uint8) | |
| else: | |
| am_keys = np.frombuffer(keys_buf, dtype=np.uint32).reshape(n_e, am.m).copy() | |
| am_vals = np.frombuffer(vals_buf, dtype=np.uint64).copy() | |
| am_counts = np.frombuffer(counts_buf, dtype=np.uint8).copy() | |
| np.savez_compressed( | |
| path, | |
| config=json.dumps(asdict(cfg)), | |
| meta=json.dumps(self.meta), | |
| fusion=json.dumps(self.fusion), | |
| log2p=self.ept.log2p, | |
| ept_params=np.array([self.ept.budget_bits, self.ept.min_len, self.ept.max_len, self.ept.alpha]), | |
| W_rec=self.core.W_rec, W_out=self.core.W_out, | |
| W_out_G=(self.core.W_outG if self.core.hier | |
| else np.zeros((0, 0), dtype=np.float32)), | |
| inhibitory=self.core.inhibitory, | |
| vocab_ids=self.vocab_ids, | |
| reg_ids=self.registry.ids, reg_offsets=self.registry.offsets, | |
| reg_counts=self.registry.counts, | |
| reg_buffer=np.frombuffer(self.registry.buffer, dtype=np.uint8), | |
| am_keys=am_keys, am_vals=am_vals, am_counts=am_counts, | |
| core_state=json.dumps({ | |
| k: (v if isinstance(v, float) else v.tolist()) | |
| for k, v in self.core.snapshot().items() | |
| }), | |
| ) | |
| def load(cls, path: str) -> "SPIRA": | |
| z = np.load(path, allow_pickle=False) | |
| cfg = SPIRAConfig(**json.loads(str(z["config"]))) | |
| m = cls(cfg) | |
| if "fusion" in z: | |
| m.fusion = json.loads(str(z["fusion"])) | |
| m.ept.log2p = z["log2p"].astype(np.float32) | |
| p = z["ept_params"] | |
| m.ept.budget_bits, m.ept.min_len, m.ept.max_len, m.ept.alpha = (float(p[0]), int(p[1]), int(p[2]), float(p[3])) | |
| m.core.W_rec = z["W_rec"].astype(np.float32) | |
| m.core.W_out = z["W_out"].astype(np.float32) | |
| m.core.inhibitory = z["inhibitory"].astype(np.bool_) | |
| w_rec_max = 3.0 * m.core.w_rec_scale | |
| m.core._rec_lo = np.where(m.core.inhibitory, -w_rec_max, 0.0).astype(np.float32) | |
| m.core._rec_hi = np.where(m.core.inhibitory, 0.0, w_rec_max).astype(np.float32) | |
| if "W_out_G" in z and cfg.hierarchy and z["W_out_G"].size > 0: | |
| m.core.W_outG = z["W_out_G"].astype(np.float32) | |
| else: | |
| # v1 checkpoint: hierarchy absent — disable the slow area | |
| cfg.hierarchy = False | |
| m.core.hier = False | |
| m.core.W_outG = None | |
| m.core.reset() | |
| if "core_state" in z: | |
| st = json.loads(str(z["core_state"])) | |
| restored = {} | |
| for k, v in st.items(): | |
| if k in ("baseline", "beta"): | |
| restored[k] = float(v) | |
| else: | |
| restored[k] = np.asarray(v, dtype=np.float32 if k not in ("refr",) else np.int8) | |
| m.core.restore(restored) | |
| m.vocab_ids = z["vocab_ids"].astype(np.uint64) | |
| m.registry.ids = z["reg_ids"].astype(np.uint64) | |
| m.registry.offsets = z["reg_offsets"].astype(np.uint64) | |
| m.registry.counts = z["reg_counts"].astype(np.uint32) | |
| m.registry.buffer = z["reg_buffer"].tobytes() | |
| m.meta = json.loads(str(z["meta"])) | |
| # rebuild hippocampus from packed entries: each entry returns to its | |
| # derivable slots (primary full-key + leave-one-out), with the same | |
| # best-effort collision policy as a live write — entries whose slots | |
| # are taken by a different key are simply not redundantly copied. | |
| keys = z["am_keys"]; vals = z["am_vals"]; counts = z["am_counts"] | |
| am = m.am | |
| for i in range(len(vals)): | |
| k, v, c = keys[i], int(vals[i]), int(counts[i]) | |
| paddr = am.primary_addr(k) | |
| if not am.used[paddr] or np.array_equal(am.keys[paddr], k): | |
| am.keys[paddr] = k | |
| am.vals[paddr] = v | |
| am.counts[paddr] = c | |
| am.used[paddr] = True | |
| for lane in range(am.m): | |
| addr = am.loo_addr(k, lane) | |
| if not am.used[addr]: | |
| am.keys[addr] = k | |
| am.vals[addr] = v | |
| am.counts[addr] = c | |
| am.used[addr] = True | |
| return m | |