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 | """Spiking core: LIF reservoir + reward-modulated (three-factor) readout. | |
| Paradigm properties — no backpropagation, no gradients, no autodiff: | |
| * Leaky integrate-and-fire (LIF) neurons with refractory periods and | |
| homeostatic adaptive thresholds; | |
| * Dale's law: each reservoir neuron is permanently excitatory or | |
| inhibitory (20% inhibitory); | |
| * Readout learning is a three-factor rule: eligibility (Hebbian pre/post | |
| coincidence = reservoir trace x outcome) gated by a dopaminergic reward | |
| prediction error d = r - baseline(r); | |
| * Recurrent synapses follow a light reward-modulated STDP (trace-based), | |
| sign-constrained to respect Dale's law; | |
| * One simulation tick = one token (event-time, not wall-clock ms). | |
| All weights are constructed deterministically from a seeded mulberry32 PRNG | |
| so the TypeScript port regenerates bit-identical connectivity. | |
| """ | |
| from __future__ import annotations | |
| import numpy as np | |
| from .hashing import Mulberry32 | |
| class SpikingCore: | |
| def __init__( | |
| self, | |
| d_in: int = 1024, | |
| k_in: int = 12, | |
| n_res: int = 512, | |
| n_out: int = 4097, # 4096 vocab + 1 UNK column | |
| seed: int = 20261022, | |
| in_fanout: int = 8, | |
| rec_fanout: int = 24, | |
| inh_frac: float = 0.50, | |
| lam: float = 0.85, # membrane leak | |
| base_thr: float = 1.0, | |
| thr_min: float = 0.4, | |
| thr_max: float = 8.0, | |
| refrac: int = 2, # ticks | |
| tau_trace: float = 0.40, # reservoir spike-trace decay | |
| w_in_scale: float = 3.0, | |
| w_rec_scale: float = 0.25, | |
| w_clip: float = 4.0, | |
| eta_delta: float = 0.5, # local Rescorla-Wagner delta rate (LMS-normalized) | |
| eta_stdp: float = 0.0015, # recurrent modulated STDP rate | |
| eta_homeo: float = 0.002, # adaptive-threshold rate | |
| target_rate: float = 0.06, # target per-neuron firing probability/tick | |
| baseline_momentum: float = 0.995, | |
| # ---- hierarchical slow area G (multi-scale cortex, v2) ---- | |
| hier: bool = True, | |
| n_g: int = 256, # slow-area neurons | |
| win: int = 8, # fast ticks per slow step | |
| w_g_scale: float = 6.0, # L -> G feedforward weight scale | |
| w_recg_scale: float = 0.15, # G recurrent scale (fixed, no STDP) | |
| lam_g: float = 0.6, # slow-area membrane leak | |
| tau_trace_g: float = 0.5, # slow-area spike-trace decay | |
| target_rate_g: float = 0.10, # target firing per slow step | |
| eta_homeo_g: float = 0.01, | |
| eta_delta_g: float = 2.0, # slow-readout delta rate (8x fewer updates) | |
| eta_beta: float = 0.02, # top-down gain learning rate | |
| beta_init: float = 0.25, | |
| beta_max: float = 2.0, | |
| # ---- v3 (Muse): multi-timescale traces + direct input pathway ---- | |
| trace_taus: tuple = (0.4,), # reservoir trace decays; readout sees ALL banks | |
| direct_input: bool = False, # readout also sees a decayed trace of the | |
| # exact input-neuron pattern (ESN-standard | |
| # [state; input] readout, omitted in v1/v2) | |
| direct_tau: float = 0.5, # decay of the direct input trace | |
| direct_gain: float = 1.0, | |
| ) -> None: | |
| self.d_in, self.k_in, self.n_res, self.n_out = d_in, k_in, n_res, n_out | |
| self.seed = seed | |
| self.in_fanout, self.rec_fanout = in_fanout, rec_fanout | |
| self.inh_frac = inh_frac | |
| self.lam, self.base_thr = lam, base_thr | |
| self.thr_min, self.thr_max = thr_min, thr_max | |
| self.refrac = refrac | |
| self.tau_trace = tau_trace | |
| self.w_in_scale, self.w_rec_scale = w_in_scale, w_rec_scale | |
| self.w_clip = w_clip | |
| self.eta_delta, self.eta_stdp = eta_delta, eta_stdp | |
| self.eta_homeo, self.target_rate = eta_homeo, target_rate | |
| self.baseline_momentum = baseline_momentum | |
| self.hier = hier | |
| self.n_g, self.win = n_g, win | |
| self.w_g_scale, self.w_recg_scale = w_g_scale, w_recg_scale | |
| self.lam_g, self.tau_trace_g = lam_g, tau_trace_g | |
| self.target_rate_g, self.eta_homeo_g = target_rate_g, eta_homeo_g | |
| self.eta_delta_g = eta_delta_g | |
| self.eta_beta, self.beta0, self.beta_max = eta_beta, beta_init, beta_max | |
| # v3 structural params | |
| self.trace_taus = tuple(float(t) for t in trace_taus) | |
| self.direct_input = bool(direct_input) | |
| self.direct_tau, self.direct_gain = float(direct_tau), float(direct_gain) | |
| self.n_banks = len(self.trace_taus) | |
| self.n_state = self.n_res * self.n_banks + (self.d_in if self.direct_input else 0) | |
| # ------------------------------------------------ deterministic wiring | |
| rng = Mulberry32(seed) | |
| # Input -> reservoir: sparse fan-out per input neuron, dense storage. | |
| self.W_in = np.zeros((d_in, n_res), dtype=np.float32) | |
| for i in range(d_in): | |
| for _ in range(in_fanout): | |
| t = int(rng.next_float() * n_res) | |
| self.W_in[i, t] += rng.next_weight(w_in_scale) | |
| # Reservoir recurrent: sparse fan-out per source, Dale's law by source. | |
| self.W_rec = np.zeros((n_res, n_res), dtype=np.float32) | |
| self.inhibitory = np.zeros(n_res, dtype=np.bool_) | |
| w_rec_max = 3.0 * w_rec_scale | |
| for j in range(n_res): | |
| is_inh = rng.next_float() < inh_frac | |
| self.inhibitory[j] = is_inh | |
| sign = -1.0 if is_inh else 1.0 | |
| for _ in range(rec_fanout): | |
| t = int(rng.next_float() * n_res) | |
| self.W_rec[j, t] += sign * abs(rng.next_weight(w_rec_scale)) | |
| # Keep Dale's law under plasticity: | |
| # excitatory rows clipped to [0, w_rec_max], inhibitory to [-w_rec_max, 0] | |
| self._rec_lo = np.where(self.inhibitory, -w_rec_max, 0.0).astype(np.float32) | |
| self._rec_hi = np.where(self.inhibitory, 0.0, w_rec_max).astype(np.float32) | |
| self.W_rec = np.clip(self.W_rec, self._rec_lo[:, None], self._rec_hi[:, None]) | |
| # Readout: small seeded init, learned by the three-factor rule. | |
| rng2 = Mulberry32(seed + 1) | |
| self.W_out = np.zeros((self.n_state, n_out), dtype=np.float32) # v3: n_state rows | |
| # sparse-ish init: each column connected from a random 25% of neurons | |
| for c in range(n_out): | |
| for _ in range(n_res // 4): | |
| r = int(rng2.next_float() * n_res) | |
| self.W_out[r, c] += rng2.next_weight(0.05) | |
| # ------------------------------------------ slow area G (hierarchy) | |
| # Wired from SEPARATE PRNG streams (seed+2 / seed+3) so the fast-area | |
| # sequences above are byte-identical to v1 — the TS port depends on | |
| # the exact call order. | |
| if hier: | |
| rng3 = Mulberry32(seed + 2) | |
| # L -> G: per fast-source fanout 2. Sign follows Dale's law of the | |
| # PRESYNAPTIC fast neuron: excitatory cells excite G, inhibitory | |
| # cells inhibit it (G sees the pooled fast spike volley). | |
| self.W_inG = np.zeros((n_res, n_g), dtype=np.float32) | |
| for i in range(n_res): | |
| sign = -1.0 if self.inhibitory[i] else 1.0 | |
| for _ in range(2): | |
| t = int(rng3.next_float() * n_g) | |
| self.W_inG[i, t] += sign * abs(rng3.next_weight(w_g_scale)) | |
| # G recurrent: sparse, Dale by source, FIXED (slower areas | |
| # plasticize slower — G plasticity lives in its readout only) | |
| self.W_recG = np.zeros((n_g, n_g), dtype=np.float32) | |
| g_rec_max = 3.0 * w_recg_scale | |
| self.inhibitory_g = np.zeros(n_g, dtype=np.bool_) | |
| for j in range(n_g): | |
| is_inh = rng3.next_float() < 0.5 | |
| self.inhibitory_g[j] = is_inh | |
| sign = -1.0 if is_inh else 1.0 | |
| for _ in range(8): | |
| t = int(rng3.next_float() * n_g) | |
| self.W_recG[j, t] += sign * abs(rng3.next_weight(w_recg_scale)) | |
| self._recg_lo = np.where(self.inhibitory_g, -g_rec_max, 0.0).astype(np.float32) | |
| self._recg_hi = np.where(self.inhibitory_g, 0.0, g_rec_max).astype(np.float32) | |
| self.W_recG = np.clip(self.W_recG, self._recg_lo[:, None], self._recg_hi[:, None]) | |
| # G readout: same seeded sparse init pattern as the fast readout | |
| rng4 = Mulberry32(seed + 3) | |
| self.W_outG = np.zeros((n_g, n_out), dtype=np.float32) | |
| for c in range(n_out): | |
| for _ in range(n_g // 4): | |
| r = int(rng4.next_float() * n_g) | |
| self.W_outG[r, c] += rng4.next_weight(0.05) | |
| self.reset() | |
| # ----------------------------------------------------------------- state | |
| def reset(self) -> None: | |
| self.v = np.zeros(self.n_res, dtype=np.float32) | |
| self.thr = np.full(self.n_res, self.base_thr, dtype=np.float32) | |
| self.rhat = np.full(self.n_res, self.target_rate, dtype=np.float32) | |
| self.refr = np.zeros(self.n_res, dtype=np.int8) | |
| self.x = np.zeros(self.n_state, dtype=np.float32) # v3: augmented readout state | |
| self.x_prev = np.zeros(self.n_res, dtype=np.float32) # reservoir trace t-1 (STDP) | |
| self.x_mean = np.zeros(self.n_state, dtype=np.float32) | |
| self.x_eff = np.zeros(self.n_state, dtype=np.float32) | |
| self._banks = np.zeros((self.n_banks, self.n_res), dtype=np.float32) | |
| self._direct = np.zeros(self.d_in if self.direct_input else 0, dtype=np.float32) | |
| self.s = np.zeros(self.n_res, dtype=np.float32) # last spike vector | |
| self.baseline = 0.0 | |
| self.tick = 0 | |
| if self.hier: | |
| self.v_G = np.zeros(self.n_g, dtype=np.float32) | |
| self.thr_G = np.full(self.n_g, self.base_thr, dtype=np.float32) | |
| self.rhat_G = np.full(self.n_g, self.target_rate_g, dtype=np.float32) | |
| self.x_G = np.zeros(self.n_g, dtype=np.float32) | |
| self.x_mean_G = np.zeros(self.n_g, dtype=np.float32) | |
| self.s_G = np.zeros(self.n_g, dtype=np.float32) | |
| self.win_spikes = np.zeros(self.n_res, dtype=np.float32) | |
| self.u_slow = np.zeros(self.n_out, dtype=np.float32) | |
| self.beta = self.beta0 | |
| # ------------------------------------------------------------- snapshots | |
| def snapshot(self) -> dict: | |
| """Full dynamical state (the homeostatic scalars are LEARNED state — | |
| thresholds and the trace mean co-adapt with the readout).""" | |
| d = { | |
| "v": self.v.copy(), "thr": self.thr.copy(), "rhat": self.rhat.copy(), | |
| "refr": self.refr.copy(), "x": self.x.copy(), "x_prev": self.x_prev.copy(), | |
| "x_mean": self.x_mean.copy(), "s": self.s.copy(), "baseline": self.baseline, | |
| "banks": self._banks.copy(), "direct": self._direct.copy(), | |
| } | |
| if self.hier: | |
| d.update({ | |
| "v_G": self.v_G.copy(), "thr_G": self.thr_G.copy(), | |
| "rhat_G": self.rhat_G.copy(), "x_G": self.x_G.copy(), | |
| "x_mean_G": self.x_mean_G.copy(), "s_G": self.s_G.copy(), | |
| "win_spikes": self.win_spikes.copy(), "u_slow": self.u_slow.copy(), | |
| "beta": self.beta, | |
| }) | |
| return d | |
| def restore(self, st: dict) -> None: | |
| self.v = st["v"].copy() | |
| self.thr = st["thr"].copy() | |
| self.rhat = st["rhat"].copy() | |
| self.refr = st["refr"].copy() | |
| self.x = st["x"].copy() | |
| self.x_prev = st["x_prev"].copy() | |
| self.x_mean = st["x_mean"].copy() | |
| self.s = st["s"].copy() | |
| self.baseline = st["baseline"] | |
| if "banks" in st: | |
| self._banks = st["banks"].copy() | |
| self._direct = st["direct"].copy() | |
| else: # reconstruct banks from the composed state vector | |
| nb = self.n_banks * self.n_res | |
| self._banks = self.x[:nb].reshape(self.n_banks, self.n_res).copy() | |
| if self.direct_input: | |
| self._direct = self.x[nb:].copy() | |
| if self.hier and "v_G" in st: | |
| self.v_G = st["v_G"].copy() | |
| self.thr_G = st["thr_G"].copy() | |
| self.rhat_G = st["rhat_G"].copy() | |
| self.x_G = st["x_G"].copy() | |
| self.x_mean_G = st["x_mean_G"].copy() | |
| self.s_G = st["s_G"].copy() | |
| self.win_spikes = st["win_spikes"].copy() | |
| self.u_slow = st["u_slow"].copy() | |
| self.beta = st["beta"] | |
| def reset_dynamics(self) -> None: | |
| """Reset transient dynamics but KEEP learned homeostatic scalars | |
| (thresholds, rate estimates, trace mean, top-down gain) — used before | |
| generation.""" | |
| thr, rhat, x_mean = self.thr, self.rhat, self.x_mean | |
| if self.hier: | |
| thr_G, rhat_G, x_mean_G, beta = self.thr_G, self.rhat_G, self.x_mean_G, self.beta | |
| self.reset() | |
| self.thr, self.rhat, self.x_mean = thr, rhat, x_mean | |
| if self.hier: | |
| self.thr_G, self.rhat_G, self.x_mean_G, self.beta = thr_G, rhat_G, x_mean_G, beta | |
| # ------------------------------------------------------------- inference | |
| def step(self, input_idx) -> np.ndarray: | |
| """Advance one tick with the given active input neurons. | |
| Returns the post-synaptic trace ``x`` (the state used for readout). | |
| """ | |
| W_in, W_rec = self.W_in, self.W_rec | |
| I = W_in[input_idx].sum(axis=0, dtype=np.float32) | |
| I += self.s @ W_rec | |
| v = self.lam * self.v + I | |
| can = self.refr == 0 | |
| v[~can] = 0.0 | |
| spikes = (v >= self.thr) & can | |
| v[spikes] = 0.0 | |
| np.subtract(self.refr, 1, out=self.refr, where=self.refr > 0) | |
| self.refr[spikes] = self.refrac | |
| self.v = v | |
| self.s = spikes.astype(np.float32) | |
| # homeostatic threshold via EMA firing-rate estimate | |
| self.rhat += 0.02 * (self.s - self.rhat) | |
| self.thr += self.eta_homeo * (self.rhat - self.target_rate) | |
| np.clip(self.thr, self.thr_min, self.thr_max, out=self.thr) | |
| # traces (v3): x_prev keeps the BANK-0 (fast) reservoir trace for STDP; | |
| # the readout state x is the concatenation of all trace banks plus the | |
| # direct input trace — context older than one tick survives in the | |
| # slow banks and in the direct trace's decayed identity history. | |
| self.x_prev = self._banks[0].copy() | |
| for i, tau in enumerate(self.trace_taus): | |
| if i == 0: | |
| self._banks[i] = tau * self._banks[i] + self.s | |
| else: | |
| # slow banks are NORMALIZED traces (decay-weighted averages): | |
| # a raw tau=0.95 sum would sit at ~20x the fast bank's scale | |
| # and swamp both the readout geometry and the LMS rate. | |
| self._banks[i] = tau * self._banks[i] + (1.0 - tau) * self.s | |
| if self.direct_input: | |
| self._direct *= np.float32(self.direct_tau) | |
| np.add.at(self._direct, input_idx, np.float32(self.direct_gain)) | |
| self.x = np.concatenate([*self._banks, self._direct]) | |
| elif self.n_banks > 1: | |
| self.x = np.concatenate([*self._banks]) | |
| else: | |
| self.x = self._banks[0] | |
| # slow EMA of the trace: homeostatic "synaptic scaling" analogue used | |
| # to mean-center the readout input (kills the shared dense background | |
| # that otherwise dominates column norms) | |
| self.x_mean = 0.999 * self.x_mean + 0.001 * self.x | |
| self.x_eff = self.x - self.x_mean | |
| self.tick += 1 | |
| if self.hier: | |
| self.win_spikes += self.s | |
| return self.x | |
| def scores(self, x: np.ndarray) -> np.ndarray: | |
| """Readout scores from a (raw) trace — centers with the EMA mean.""" | |
| return (x - self.x_mean) @ self.W_out | |
| # ------------------------------------------------- slow area (hierarchy) | |
| def step_G(self) -> None: | |
| """Advance the slow area one slow step (one window of ``win`` fast | |
| ticks). | |
| Input: the pooled fast-area spike volley over the window (bottom-up, | |
| sign-preserving). Output: the cached slow readout ``u_slow``, used as | |
| a top-down prior fused into the fast prediction with gain ``beta``. | |
| Pure inference dynamics (homeostasis included, same regime as the | |
| fast area); learning happens in ``learn_G``. | |
| """ | |
| c = self.win_spikes * np.float32(1.0 / self.win) | |
| I = c @ self.W_inG | |
| I += self.s_G @ self.W_recG | |
| v = self.lam_g * self.v_G + I | |
| spikes = v >= self.thr_G | |
| v[spikes] = 0.0 | |
| self.v_G = v | |
| self.s_G = spikes.astype(np.float32) | |
| self.rhat_G += 0.05 * (self.s_G - self.rhat_G) | |
| self.thr_G += self.eta_homeo_g * (self.rhat_G - self.target_rate_g) | |
| np.clip(self.thr_G, self.thr_min, self.thr_max, out=self.thr_G) | |
| self.x_G = self.tau_trace_g * self.x_G + self.s_G | |
| self.x_mean_G = 0.999 * self.x_mean_G + 0.001 * self.x_G | |
| self.win_spikes[:] = 0.0 | |
| self.u_slow = (self.x_G - self.x_mean_G) @ self.W_outG | |
| def learn_G(self, target: int, d: float) -> int: | |
| """Delta-rule update on the SLOW readout, gated by the same dopaminergic | |
| RPE that gates the fast one (a global neuromodulator broadcasts to both | |
| areas — each readout computes its own local error from its own | |
| activity). ``target`` is the majority vocabulary column of the NEXT | |
| window. Returns 1 if the slow readout's own top prediction matches it. | |
| """ | |
| x_effG = self.x_G - self.x_mean_G | |
| u = self.u_slow | |
| V = self.n_out - 1 | |
| own_pred = int(np.argmax(u[:V])) | |
| correct = 1 if own_pred == target else 0 | |
| nx = 1.0 / (float(x_effG @ x_effG) + 1e-6) | |
| k = min(16, V) | |
| act = np.argpartition(u[:V], -k)[-k:] | |
| if target < V: | |
| cols = np.concatenate([act, np.array([target], dtype=np.int64)]) | |
| desired = np.zeros(len(cols), dtype=np.float32) | |
| desired[-1] = 1.0 | |
| else: | |
| cols = act | |
| desired = np.zeros(len(cols), dtype=np.float32) | |
| errs = u[cols] - desired | |
| gate = float(np.clip(1.0 + d, 0.25, 2.0)) | |
| eta = np.float32(self.eta_delta_g * gate * nx) | |
| W = self.W_outG | |
| W[:, cols] -= eta * errs[None, :] * x_effG[:, None] | |
| W[:, cols] = np.clip(W[:, cols], -self.w_clip, self.w_clip) | |
| return correct | |
| # ------------------------------------------------------------- learning | |
| def learn(self, pred: int, target: int, x: np.ndarray, | |
| u: np.ndarray | None = None, active_k: int = 16) -> tuple[int, float, float]: | |
| """Local delta-rule readout update (Rescorla-Wagner) gated by a | |
| dopaminergic reward-prediction error, + modulated recurrent STDP. | |
| Each readout unit computes its OWN prediction error from its own | |
| activity (no gradients, no chain rule, no backprop): units in the | |
| active top-k set and the instructive target unit move their weights | |
| to reduce (score - desired). The learning rate is gated by dopamine | |
| d = reward - baseline (RPE), the classical Rescorla-Wagner / TD | |
| account of neuromodulation. | |
| Returns (correct, reward, dopamine). | |
| """ | |
| correct = 1 if pred == target else 0 | |
| reward = float(correct) | |
| self.baseline = self.baseline_momentum * self.baseline + (1 - self.baseline_momentum) * reward | |
| d = reward - self.baseline # dopaminergic RPE | |
| W_out = self.W_out | |
| x_eff = x - self.x_mean | |
| nx = 1.0 / (float(x_eff @ x_eff) + 1e-6) # LMS normalization | |
| V = self.n_out - 1 | |
| if u is None: | |
| u = x_eff @ W_out | |
| # active set: top-k scored columns (+ target if in vocabulary) | |
| k = min(active_k, V) | |
| act = np.argpartition(u[:V], -k)[-k:] | |
| if target < V: | |
| cols = np.concatenate([act, np.array([target], dtype=np.int64)]) | |
| else: | |
| cols = act | |
| desired = np.zeros(len(cols), dtype=np.float32) | |
| if target < V: | |
| desired[-1] = 1.0 # one-hot with margin 1 | |
| errs = u[cols] - desired | |
| # dopamine gates plasticity amplitude (dips shrink, bursts amplify) | |
| gate = float(np.clip(1.0 + d, 0.25, 2.0)) | |
| eta = np.float32(self.eta_delta * gate * nx) | |
| W_out[:, cols] -= eta * errs[None, :] * x_eff[:, None] | |
| # clip in place (fancy-indexed views need explicit writeback) | |
| W_out[:, cols] = np.clip(W_out[:, cols], -self.w_clip, self.w_clip) | |
| # light reward-modulated STDP on recurrent synapses of spiked neurons | |
| if self.s.any(): | |
| spiked = np.nonzero(self.s)[0] | |
| W_rec = self.W_rec | |
| delta = np.float32(self.eta_stdp * d) | |
| # pre-trace (x_prev) paired with post spikes, gated by dopamine | |
| W_rec[:, spiked] += delta * self.x_prev[:, None] | |
| np.clip(W_rec[:, spiked], self._rec_lo[:, None], self._rec_hi[:, None], | |
| out=W_rec[:, spiked]) | |
| # top-down gain: TRUST dynamics, gated by the same RPE (a | |
| # noradrenergic gain-modulation analogue) with counterfactual credit: | |
| # fused right + fast wrong -> the slow prior SAVED the choice -> trust builds (additive) | |
| # fused wrong + fast right -> the slow prior CORRUPTED it -> trust decays (multiplicative) | |
| # Trust building slowly and decaying fast (loss aversion) lets the | |
| # system keep a genuinely informative slow area while driving a | |
| # noise- or frequency-prior-dominated one to beta ~ 0. | |
| if self.hier and target < V: | |
| pred_fast = int(np.argmax(u[:V])) | |
| if pred_fast != target and correct: # slow saved it | |
| self.beta = min(self.beta + self.eta_beta * gate, self.beta_max) | |
| elif pred_fast == target and not correct: # slow corrupted it | |
| self.beta *= float(0.98) ** gate | |
| if self.beta < 1e-4: | |
| self.beta = 0.0 | |
| else: | |
| us_t = float(self.u_slow[target]) | |
| us_p = float(self.u_slow[pred_fast]) | |
| norm = abs(us_t) + abs(us_p) + 1e-6 | |
| margin = float(np.tanh(3.0 * (us_t - us_p) / norm)) | |
| self.beta = float(np.clip( | |
| self.beta + 0.25 * self.eta_beta * gate * margin, | |
| 0.0, self.beta_max)) | |
| return correct, reward, d | |
| def learn_reservoir(self, pred: int, target: int, x: np.ndarray, | |
| u: np.ndarray | None = None) -> tuple[int, float, float]: | |
| """Training-regime dynamics WITHOUT readout plasticity: same reward, | |
| same dopaminergic RPE, same recurrent STDP — only the W_out delta | |
| update is skipped. Used at inference so the reservoir runs in its | |
| training operating regime (continual plasticity, frozen readout).""" | |
| correct = 1 if pred == target else 0 | |
| reward = float(correct) | |
| self.baseline = self.baseline_momentum * self.baseline + (1 - self.baseline_momentum) * reward | |
| d = reward - self.baseline | |
| if self.s.any(): | |
| spiked = np.nonzero(self.s)[0] | |
| W_rec = self.W_rec | |
| delta = np.float32(self.eta_stdp * d) | |
| W_rec[:, spiked] += delta * self.x_prev[:, None] | |
| W_rec[:, spiked] = np.clip(W_rec[:, spiked], self._rec_lo[:, None], self._rec_hi[:, None]) | |
| return correct, reward, d | |
| # ------------------------------------------------------- training driver | |
| def train_pass( | |
| self, | |
| neurons: np.ndarray, # (T, k_in) int32 active input neurons per tick | |
| cols: np.ndarray, # (T,) int64 target column per tick (col of token t+1) | |
| log_every: int = 5000, | |
| log_fn=None, | |
| decay_every: int = 1024, | |
| decay: float = 1e-4, | |
| record_raster_from: int | None = None, | |
| ) -> dict: | |
| """One epoch over a token stream (teacher forcing on the next token). | |
| ``neurons[t]`` drives the reservoir at tick t; the prediction made | |
| from the resulting state is scored against ``cols[t]``. | |
| """ | |
| self.reset() | |
| T = len(cols) | |
| acc_ema = None | |
| d_ema = 0.0 | |
| n_correct = 0 | |
| n_correct_fast = 0 | |
| raster: list[tuple[int, np.ndarray]] = [] | |
| V_vocab = self.n_out - 1 # UNK occupies the last column | |
| import time | |
| t0 = time.perf_counter() | |
| for t in range(T): | |
| x = self.step(neurons[t]) | |
| u = (x - self.x_mean) @ self.W_out | |
| if self.hier: | |
| # behavioural prediction = fast readout + top-down slow prior | |
| fused = u + np.float32(self.beta) * self.u_slow | |
| pred = int(np.argmax(fused[:V_vocab])) | |
| else: | |
| pred = int(np.argmax(u[:V_vocab])) # vocabulary-only prediction | |
| correct, reward, d = self.learn(pred, int(cols[t]), x, u=u) | |
| n_correct += correct | |
| if self.hier: | |
| pred_fast = int(np.argmax(u[:V_vocab])) | |
| if pred_fast == int(cols[t]): | |
| n_correct_fast += 1 | |
| if (t + 1) % self.win == 0: | |
| self.step_G() | |
| if t + 1 < T: | |
| seg = cols[t + 1 : min(t + 1 + self.win, T)] | |
| g_target = int(np.bincount(seg).argmax()) | |
| self.learn_G(g_target, d) | |
| acc_ema = reward if acc_ema is None else 0.999 * acc_ema + 0.001 * reward | |
| d_ema = 0.999 * d_ema + 0.001 * d | |
| if record_raster_from is not None and t >= record_raster_from: | |
| raster.append((t, np.nonzero(self.s)[0].astype(np.int16))) | |
| if (t + 1) % log_every == 0 and log_fn is not None: | |
| dt = time.perf_counter() - t0 | |
| log_fn({ | |
| "tick": t + 1, "total": T, | |
| "acc_ema": round(acc_ema, 4), | |
| "acc_fast_ema": (round(0.999 * (n_correct_fast / max(t + 1, 1)), 4) | |
| if self.hier else None), | |
| "beta": (round(self.beta, 3) if self.hier else None), | |
| "dopamine_ema": round(d_ema, 4), | |
| "ticks_per_sec": round((t + 1) / max(dt, 1e-9), 1), | |
| "elapsed_s": round(dt, 1), | |
| }) | |
| if (t + 1) % decay_every == 0: | |
| self.W_out *= np.float32(1.0 - decay) | |
| if self.hier: | |
| self.W_outG *= np.float32(1.0 - decay) | |
| dt = time.perf_counter() - t0 | |
| return { | |
| "ticks": T, | |
| "accuracy": n_correct / max(T, 1), | |
| "accuracy_fast_only": (n_correct_fast / max(T, 1)) if self.hier else None, | |
| "beta": float(self.beta) if self.hier else None, | |
| "acc_ema_final": acc_ema, | |
| "elapsed_s": dt, | |
| "ticks_per_sec": T / max(dt, 1e-9), | |
| "raster": raster, | |
| } | |
| # ------------------------------------------------------------ (de)hydrate | |
| def state_dict(self) -> dict: | |
| sd = { | |
| "config": { | |
| "d_in": self.d_in, "k_in": self.k_in, "n_res": self.n_res, | |
| "n_out": self.n_out, "seed": self.seed, | |
| "in_fanout": self.in_fanout, "rec_fanout": self.rec_fanout, | |
| "inh_frac": self.inh_frac, "lam": self.lam, | |
| "base_thr": self.base_thr, "refrac": self.refrac, | |
| "tau_trace": self.tau_trace, | |
| "w_in_scale": self.w_in_scale, "w_rec_scale": self.w_rec_scale, | |
| "w_clip": self.w_clip, "eta_delta": self.eta_delta, "eta_stdp": self.eta_stdp, | |
| "eta_homeo": self.eta_homeo, | |
| "target_rate": self.target_rate, | |
| "hier": self.hier, "n_g": self.n_g, "win": self.win, | |
| "w_g_scale": self.w_g_scale, "w_recg_scale": self.w_recg_scale, | |
| "lam_g": self.lam_g, "tau_trace_g": self.tau_trace_g, | |
| "target_rate_g": self.target_rate_g, "eta_homeo_g": self.eta_homeo_g, | |
| "eta_delta_g": self.eta_delta_g, | |
| "eta_beta": self.eta_beta, "beta_init": self.beta0, "beta_max": self.beta_max, | |
| "trace_taus": list(self.trace_taus), "direct_input": self.direct_input, | |
| "direct_tau": self.direct_tau, "direct_gain": self.direct_gain, | |
| }, | |
| "W_rec": self.W_rec.astype(np.float32), | |
| "W_out": self.W_out.astype(np.float32), | |
| "inhibitory": self.inhibitory, | |
| } | |
| if self.hier: | |
| sd["W_out_G"] = self.W_outG.astype(np.float32) | |
| return sd | |
| def load_state_dict(self, sd: dict) -> None: | |
| cfg = sd["config"] | |
| self.__dict__.update(cfg) | |
| # v3 derived structure (defaults reproduce v2 for old checkpoints) | |
| self.trace_taus = tuple(float(t) for t in getattr(self, "trace_taus", (self.tau_trace,))) | |
| self.direct_input = bool(getattr(self, "direct_input", False)) | |
| self.n_banks = len(self.trace_taus) | |
| self.n_state = int(sd["W_out"].shape[0]) | |
| self.W_rec = np.asarray(sd["W_rec"], dtype=np.float32) | |
| self.W_out = np.asarray(sd["W_out"], dtype=np.float32) | |
| self.inhibitory = np.asarray(sd["inhibitory"], dtype=np.bool_) | |
| w_rec_max = 3.0 * self.w_rec_scale | |
| self._rec_lo = np.where(self.inhibitory, -w_rec_max, 0.0).astype(np.float32) | |
| self._rec_hi = np.where(self.inhibitory, 0.0, w_rec_max).astype(np.float32) | |
| if "W_out_G" in sd and self.hier: | |
| self.W_outG = np.asarray(sd["W_out_G"], dtype=np.float32) | |
| else: | |
| # v1 checkpoint without the slow area | |
| self.hier = False | |
| self.reset() | |