Add cortex/memory_bank.py

cortex/memory_bank.py

@@ -0,0 +1,170 @@
+"""
+MemoryBank: Persistent episodic memory via gated cross-attention.
+
+Inspired by LM2 (Large Memory Models, Kang et al. 2025) and WISE (Xia et al. 2024).
+
+Architecture:
+    - A learnable memory matrix M ∈ R^{N×D} per injection layer
+    - Cross-attention: query = hidden_states, key/value = memory
+    - LSTM-style gating: forget gate, input gate, output gate
+    - Memory is updated per forward pass (recurrent across tokens)
+    
+Failure mode addressed: 
+    - Short context window / lost long-range dependencies
+    - Catastrophic forgetting during conversation
+    - Inability to maintain working memory for multi-step reasoning
+    
+Injection point: POST_ATTENTION (between attention and FFN)
+    - Rationale: Memory should augment what attention retrieved, before 
+      the FFN processes it for next-token prediction
+"""
+
+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 MemoryBank(CortexModule):
+    """
+    Gated cross-attention memory bank.
+    
+    The memory persists across forward passes (useful for multi-turn conversations
+    or long documents processed in chunks). It learns to read and write information
+    via gated cross-attention.
+    
+    Args:
+        hidden_dim: Model's hidden dimension
+        num_slots: Number of memory slots (like "memory addresses")
+        num_heads: Number of attention heads for cross-attention
+        dropout: Dropout rate
+        gate_bias_init: Initial bias for the output gate (negative = conservative writing)
+    """
+    
+    def __init__(
+        self,
+        hidden_dim: int,
+        num_slots: int = 64,
+        num_heads: int = 8,
+        dropout: float = 0.0,
+        gate_bias_init: float = -2.0,
+        target_layers: Union[List[int], str] = "middle",
+    ):
+        super().__init__(InjectionPoint.POST_ATTENTION, target_layers)
+        
+        self.hidden_dim = hidden_dim
+        self.num_slots = num_slots
+        self.num_heads = num_heads
+        self.head_dim = hidden_dim // num_heads
+        
+        # Memory bank: learnable initial state
+        self.memory_init = nn.Parameter(torch.randn(1, num_slots, hidden_dim) * 0.02)
+        
+        # Cross-attention: query from hidden states, key/value from memory
+        self.q_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
+        self.k_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
+        self.v_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
+        self.o_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
+        
+        # Gating mechanism (LSTM-style)
+        self.forget_gate = nn.Linear(hidden_dim, hidden_dim)
+        self.input_gate = nn.Linear(hidden_dim, hidden_dim)
+        self.output_gate = nn.Linear(hidden_dim, hidden_dim)
+        
+        # Output gate bias: negative means "default to not modifying hidden states"
+        # This is critical for stable injection into a pretrained model
+        nn.init.constant_(self.output_gate.bias, gate_bias_init)
+        
+        # Layer norm for memory read output
+        self.mem_norm = nn.LayerNorm(hidden_dim)
+        
+        # Dropout
+        self.dropout = nn.Dropout(dropout)
+        
+        # Running memory state (not a parameter — persistent buffer)
+        self.register_buffer("_memory_state", None, persistent=False)
+    
+    @property
+    def memory(self) -> torch.Tensor:
+        """Current memory state. Initialized from memory_init if not set."""
+        if self._memory_state is None:
+            return self.memory_init
+        return self._memory_state
+    
+    def reset_memory(self, batch_size: int = 1):
+        """Reset memory to initial state."""
+        self._memory_state = self.memory_init.expand(batch_size, -1, -1).clone()
+    
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        layer_idx: int,
+        attention_mask: Optional[torch.Tensor] = None,
+        **kwargs
+    ) -> torch.Tensor:
+        """
+        Read from and write to memory.
+        
+        Args:
+            hidden_states: [batch, seq_len, hidden_dim]
+            layer_idx: Current layer index
+            
+        Returns:
+            Modified hidden_states with memory-augmented information
+        """
+        batch_size, seq_len, _ = hidden_states.shape
+        
+        # Ensure memory is batch-expanded
+        memory = self.memory
+        if memory.shape[0] == 1 and batch_size > 1:
+            memory = memory.expand(batch_size, -1, -1)
+        if memory.shape[0] != batch_size:
+            memory = self.memory_init.expand(batch_size, -1, -1)
+        
+        memory = memory.to(hidden_states.device)
+        
+        # === MEMORY READ (cross-attention) ===
+        Q = self.q_proj(hidden_states)  # [B, T, D]
+        K = self.k_proj(memory)          # [B, N, D]
+        V = self.v_proj(memory)          # [B, N, D]
+        
+        # Reshape for multi-head attention
+        Q = Q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+        K = K.view(batch_size, self.num_slots, self.num_heads, self.head_dim).transpose(1, 2)
+        V = V.view(batch_size, self.num_slots, self.num_heads, self.head_dim).transpose(1, 2)
+        
+        # Scaled dot-product attention
+        attn_weights = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.head_dim)
+        attn_weights = F.softmax(attn_weights, dim=-1)
+        attn_weights = self.dropout(attn_weights)
+        
+        mem_read = torch.matmul(attn_weights, V)  # [B, H, T, D_h]
+        mem_read = mem_read.transpose(1, 2).contiguous().view(batch_size, seq_len, self.hidden_dim)
+        mem_read = self.o_proj(mem_read)
+        mem_read = self.mem_norm(mem_read)
+        
+        # === OUTPUT GATING ===
+        g_out = torch.sigmoid(self.output_gate(mem_read))
+        enhanced = hidden_states + g_out * mem_read
+        
+        # === MEMORY WRITE (update memory state) ===
+        write_signal = hidden_states[:, -1:, :]  # [B, 1, D]
+        write_signal_expanded = write_signal.expand_as(memory)
+        
+        g_forget = torch.sigmoid(self.forget_gate(write_signal_expanded))
+        g_input = torch.sigmoid(self.input_gate(write_signal_expanded))
+        
+        new_memory = g_forget * memory + g_input * torch.tanh(write_signal_expanded)
+        self._memory_state = new_memory.detach()
+        
+        return enhanced
+    
+    def on_inject(self, model: nn.Module, layer_idx: int, hidden_dim: int):
+        """Initialize projection weights to be near-zero for stable injection."""
+        nn.init.zeros_(self.o_proj.weight)
+        
+    def extra_repr(self):
+        return (f"hidden_dim={self.hidden_dim}, num_slots={self.num_slots}, "
+                f"num_heads={self.num_heads}, {super().extra_repr()}")