""" Direction G: Self-Improving Retrieval for PC-SHO-DLM + MSA A retrieval system that improves with every query -- no explicit retraining. Each query now uses the repaired unified path: 1. Settle hidden states toward a low-energy solution (fast timescale) 2. Apply model updates once from the settled state (slow timescale) 3. Update router parameters from the settled retrieval signal Convergence guarantee (Borkar 2008 two-timescale + PC contraction): E[||W_QR^N - W_QR*||^2] = O(1 / sqrt(N)) Key insight: predictive-coding settling supplies the local error signals, but the stable training rule is to update model parameters from settled states rather than from transient microsteps. Router weights still adapt online from the retrieval signal within the query. Safety mechanisms: - Elastic regularization: prevents catastrophic drift from initial weights - Snapshot/rollback: revert if quality degrades - Drift monitoring: ||theta_n - theta_0|| / ||theta_0|| tracked per query """ import copy import math from dataclasses import dataclass, field from typing import Optional, Tuple, List, Dict import torch import torch.nn as nn import torch.nn.functional as F from model import PCSHODLM, PCSHOConfig, InferenceUpdater from msa import ( MSAConfig, MSALayer, MemoryBank, MemoryEncoder, RouterProjector, create_msa_layers, chunk_mean_pool, compute_routing_aux_loss, ) # ============================================================================ # Configuration # ============================================================================ @dataclass class SelfImprovingConfig: """Configuration for the self-improving retrieval system.""" # Elastic regularization elastic_lambda: float = 0.01 # L_elastic = lambda * ||theta - theta_0||^2 drift_threshold: float = 0.10 # activate elastic reg when drift > 10% drift_hard_cap: float = 0.30 # force rollback if drift > 30% # Unified settling for retrieval queries n_settling_steps: int = 6 # inner loop iterations per query param_lr_scale: float = 0.01 # slow timescale for parameters # Quality tracking quality_ema_alpha: float = 0.1 # exponential moving average smoothing quality_window: int = 10 # window for rolling average # Snapshot policy snapshot_every: int = 10 # save snapshot every N queries max_snapshots: int = 5 # keep at most this many snapshots # Router-specific learning rate scaling router_lr_boost: float = 2.0 # router params get boosted LR readout_lr_scale: float = 0.5 # readout params get reduced LR # ============================================================================ # Retrieval Quality Metric # ============================================================================ class RetrievalQualityTracker: """Tracks Q_N = retrieval quality at query N. Quality is measured as a composite of: - Router confidence: max routing score for selected documents - Settling energy reduction: E_final / E_initial (lower is better) - Answer coherence: negative entropy of output distribution All three are normalized to [0, 1] and combined. """ def __init__(self, ema_alpha: float = 0.1): self.ema_alpha = ema_alpha self.history: List[float] = [] self.components: List[Dict[str, float]] = [] self._ema = 0.0 self._initialized = False def record(self, router_confidence: float, energy_ratio: float, answer_coherence: float) -> float: """Record quality for one query and return composite Q_N.""" # Router confidence: already in [0, 1] (cosine similarity based) q_router = max(0.0, min(1.0, router_confidence)) # Energy ratio: E_final / E_initial. Lower = better settling. # Map to [0, 1] where 1 = perfect settling (ratio -> 0) q_energy = max(0.0, min(1.0, 1.0 - energy_ratio)) # Answer coherence: negative entropy normalized by log(vocab_size) # Higher coherence (lower entropy) = better. Already in [0, 1]. q_coherence = max(0.0, min(1.0, answer_coherence)) # Composite: weighted average q_n = 0.4 * q_router + 0.3 * q_energy + 0.3 * q_coherence self.history.append(q_n) self.components.append({ "router_confidence": q_router, "energy_reduction": q_energy, "answer_coherence": q_coherence, "composite": q_n, }) # Update EMA if not self._initialized: self._ema = q_n self._initialized = True else: self._ema = self.ema_alpha * q_n + (1 - self.ema_alpha) * self._ema return q_n @property def current_quality(self) -> float: return self._ema if self._initialized else 0.0 @property def n_queries(self) -> int: return len(self.history) def get_improvement_curve(self) -> List[float]: """Return the full Q_N sequence.""" return list(self.history) def get_rolling_average(self, window: int = 10) -> List[float]: """Return rolling average of quality for smoother visualization.""" if len(self.history) < window: return list(self.history) result = [] for i in range(len(self.history)): start = max(0, i - window + 1) result.append(sum(self.history[start:i + 1]) / (i - start + 1)) return result # ============================================================================ # Drift Monitor # ============================================================================ class DriftMonitor: """Monitors parameter drift: drift_n = ||theta_n - theta_0|| / ||theta_0||. Provides per-component drift (router, forward blocks, feedback, readout) and aggregate drift for the elastic regularization trigger. """ def __init__(self): self._theta_0: Optional[Dict[str, torch.Tensor]] = None self._theta_0_norm: float = 0.0 self.history: List[float] = [] self.component_history: List[Dict[str, float]] = [] def set_baseline(self, model: nn.Module, msa_layers: nn.ModuleList) -> None: """Snapshot initial parameters as theta_0.""" self._theta_0 = {} total_norm_sq = 0.0 for name, p in model.named_parameters(): self._theta_0[f"model.{name}"] = p.data.clone() total_norm_sq += p.data.norm().item() ** 2 for name, p in msa_layers.named_parameters(): self._theta_0[f"msa.{name}"] = p.data.clone() total_norm_sq += p.data.norm().item() ** 2 self._theta_0_norm = math.sqrt(total_norm_sq) def compute_drift(self, model: nn.Module, msa_layers: nn.ModuleList) -> float: """Compute current drift from baseline. Returns scalar drift ratio.""" if self._theta_0 is None: return 0.0 delta_sq = 0.0 component_deltas = {"router": 0.0, "forward": 0.0, "feedback": 0.0, "other": 0.0} for name, p in model.named_parameters(): key = f"model.{name}" if key in self._theta_0: d = (p.data - self._theta_0[key].to(p.device)).norm().item() ** 2 delta_sq += d if "forward_blocks" in name: component_deltas["forward"] += d elif "feedback_blocks" in name: component_deltas["feedback"] += d else: component_deltas["other"] += d for name, p in msa_layers.named_parameters(): key = f"msa.{name}" if key in self._theta_0: d = (p.data - self._theta_0[key].to(p.device)).norm().item() ** 2 delta_sq += d if "router" in name: component_deltas["router"] += d else: component_deltas["other"] += d drift = math.sqrt(delta_sq) / max(self._theta_0_norm, 1e-10) self.history.append(drift) # Normalize component deltas component_drift = { k: math.sqrt(v) / max(self._theta_0_norm, 1e-10) for k, v in component_deltas.items() } self.component_history.append(component_drift) return drift def get_elastic_penalty(self, model: nn.Module, msa_layers: nn.ModuleList, lam: float) -> torch.Tensor: """Compute L_elastic = lambda * ||theta - theta_0||^2. Returns a differentiable scalar loss to be added to the energy. """ if self._theta_0 is None: return torch.tensor(0.0) penalty = torch.tensor(0.0, device=next(model.parameters()).device) for name, p in model.named_parameters(): key = f"model.{name}" if key in self._theta_0: penalty = penalty + (p - self._theta_0[key].to(p.device)).pow(2).sum() for name, p in msa_layers.named_parameters(): key = f"msa.{name}" if key in self._theta_0: penalty = penalty + (p - self._theta_0[key].to(p.device)).pow(2).sum() return lam * penalty # ============================================================================ # Self-Improving Retriever # ============================================================================ class SelfImprovingRetriever: """A retrieval system that improves with every query. Wraps PC-SHO-DLM model + MSA layers + MemoryBank into a unified retrieval engine where each query triggers settling, then applies post-settle model updates and router adaptation. Convergence bound: E[||W_QR^N - W_QR*||^2] = O(1 / sqrt(N)) This follows from Borkar (2008) two-timescale stochastic approximation: the fast process (hidden state settling) converges at rate O(1/K) per query, while the slow process (parameter updates) converges at rate O(1/sqrt(N)) over queries, because the effective noise variance is bounded by the settling residual which contracts geometrically. Usage: retriever = SelfImprovingRetriever(model, msa_layers, memory_bank, config) for text in queries: answer = retriever.query(text) curve = retriever.get_improvement_curve() """ def __init__( self, model: PCSHODLM, msa_layers: nn.ModuleList, memory_bank: MemoryBank, config: Optional[SelfImprovingConfig] = None, device: str = "cpu", ): self.model = model self.msa_layers = msa_layers self.memory_bank = memory_bank self.config = config or SelfImprovingConfig() self.device = device # Core tracking self.quality_tracker = RetrievalQualityTracker( ema_alpha=self.config.quality_ema_alpha ) self.drift_monitor = DriftMonitor() self.drift_monitor.set_baseline(model, msa_layers) # Query counter self._query_count = 0 # Snapshot management self._snapshots: List[Dict[str, torch.Tensor]] = [] self._snapshot_queries: List[int] = [] self._save_snapshot() # initial snapshot # Energy history per query (for diagnostics) self.energy_traces: List[List[float]] = [] # ------------------------------------------------------------------ # Snapshot / Rollback # ------------------------------------------------------------------ def _save_snapshot(self) -> None: """Save current parameters as a snapshot.""" snapshot = {} for name, p in self.model.named_parameters(): snapshot[f"model.{name}"] = p.data.clone() for name, p in self.msa_layers.named_parameters(): snapshot[f"msa.{name}"] = p.data.clone() self._snapshots.append(snapshot) self._snapshot_queries.append(self._query_count) # Prune old snapshots while len(self._snapshots) > self.config.max_snapshots: self._snapshots.pop(0) self._snapshot_queries.pop(0) def _restore_snapshot(self, idx: int = -1) -> None: """Restore parameters from a snapshot.""" snapshot = self._snapshots[idx] for name, p in self.model.named_parameters(): key = f"model.{name}" if key in snapshot: p.data.copy_(snapshot[key]) for name, p in self.msa_layers.named_parameters(): key = f"msa.{name}" if key in snapshot: p.data.copy_(snapshot[key]) def reset_to_original(self) -> None: """Reset all parameters to the original (query 0) state.""" self._restore_snapshot(0) self._query_count = 0 self.quality_tracker = RetrievalQualityTracker( ema_alpha=self.config.quality_ema_alpha ) self.drift_monitor = DriftMonitor() self.drift_monitor.set_baseline(self.model, self.msa_layers) self._snapshots = self._snapshots[:1] self._snapshot_queries = self._snapshot_queries[:1] self.energy_traces = [] def save_state(self, path: str) -> None: """Save full retriever state to disk.""" state = { "model_state": self.model.state_dict(), "msa_state": self.msa_layers.state_dict(), "quality_history": self.quality_tracker.history, "quality_components": self.quality_tracker.components, "drift_history": self.drift_monitor.history, "drift_components": self.drift_monitor.component_history, "query_count": self._query_count, "energy_traces": self.energy_traces, "config": self.config, } torch.save(state, path) def load_state(self, path: str) -> None: """Load retriever state from disk.""" state = torch.load(path, map_location=self.device, weights_only=False) self.model.load_state_dict(state["model_state"]) self.msa_layers.load_state_dict(state["msa_state"]) self.quality_tracker.history = state["quality_history"] self.quality_tracker.components = state["quality_components"] self.drift_monitor.history = state["drift_history"] self.drift_monitor.component_history = state["drift_components"] self._query_count = state["query_count"] self.energy_traces = state["energy_traces"] if "config" in state: self.config = state["config"] # ------------------------------------------------------------------ # Core: Query Processing with Self-Improvement # ------------------------------------------------------------------ def _tokenize(self, text: str) -> torch.Tensor: """Simple byte-level tokenization (matches MemoryEncoder).""" tokens = [min(b + 1, self.model.config.vocab_size - 1) for b in text.encode("utf-8")[:self.model.config.max_seq_len]] t = torch.tensor(tokens, dtype=torch.long, device=self.device).unsqueeze(0) if t.shape[1] < self.model.config.max_seq_len: t = F.pad(t, (0, self.model.config.max_seq_len - t.shape[1])) return t def _retrieve_documents(self, h_query: torch.Tensor, layer_idx: int ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], float, List[str]]: """Route query through MSA to retrieve relevant documents. Returns: memory_k: compressed keys from top-k docs (or None) memory_v: compressed values from top-k docs (or None) router_confidence: max routing score (for quality tracking) selected_ids: IDs of selected documents """ if len(self.memory_bank) == 0: return None, None, 0.0, [] msa_start = len(self.model.forward_blocks) // 2 msa_idx = layer_idx - msa_start if msa_idx < 0 or msa_idx >= len(self.msa_layers): return None, None, 0.0, [] msa_layer = self.msa_layers[msa_idx] # Get all routing keys from memory bank routing_keys, chunk_doc_ids = self.memory_bank.get_routing_keys(layer_idx) if routing_keys is None: return None, None, 0.0, [] routing_keys = routing_keys.to(self.device) # Compute routing scores scores = msa_layer.compute_routing_scores(h_query, routing_keys) # (1, N_chunks) # Top-k selection k = min(self.msa_layers[0].msa_config.top_k, scores.shape[1]) top_scores, top_indices = scores.topk(k, dim=1) router_confidence = top_scores.max().item() # Map chunk indices to document IDs selected_doc_ids = list(set( chunk_doc_ids[idx.item()] for idx in top_indices[0] if idx.item() < len(chunk_doc_ids) )) # Load compressed KV for selected documents memory_k, memory_v = self.memory_bank.get_kv(selected_doc_ids, layer_idx) if memory_k is not None: memory_k = memory_k.unsqueeze(0).to(self.device) # (1, S_mem, D) memory_v = memory_v.unsqueeze(0).to(self.device) return memory_k, memory_v, router_confidence, selected_doc_ids def query(self, text: str) -> Dict: """Process a query with self-improving retrieval. This is the main entry point. Each call: 1. Tokenizes the query 2. Runs the forward pass through lower layers 3. At MSA layers, routes to memory bank and retrieves documents 4. Runs shared settling, then post-settle model and router updates 5. Decodes the answer 6. Tracks quality, drift, and applies elastic regularization if needed Args: text: query text Returns: dict with keys: answer_logits, answer_tokens, quality, drift, energy_trace, retrieved_docs """ self._query_count += 1 config = self.config model = self.model mc = model.config # Tokenize tokens = self._tokenize(text) B, S = tokens.shape # Create a partial mask: treat last 25% of non-padding tokens as # "to predict" (simulates the query -> answer pattern) non_pad = (tokens != 0).sum(dim=1).item() mask_start = max(1, int(non_pad * 0.75)) mask = torch.zeros(B, S, dtype=torch.bool, device=self.device) mask[0, mask_start:non_pad] = True # If no tokens to predict, mask the last token if not mask.any(): mask[0, max(0, non_pad - 1)] = True # Timestep (low noise -- we want mostly-clean settling) t = torch.ones(B, dtype=torch.long, device=self.device) # Embed and forward through lower layers h_0 = model.embed_input(tokens, t) h = [h_0] current = h_0 msa_start = len(model.forward_blocks) // 2 max_router_conf = 0.0 all_retrieved_docs = [] # Forward pass with MSA retrieval at upper layers for l, block in enumerate(model.forward_blocks): current = block(current) # At MSA layers: retrieve from memory if l >= msa_start: mem_k, mem_v, conf, doc_ids = self._retrieve_documents(current, l) max_router_conf = max(max_router_conf, conf) all_retrieved_docs.extend(doc_ids) # Apply sparse attention from MSA layer if we retrieved docs if mem_k is not None: msa_idx = l - msa_start if msa_idx < len(self.msa_layers): current = self.msa_layers[msa_idx](current, mem_k, mem_v) h.append(current) # Store h_init for settling h_init = [hi.detach() for hi in h] # --- Settling followed by canonical post-settle learning --- L = model.n_active_layers v = [torch.zeros_like(h_init[l + 1]) for l in range(L)] energies = [] # Check drift before settling to decide on elastic regularization current_drift = self.drift_monitor.compute_drift(model, self.msa_layers) use_elastic = current_drift > config.drift_threshold # Hard cap: rollback if drift is too large if current_drift > config.drift_hard_cap and len(self._snapshots) > 1: self._restore_snapshot(-1) current_drift = self.drift_monitor.compute_drift(model, self.msa_layers) for k in range(config.n_settling_steps): # Adaptive active tokens after first step if k > 0: with torch.no_grad(): uncertainty = model.compute_token_uncertainty(h) active_tokens = torch.sigmoid( (uncertainty - mc.settling_threshold) / mc.settling_temperature ) else: active_tokens = None h, v, energy = model.settling_step( h, v, h_init, tokens, mask, t, active_tokens=active_tokens, ) energies.append(energy) model.post_settle_update( h, x_input=tokens, x_0=tokens, mask=mask, t=t, param_lr_scale=config.param_lr_scale, energies=energies, ) # Apply elastic regularization after the canonical model update. if use_elastic: penalty = self.drift_monitor.get_elastic_penalty( model, self.msa_layers, config.elastic_lambda ) if penalty.requires_grad: model.zero_grad(set_to_none=True) self.msa_layers.zero_grad(set_to_none=True) penalty.backward() with torch.no_grad(): lr = config.param_lr_scale for p in list(model.parameters()) + list(self.msa_layers.parameters()): if p.grad is not None: p.data -= lr * p.grad p.grad.zero_() # Also update MSA router parameters from routing errors self._update_routers(h, tokens, mask, t) self.energy_traces.append(energies) # --- Decode answer --- with torch.no_grad(): logits = model.readout(model.readout_norm(h[-1])) probs = F.softmax(logits, dim=-1) answer_tokens = logits[0, mask_start:non_pad].argmax(dim=-1) # Compute quality components energy_ratio = energies[-1] / max(energies[0], 1e-8) if energies else 1.0 answer_probs = probs[0, mask_start:non_pad] entropy = -(answer_probs * (answer_probs + 1e-10).log()).sum(dim=-1) max_entropy = math.log(mc.vocab_size) answer_coherence = 1.0 - (entropy.mean().item() / max_entropy) # Record quality q_n = self.quality_tracker.record( router_confidence=max_router_conf, energy_ratio=max(0.0, min(1.0, energy_ratio)), answer_coherence=answer_coherence, ) # Periodic snapshot if self._query_count % config.snapshot_every == 0: self._save_snapshot() return { "answer_logits": logits, "answer_tokens": answer_tokens, "quality": q_n, "drift": current_drift, "energy_trace": energies, "retrieved_docs": list(set(all_retrieved_docs)), "query_number": self._query_count, } def _update_routers(self, h_settled: list, tokens: torch.Tensor, mask: torch.Tensor, t: torch.Tensor) -> None: """Update MSA router parameters using settled hidden states. The router projectors (W_QR, W_KR) are updated via the contrastive routing loss, using the settled states as signal for what the "correct" routing should have been. """ msa_start = len(self.model.forward_blocks) // 2 for msa_idx, msa_layer in enumerate(self.msa_layers): layer_idx = msa_start + msa_idx if layer_idx + 1 >= len(h_settled): continue h_at_layer = h_settled[layer_idx + 1].detach() # Get routing keys from memory routing_keys, _ = self.memory_bank.get_routing_keys(layer_idx) if routing_keys is None: continue routing_keys = routing_keys.to(self.device) # Compute current routing scores scores = msa_layer.compute_routing_scores(h_at_layer, routing_keys) # Self-supervised signal: top-scored docs are "positive", # bottom-scored are "negative" k = min(self.msa_layers[0].msa_config.top_k, scores.shape[1]) if scores.shape[1] <= k: continue _, top_idx = scores.topk(k, dim=1) _, bot_idx = scores.topk(scores.shape[1] - k, dim=1, largest=False) scores_pos = scores.gather(1, top_idx) scores_neg = scores.gather(1, bot_idx) # Contrastive loss for router router_loss = compute_routing_aux_loss( scores_pos, scores_neg, temperature=msa_layer.msa_config.aux_temperature, ) if router_loss.requires_grad: router_loss.backward() lr = self.config.param_lr_scale * self.config.router_lr_boost with torch.no_grad(): nn.utils.clip_grad_norm_(msa_layer.router.parameters(), 1.0) for p in msa_layer.router.parameters(): if p.grad is not None: p.data -= lr * p.grad p.grad.zero_() # ------------------------------------------------------------------ # Diagnostics # ------------------------------------------------------------------ def get_improvement_curve(self) -> List[float]: """Return Q_1, Q_2, ..., Q_N quality sequence.""" return self.quality_tracker.get_improvement_curve() def get_drift_curve(self) -> List[float]: """Return drift_1, drift_2, ..., drift_N.""" return self.drift_monitor.history def get_diagnostics(self) -> Dict: """Return comprehensive diagnostics.""" return { "n_queries": self._query_count, "current_quality": self.quality_tracker.current_quality, "quality_curve": self.get_improvement_curve(), "quality_rolling": self.quality_tracker.get_rolling_average( self.config.quality_window ), "drift_curve": self.get_drift_curve(), "drift_components": self.drift_monitor.component_history, "energy_traces": self.energy_traces, "n_snapshots": len(self._snapshots), "memory_bank_size": len(self.memory_bank), } def theoretical_bound(self, N: int) -> float: """Compute the theoretical convergence bound at query N. E[||W_QR^N - W_QR*||^2] = C / sqrt(N) The constant C depends on the settling contraction rate rho and the noise variance sigma^2 of the stochastic gradient: C = sigma^2 / (1 - rho^K) where K = n_settling_steps and rho < 1 is the SHO contraction rate. We estimate C from the empirical quality curve. """ if N == 0: return float("inf") # Estimate C from the first few queries if len(self.quality_tracker.history) >= 2: q1 = 1.0 - self.quality_tracker.history[0] c_est = q1 # rough: error at N=1 should be ~C/1 else: c_est = 1.0 return c_est / math.sqrt(N) # ============================================================================ # Simulation: demonstrate self-improvement over 50 queries # ============================================================================ def run_simulation(n_queries: int = 50, device: str = "cpu") -> Dict: """Run a self-improving retrieval simulation. Creates a small model, populates a memory bank with synthetic documents, and issues a sequence of queries. Each query triggers unified settling that updates both hidden states and retrieval parameters. Returns: Dict with quality curve, drift curve, energy traces, and diagnostics. """ print("=" * 70) print("Direction G: Self-Improving Retrieval Simulation") print("=" * 70) # --- Setup --- model_config = PCSHOConfig( vocab_size=300, max_seq_len=128, d_model=128, n_heads=4, n_layers=4, d_ff=256, n_diffusion_steps=50, n_settling_steps=4, eta_base=0.05, online_learn_lr=1e-4, ) msa_config = MSAConfig( chunk_size=32, top_k=4, router_dim=64, n_router_heads=4, apply_from_layer=2, ) si_config = SelfImprovingConfig( elastic_lambda=0.005, drift_threshold=0.15, drift_hard_cap=0.40, n_settling_steps=4, param_lr_scale=0.02, quality_ema_alpha=0.15, snapshot_every=10, ) print(f"\nModel: d={model_config.d_model}, L={model_config.n_layers}, " f"H={model_config.n_heads}") print(f"MSA: top_k={msa_config.top_k}, router_dim={msa_config.router_dim}") print(f"Settling steps per query: {si_config.n_settling_steps}") # Create model and MSA layers model = PCSHODLM(model_config).to(device) msa_layers = create_msa_layers(model_config, msa_config).to(device) memory_bank = MemoryBank(chunk_size=msa_config.chunk_size) param_count = sum(p.numel() for p in model.parameters()) msa_param_count = sum(p.numel() for p in msa_layers.parameters()) print(f"Parameters: model={param_count:,}, MSA={msa_param_count:,}") # --- Populate memory bank with synthetic documents --- documents = [ "The speed of light in vacuum is approximately 299792458 meters per second.", "Photosynthesis converts carbon dioxide and water into glucose and oxygen.", "The Pythagorean theorem states that a squared plus b squared equals c squared.", "DNA stores genetic information using four nucleotide bases: A T G and C.", "Gravity is the force of attraction between objects with mass.", "Water freezes at zero degrees Celsius and boils at one hundred degrees.", "The mitochondria are the powerhouse of the cell.", "Newtons first law states an object in motion stays in motion.", "The periodic table organizes elements by atomic number and properties.", "Evolution by natural selection drives adaptation in populations.", "Quantum mechanics describes behavior of matter at atomic scales.", "The human genome contains approximately three billion base pairs.", "Plate tectonics explains the movement of Earths lithospheric plates.", "Entropy always increases in an isolated system.", "General relativity describes gravity as curvature of spacetime.", ] print(f"\nEncoding {len(documents)} documents into memory bank...") for i, doc in enumerate(documents): MemoryEncoder.encode_document( model, doc, f"doc_{i}", memory_bank, msa_layers, chunk_size=msa_config.chunk_size, device=device, ) print(f"Memory bank: {len(memory_bank)} documents") # --- Build retriever --- retriever = SelfImprovingRetriever( model=model, msa_layers=msa_layers, memory_bank=memory_bank, config=si_config, device=device, ) # --- Query sequence --- queries = [ "What is the speed of light?", "How do plants make food?", "What is the Pythagorean theorem?", "What are the bases of DNA?", "Why do objects fall?", "At what temperature does water freeze?", "What produces energy in cells?", "What happens to moving objects?", "How are chemical elements organized?", "What drives evolution?", "How do atoms behave?", "How large is the human genome?", "What moves the continents?", "Does entropy increase or decrease?", "How does gravity work in general relativity?", # Repeat with variations to show learning "Tell me about light speed.", "Explain photosynthesis.", "Describe the Pythagorean relationship.", "What nucleotides make up DNA?", "Why is there gravity?", "When does water boil?", "Where is energy made in a cell?", "Do objects keep moving?", "What is the periodic table?", "How does natural selection work?", "What is quantum mechanics about?", "How many base pairs in human DNA?", "What are tectonic plates?", "Explain the second law of thermodynamics.", "Describe spacetime curvature.", # More variations "Light travels at what speed?", "CO2 and water become what in plants?", "Right triangles follow what rule?", "Adenine thymine guanine cytosine are what?", "Mass attracts mass through what force?", "Zero degrees Celsius is the freezing point of what?", "Mitochondria function is what?", "Inertia means what?", "Elements are ordered by what?", "Survival of the fittest is part of what?", "Subatomic particles follow what physics?", "Three billion base pairs are in what?", "Continental drift is caused by what?", "Isolated systems and entropy?", "Einstein described gravity as what?", # Final batch "Speed of electromagnetic radiation in vacuum?", "Chloroplasts perform what process?", "a^2 + b^2 = c^2 is called what?", "The double helix stores information using what?", "What bends spacetime?", ] queries = queries[:n_queries] print(f"\nRunning {len(queries)} queries with self-improving retrieval...\n") print(f"{'Query':>5} | {'Q_N':>6} | {'Drift':>7} | {'E_ratio':>8} | {'Retrieved':>9} | Text") print("-" * 90) for i, q in enumerate(queries): result = retriever.query(q) # Energy ratio for display etrace = result["energy_trace"] e_ratio = etrace[-1] / max(etrace[0], 1e-8) if len(etrace) >= 2 else 1.0 print(f"{result['query_number']:>5} | {result['quality']:>6.3f} | " f"{result['drift']:>7.4f} | {e_ratio:>8.4f} | " f"{len(result['retrieved_docs']):>9} | {q[:40]}") # --- Summary --- diagnostics = retriever.get_diagnostics() curve = diagnostics["quality_curve"] drift = diagnostics["drift_curve"] print("\n" + "=" * 70) print("RESULTS SUMMARY") print("=" * 70) # Quality improvement first_5 = sum(curve[:5]) / min(5, len(curve)) last_5 = sum(curve[-5:]) / min(5, len(curve)) print(f"\nRetrieval Quality (Q_N):") print(f" First 5 queries (avg): {first_5:.4f}") print(f" Last 5 queries (avg): {last_5:.4f}") print(f" Improvement: {last_5 - first_5:+.4f} ({(last_5/max(first_5,1e-8) - 1)*100:+.1f}%)") print(f" Final EMA quality: {diagnostics['current_quality']:.4f}") # Drift if drift: print(f"\nParameter Drift:") print(f" Final drift: {drift[-1]:.4f}") print(f" Max drift: {max(drift):.4f}") print(f" Elastic reg activated: {sum(1 for d in drift if d > si_config.drift_threshold)} times") # Convergence bound print(f"\nConvergence Bound E[||W_QR^N - W_QR*||^2] = O(1/sqrt(N)):") for n in [1, 10, 25, 50]: if n <= n_queries: bound = retriever.theoretical_bound(n) actual = 1.0 - (curve[n - 1] if n <= len(curve) else curve[-1]) print(f" N={n:>3}: bound={bound:.4f}, actual_error={actual:.4f}") print(f"\nSnapshots saved: {diagnostics['n_snapshots']}") print(f"Memory bank: {diagnostics['memory_bank_size']} documents") # ASCII quality curve print(f"\nQuality Curve (Q_N over queries):") rolling = diagnostics["quality_rolling"] if rolling: max_q = max(rolling) if max(rolling) > 0 else 1.0 min_q = min(rolling) bar_width = 40 for i, q in enumerate(rolling): if i % max(1, len(rolling) // 20) == 0 or i == len(rolling) - 1: normalized = (q - min_q) / max(max_q - min_q, 1e-8) bar = "#" * int(normalized * bar_width) print(f" Q_{i+1:>3}: {q:.3f} |{bar}") return diagnostics # ============================================================================ # Entry point # ============================================================================ if __name__ == "__main__": diagnostics = run_simulation(n_queries=50, device="cpu")