SPIRA v2: entropy-patch tokenization + hierarchical dopamine-gated spiking cortex + hippocampal one-shot memory with replay consolidation
032a961 verified | """Hippocampal one-shot associative memory. | |
| Biological motivation (complementary learning systems): a fast, episodic, | |
| one-shot memory next to a slow statistical cortex (the spiking core). | |
| Key design — hybrid addressing (dentate-gyrus separation + CA3 completion): | |
| * The key of a memory is the last ``m`` tokens (patches), one 32-bit lane | |
| per context slot. | |
| * PRIMARY address: a hash of the FULL key (pattern separation — every | |
| distinct context gets its own slot, so capacity is real). | |
| * SECONDARY addresses: one per lane (auto-associative backup probes — a | |
| query with a corrupted context still reaches the entry through its | |
| intact lanes). | |
| * Retrieval is two-tier: a full-key match at the primary address | |
| dominates; otherwise candidates matching >= min_lanes lanes vote by | |
| lane-match count x log evidence count. | |
| This yields both capacity AND graceful degradation: corrupting 1 of 4 | |
| context tokens still retrieves the original continuation through the | |
| intact lanes. | |
| """ | |
| from __future__ import annotations | |
| import numpy as np | |
| from .hashing import LANE_SALTS, splitmix64 | |
| class Hippocampus: | |
| def __init__( | |
| self, | |
| m: int = 4, | |
| size: int = 1 << 23, | |
| min_lanes: int = 3, | |
| ) -> None: | |
| self.m = m | |
| self.size = size | |
| self.min_lanes = min_lanes | |
| # slot storage (vals are stable uint64 PATCH IDS, never indices) | |
| self.keys = np.zeros((size, m), dtype=np.uint32) | |
| self.vals = np.zeros(size, dtype=np.uint64) | |
| self.counts = np.zeros(size, dtype=np.uint8) | |
| self.used = np.zeros(size, dtype=np.bool_) | |
| self.n_writes = 0 | |
| # ---------------------------------------------------------------- utils | |
| def primary_addr(self, lane_keys: np.ndarray) -> int: | |
| """Pattern-separation hash of the FULL key -> primary slot | |
| (dentate-gyrus style: every distinct context gets its own address).""" | |
| z = (int(lane_keys[0]) | |
| ^ ((int(lane_keys[1]) << 8 | int(lane_keys[1]) >> 24) & 0xFFFFFFFF) | |
| ^ ((int(lane_keys[2]) << 16 | int(lane_keys[2]) >> 16) & 0xFFFFFFFF) | |
| ^ ((int(lane_keys[3]) << 24 | int(lane_keys[3]) >> 8) & 0xFFFFFFFF)) | |
| return splitmix64(z) % self.size | |
| def loo_addr(self, lane_keys: np.ndarray, l: int) -> int: | |
| """Leave-one-out hash of the key minus lane l -> CA3 completion slot. | |
| A query whose lane l is corrupted can still compute this address | |
| from its intact lanes and reach the stored entry.""" | |
| a = [int(lane_keys[i]) for i in range(self.m) if i != l] | |
| z = a[0] ^ (((a[1] << 13) | (a[1] >> 19)) & 0xFFFFFFFF) \ | |
| ^ (((a[2] << 26) | (a[2] >> 6)) & 0xFFFFFFFF) | |
| return splitmix64(z ^ (0x9E3779B9 + l * 0x85EBCA6B)) % self.size | |
| def lane_keys(patch_ids) -> np.ndarray: | |
| """(m,) uint32 lane keys for a context of m patch ids.""" | |
| out = np.empty(len(patch_ids), dtype=np.uint32) | |
| for lane, pid in enumerate(patch_ids): | |
| out[lane] = splitmix64(int(pid) ^ LANE_SALTS[lane]) >> 32 | |
| return out | |
| def lane_keys_stream(self, patch_ids: np.ndarray) -> np.ndarray: | |
| """Vectorized lane keys for a whole token stream: (T, m) uint32.""" | |
| T = len(patch_ids) | |
| out = np.empty((T, self.m), dtype=np.uint32) | |
| for lane in range(self.m): | |
| # splitmix64 on an array of uint64 ids, take high 32 bits | |
| z = patch_ids.astype(np.uint64) ^ np.uint64(LANE_SALTS[lane]) | |
| 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[:, lane] = (z >> np.uint64(32)).astype(np.uint32) | |
| return out | |
| # ---------------------------------------------------------------- write | |
| def write(self, lane_keys: np.ndarray, val: int) -> None: | |
| """One-shot write of (context -> continuation patch id). | |
| Stores the entry at its PRIMARY (full-key hash) slot — pattern | |
| separation gives every distinct context its own address — plus | |
| best-effort copies at the m lane slots for corrupted-cue retrieval. | |
| """ | |
| # primary slot: exact-context capacity | |
| paddr = self.primary_addr(lane_keys) | |
| if not self.used[paddr]: | |
| self.keys[paddr] = lane_keys | |
| self.vals[paddr] = val | |
| self.counts[paddr] = 1 | |
| self.used[paddr] = True | |
| elif np.array_equal(self.keys[paddr], lane_keys): | |
| if self.vals[paddr] == val: | |
| if self.counts[paddr] < 255: | |
| self.counts[paddr] += 1 | |
| else: | |
| c = int(self.counts[paddr]) | |
| if c > 1: | |
| self.counts[paddr] = c - 1 | |
| else: | |
| self.vals[paddr] = val | |
| self.counts[paddr] = 1 | |
| elif int(self.counts[paddr]) <= 1: | |
| # evict unconsolidated memories only (keep consolidated ones) | |
| self.keys[paddr] = lane_keys | |
| self.vals[paddr] = val | |
| self.counts[paddr] = 1 | |
| # leave-one-out completion slots: robust to one corrupted lane | |
| for l in range(self.m): | |
| addr = self.loo_addr(lane_keys, l) | |
| if not self.used[addr]: | |
| self.keys[addr] = lane_keys | |
| self.vals[addr] = val | |
| self.counts[addr] = 1 | |
| self.used[addr] = True | |
| elif np.array_equal(self.keys[addr], lane_keys): | |
| if self.vals[addr] == val: | |
| if self.counts[addr] < 255: | |
| self.counts[addr] += 1 | |
| # else: a different key occupies this completion slot — keep it | |
| self.n_writes += 1 | |
| def write_stream(self, lane_stream: np.ndarray, vals: np.ndarray) -> None: | |
| """Bulk one-shot write over a stream (val[t] follows context of | |
| lane_stream[t]).""" | |
| for t in range(len(vals)): | |
| self.write(lane_stream[t], int(vals[t])) | |
| # ---------------------------------------------------------------- query | |
| def query(self, lane_keys: np.ndarray): | |
| """Query with the current context. Returns (patch_id, weight) or | |
| (None, 0.0). | |
| Two-tier voting: a candidate matching ALL m lanes dominates any | |
| partial match (a full-key entry must never lose to a partial match | |
| with a higher count); within the winning tier, votes scale with | |
| lane-match count x log evidence count. | |
| """ | |
| votes: dict[int, float] = {} | |
| max_match: dict[int, int] = {} | |
| # primary probe: exact full-key retrieval | |
| paddr = self.primary_addr(lane_keys) | |
| if self.used[paddr] and np.array_equal(self.keys[paddr], lane_keys): | |
| v = int(self.vals[paddr]) | |
| c = int(self.counts[paddr]) | |
| votes[v] = votes.get(v, 0.0) + self.m * (1.0 + np.log1p(c)) | |
| max_match[v] = self.m | |
| # leave-one-out probes: corrupted-cue completion (CA3) | |
| for l in range(self.m): | |
| addr = self.loo_addr(lane_keys, l) | |
| if not self.used[addr]: | |
| continue | |
| match = int(np.count_nonzero(self.keys[addr] == lane_keys)) | |
| if match >= self.min_lanes: | |
| v = int(self.vals[addr]) | |
| c = int(self.counts[addr]) | |
| w = match * (1.0 + np.log1p(c)) | |
| votes[v] = votes.get(v, 0.0) + w | |
| if match > max_match.get(v, 0): | |
| max_match[v] = match | |
| if not votes: | |
| return None, 0.0 | |
| if max(max_match.values()) == self.m: | |
| votes = {v: w for v, w in votes.items() if max_match[v] == self.m} | |
| best_v, best_w = max(votes.items(), key=lambda kv: kv[1]) | |
| return best_v, best_w | |
| # ---------------------------------------------------------------- stats | |
| def stats(self) -> dict: | |
| n_used = int(self.used.sum()) | |
| return { | |
| "slots": self.size, | |
| "used": n_used, | |
| "load_factor": n_used / self.size, | |
| "n_writes": self.n_writes, | |
| } | |
| def _slots_of(self, lane_keys: np.ndarray) -> list[int]: | |
| """All derivable slot addresses of a key: primary + leave-one-out.""" | |
| out = [self.primary_addr(lane_keys)] | |
| for l in range(self.m): | |
| out.append(self.loo_addr(lane_keys, l)) | |
| return out | |
| def top_entries(self, registry=None, min_count: int = 2, cap: int | None = None): | |
| """Deduplicated (lane_keys, val, count) list for export/eval, sorted by | |
| count descending. | |
| Dedup is EXACT: every slot an entry occupies is derivable from its | |
| key (the primary full-key hash plus the m leave-one-out hashes), so | |
| we mark all slots that actually hold this key. (Collisions may have | |
| blocked some redundant writes — those slots hold a different key and | |
| are not marked.) | |
| """ | |
| visited = np.zeros(self.size, dtype=np.bool_) | |
| items = [] | |
| idx = np.nonzero(self.used)[0] | |
| for addr in idx: | |
| if visited[addr]: | |
| continue | |
| k = self.keys[addr] | |
| c = int(self.counts[addr]) | |
| for a in self._slots_of(k): | |
| if self.used[a] and np.array_equal(self.keys[a], k): | |
| visited[a] = True | |
| if c >= min_count: | |
| items.append((k.copy(), int(self.vals[addr]), c)) | |
| items.sort(key=lambda it: it[2], reverse=True) | |
| if cap is not None: | |
| items = items[:cap] | |
| return items | |