Upload model_architecture.py

model_architecture.py

@@ -0,0 +1,616 @@
+"""
+LMCODE: Language Model with Memory CODE
+=======================================
+
+A memory-augmented language model with dual memory systems:
+- Short-term memory: Working memory for immediate context (like Transformer KV cache)
+- Long-term memory: Persistent storage for knowledge and experiences
+
+Inspired by:
+- LongMem (2023): Augmenting LLMs with Long-Term Memory
+- MemoRAG (2024): Dual-system RAG with Global and Local Memory
+- CAMELoT (2024): Training-free Consolidated Associative Memory
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Tuple, Dict, List
+import math
+
+
+class ShortTermMemory(nn.Module):
+    """
+    Short-term memory module using recurrent state updates.
+    Acts as working memory for immediate context (similar to Transformer KV cache).
+    """
+    
+    def __init__(self, hidden_size: int, memory_size: int = 512):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.memory_size = memory_size
+        
+        # Memory state
+        self.memory = nn.Parameter(torch.zeros(1, memory_size, hidden_size))
+        
+        # Memory update gate
+        self.update_gate = nn.Linear(hidden_size, 1)
+        
+        # Memory read/write projections
+        self.read_proj = nn.Linear(hidden_size, hidden_size)
+        self.write_proj = nn.Linear(hidden_size, hidden_size)
+        
+        # Initialize memory
+        nn.init.normal_(self.memory, mean=0, std=0.02)
+    
+    def forward(self, hidden_states: torch.Tensor, 
+                attention_mask: Optional[torch.Tensor] = None) -> Dict[str, torch.Tensor]:
+        """
+        Update and read from short-term memory.
+        
+        Args:
+            hidden_states: Current hidden states [batch_size, seq_len, hidden_size]
+            attention_mask: Optional attention mask
+            
+        Returns:
+            Dictionary with memory outputs
+        """
+        batch_size, seq_len, _ = hidden_states.shape
+        
+        # Expand memory for batch
+        memory = self.memory.expand(batch_size, -1, -1)  # [batch, memory_size, hidden]
+        
+        # Compute update gates
+        update_scores = torch.sigmoid(self.update_gate(hidden_states))  # [batch, seq, 1]
+        
+        # Read from memory
+        read_query = self.read_proj(hidden_states)  # [batch, seq, hidden]
+        read_scores = torch.matmul(read_query, memory.transpose(-1, -2))  # [batch, seq, memory_size]
+        read_scores = read_scores / math.sqrt(self.hidden_size)
+        
+        if attention_mask is not None:
+            read_scores = read_scores.masked_fill(~attention_mask.unsqueeze(-1).bool(), -1e9)
+        
+        read_weights = F.softmax(read_scores, dim=-1)  # [batch, seq, memory_size]
+        memory_read = torch.matmul(read_weights, memory)  # [batch, seq, hidden]
+        
+        # Write to memory (soft update)
+        write_values = self.write_proj(hidden_states)  # [batch, seq, hidden]
+        write_weights = update_scores.expand(-1, -1, self.memory_size)  # [batch, seq, memory_size]
+        
+        # Update memory (detached to prevent gradient flow through time)
+        memory_delta = torch.matmul(write_weights.transpose(-1, -2), write_values)  # [batch, memory_size, hidden]
+        memory = memory + 0.1 * memory_delta  # Small learning rate for memory update
+        
+        # Combine hidden states with memory read
+        output = hidden_states + 0.5 * memory_read  # Residual connection
+        
+        return {
+            'output': output,
+            'memory_read': memory_read,
+            'read_weights': read_weights,
+            'updated_memory': memory
+        }
+
+
+class LongTermMemory(nn.Module):
+    """
+    Long-term memory module for persistent storage of knowledge.
+    Uses a key-value store with retrieval mechanism.
+    
+    Inspired by:
+    - LongMem's adaptive residual side-network
+    - MemoRAG's global memory for overview
+    """
+    
+    def __init__(self, hidden_size: int, num_slots: int = 10000):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_slots = num_slots
+        
+        # Memory storage (key-value pairs)
+        self.memory_keys = nn.Parameter(torch.zeros(num_slots, hidden_size))
+        self.memory_values = nn.Parameter(torch.zeros(num_slots, hidden_size))
+        
+        # Memory metadata (importance, timestamps, etc.)
+        self.memory_importance = nn.Parameter(torch.ones(num_slots))
+        
+        # Query projection
+        self.query_proj = nn.Linear(hidden_size, hidden_size)
+        self.key_proj = nn.Linear(hidden_size, hidden_size)
+        self.value_proj = nn.Linear(hidden_size, hidden_size)
+        
+        # Memory consolidation (for summarizing old memories)
+        self.consolidation_layer = nn.Sequential(
+            nn.Linear(hidden_size * 2, hidden_size),
+            nn.ReLU(),
+            nn.Linear(hidden_size, hidden_size)
+        )
+        
+        # Initialize
+        nn.init.normal_(self.memory_keys, mean=0, std=0.02)
+        nn.init.normal_(self.memory_values, mean=0, std=0.02)
+        nn.init.constant_(self.memory_importance, 1.0)
+    
+    def forward(self, hidden_states: torch.Tensor,
+                store_new: bool = False,
+                new_content: Optional[torch.Tensor] = None) -> Dict[str, torch.Tensor]:
+        """
+        Retrieve from or store to long-term memory.
+        
+        Args:
+            hidden_states: Current hidden states [batch_size, seq_len, hidden_size]
+            store_new: Whether to store new content
+            new_content: New content to store [batch_size, content_len, hidden_size]
+            
+        Returns:
+            Dictionary with memory outputs
+        """
+        batch_size, seq_len, _ = hidden_states.shape
+        
+        # Project queries
+        queries = self.query_proj(hidden_states)  # [batch, seq, hidden]
+        
+        # Project memory keys
+        keys = self.key_proj(self.memory_keys)  # [num_slots, hidden]
+        
+        # Compute retrieval scores
+        scores = torch.matmul(queries, keys.transpose(-1, -2))  # [batch, seq, num_slots]
+        scores = scores / math.sqrt(self.hidden_size)
+        
+        # Apply importance weighting
+        importance_weights = torch.sigmoid(self.memory_importance)  # [num_slots]
+        scores = scores * importance_weights.unsqueeze(0).unsqueeze(0)  # [batch, seq, num_slots]
+        
+        # Get top-k memories
+        top_k = min(10, self.num_slots)
+        top_scores, top_indices = torch.topk(scores, k=top_k, dim=-1)  # [batch, seq, top_k]
+        
+        # Softmax over top-k
+        retrieval_weights = F.softmax(top_scores, dim=-1)  # [batch, seq, top_k]
+        
+        # Gather retrieved memories
+        retrieved_values = self.memory_values[top_indices]  # [batch, seq, top_k, hidden]
+        
+        # Weighted sum
+        memory_output = torch.sum(retrieved_values * retrieval_weights.unsqueeze(-1), dim=-2)  # [batch, seq, hidden]
+        
+        # Project memory output
+        memory_output = self.value_proj(memory_output)
+        
+        # Store new memories if requested
+        if store_new and new_content is not None:
+            self._store_new_memories(new_content)
+        
+        return {
+            'output': memory_output,
+            'retrieval_weights': retrieval_weights,
+            'top_indices': top_indices,
+            'retrieval_scores': top_scores
+        }
+    
+    def _store_new_memories(self, content: torch.Tensor):
+        """
+        Store new content in long-term memory.
+        Uses a simple FIFO with consolidation strategy.
+        
+        Args:
+            content: Content to store [batch_size, content_len, hidden_size]
+        """
+        with torch.no_grad():
+            # Average content across batch and sequence
+            content_repr = content.mean(dim=(0, 1))  # [hidden_size]
+            
+            # Find least important slot
+            _, least_important_idx = torch.min(self.memory_importance, dim=0)
+            
+            # Update memory (with small learning rate)
+            lr = 0.01
+            self.memory_keys.data[least_important_idx] = \
+                (1 - lr) * self.memory_keys.data[least_important_idx] + lr * content_repr
+            self.memory_values.data[least_important_idx] = \
+                (1 - lr) * self.memory_values.data[least_important_idx] + lr * content_repr
+            
+            # Reset importance (new memories start neutral)
+            self.memory_importance.data[least_important_idx] = 1.0
+    
+    def consolidate_memories(self, threshold: float = 0.1):
+        """
+        Consolidate similar memories to prevent redundancy.
+        
+        Args:
+            threshold: Similarity threshold for consolidation
+        """
+        with torch.no_grad():
+            # Compute pairwise similarity
+            keys_norm = F.normalize(self.memory_keys, dim=-1)
+            similarity = torch.matmul(keys_norm, keys_norm.transpose(-1, -2))
+            
+            # Find highly similar pairs
+            mask = torch.triu(similarity > threshold, diagonal=1)
+            
+            if mask.any():
+                # Consolidate: merge similar memories
+                indices = torch.where(mask)
+                
+                for i, j in zip(indices[0], indices[1]):
+                    # Merge j into i (weighted by importance)
+                    imp_i = torch.sigmoid(self.memory_importance[i])
+                    imp_j = torch.sigmoid(self.memory_importance[j])
+                    total_imp = imp_i + imp_j
+                    
+                    if total_imp > 0:
+                        self.memory_keys.data[i] = \
+                            (imp_i * self.memory_keys.data[i] + imp_j * self.memory_keys.data[j]) / total_imp
+                        self.memory_values.data[i] = \
+                            (imp_i * self.memory_values.data[i] + imp_j * self.memory_values.data[j]) / total_imp
+                        self.memory_importance.data[i] = torch.log(torch.exp(self.memory_importance.data[i]) + 
+                                                                 torch.exp(self.memory_importance.data[j]))
+                        
+                        # Zero out merged memory
+                        self.memory_keys.data[j] = 0
+                        self.memory_values.data[j] = 0
+                        self.memory_importance.data[j] = -10  # Mark as inactive
+
+
+class MemoryAugmentedLayer(nn.Module):
+    """
+    Transformer layer augmented with both short-term and long-term memory.
+    """
+    
+    def __init__(self, hidden_size: int, num_heads: int, 
+                 short_term_memory_size: int = 512,
+                 long_term_memory_slots: int = 10000):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        
+        # Self-attention
+        self.attention = nn.MultiheadAttention(hidden_size, num_heads, batch_first=True)
+        self.attention_norm = nn.LayerNorm(hidden_size)
+        
+        # Feed-forward
+        self.ffn = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size * 4),
+            nn.GELU(),
+            nn.Linear(hidden_size * 4, hidden_size)
+        )
+        self.ffn_norm = nn.LayerNorm(hidden_size)
+        
+        # Memory modules
+        self.short_term_memory = ShortTermMemory(hidden_size, short_term_memory_size)
+        self.long_term_memory = LongTermMemory(hidden_size, long_term_memory_slots)
+        
+        # Memory gating
+        self.memory_gate = nn.Linear(hidden_size * 2, hidden_size)
+    
+    def forward(self, hidden_states: torch.Tensor,
+                attention_mask: Optional[torch.Tensor] = None,
+                use_long_term_memory: bool = True,
+                store_long_term: bool = False) -> Dict[str, torch.Tensor]:
+        """
+        Forward pass through memory-augmented layer.
+        
+        Args:
+            hidden_states: Input hidden states [batch_size, seq_len, hidden_size]
+            attention_mask: Optional attention mask
+            use_long_term_memory: Whether to use long-term memory
+            store_long_term: Whether to store in long-term memory
+            
+        Returns:
+            Dictionary with layer outputs
+        """
+        residual = hidden_states
+        
+        # Self-attention
+        attn_output, attn_weights = self.attention(
+            self.attention_norm(hidden_states),
+            self.attention_norm(hidden_states),
+            self.attention_norm(hidden_states),
+            key_padding_mask=attention_mask
+        )
+        hidden_states = residual + attn_output
+        
+        # Short-term memory
+        memory_output = self.short_term_memory(hidden_states, attention_mask)
+        hidden_states = memory_output['output']
+        
+        # Long-term memory (optional)
+        long_term_output = None
+        if use_long_term_memory:
+            long_term_output = self.long_term_memory(
+                hidden_states,
+                store_new=store_long_term,
+                new_content=residual  # Store original input
+            )
+            # Gate between attention and memory
+            memory_gate_input = torch.cat([hidden_states, long_term_output['output']], dim=-1)
+            gate = torch.sigmoid(self.memory_gate(memory_gate_input))
+            hidden_states = gate * hidden_states + (1 - gate) * long_term_output['output']
+        
+        # Feed-forward
+        ffn_output = self.ffn(self.ffn_norm(hidden_states))
+        hidden_states = hidden_states + ffn_output
+        
+        return {
+            'hidden_states': hidden_states,
+            'attention_weights': attn_weights,
+            'short_term_output': memory_output,
+            'long_term_output': long_term_output
+        }
+
+
+class LMCODE(nn.Module):
+    """
+    Language Model with Memory CODE (LMCODE).
+    
+    A transformer-based language model augmented with:
+    1. Short-term memory: Working memory for immediate context
+    2. Long-term memory: Persistent storage for knowledge and experiences
+    
+    Architecture:
+    - Embedding layer
+    - Multiple memory-augmented transformer layers
+    - Language model head
+    
+    Memory Flow:
+    - Short-term: Updated every forward pass, decays over time
+    - Long-term: Consolidated periodically, stores important experiences
+    """
+    
+    def __init__(self, config: Optional['LMCODEConfig'] = None):
+        super().__init__()
+        
+        if config is None:
+            config = LMCODEConfig()
+        
+        self.config = config
+        
+        # Embeddings
+        self.token_embeddings = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.position_embeddings = nn.Embedding(2048, config.hidden_size)
+        
+        # Memory-augmented layers
+        self.layers = nn.ModuleList([
+            MemoryAugmentedLayer(
+                hidden_size=config.hidden_size,
+                num_heads=config.num_heads,
+                short_term_memory_size=config.short_term_memory_size,
+                long_term_memory_slots=config.long_term_memory_slots
+            )
+            for _ in range(config.num_layers)
+        ])
+        
+        # Output normalization
+        self.final_norm = nn.LayerNorm(config.hidden_size)
+        
+        # Language model head
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        
+        # Initialize weights
+        self.apply(self._init_weights)
+    
+    def _init_weights(self, module):
+        """Initialize model weights."""
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0, std=0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0, std=0.02)
+        elif isinstance(module, nn.LayerNorm):
+            nn.init.zeros_(module.bias)
+            nn.init.ones_(module.weight)
+    
+    def forward(self, input_ids: torch.Tensor,
+                attention_mask: Optional[torch.Tensor] = None,
+                use_long_term_memory: bool = True,
+                store_long_term: bool = False,
+                labels: Optional[torch.Tensor] = None) -> Dict[str, torch.Tensor]:
+        """
+        Forward pass through LMCODE model.
+        
+        Args:
+            input_ids: Input token IDs [batch_size, seq_len]
+            attention_mask: Optional attention mask [batch_size, seq_len]
+            use_long_term_memory: Whether to use long-term memory
+            store_long_term: Whether to store in long-term memory
+            labels: Optional labels for computing loss
+            
+        Returns:
+            Dictionary with model outputs
+        """
+        batch_size, seq_len = input_ids.shape
+        
+        # Create position IDs
+        position_ids = torch.arange(seq_len, device=input_ids.device).unsqueeze(0).expand(batch_size, -1)
+        
+        # Embeddings
+        token_embeds = self.token_embeddings(input_ids)
+        position_embeds = self.position_embeddings(position_ids)
+        hidden_states = token_embeds + position_embeds
+        
+        # Apply dropout
+        hidden_states = F.dropout(hidden_states, p=0.1, training=self.training)
+        
+        # Track intermediate outputs
+        all_hidden_states = []
+        all_attention_weights = []
+        all_short_term_outputs = []
+        all_long_term_outputs = []
+        
+        # Pass through layers
+        for layer in self.layers:
+            layer_output = layer(
+                hidden_states,
+                attention_mask=attention_mask,
+                use_long_term_memory=use_long_term_memory,
+                store_long_term=store_long_term
+            )
+            
+            hidden_states = layer_output['hidden_states']
+            all_hidden_states.append(hidden_states)
+            all_attention_weights.append(layer_output['attention_weights'])
+            all_short_term_outputs.append(layer_output['short_term_output'])
+            all_long_term_outputs.append(layer_output['long_term_output'])
+        
+        # Final normalization
+        hidden_states = self.final_norm(hidden_states)
+        
+        # Language model logits
+        logits = self.lm_head(hidden_states)
+        
+        # Compute loss if labels provided
+        loss = None
+        if labels is not None:
+            loss = F.cross_entropy(
+                logits.view(-1, self.config.vocab_size),
+                labels.view(-1),
+                ignore_index=-100
+            )
+        
+        return {
+            'logits': logits,
+            'loss': loss,
+            'hidden_states': all_hidden_states,
+            'attention_weights': all_attention_weights,
+            'short_term_outputs': all_short_term_outputs,
+            'long_term_outputs': all_long_term_outputs
+        }
+    
+    def generate(self, input_ids: torch.Tensor,
+                 max_length: int = 100,
+                 temperature: float = 1.0,
+                 top_k: int = 50,
+                 top_p: float = 0.9,
+                 use_long_term_memory: bool = True,
+                 store_long_term: bool = True) -> torch.Tensor:
+        """
+        Generate text autoregressively.
+        
+        Args:
+            input_ids: Initial token IDs [batch_size, seq_len]
+            max_length: Maximum generation length
+            temperature: Sampling temperature
+            top_k: Top-k sampling
+            top_p: Nucleus sampling
+            use_long_term_memory: Whether to use long-term memory during generation
+            store_long_term: Whether to store generated text in long-term memory
+            
+        Returns:
+            Generated token IDs [batch_size, total_length]
+        """
+        batch_size = input_ids.shape[0]
+        
+        for _ in range(max_length):
+            # Forward pass
+            with torch.no_grad():
+                outputs = self(
+                    input_ids,
+                    use_long_term_memory=use_long_term_memory,
+                    store_long_term=store_long_term
+                )
+                
+                logits = outputs['logits'][:, -1, :] / temperature
+                
+                # Top-k filtering
+                if top_k > 0:
+                    indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
+                    logits[indices_to_remove] = -float('Inf')
+                
+                # Top-p (nucleus) filtering
+                if top_p < 1.0:
+                    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+                    cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+                    
+                    sorted_indices_to_remove = cumulative_probs > top_p
+                    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+                    sorted_indices_to_remove[..., 0] = 0
+                    
+                    indices_to_remove = sorted_indices_to_remove.scatter(
+                        1, sorted_indices, sorted_indices_to_remove
+                    )
+                    logits[indices_to_remove] = -float('Inf')
+                
+                # Sample next token
+                probs = F.softmax(logits, dim=-1)
+                next_token = torch.multinomial(probs, num_samples=1)
+                
+                # Append to sequence
+                input_ids = torch.cat([input_ids, next_token], dim=-1)
+        
+        return input_ids
+    
+    def store_experience(self, text: str, metadata: Optional[Dict] = None):
+        """
+        Store a text experience in long-term memory.
+        
+        Args:
+            text: Text to store
+            metadata: Optional metadata
+        """
+        # Tokenize text (simplified - in practice, use proper tokenizer)
+        # For demonstration, we'll create dummy embeddings
+        # In production, use: tokens = tokenizer.encode(text)
+        
+        # Create a representation (in practice, pass through model)
+        dummy_input = torch.zeros(1, 10, self.config.hidden_size)
+        
+        with torch.no_grad():
+            # Get representation from model
+            representation = self.forward(dummy_input, use_long_term_memory=False)
+            # Store in long-term memory
+            self.layers[0].long_term_memory._store_new_memories(representation['hidden_states'][-1])
+    
+    def query_memory(self, query: str, top_k: int = 5) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Query long-term memory for relevant information.
+        
+        Args:
+            query: Query text
+            top_k: Number of results to retrieve
+            
+        Returns:
+            retrieved_memories: Retrieved memory representations
+            indices: Indices of retrieved memories
+        """
+        # Create query representation
+        dummy_input = torch.zeros(1, 10, self.config.hidden_size)
+        
+        with torch.no_grad():
+            query_representation = self.forward(dummy_input, use_long_term_memory=False)
+            query_repr = query_representation['hidden_states'][-1].mean(dim=1)
+            
+            # Retrieve from long-term memory
+            retrieved, indices = self.layers[0].long_term_memory.retrieve(query_repr, top_k=top_k)
+        
+        return retrieved, indices
+
+
+class LMCODEConfig:
+    """Configuration for LMCODE model."""
+    
+    def __init__(self, vocab_size: int = 50257, hidden_size: int = 512, 
+                 num_layers: int = 6, num_heads: int = 8,
+                 short_term_memory_size: int = 512,
+                 long_term_memory_slots: int = 10000):
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_layers = num_layers
+        self.num_heads = num_heads
+        self.short_term_memory_size = short_term_memory_size
+        self.long_term_memory_slots = long_term_memory_slots
+    
+    def to_dict(self) -> Dict:
+        return {
+            'vocab_size': self.vocab_size,
+            'hidden_size': self.hidden_size,
+            'num_layers': self.num_layers,
+            'num_heads': self.num_heads,
+            'short_term_memory_size': self.short_term_memory_size,
+            'long_term_memory_slots': self.long_term_memory_slots
+        }
+    
+    @classmethod
+    def from_dict(cls, config_dict: Dict) -> 'LMCODEConfig':
+        return cls(**config_dict)