Neural Network Memory Systems

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

class BiologicalMemoryNetwork(nn.Module):
    def __init__(self, input_size, hidden_size, memory_slots=256):
        super().__init__()
        
        # Hippocampal-inspired short-term memory
        self.stm = nn.LSTM(input_size, hidden_size, batch_first=True)
        
        # Cortical-inspired long-term memory
        self.ltm_keys = nn.Parameter(torch.randn(memory_slots, hidden_size))
        self.ltm_values = nn.Parameter(torch.randn(memory_slots, hidden_size))
        
        # Attention mechanism for memory retrieval
        self.attention = nn.MultiheadAttention(hidden_size, num_heads=8)
        
        # Memory consolidation network
        self.consolidation = nn.Sequential(
            nn.Linear(hidden_size * 2, hidden_size),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(hidden_size, hidden_size)
        )
        
        # Forgetting gate (synaptic plasticity)
        self.forget_gate = nn.Sigmoid()
        self.decay_rate = 0.99
        
    def forward(self, x, memory_state=None):
        batch_size, seq_len, _ = x.shape
        
        # Process through short-term memory
        stm_out, (h_n, c_n) = self.stm(x, memory_state)
        
        # Retrieve from long-term memory using attention
        ltm_query = stm_out.view(-1, stm_out.size(-1))
        similarity = torch.matmul(ltm_query, self.ltm_keys.t())
        attention_weights = F.softmax(similarity / np.sqrt(self.ltm_keys.size(1)), dim=-1)
        retrieved_memory = torch.matmul(attention_weights, self.ltm_values)
        retrieved_memory = retrieved_memory.view(batch_size, seq_len, -1)
        
        # Consolidate memories
        consolidated = self.consolidation(
            torch.cat([stm_out, retrieved_memory], dim=-1)
        )
        
        # Apply forgetting mechanism
        forget_mask = self.forget_gate(consolidated)
        output = consolidated * forget_mask
        
        # Update long-term memory (hebbian learning)
        self.update_ltm(stm_out.detach())
        
        return output, (h_n, c_n)
    
    def update_ltm(self, new_memories):
        """Hebbian-inspired memory update rule"""
        with torch.no_grad():
            # Decay existing memories
            self.ltm_values *= self.decay_rate
            
            # Find least used memory slots
            usage = torch.norm(self.ltm_values, dim=1)
            _, indices = torch.topk(usage, k=new_memories.size(1), largest=False)
            
            # Store new memories in least used slots
            new_memories_flat = new_memories.view(-1, new_memories.size(-1))
            for i, idx in enumerate(indices[:new_memories_flat.size(0)]):
                self.ltm_keys[idx] = new_memories_flat[i]
                self.ltm_values[idx] = new_memories_flat[i]
    
    def recall_memory(self, cue, top_k=5):
        """Recall memories similar to given cue"""
        similarity = torch.matmul(cue, self.ltm_keys.t())
        values, indices = torch.topk(similarity, k=top_k)
        
        recalled_memories = self.ltm_values[indices]
        confidence_scores = F.softmax(values, dim=-1)
        
        return recalled_memories, confidence_scores

# Working Memory Module (Prefrontal Cortex inspired)
class WorkingMemory(nn.Module):
    def __init__(self, capacity=7, feature_dim=256):
        super().__init__()
        self.capacity = capacity  # Miller's magical number 7±2
        self.feature_dim = feature_dim
        
        # Memory slots
        self.memory_slots = nn.Parameter(torch.zeros(capacity, feature_dim))
        
        # Gating mechanisms
        self.input_gate = nn.Linear(feature_dim, capacity)
        self.maintenance_gate = nn.Linear(feature_dim, capacity)
        self.output_gate = nn.Linear(feature_dim, capacity)
        
    def forward(self, input_features, query):
        # Determine which slots to update
        input_weights = torch.sigmoid(self.input_gate(input_features))
        
        # Update memory slots
        for i in range(self.capacity):
            update_gate = input_weights[:, i:i+1]
            self.memory_slots.data[i] = (
                update_gate * input_features + 
                (1 - update_gate) * self.memory_slots[i]
            )
        
        # Retrieve from working memory
        attention_scores = torch.matmul(query, self.memory_slots.t())
        attention_weights = F.softmax(attention_scores, dim=-1)
        retrieved = torch.matmul(attention_weights, self.memory_slots)
        
        return retrieved