Memory Augmentation › Memory Writing and Updating
Language AI HandbookAll chapters Memory Augmentation › Memory Writing and UpdatingOne section · Sections in this chapter5 / 12 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Reading is only half the story. The model must also decide what to write to memory and when to update existing entries. Several strategies exist.
Append-Only Memory
The simplest approach appends new key-value pairs without modifying existing entries. This preserves all historical information but causes unbounded memory growth.
class AppendOnlyMemory:
"""
Memory that grows by appending new entries.
Simple but unbounded in size.
"""
def __init__(self, key_dim, value_dim):
self.key_dim = key_dim
self.value_dim = value_dim
self.keys = []
self.values = []
def write(self, key, value):
"""Add a new entry to memory."""
self.keys.append(key)
self.values.append(value)
def read(self, query, k=10):
"""Read from memory using top-k retrieval."""
if not self.keys:
return None, None
keys_tensor = torch.stack(self.keys)
values_tensor = torch.stack(self.values)
# Compute similarities
scores = torch.matmul(query, keys_tensor.T)
k = min(k, len(self.keys))
topk_scores, topk_indices = torch.topk(scores.squeeze(), k)
# Retrieve
topk_values = values_tensor[topk_indices]
attention = F.softmax(topk_scores, dim=-1)
retrieved = torch.sum(attention.unsqueeze(-1) * topk_values, dim=0)
return retrieved, topk_indices
def size(self):
return len(self.keys)
# Simulate memory growth
memory = AppendOnlyMemory(key_dim=64, value_dim=128)
tokens_processed = [100, 500, 1000, 5000, 10000]
for n in tokens_processed:
# Simulate adding entries
for _ in range(n - memory.size()):
key = torch.randn(64)
value = torch.randn(128)
memory.write(key, value)Append-Only Memory Growth:
Tokens Processed Memory Size Memory (MB, est.)
-------------------------------------------------------
100 100 0.07
500 500 0.37
1,000 1,000 0.73
5,000 5,000 3.66
10,000 10,000 7.32Memory usage grows linearly with tokens processed. After processing 10,000 tokens, the memory consumes approximately 0.73 MB. While this seems manageable, extrapolating to 1 million tokens would require 73 MB, and to 100 million tokens would require 7.3 GB. For long-running applications like continuous conversation or document streaming, this unbounded growth eventually exhausts available memory.
Bounded Memory with Eviction
To bound memory size, older or less-used entries must be evicted. Common strategies include FIFO (first-in, first-out), LRU (least recently used), and importance-weighted eviction.
class BoundedMemory:
"""
Memory with fixed maximum size.
Evicts entries when capacity is reached.
"""
def __init__(self, max_size, key_dim, value_dim, eviction="fifo"):
self.max_size = max_size
self.key_dim = key_dim
self.value_dim = value_dim
self.eviction = eviction
# Storage
self.keys = torch.zeros(max_size, key_dim)
self.values = torch.zeros(max_size, value_dim)
self.write_ptr = 0
self.filled = 0
# For LRU tracking
self.access_times = torch.zeros(max_size)
self.current_time = 0
def write(self, key, value):
"""Write to memory, evicting if necessary."""
if self.eviction == "fifo":
# Write at current pointer, wrapping around
idx = self.write_ptr
self.write_ptr = (self.write_ptr + 1) % self.max_size
elif self.eviction == "lru":
if self.filled < self.max_size:
idx = self.filled
else:
# Find least recently accessed
idx = torch.argmin(self.access_times[: self.filled]).item()
else:
raise ValueError(f"Unknown eviction policy: {self.eviction}")
self.keys[idx] = key
self.values[idx] = value
self.filled = min(self.filled + 1, self.max_size)
self.access_times[idx] = self.current_time
self.current_time += 1
def read(self, query, k=10):
"""Read with top-k retrieval, updating access times for LRU."""
if self.filled == 0:
return None, None
active_keys = self.keys[: self.filled]
active_values = self.values[: self.filled]
scores = torch.matmul(query, active_keys.T)
k = min(k, self.filled)
topk_scores, topk_indices = torch.topk(scores.squeeze(), k)
# Update access times for LRU
self.access_times[topk_indices] = self.current_time
self.current_time += 1
topk_values = active_values[topk_indices]
attention = F.softmax(topk_scores, dim=-1)
retrieved = torch.sum(attention.unsqueeze(-1) * topk_values, dim=0)
return retrieved, topk_indices
# Compare eviction policies
fifo_mem = BoundedMemory(
max_size=100, key_dim=64, value_dim=128, eviction="fifo"
)
lru_mem = BoundedMemory(max_size=100, key_dim=64, value_dim=128, eviction="lru")Bounded Memory with Different Eviction Policies: Maximum size: 100 slots FIFO (First-In, First-Out): - Evicts oldest entries regardless of usage - Simple, predictable behavior - Good when all information has equal importance LRU (Least Recently Used): - Evicts entries not accessed recently - Keeps frequently-needed information - Better for repeated retrieval patterns
The choice of eviction policy depends on the access patterns of your application. FIFO works well when information has a natural temporal ordering and older entries become less relevant over time. LRU is better suited for applications with "hot" entries that are accessed repeatedly, as it keeps frequently-used information available even if it was written long ago.
Learnable Writing Decisions
More sophisticated approaches learn when and what to write. A gating mechanism can control write operations based on content importance.
class GatedMemoryWriter(nn.Module):
"""
Learns to gate write operations based on content importance.
A write gate determines whether new content should be stored.
This prevents memory from filling with redundant information.
"""
def __init__(self, input_dim, key_dim, value_dim, memory_size):
super().__init__()
self.memory_size = memory_size
# Projections for key and value
self.key_proj = nn.Linear(input_dim, key_dim)
self.value_proj = nn.Linear(input_dim, value_dim)
# Write gate: determines importance
self.write_gate = nn.Sequential(
nn.Linear(input_dim, 64), nn.ReLU(), nn.Linear(64, 1), nn.Sigmoid()
)
# Memory storage
self.keys = nn.Parameter(
torch.zeros(memory_size, key_dim), requires_grad=False
)
self.values = nn.Parameter(
torch.zeros(memory_size, value_dim), requires_grad=False
)
self.write_ptr = 0
def forward(self, x, threshold=0.5):
"""
Conditionally write to memory based on learned gate.
Args:
x: Input representation to potentially store
threshold: Write gate threshold
Returns:
wrote: Whether the write occurred
gate_value: The gate's output (importance score)
"""
# Compute write gate
gate_value = self.write_gate(x).squeeze()
if gate_value > threshold:
key = self.key_proj(x)
value = self.value_proj(x)
# Write to memory
idx = self.write_ptr % self.memory_size
self.keys.data[idx] = key.detach()
self.values.data[idx] = value.detach()
self.write_ptr += 1
return True, gate_value.item()
return False, gate_value.item()
# Demonstrate gated writing
writer = GatedMemoryWriter(
input_dim=256, key_dim=64, value_dim=128, memory_size=100
)
# Simulate different input "importance" levels
inputs = [
torch.randn(256) * 0.5, # Low magnitude (likely unimportant)
torch.randn(256) * 2.0, # High magnitude (likely important)
torch.randn(256) * 1.0, # Medium magnitude
]
write_results = []
for i, inp in enumerate(inputs):
wrote, gate = writer(inp)
write_results.append({"input_idx": i, "gate_value": gate, "wrote": wrote})Gated Memory Writing:
Input Gate Value Wrote to Memory
------------------------------------------
0 0.451 No
1 0.426 No
2 0.477 NoThe gate values demonstrate learned selectivity. With the default threshold of 0.5, inputs with gate values above this threshold get written to memory while others are discarded. In this example with random initialization, the gate produces varying importance scores. After training, the gate would learn to assign high scores to information that proves useful for downstream tasks and low scores to redundant or irrelevant content.
The write gate learns to recognize important information that's worth storing. During training, the model receives gradients based on whether stored information proved useful for downstream tasks. This naturally filters out redundant or irrelevant content.