""" Direction E: Neuromorphic Cortical-Hippocampal Loop for PC-SHO-DLM + MSA Maps the PC-SHO-DLM + MSA architecture onto biologically-named circuits: Cortical columns -> forward/feedback blocks (bottom-up / top-down) Hippocampal memory -> MSA MemoryBank (compressed KV store) Entorhinal routing -> MSA RouterProjector (query-key relevance) Reconsolidation -> post-settle param update from settled errors This is a *wrapper* over the existing architecture -- no reimplementation. Every public method delegates to real PCSHODLM / MSA methods and logs a structured timeline of (timestamp, event_type, magnitude) events for downstream analysis and prediction validation. """ import time from dataclasses import dataclass, field from typing import Dict, List, Optional, Tuple import torch import torch.nn as nn import torch.nn.functional as F from src.model import PCSHODLM, PCSHOConfig from src.msa import ( MSAConfig, MSALayer, MemoryBank, MemoryEncoder, RouterProjector, chunk_mean_pool, create_msa_layers, ) # ------------------------------------------------------------------------- # Timeline event record # ------------------------------------------------------------------------- @dataclass class TimelineEvent: """One observable event in the cortical-hippocampal loop.""" timestamp: float # monotonic wall-clock seconds since loop creation event_type: str # "cortical_error" | "hippocampal_retrieval" | "reconsolidation" magnitude: float # scalar summary of the event's intensity metadata: dict = field(default_factory=dict) # ------------------------------------------------------------------------- # CorticalHippocampalLoop # ------------------------------------------------------------------------- class CorticalHippocampalLoop(nn.Module): """Biologically-named wrapper mapping PC-SHO-DLM + MSA to brain circuits. Lifecycle for one stimulus (sequence of tokens): 1. cortex_process(stimulus) -- bottom-up forward pass 2. hippocampus_query(cortical_state) -- entorhinal routing scores 3. hippocampus_retrieve(indices) -- load compressed KV 4. cortex_settle(stimulus, memory) -- PC settling with context 5. reconsolidate(settled_state) -- unified parameter updates Every step is logged to ``self.timeline``. """ def __init__( self, model: PCSHODLM, msa_layers: nn.ModuleList, memory_bank: MemoryBank, msa_config: Optional[MSAConfig] = None, ): super().__init__() self.model = model self.msa_layers = msa_layers self.memory_bank = memory_bank self.msa_config = msa_config or MSAConfig() self._last_x_t: Optional[torch.Tensor] = None # Which model layers correspond to MSA-augmented layers n_layers = model.config.n_layers self._msa_start = max(0, n_layers - len(msa_layers)) # Observable timeline self.timeline: List[TimelineEvent] = [] self._t0 = time.monotonic() # -- internal helpers -------------------------------------------------- def _ts(self) -> float: """Seconds since construction.""" return time.monotonic() - self._t0 def _log(self, event_type: str, magnitude: float, **meta) -> None: self.timeline.append( TimelineEvent( timestamp=self._ts(), event_type=event_type, magnitude=magnitude, metadata=meta, ) ) # ===================================================================== # 1. cortex_process (cortical columns -> forward pass) # ===================================================================== def cortex_process( self, stimulus: torch.Tensor, t: Optional[torch.Tensor] = None, ) -> Tuple[list, torch.Tensor, torch.Tensor]: """Bottom-up forward pass through cortical columns. Args: stimulus: (B, S) token ids (possibly partially masked). t: (B,) diffusion timesteps. Defaults to t=1 (near-clean). Returns: h_init: list of hidden states [h_0 ... h_L] from amortized pass. mask: (B, S) boolean mask of corrupted positions. t: timestep tensor used. """ B, S = stimulus.shape device = stimulus.device if t is None: t = torch.ones(B, dtype=torch.long, device=device) # Corrupt at the requested timestep x_t, mask = self.model.schedule.corrupt(stimulus, t, self.model.config.mask_token_id) self._last_x_t = x_t.detach().clone() # Embed and run amortized forward pass (bottom-up cortical sweep) h_0 = self.model.embed_input(x_t, t) h_init = self.model.amortized_forward_pass(h_0) # Log cortical prediction error (mean L2 across layers) with torch.no_grad(): mu_up, mu_down = self.model.compute_predictions(h_init) L = self.model.config.n_layers eps_up = [h_init[l + 1] - mu_up[l] for l in range(L)] error_mag = sum(e.norm().item() for e in eps_up) / L self._log("cortical_error", error_mag, stage="process") return h_init, mask, t # ===================================================================== # 2. hippocampus_query (entorhinal routing -> MSA router scores) # ===================================================================== def hippocampus_query( self, cortical_state: list, ) -> Tuple[Dict[int, torch.Tensor], Dict[int, List[str]]]: """Route cortical state through entorhinal projectors to score memories. Args: cortical_state: hidden states [h_0 ... h_L]. Returns: all_scores: {layer_idx: (B, N_chunks) routing scores} all_chunk_ids: {layer_idx: list of doc-chunk identifiers} """ all_scores: Dict[int, torch.Tensor] = {} all_chunk_ids: Dict[int, List[str]] = {} if len(self.memory_bank) == 0: return all_scores, all_chunk_ids for msa_idx, msa_layer in enumerate(self.msa_layers): layer_idx = self._msa_start + msa_idx if layer_idx >= len(cortical_state) - 1: continue routing_keys, chunk_ids = self.memory_bank.get_routing_keys(layer_idx) if routing_keys is None: continue h_query = cortical_state[layer_idx + 1] # hidden state at this layer scores = msa_layer.compute_routing_scores(h_query, routing_keys) all_scores[layer_idx] = scores all_chunk_ids[layer_idx] = chunk_ids # Log retrieval event: mean top-1 score across layers if all_scores: top1_vals = [s.max().item() for s in all_scores.values()] mean_top1 = sum(top1_vals) / len(top1_vals) else: mean_top1 = 0.0 self._log("hippocampal_retrieval", mean_top1, stage="query") return all_scores, all_chunk_ids # ===================================================================== # 3. hippocampus_retrieve (load compressed KV for winning docs) # ===================================================================== def hippocampus_retrieve( self, scores: Dict[int, torch.Tensor], chunk_ids: Dict[int, List[str]], top_k: Optional[int] = None, ) -> Dict[int, Tuple[torch.Tensor, torch.Tensor]]: """Select top-k documents and fetch their compressed KV pairs. Args: scores: per-layer routing scores from hippocampus_query. chunk_ids: per-layer chunk doc-id lists. top_k: override for msa_config.top_k. Returns: memory_kv: {layer_idx: (K_compressed, V_compressed)} both (B, S_mem, D). """ top_k = top_k or self.msa_config.top_k memory_kv: Dict[int, Tuple[torch.Tensor, torch.Tensor]] = {} for layer_idx, layer_scores in scores.items(): cids = chunk_ids[layer_idx] B = layer_scores.shape[0] # Pick top-k chunks per batch element (use first batch element for doc selection) k = min(top_k, layer_scores.shape[1]) _, topk_idx = layer_scores[0].topk(k) selected_docs = list({cids[i.item()] for i in topk_idx}) K_sel, V_sel = self.memory_bank.get_kv(selected_docs, layer_idx) if K_sel is not None: # Expand to batch dimension K_sel = K_sel.unsqueeze(0).expand(B, -1, -1) V_sel = V_sel.unsqueeze(0).expand(B, -1, -1) memory_kv[layer_idx] = (K_sel, V_sel) # Log n_retrieved = sum(kv[0].shape[1] for kv in memory_kv.values()) if memory_kv else 0 self._log("hippocampal_retrieval", float(n_retrieved), stage="retrieve") return memory_kv # ===================================================================== # 4. cortex_settle (PC settling with retrieved hippocampal context) # ===================================================================== def cortex_settle( self, stimulus: torch.Tensor, h_init: list, mask: torch.Tensor, t: torch.Tensor, memory_kv: Optional[Dict[int, Tuple[torch.Tensor, torch.Tensor]]] = None, ) -> Tuple[list, list, list]: """Run PC-SHO settling, optionally injecting hippocampal memory. If ``memory_kv`` is provided, the retrieved context is injected into the corresponding MSA layers before each settling iteration by running an MSA forward pass that blends local representations with retrieved KV pairs. Args: stimulus: (B, S) original (uncorrupted) token ids -- used as x_0. h_init: amortized hidden states from cortex_process. mask: corruption mask. t: timestep tensor. memory_kv: retrieved memory from hippocampus_retrieve. Returns: h_settled: settled hidden states. v_final: final velocities. energies: energy trace. """ # Inject memory into h_init before settling: for each MSA layer, # run a single sparse-attention pass that conditions h on memory. if memory_kv: for layer_idx, (K_mem, V_mem) in memory_kv.items(): msa_idx = layer_idx - self._msa_start if 0 <= msa_idx < len(self.msa_layers): h_layer = h_init[layer_idx + 1] h_init[layer_idx + 1] = self.msa_layers[msa_idx]( h_layer, memory_k=K_mem, memory_v=V_mem ) # Delegate to the model's settle method h_settled, v_final, energies, _, _ = self.model.settle( h_init, stimulus, mask, t, ) # Log cortical error after settling with torch.no_grad(): mu_up, _ = self.model.compute_predictions(h_settled) L = self.model.config.n_layers eps_up = [h_settled[l + 1] - mu_up[l] for l in range(L)] error_mag = sum(e.norm().item() for e in eps_up) / L self._log("cortical_error", error_mag, stage="settle") return h_settled, v_final, energies # ===================================================================== # 5. reconsolidate (unified parameter updates) # ===================================================================== def reconsolidate( self, settled_state: list, stimulus: torch.Tensor, mask: torch.Tensor, t: torch.Tensor, param_lr_scale: float = 0.01, ) -> float: """Update parameters from settled prediction errors (reconsolidation). Uses the model's canonical post-settle updater so the parameter change is computed from the converged cortical state rather than from a transient microstep. Args: settled_state: hidden states after settling. stimulus: clean token ids (x_0). mask: corruption mask. t: timesteps. param_lr_scale: learning-rate multiplier for the slow (parameter) timescale. Returns: energy: energy after the reconsolidation step. """ # Snapshot parameter norms before update (for magnitude logging) with torch.no_grad(): pre_norms = { name: p.data.norm().item() for name, p in self.model.named_parameters() } # h_init = settled_state (reconsolidation starts from the settled point) h_init = [s.detach() for s in settled_state] update = self.model.post_settle_update( h_init, x_input=self._last_x_t if self._last_x_t is not None else stimulus, x_0=stimulus, mask=mask, t=t, param_lr_scale=param_lr_scale, energies=[ float(self.model.measure_energy(h_init, h_init, stimulus, mask, t)) ], ) energy = float(update["energy"]) # Compute parameter change magnitude with torch.no_grad(): delta = 0.0 count = 0 for name, p in self.model.named_parameters(): if name in pre_norms: delta += abs(p.data.norm().item() - pre_norms[name]) count += 1 mean_delta = delta / max(count, 1) self._log("reconsolidation", mean_delta, energy=energy) return energy # ===================================================================== # Full loop: stimulus -> ... -> reconsolidate # ===================================================================== def full_loop( self, stimulus: torch.Tensor, t: Optional[torch.Tensor] = None, n_hops: int = 1, reconsolidate: bool = True, param_lr_scale: float = 0.01, ) -> dict: """Run the complete cortical-hippocampal loop. Args: stimulus: (B, S) token ids. t: diffusion timesteps. n_hops: number of query-retrieve-settle iterations (multi-hop). reconsolidate: whether to run the reconsolidation step. param_lr_scale: LR scale for reconsolidation. Returns: Dictionary with h_settled, energies, and the full timeline. """ # 1. Cortical forward pass h_init, mask, t = self.cortex_process(stimulus, t) h_current = h_init all_energies: List[float] = [] for hop in range(n_hops): # 2. Hippocampal query scores, chunk_ids = self.hippocampus_query(h_current) # 3. Hippocampal retrieve memory_kv = self.hippocampus_retrieve(scores, chunk_ids) # 4. Cortex settle with memory h_settled, v_final, energies = self.cortex_settle( stimulus, [h.clone() for h in h_current], mask, t, memory_kv ) h_current = h_settled all_energies.extend(energies) # 5. Reconsolidate recon_energy = None if reconsolidate: recon_energy = self.reconsolidate( h_current, stimulus, mask, t, param_lr_scale ) return { "h_settled": h_current, "energies": all_energies, "reconsolidation_energy": recon_energy, "timeline": list(self.timeline), # copy } # ===================================================================== # 6/7. Validation predicates # ===================================================================== def validate_predictions( self, retrieval_tolerance: float = 0.0, ) -> Dict[str, bool]: """Check three neuroscience-inspired predictions against the timeline. (a) Cortical errors precede hippocampal retrieval: every retrieval event is preceded by a cortical_error event. (b) Parameters change during recall: at least one reconsolidation event has magnitude > 0. (c) Retrieval precision improves across hops: hippocampal_retrieval magnitudes at 'query' stage are non-decreasing across successive hops (within tolerance). For untrained models use ``retrieval_tolerance > 0`` to account for random-weight noise. Args: retrieval_tolerance: allowed fractional decrease per hop for (c). 0.0 = strict non-decreasing, 0.1 = allow 10% drop, etc. Returns: {"errors_precede_retrieval": bool, "params_change_during_recall": bool, "retrieval_improves_across_hops": bool} """ # (a) Errors precede retrieval errors_precede = True last_cortical_ts = None for evt in self.timeline: if evt.event_type == "cortical_error": last_cortical_ts = evt.timestamp elif evt.event_type == "hippocampal_retrieval": if last_cortical_ts is None: errors_precede = False break # (b) Parameters change during recall recon_events = [e for e in self.timeline if e.event_type == "reconsolidation"] params_change = any(e.magnitude > 0.0 for e in recon_events) # (c) Retrieval precision improves across hops query_scores = [ e.magnitude for e in self.timeline if e.event_type == "hippocampal_retrieval" and e.metadata.get("stage") == "query" ] if len(query_scores) >= 2: retrieval_improves = all( query_scores[i + 1] >= query_scores[i] * (1.0 - retrieval_tolerance) - 1e-6 for i in range(len(query_scores) - 1) ) else: # Vacuously true with < 2 hops retrieval_improves = True return { "errors_precede_retrieval": errors_precede, "params_change_during_recall": params_change, "retrieval_improves_across_hops": retrieval_improves, } # ===================================================================== # 8. plot_dynamics (3-panel visualization) # ===================================================================== def plot_dynamics(self, save_path: Optional[str] = None) -> None: """Create a 3-panel dynamics visualization. Panel 1: Cortical error magnitude over time. Panel 2: Hippocampal retrieval scores over time. Panel 3: Reconsolidation parameter-change magnitude over time. Falls back to a text table if matplotlib is not available. """ # Separate timeline events by type cortical = [(e.timestamp, e.magnitude) for e in self.timeline if e.event_type == "cortical_error"] retrieval = [(e.timestamp, e.magnitude) for e in self.timeline if e.event_type == "hippocampal_retrieval"] recon = [(e.timestamp, e.magnitude) for e in self.timeline if e.event_type == "reconsolidation"] try: import matplotlib.pyplot as plt fig, axes = plt.subplots(1, 3, figsize=(15, 4)) if cortical: ts, mag = zip(*cortical) axes[0].plot(ts, mag, "o-", color="tab:blue") axes[0].set_title("Cortical Error") axes[0].set_xlabel("Time (s)") axes[0].set_ylabel("Mean L2 Error") if retrieval: ts, mag = zip(*retrieval) axes[1].plot(ts, mag, "s-", color="tab:orange") axes[1].set_title("Hippocampal Retrieval") axes[1].set_xlabel("Time (s)") axes[1].set_ylabel("Score / Count") if recon: ts, mag = zip(*recon) axes[2].plot(ts, mag, "^-", color="tab:green") axes[2].set_title("Reconsolidation") axes[2].set_xlabel("Time (s)") axes[2].set_ylabel("Param Delta") plt.tight_layout() if save_path: plt.savefig(save_path, dpi=150) print(f"Saved plot to {save_path}") else: plt.show() except ImportError: # Fallback: text table print("=" * 72) print("CORTICAL-HIPPOCAMPAL LOOP DYNAMICS") print("=" * 72) print("\n--- Panel 1: Cortical Error ---") for ts, mag in cortical: print(f" t={ts:.4f}s error={mag:.6f}") print("\n--- Panel 2: Hippocampal Retrieval ---") for ts, mag in retrieval: print(f" t={ts:.4f}s score={mag:.6f}") print("\n--- Panel 3: Reconsolidation ---") for ts, mag in recon: print(f" t={ts:.4f}s delta={mag:.8f}") print("=" * 72) # ===================================================================== # Utility: reset timeline for a fresh trial # ===================================================================== def reset_timeline(self) -> None: """Clear all recorded events and reset the clock.""" self.timeline.clear() self._t0 = time.monotonic()