Add cortex/hallucination_gate.py

cortex/hallucination_gate.py

@@ -0,0 +1,151 @@
+"""
+HallucinationGate: Per-token confidence estimation and gated suppression.
+
+Inspired by InternalInspector (Chen et al. 2024) and "The Map of Misbelief" (2025).
+
+Architecture:
+    - A lightweight probe network that reads the residual stream at each layer
+    - Outputs a per-token confidence score in [0, 1]
+    - When confidence is low, the gate suppresses the residual update 
+      (pulling toward the safe "prior" — i.e., the identity/previous representation)
+    - Multi-layer aggregation: confidence from different layers is ensembled
+    
+Failure mode addressed:
+    - Hallucination: model generates confident-sounding but factually wrong text
+    - By detecting low-confidence internal states BEFORE they reach the output,
+      we can suppress risky updates to the residual stream
+    - The model effectively "hedges" — when unsure, it preserves its prior state
+      rather than committing to a potentially wrong update
+    
+Injection point: POST_FFN (after each transformer block)
+    - Rationale: The full block output contains the richest signal about confidence.
+      We gate the entire block's contribution to the residual stream.
+
+Key insight from InternalInspector: Internal states (especially attention outputs and 
+FFN outputs) contain significantly more information about correctness than the final 
+output distribution. I(Θ; K(X)|X) - I(Y; K(X)|X) ≥ Δ, where Δ is large.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Union, List
+from cortex.core import CortexModule, InjectionPoint
+
+
+class ConfidenceProbe(nn.Module):
+    """
+    Lightweight MLP that estimates per-token confidence from hidden states.
+    
+    Uses a bottleneck architecture to keep parameter count low.
+    Trained with contrastive loss on correct/incorrect generations.
+    """
+    
+    def __init__(self, hidden_dim: int, bottleneck_dim: int = 64):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(hidden_dim, bottleneck_dim),
+            nn.GELU(),
+            nn.Linear(bottleneck_dim, bottleneck_dim),
+            nn.GELU(),
+            nn.Linear(bottleneck_dim, 1),
+        )
+        # Initialize to output ~0.5 (neutral confidence) — sigmoid(0) = 0.5
+        nn.init.zeros_(self.net[-1].weight)
+        nn.init.zeros_(self.net[-1].bias)
+    
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            hidden_states: [batch, seq_len, hidden_dim]
+        Returns:
+            confidence: [batch, seq_len, 1] in (0, 1)
+        """
+        return torch.sigmoid(self.net(hidden_states))
+
+
+class HallucinationGate(CortexModule):
+    """
+    Gates the residual stream update based on internal confidence.
+    
+    When the model is confident (confidence → 1), the layer's output passes through 
+    normally. When uncertain (confidence → 0), the gate suppresses the layer's update,
+    effectively making it a near-identity operation.
+    
+    This creates an implicit "I don't know" signal that propagates through the model,
+    causing it to rely more on earlier, safer representations.
+    
+    Args:
+        hidden_dim: Model hidden dimension
+        bottleneck_dim: Confidence probe bottleneck size
+        suppression_strength: How aggressively to suppress uncertain updates (0-1)
+        confidence_threshold: Below this confidence, suppression kicks in hard
+    """
+    
+    def __init__(
+        self,
+        hidden_dim: int,
+        bottleneck_dim: int = 64,
+        suppression_strength: float = 0.8,
+        confidence_threshold: float = 0.3,
+        target_layers: Union[List[int], str] = "deep",
+    ):
+        super().__init__(InjectionPoint.POST_FFN, target_layers)
+        
+        self.hidden_dim = hidden_dim
+        self.suppression_strength = suppression_strength
+        self.confidence_threshold = confidence_threshold
+        
+        # Confidence probe
+        self.probe = ConfidenceProbe(hidden_dim, bottleneck_dim)
+        
+        # Learnable gate bias per dimension — allows the model to learn which 
+        # dimensions are safe to suppress vs which should always pass through
+        self.dim_gate = nn.Parameter(torch.zeros(1, 1, hidden_dim))
+        
+        # Running confidence for monitoring/logging
+        self.register_buffer("_last_confidence", torch.tensor(0.5), persistent=False)
+        self.register_buffer("_running_input", None, persistent=False)
+    
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        layer_idx: int,
+        **kwargs
+    ) -> torch.Tensor:
+        """
+        Gate the hidden states based on confidence.
+        
+        The gating formula is:
+            output = hidden_states * (1 - suppression * (1 - confidence))
+        
+        When confidence = 1: output = hidden_states (pass through)
+        When confidence = 0: output = hidden_states * (1 - suppression)
+        
+        This is applied as a soft gate on the RESIDUAL UPDATE, not the full hidden state.
+        """
+        # Estimate confidence
+        confidence = self.probe(hidden_states)  # [B, T, 1]
+        self._last_confidence = confidence.mean().detach()
+        
+        # Compute per-dimension gate
+        dim_bias = torch.sigmoid(self.dim_gate)  # [1, 1, D], in (0,1)
+        
+        # Effective gate: combines token-level confidence with dimension-level bias
+        # High dim_bias = dimension always passes through
+        # Low dim_bias = dimension is gated by confidence
+        gate = 1.0 - self.suppression_strength * (1.0 - confidence) * (1.0 - dim_bias)
+        
+        # Apply gate
+        gated_output = hidden_states * gate
+        
+        return gated_output
+    
+    def get_confidence(self) -> float:
+        """Return the last computed average confidence score."""
+        return self._last_confidence.item()
+    
+    def extra_repr(self):
+        return (f"hidden_dim={self.hidden_dim}, "
+                f"suppression_strength={self.suppression_strength}, "
+                f"{super().extra_repr()}")