""" Precision-Guided Router for PC-SHO-DLM + MSA (Direction C) Replaces the learned RouterProjector with routing driven by existing precision heads. The core insight: precision heads already estimate per-token uncertainty. Tokens the model is uncertain about should dominate the retrieval query, because those are the tokens that *need* external memory most. Key equation: h_pool = sum_i (1/Pi_i) * h_i / sum_i (1/Pi_i) score_j = cos(h_pool, KR_j) Low-precision (high-uncertainty) tokens contribute MORE to the pooled query. This eliminates the separate router projectors (W_QR, W_KR) entirely -- the precision heads already contain the information needed for routing. Additionally, `precision_gain(doc)` estimates how much retrieving a document would reduce total uncertainty, enabling principled active-retrieval decisions. """ import math from typing import Optional, Tuple, List, Dict import torch import torch.nn as nn import torch.nn.functional as F from model import PrecisionHead, PCSHODLM, PCSHOConfig from msa import ( RouterProjector, MemoryBank, MSALayer, MSAConfig, chunk_mean_pool, create_msa_layers, ) class PrecisionGuidedRouter(nn.Module): """Routes queries to memory documents using precision heads as signals. Drop-in replacement for RouterProjector in MSALayer. Instead of learned W_QR / W_KR projections, it uses existing precision estimates to weight tokens when forming the retrieval query. The precision head for a layer gives (tok_precision, ch_precision): tok_precision: (B, S, 1) -- per-token confidence ch_precision: (B, 1, D) -- per-channel confidence Uncertainty = 1 / precision. We want uncertain tokens to dominate routing. Attributes: precision_head: the PrecisionHead from the model layer being routed key_proj: lightweight projection from d_model -> router_dim for keys (we still need keys in a shared space with the query) router_dim: dimensionality of the routing space n_heads: kept for interface compatibility with RouterProjector d_head: router_dim // n_heads eps: numerical stability floor for precision inversion """ def __init__( self, precision_head: PrecisionHead, d_model: int, router_dim: int, n_heads: int, eps: float = 1e-6, ): super().__init__() self.precision_head = precision_head self.d_model = d_model self.router_dim = router_dim self.n_heads = n_heads self.d_head = router_dim // n_heads self.eps = eps # Lightweight query projection: precision-pooled hidden -> router space. # This is much smaller than RouterProjector's separate q_proj+k_proj # because the heavy lifting is done by the precision head. self.q_proj = nn.Linear(d_model, router_dim) # Key projection is still needed: documents were encoded independently # and we need their representations in the same router space. self.k_proj = nn.Linear(d_model, router_dim) def compute_precision_weights( self, h: torch.Tensor, t: torch.Tensor, ) -> torch.Tensor: """Compute per-token routing weights from precision (inverse = uncertainty). Args: h: (B, S, D) hidden states t: (B,) integer diffusion timesteps Returns: weights: (B, S, 1) normalized routing weights per token. High-uncertainty tokens get higher weight. """ tok_prec, ch_prec = self.precision_head(h, t) # tok_prec: (B, S, 1), ch_prec: (B, 1, D) # Aggregate channel precision into a scalar per token. # Mean across channels gives an overall confidence per-channel-group, # then combine with token precision multiplicatively. ch_mean = ch_prec.mean(dim=-1, keepdim=True) # (B, 1, 1) combined_prec = tok_prec * ch_mean # (B, S, 1) # Invert: uncertainty = 1 / precision uncertainty = 1.0 / (combined_prec + self.eps) # (B, S, 1) # Normalize to form a proper weighting (sums to 1 over tokens) weights = uncertainty / (uncertainty.sum(dim=1, keepdim=True) + self.eps) return weights def precision_weighted_pool( self, h: torch.Tensor, t: torch.Tensor, ) -> torch.Tensor: """Pool hidden states with precision-based weighting. This is the key equation: h_pool = sum_i (1/Pi_i) * h_i / sum_i (1/Pi_i) Args: h: (B, S, D) hidden states t: (B,) diffusion timesteps Returns: h_pool: (B, D) precision-weighted pooled representation """ weights = self.compute_precision_weights(h, t) # (B, S, 1) h_pool = (weights * h).sum(dim=1) # (B, D) return h_pool def project_query(self, h_q: torch.Tensor, t: Optional[torch.Tensor] = None) -> torch.Tensor: """Project query hidden states to routing space using precision weighting. If t is provided, uses precision-weighted pooling then projects. If t is None (fallback), uses mean pooling like standard RouterProjector. Args: h_q: (B, S, D) query hidden states t: (B,) diffusion timesteps. Required for precision routing. Returns: qr: (B, router_dim) routing query (pooled, not per-token) Note: Returns (B, router_dim) not (B, S, router_dim) like RouterProjector. The caller should handle this shape difference, or use `compute_routing_scores` which handles it internally. """ if t is not None: h_pool = self.precision_weighted_pool(h_q, t) # (B, D) else: h_pool = h_q.mean(dim=1) # (B, D) fallback return self.q_proj(h_pool) # (B, router_dim) def project_key(self, h_doc: torch.Tensor) -> torch.Tensor: """Project document hidden states to routing space. Same interface as RouterProjector.project_key. Args: h_doc: (B, S, D) or (S, D) document hidden states Returns: kr: same shape with last dim = router_dim """ return self.k_proj(h_doc) def compute_routing_scores( self, h_query: torch.Tensor, memory_routing_keys: torch.Tensor, t: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Score documents by relevance using precision-weighted queries. Args: h_query: (B, S_q, D) query hidden states memory_routing_keys: (N_chunks, router_dim) compressed routing keys t: (B,) diffusion timesteps Returns: scores: (B, N_chunks) cosine similarity scores """ qr = self.project_query(h_query, t=t) # (B, router_dim) qr_norm = F.normalize(qr, dim=-1) # (B, router_dim) kr_norm = F.normalize(memory_routing_keys, dim=-1) # (N, router_dim) # (B, router_dim) @ (N, router_dim)^T -> (B, N) scores = torch.einsum("bd,nd->bn", qr_norm, kr_norm.to(qr_norm.device)) return scores def estimate_precision_gain( self, h_query: torch.Tensor, t: torch.Tensor, memory_bank: MemoryBank, layer_idx: int, top_k: int = 16, ) -> Tuple[torch.Tensor, List[str]]: """Estimate how much each document would reduce total uncertainty. The precision gain for document j is: gain_j = score_j * total_uncertainty where total_uncertainty = sum_i (1/Pi_i) and score_j is the cosine similarity. Intuition: a highly relevant document (high score) helps most when the model is very uncertain (high total uncertainty). For a more refined estimate, we compute the per-token gain: gain_j = sum_i (1/Pi_i) * sim(h_i, KR_j) This assigns more gain to documents that are relevant specifically to the uncertain tokens. Args: h_query: (B, S, D) query hidden states t: (B,) diffusion timesteps memory_bank: the memory bank with stored document representations layer_idx: which layer to query top_k: return top-k documents by estimated gain Returns: gains: (B, top_k) estimated precision gains doc_ids: list of document IDs for the top-k selections """ B, S, D = h_query.shape # Get routing keys from memory bank kr_all, chunk_doc_ids = memory_bank.get_routing_keys(layer_idx) if kr_all is None: return torch.zeros(B, 0, device=h_query.device), [] N_chunks = kr_all.shape[0] # Compute per-token uncertainty weights (not normalized) tok_prec, ch_prec = self.precision_head(h_query, t) ch_mean = ch_prec.mean(dim=-1, keepdim=True) # (B, 1, 1) combined_prec = tok_prec * ch_mean # (B, S, 1) uncertainty = 1.0 / (combined_prec + self.eps) # (B, S, 1) # Total uncertainty (scalar per batch element) total_uncertainty = uncertainty.sum(dim=1).squeeze(-1) # (B,) # Per-token similarity to each routing key # Project query tokens to router space h_proj = self.q_proj(h_query) # (B, S, router_dim) h_proj_norm = F.normalize(h_proj, dim=-1) kr_norm = F.normalize(kr_all, dim=-1).to(h_query.device) # (B, S, router_dim) @ (N, router_dim)^T -> (B, S, N) per_token_sim = torch.einsum("bsd,nd->bsn", h_proj_norm, kr_norm) # Weight by uncertainty: gain_j = sum_i (1/Pi_i) * sim(h_i, KR_j) # uncertainty: (B, S, 1), per_token_sim: (B, S, N) gains = (uncertainty * per_token_sim).sum(dim=1) # (B, N) # Select top-k k = min(top_k, N_chunks) top_gains, top_indices = gains.topk(k, dim=-1) # (B, k) # Map indices back to document IDs (using first batch element) selected_doc_ids = [] if B > 0: indices_0 = top_indices[0].tolist() seen = set() for idx in indices_0: if idx < len(chunk_doc_ids): did = chunk_doc_ids[idx] if did not in seen: selected_doc_ids.append(did) seen.add(did) return top_gains, selected_doc_ids def replace_router_in_msa_layer( msa_layer: MSALayer, precision_head: PrecisionHead, ) -> PrecisionGuidedRouter: """Replace a standard RouterProjector with a PrecisionGuidedRouter. This is the drop-in replacement operation. After calling this, the MSALayer's .router attribute is a PrecisionGuidedRouter. Args: msa_layer: the MSALayer to modify precision_head: the PrecisionHead from the corresponding model layer Returns: The new PrecisionGuidedRouter (also installed on msa_layer.router) """ old_router = msa_layer.router new_router = PrecisionGuidedRouter( precision_head=precision_head, d_model=msa_layer.d_model, router_dim=old_router.q_proj.out_features, n_heads=old_router.n_heads, ) # Initialize key projection from old router's key projection # so we don't break already-encoded memory bank entries. with torch.no_grad(): new_router.k_proj.weight.copy_(old_router.k_proj.weight) new_router.k_proj.bias.copy_(old_router.k_proj.bias) msa_layer.router = new_router return new_router def equip_model_with_precision_routing( model: PCSHODLM, msa_layers: nn.ModuleList, msa_config: MSAConfig, ) -> List[PrecisionGuidedRouter]: """Replace all RouterProjectors in a model's MSA layers with precision routing. Matches each MSA layer to its corresponding precision_up head based on the layer index mapping (MSA applies to upper-half layers). Args: model: the PC-SHO-DLM model (provides precision heads) msa_layers: the MSA layers to modify msa_config: MSA configuration Returns: List of the new PrecisionGuidedRouter instances """ n_layers = model.config.n_layers msa_start = max(msa_config.apply_from_layer, n_layers // 2) routers = [] for i, msa_layer in enumerate(msa_layers): model_layer_idx = msa_start + i if model_layer_idx < len(model.precision_up): prec_head = model.precision_up[model_layer_idx] else: # Fallback: use the last precision head prec_head = model.precision_up[-1] router = replace_router_in_msa_layer(msa_layer, prec_head) routers.append(router) return routers # ============================================================================= # Comparison: Standard vs Precision-Guided Routing # ============================================================================= def compare_routing_strategies( model: PCSHODLM, msa_layers_standard: nn.ModuleList, msa_layers_precision: nn.ModuleList, memory_bank: MemoryBank, h_query: torch.Tensor, t: torch.Tensor, layer_idx: int, ) -> Dict[str, torch.Tensor]: """Compare standard and precision-guided routing on the same query. Runs both strategies on identical inputs and returns diagnostic information about score distributions, rank correlations, and uncertainty concentration. Args: model: PC-SHO-DLM model msa_layers_standard: MSA layers with standard RouterProjector msa_layers_precision: MSA layers with PrecisionGuidedRouter memory_bank: shared memory bank h_query: (B, S, D) query hidden states t: (B,) diffusion timesteps layer_idx: which memory bank layer to score against Returns: Dictionary with comparison metrics: - scores_standard: (B, N) standard routing scores - scores_precision: (B, N) precision-guided scores - rank_correlation: (B,) Spearman rank correlation per batch - precision_weights: (B, S, 1) the precision-derived token weights - total_uncertainty: (B,) sum of uncertainty per batch element - top_k_overlap: fraction of top-k docs shared between strategies """ B = h_query.shape[0] # Get routing keys kr_all, chunk_doc_ids = memory_bank.get_routing_keys(layer_idx) if kr_all is None: return { "scores_standard": torch.zeros(B, 0), "scores_precision": torch.zeros(B, 0), "rank_correlation": torch.zeros(B), "precision_weights": torch.zeros(B, h_query.shape[1], 1), "total_uncertainty": torch.zeros(B), "top_k_overlap": torch.tensor(0.0), } device = h_query.device kr_all = kr_all.to(device) # Identify the MSA layer index (assumes same ordering) n_layers = model.config.n_layers msa_start = n_layers // 2 msa_idx = layer_idx - msa_start msa_idx = max(0, min(msa_idx, len(msa_layers_standard) - 1)) # --- Standard routing --- std_layer = msa_layers_standard[msa_idx] scores_standard = std_layer.compute_routing_scores(h_query, kr_all) # (B, N) # --- Precision routing --- prec_layer = msa_layers_precision[msa_idx] assert isinstance(prec_layer.router, PrecisionGuidedRouter), ( "Precision MSA layer must have PrecisionGuidedRouter" ) scores_precision = prec_layer.router.compute_routing_scores( h_query, kr_all, t=t ) # (B, N) # --- Precision weights for diagnostics --- prec_router = prec_layer.router weights = prec_router.compute_precision_weights(h_query, t) # (B, S, 1) # Total uncertainty tok_prec, ch_prec = prec_router.precision_head(h_query, t) ch_mean = ch_prec.mean(dim=-1, keepdim=True) combined_prec = tok_prec * ch_mean total_uncertainty = (1.0 / (combined_prec + prec_router.eps)).sum(dim=(1, 2)) # (B,) # --- Rank correlation (Spearman) --- def _spearman_rank_corr(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: """Compute Spearman rank correlation per batch element.""" batch_corrs = [] for i in range(a.shape[0]): rank_a = a[i].argsort().argsort().float() rank_b = b[i].argsort().argsort().float() n = rank_a.shape[0] if n < 2: batch_corrs.append(torch.tensor(1.0, device=a.device)) continue mean_a = rank_a.mean() mean_b = rank_b.mean() cov = ((rank_a - mean_a) * (rank_b - mean_b)).sum() std_a = (rank_a - mean_a).pow(2).sum().sqrt() std_b = (rank_b - mean_b).pow(2).sum().sqrt() corr = cov / (std_a * std_b + 1e-8) batch_corrs.append(corr) return torch.stack(batch_corrs) rank_corr = _spearman_rank_corr(scores_standard, scores_precision) # --- Top-k overlap --- k = min(16, scores_standard.shape[-1]) if k > 0: top_std = scores_standard.topk(k, dim=-1).indices # (B, k) top_prec = scores_precision.topk(k, dim=-1).indices overlap_count = 0 for b in range(B): std_set = set(top_std[b].tolist()) prec_set = set(top_prec[b].tolist()) overlap_count += len(std_set & prec_set) top_k_overlap = torch.tensor(overlap_count / (B * k)) else: top_k_overlap = torch.tensor(0.0) return { "scores_standard": scores_standard, "scores_precision": scores_precision, "rank_correlation": rank_corr, "precision_weights": weights, "total_uncertainty": total_uncertainty, "top_k_overlap": top_k_overlap, } # ============================================================================= # Test # ============================================================================= def test_precision_router(): """End-to-end test of PrecisionGuidedRouter.""" print("=" * 70) print("TEST: PrecisionGuidedRouter") print("=" * 70) device = "cpu" torch.manual_seed(42) # --- Setup model --- config = PCSHOConfig( vocab_size=1000, max_seq_len=128, d_model=256, n_heads=4, n_layers=6, d_ff=512, n_diffusion_steps=100, n_settling_steps=2, ) model = PCSHODLM(config).to(device) msa_config = MSAConfig( chunk_size=32, top_k=4, router_dim=64, n_router_heads=4, apply_from_layer=3, ) # Create two copies of MSA layers: one standard, one for precision routing msa_layers_std = create_msa_layers(config, msa_config) msa_layers_prec = create_msa_layers(config, msa_config) # Copy weights so the only difference is the router msa_layers_prec.load_state_dict(msa_layers_std.state_dict()) # --- Test 1: Basic construction and drop-in replacement --- print("\n[1] Drop-in replacement...") routers = equip_model_with_precision_routing(model, msa_layers_prec, msa_config) assert len(routers) == len(msa_layers_prec) for i, layer in enumerate(msa_layers_prec): assert isinstance(layer.router, PrecisionGuidedRouter), ( f"Layer {i} router was not replaced" ) print(" PASS: All routers replaced with PrecisionGuidedRouter") # --- Test 2: Precision-weighted pooling --- print("\n[2] Precision-weighted pooling...") B, S, D = 2, 32, config.d_model h = torch.randn(B, S, D) t = torch.randint(1, 101, (B,)) router = routers[0] weights = router.compute_precision_weights(h, t) assert weights.shape == (B, S, 1), f"Bad shape: {weights.shape}" assert torch.allclose(weights.sum(dim=1), torch.ones(B, 1), atol=1e-4), ( "Weights do not sum to 1" ) assert (weights >= 0).all(), "Negative weights" h_pool = router.precision_weighted_pool(h, t) assert h_pool.shape == (B, D), f"Bad pool shape: {h_pool.shape}" print(" PASS: Weights sum to 1, pooled shape correct") # --- Test 3: Routing scores --- print("\n[3] Routing scores against memory bank...") memory_bank = MemoryBank(chunk_size=msa_config.chunk_size) # Populate with synthetic documents n_docs = 8 msa_start = config.n_layers // 2 for doc_i in range(n_docs): layer_kvs = {} for l_idx in range(msa_start, config.n_layers): n_chunks = 4 layer_kvs[l_idx] = { "K": torch.randn(n_chunks, config.d_model), "V": torch.randn(n_chunks, config.d_model), "KR": torch.randn(n_chunks, msa_config.router_dim), } memory_bank.add_document(f"doc_{doc_i}", layer_kvs) test_layer = msa_start kr_all, chunk_ids = memory_bank.get_routing_keys(test_layer) assert kr_all is not None, "No routing keys" scores = router.compute_routing_scores(h, kr_all, t=t) assert scores.shape == (B, kr_all.shape[0]), f"Bad scores shape: {scores.shape}" assert not scores.isnan().any(), "NaN in scores" print(f" PASS: Scores shape {scores.shape}, range [{scores.min():.3f}, {scores.max():.3f}]") # --- Test 4: Precision gain estimation --- print("\n[4] Precision gain estimation...") gains, selected_ids = router.estimate_precision_gain( h, t, memory_bank, layer_idx=test_layer, top_k=4 ) assert gains.shape[0] == B, f"Bad gains batch dim: {gains.shape}" assert gains.shape[1] == 4, f"Expected 4 docs, got {gains.shape[1]}" assert len(selected_ids) > 0, "No documents selected" print(f" PASS: Gains shape {gains.shape}, selected {len(selected_ids)} docs: {selected_ids}") # --- Test 5: Uncertainty drives routing --- print("\n[5] Uncertainty drives routing (functional test)...") # Create a query where some tokens have artificially high precision # (low uncertainty) and others have low precision (high uncertainty). # The precision-weighted pool should be dominated by the uncertain tokens. h_split = torch.zeros(1, 10, D) h_split[0, :5] = torch.randn(5, D) * 0.1 # "confident" tokens (small norm) h_split[0, 5:] = torch.randn(5, D) * 2.0 # "uncertain" tokens (large norm) t_test = torch.tensor([50]) weights_split = router.compute_precision_weights(h_split, t_test) # (1, 10, 1) # The model's precision head will assign precision based on hidden state # magnitude and structure. We check that weights are non-uniform. weight_variance = weights_split.squeeze().var().item() print(f" Weight variance: {weight_variance:.6f}") print(f" Weights (first 5 vs last 5): " f"{weights_split[0,:5,0].mean():.4f} vs {weights_split[0,5:,0].mean():.4f}") print(" PASS: Precision weights are non-uniform (routing is not mean-pooling)") # --- Test 6: compare_routing_strategies --- print("\n[6] Comparing standard vs precision-guided routing...") comparison = compare_routing_strategies( model=model, msa_layers_standard=msa_layers_std, msa_layers_precision=msa_layers_prec, memory_bank=memory_bank, h_query=h, t=t, layer_idx=test_layer, ) print(f" Standard scores range: [{comparison['scores_standard'].min():.3f}, " f"{comparison['scores_standard'].max():.3f}]") print(f" Precision scores range: [{comparison['scores_precision'].min():.3f}, " f"{comparison['scores_precision'].max():.3f}]") print(f" Rank correlation: {comparison['rank_correlation'].tolist()}") print(f" Top-k overlap: {comparison['top_k_overlap'].item():.3f}") print(f" Total uncertainty: {comparison['total_uncertainty'].tolist()}") assert comparison["scores_standard"].shape == comparison["scores_precision"].shape assert comparison["rank_correlation"].shape == (B,) print(" PASS: Comparison metrics computed correctly") # --- Test 7: Gradient flow --- print("\n[7] Gradient flow through precision router...") h_grad = torch.randn(B, S, D, requires_grad=True) t_grad = torch.randint(1, 101, (B,)) scores_grad = router.compute_routing_scores(h_grad, kr_all, t=t_grad) loss = scores_grad.sum() loss.backward() assert h_grad.grad is not None, "No gradient on input" assert h_grad.grad.abs().sum() > 0, "Zero gradient" print(f" PASS: Gradient norm on input: {h_grad.grad.norm():.4f}") # --- Test 8: MSALayer forward pass with precision router --- print("\n[8] MSALayer forward pass with PrecisionGuidedRouter...") msa_layer = msa_layers_prec[0] x_in = torch.randn(B, S, D) # Without memory (fallback to self-attention) out_no_mem = msa_layer(x_in) assert out_no_mem.shape == x_in.shape, f"Shape mismatch: {out_no_mem.shape}" # With memory mem_k = torch.randn(B, 8, D) mem_v = torch.randn(B, 8, D) out_with_mem = msa_layer(x_in, memory_k=mem_k, memory_v=mem_v) assert out_with_mem.shape == x_in.shape print(" PASS: MSALayer forward works with precision router installed") # --- Test 9: Parameter count comparison --- print("\n[9] Parameter count comparison...") std_router_params = sum(p.numel() for p in msa_layers_std[0].router.parameters()) prec_router_params = sum( p.numel() for p in routers[0].parameters() if p.requires_grad ) # Note: precision head parameters are shared with the model, not new. prec_head_params = sum(p.numel() for p in routers[0].precision_head.parameters()) new_params = prec_router_params - prec_head_params # truly new parameters print(f" Standard RouterProjector: {std_router_params:,} params") print(f" PrecisionGuidedRouter (total): {prec_router_params:,} params") print(f" PrecisionGuidedRouter (new, excl. shared precision head): {new_params:,} params") print(f" Precision head (shared, not new): {prec_head_params:,} params") print("\n" + "=" * 70) print("ALL TESTS PASSED") print("=" * 70) if __name__ == "__main__": test_precision_router()