File size: 3,648 Bytes
ecc111a 0bf798e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 | """
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
} |