Update memory.py

memory.py

@@ -1,61 +1,55 @@
+# cognitive_net/memory.py
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from collections import deque
-from typing import Dict, List, Optional, Tuple
+from typing import Deque, Dict, Any
 
 class CognitiveMemory(nn.Module):
-    """Differentiable memory system with consolidation and retrieval"""
+    """Differentiable memory system with biological consolidation mechanisms"""
     def __init__(self, context_size: int, capacity: int = 100):
         super().__init__()
         self.context_size = context_size
         self.capacity = capacity
-        self.memory_queue = deque(maxlen=capacity)
+        self.memory_queue: Deque[Dict[str, Any]] = deque(maxlen=capacity)
         
-        # Memory importance parameters
-        self.importance_decay = nn.Parameter(torch.tensor(0.95))
-        self.consolidation_threshold = 0.7
-        
-        # Memory projection layers
+        # Memory projection layers with adaptive scaling
         self.key_proj = nn.Linear(context_size, 64)
         self.value_proj = nn.Linear(context_size, 64)
+        self.importance_decay = nn.Parameter(torch.tensor(0.95))
+        
+        # Consolidation parameters
+        self.consolidation_threshold = 0.7
+        self.age_decay = 0.1
 
     def add_memory(self, context: torch.Tensor, activation: float):
-        """Store new memory with adaptive importance"""
-        # Ensure context is 1D tensor with single value
-        context = context.reshape(-1)
+        """Store memory with dynamic importance weighting"""
         importance = torch.sigmoid(torch.tensor(activation * 0.5 + 0.2))
         self.memory_queue.append({
-            'context': context.detach(),
+            'context': context.detach().clone(),
             'importance': importance,
-            'age': 0.0
+            'age': torch.tensor(0.0)
         })
-        
+
     def consolidate_memories(self):
-        """Memory consolidation through importance reweighting"""
+        """Memory optimization through importance-based pruning"""
+        new_queue = deque(maxlen=self.capacity)
         for mem in self.memory_queue:
             mem['importance'] *= self.importance_decay
-            mem['age'] += 0.1
-            
-        # Remove unimportant memories
-        self.memory_queue = deque(
-            [m for m in self.memory_queue if m['importance'] > 0.2],
-            maxlen=self.capacity
-        )
-        
+            mem['age'] += self.age_decay
+            if mem['importance'] > 0.2:
+                new_queue.append(mem)
+        self.memory_queue = new_queue
+
     def retrieve(self, query: torch.Tensor) -> torch.Tensor:
-        """Attention-based memory retrieval"""
+        """Content-based memory retrieval with attention"""
         if not self.memory_queue:
-            return torch.zeros_like(query)
+            return torch.zeros(64, device=query.device)
             
-        # Ensure query is 1D tensor with single value
-        query = query.reshape(1, 1)
-        memories = torch.stack([m['context'].reshape(1, 1) for m in self.memory_queue])
-        
-        keys = self.key_proj(memories)
-        values = self.value_proj(memories)
-        query_proj = self.key_proj(query)
+        contexts = torch.stack([m['context'] for m in self.memory_queue])
+        keys = self.key_proj(contexts)
+        values = self.value_proj(contexts)
+        query_proj = self.key_proj(query.unsqueeze(0))
         
-        scores = F.softmax(torch.matmul(keys, query_proj.transpose(0, 1)), dim=0)
-        retrieved = torch.matmul(scores.transpose(0, 1), values)
-        return retrieved.squeeze(0)
+        scores = F.softmax(keys @ query_proj.T, dim=0)
+        return (scores * values).sum(dim=0)