Add cortex/pause_and_think.py

cortex/pause_and_think.py

@@ -0,0 +1,153 @@
+"""
+PauseAndThink: Learnable latent computation tokens for extended reasoning.
+
+Inspired by "Think before you speak: Training LMs with Pause Tokens" (Goyal et al. 2023)
+and "Thoughtbubbles" (2025).
+
+Architecture:
+    - Injects K learnable "thinking" embeddings into the residual stream
+    - These embeddings participate in self-attention (attend to all real tokens)
+    - But their outputs are NOT projected to vocabulary — they are computation-only
+    - After processing, the thinking tokens' representations are compressed 
+      back into the real token positions via cross-attention
+    - Effectively gives the model K extra "thinking steps" per layer
+    
+Failure mode addressed:
+    - Limited compute per token: standard transformers do fixed computation per token.
+      Some tokens (e.g., the final answer token) need more computation.
+    - Shallow reasoning: complex multi-step reasoning requires iterating over information.
+      Thinking tokens provide extra "scratch space" for intermediate computation.
+    - CoT without output tokens: like chain-of-thought, but in latent space —
+      no extra output tokens are needed.
+
+Injection point: RESIDUAL_STREAM (wraps full block)
+    - Rationale: The thinking tokens need to participate in attention.
+      By operating at the residual stream level, they can read from and write back
+      to the main representation pathway.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+from typing import Optional, Union, List
+from cortex.core import CortexModule, InjectionPoint
+
+
+class PauseAndThink(CortexModule):
+    """
+    Injects learnable computation tokens into the hidden state sequence.
+    
+    These tokens expand the sequence temporarily, participate in attention,
+    then their information is compressed back into the original positions.
+    
+    Args:
+        hidden_dim: Model hidden dimension
+        num_think_tokens: How many thinking tokens to inject
+        compress_method: How to fold thinking token info back ("cross_attn" or "gate")
+        init_scale: Scale for thinking token initialization
+    """
+    
+    def __init__(
+        self,
+        hidden_dim: int,
+        num_think_tokens: int = 8,
+        compress_method: str = "gate",
+        init_scale: float = 0.01,
+        target_layers: Union[List[int], str] = "middle",
+    ):
+        super().__init__(InjectionPoint.RESIDUAL_STREAM, target_layers)
+        
+        self.hidden_dim = hidden_dim
+        self.num_think_tokens = num_think_tokens
+        self.compress_method = compress_method
+        
+        # Learnable thinking token embeddings
+        self.think_tokens = nn.Parameter(
+            torch.randn(1, num_think_tokens, hidden_dim) * init_scale
+        )
+        
+        # Projection to adapt thinking tokens to the current context
+        self.context_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
+        nn.init.eye_(self.context_proj.weight)  # Initialize as identity
+        
+        if compress_method == "cross_attn":
+            # Cross-attention to compress thinking tokens back
+            self.compress_q = nn.Linear(hidden_dim, hidden_dim, bias=False)
+            self.compress_k = nn.Linear(hidden_dim, hidden_dim, bias=False)
+            self.compress_v = nn.Linear(hidden_dim, hidden_dim, bias=False)
+            self.compress_out = nn.Linear(hidden_dim, hidden_dim, bias=False)
+            nn.init.zeros_(self.compress_out.weight)  # Start as no-op
+        elif compress_method == "gate":
+            # Simpler: gated addition from mean-pooled thinking tokens
+            self.compress_gate = nn.Sequential(
+                nn.Linear(hidden_dim * 2, hidden_dim),
+                nn.Sigmoid()
+            )
+            self.compress_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
+            nn.init.zeros_(self.compress_proj.weight)  # Start as no-op
+            # Initialize gate bias to make it ~0 initially (conservative)
+            nn.init.constant_(self.compress_gate[0].bias, -3.0)
+        
+        # Layer norm
+        self.think_norm = nn.LayerNorm(hidden_dim)
+    
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        layer_idx: int,
+        attention_mask: Optional[torch.Tensor] = None,
+        **kwargs
+    ) -> torch.Tensor:
+        """
+        Inject thinking tokens, let them process via the original layer,
+        then compress their information back.
+        
+        Since we operate at RESIDUAL_STREAM level (post full block), we use the
+        thinking tokens as a read-process-write cycle on the hidden states.
+        """
+        batch_size, seq_len, _ = hidden_states.shape
+        
+        # === GENERATE CONTEXT-AWARE THINKING TOKENS ===
+        context = hidden_states.mean(dim=1, keepdim=True)  # [B, 1, D]
+        
+        think = self.think_tokens.expand(batch_size, -1, -1)  # [B, K, D]
+        think = think + self.context_proj(context)  # Context-conditioned
+        think = self.think_norm(think)
+        
+        # === THINK: Let thinking tokens interact with hidden states ===
+        Q = think  # [B, K, D]
+        K = hidden_states  # [B, T, D]
+        V = hidden_states  # [B, T, D]
+        
+        attn_scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.hidden_dim)
+        attn_weights = F.softmax(attn_scores, dim=-1)
+        think_output = torch.matmul(attn_weights, V)  # [B, K, D]
+        
+        # === COMPRESS: Fold thinking token info back into hidden states ===
+        if self.compress_method == "cross_attn":
+            cQ = self.compress_q(hidden_states)   # [B, T, D]
+            cK = self.compress_k(think_output)    # [B, K, D]
+            cV = self.compress_v(think_output)    # [B, K, D]
+            
+            c_attn = torch.matmul(cQ, cK.transpose(-2, -1)) / math.sqrt(self.hidden_dim)
+            c_weights = F.softmax(c_attn, dim=-1)
+            compressed = torch.matmul(c_weights, cV)  # [B, T, D]
+            compressed = self.compress_out(compressed)
+            
+            return hidden_states + compressed
+            
+        elif self.compress_method == "gate":
+            think_summary = think_output.mean(dim=1, keepdim=True)  # [B, 1, D]
+            think_summary = think_summary.expand_as(hidden_states)   # [B, T, D]
+            
+            gate_input = torch.cat([hidden_states, think_summary], dim=-1)
+            gate = self.compress_gate(gate_input)  # [B, T, D]
+            
+            projected = self.compress_proj(think_summary)
+            
+            return hidden_states + gate * projected
+    
+    def extra_repr(self):
+        return (f"hidden_dim={self.hidden_dim}, num_think_tokens={self.num_think_tokens}, "
+                f"compress={self.compress_method}, {super().extra_repr()}")