Memory Augmentation › Memorizing Transformers
Language AI HandbookAll chapters Memory Augmentation › Memorizing TransformersOne section · Sections in this chapter6 / 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
The Memorizing Transformer (Wu et al., 2022) represents a practical integration of external memory with transformer architecture. It extends a standard transformer by adding a -nearest neighbors (NN) retrieval mechanism that queries a large external memory during the attention computation.
In NN retrieval, given a query vector, we find the memory entries whose keys are most similar to the query (typically measured by dot product or cosine similarity). Only these entries participate in the subsequent attention computation, making retrieval efficient regardless of total memory size.
The key insight is that you can store key-value pairs from previous forward passes and retrieve them during current processing. This effectively extends context far beyond the model's native window without increasing the attention computation cost.
class MemorizingAttention(nn.Module):
"""
Simplified Memorizing Transformer attention layer.
Combines local self-attention with kNN retrieval from external memory.
"""
def __init__(self, d_model, num_heads, memory_size, k_neighbors=32):
super().__init__()
self.d_model = d_model
self.num_heads = num_heads
self.head_dim = d_model // num_heads
self.k_neighbors = k_neighbors
# Standard attention projections
self.q_proj = nn.Linear(d_model, d_model)
self.k_proj = nn.Linear(d_model, d_model)
self.v_proj = nn.Linear(d_model, d_model)
self.out_proj = nn.Linear(d_model, d_model)
# External memory (simplified: stored as tensors)
self.memory_keys = None
self.memory_values = None
self.memory_size = memory_size
# Gate for combining local and memory attention
self.memory_gate = nn.Linear(d_model, 1)
def update_memory(self, keys, values):
"""Add new key-value pairs to external memory."""
if self.memory_keys is None:
self.memory_keys = keys.detach()
self.memory_values = values.detach()
else:
self.memory_keys = torch.cat(
[self.memory_keys, keys.detach()], dim=0
)
self.memory_values = torch.cat(
[self.memory_values, values.detach()], dim=0
)
# Truncate if exceeds size
if self.memory_keys.shape[0] > self.memory_size:
self.memory_keys = self.memory_keys[-self.memory_size :]
self.memory_values = self.memory_values[-self.memory_size :]
def retrieve_from_memory(self, queries):
"""
Retrieve top-k most similar key-value pairs from memory.
Args:
queries: (batch, seq_len, d_model)
Returns:
retrieved_values: (batch, seq_len, d_model)
"""
if self.memory_keys is None or self.memory_keys.shape[0] == 0:
return torch.zeros_like(queries)
batch, seq_len, _ = queries.shape
# Flatten queries for retrieval
q_flat = queries.view(-1, self.d_model) # (batch * seq_len, d_model)
# Compute similarities
scores = torch.matmul(
q_flat, self.memory_keys.T
) # (batch * seq_len, memory_size)
scores = scores / np.sqrt(self.head_dim)
# Top-k retrieval
k = min(self.k_neighbors, self.memory_keys.shape[0])
topk_scores, topk_indices = torch.topk(scores, k, dim=-1)
# Gather values
topk_values = self.memory_values[
topk_indices
] # (batch * seq_len, k, d_model)
# Attention over retrieved values
attn = F.softmax(topk_scores, dim=-1) # (batch * seq_len, k)
retrieved = torch.einsum("bk,bkd->bd", attn, topk_values)
return retrieved.view(batch, seq_len, self.d_model)
def forward(self, x, store_to_memory=True):
"""
Forward pass with local attention and memory retrieval.
Args:
x: (batch, seq_len, d_model)
store_to_memory: Whether to store this chunk's KV to memory
Returns:
output: (batch, seq_len, d_model)
"""
batch, seq_len, _ = x.shape
# Compute Q, K, V for local attention (keep original for memory ops)
Q_orig = self.q_proj(x)
K_orig = self.k_proj(x)
V_orig = self.v_proj(x)
# Reshape for multi-head attention
Q = Q_orig.view(
batch, seq_len, self.num_heads, self.head_dim
).transpose(1, 2)
K = K_orig.view(
batch, seq_len, self.num_heads, self.head_dim
).transpose(1, 2)
V = V_orig.view(
batch, seq_len, self.num_heads, self.head_dim
).transpose(1, 2)
# Local attention
scores = torch.matmul(Q, K.transpose(-2, -1)) / np.sqrt(self.head_dim)
local_attn = F.softmax(scores, dim=-1)
local_out = torch.matmul(local_attn, V)
local_out = (
local_out.transpose(1, 2)
.contiguous()
.view(batch, seq_len, self.d_model)
)
# Memory retrieval (use original Q before multi-head reshape)
memory_out = self.retrieve_from_memory(Q_orig)
# Gate between local and memory
gate = torch.sigmoid(self.memory_gate(x)) # (batch, seq_len, 1)
combined = gate * local_out + (1 - gate) * memory_out
# Store current KV to memory for future retrieval
if store_to_memory:
# Use original K, V reshaped for flat storage
self.update_memory(
K_orig.reshape(-1, self.d_model),
V_orig.reshape(-1, self.d_model),
)
return self.out_proj(combined)
# Demonstrate memorizing attention
mem_attn = MemorizingAttention(
d_model=256, num_heads=8, memory_size=10000, k_neighbors=32
)# Simulate processing multiple chunks
chunk_size = 512
num_chunks = 10
memory_sizes = []
for i in range(num_chunks):
x = torch.randn(2, chunk_size, 256) # batch=2, seq=512, d=256
output = mem_attn(x, store_to_memory=True)
if mem_attn.memory_keys is not None:
memory_sizes.append(mem_attn.memory_keys.shape[0])
else:
memory_sizes.append(0)Memorizing Attention: Processing Chunks
Chunk Tokens Seen Memory Size Effective Context
------------------------------------------------------------
1 1,024 1,024 1,536
2 2,048 2,048 2,560
3 3,072 3,072 3,584
4 4,096 4,096 4,608
5 5,120 5,120 5,632
6 6,144 6,144 6,656
7 7,168 7,168 7,680
8 8,192 8,192 8,704
9 9,216 9,216 9,728
10 10,240 10,000 10,512The memory grows with each processed chunk, accumulating key-value pairs from all previous tokens. By chunk 10, the model has seen over 10,000 tokens but only needs to store 8,192 keys in memory (bounded by the configured memory size). The effective context, which is the combination of the local window (512 tokens) plus the memory (8,192 entries), reaches 8,704 positions that the model can attend to when generating the next token.
The Memorizing Transformer demonstrates how external memory scales context dramatically. After processing just 10 chunks, the model can access information from 10,000+ tokens ago while keeping each attention operation limited to a small local window plus retrieved neighbors.
Memory Compression for Long-Term Storage
Raw storage of all past key-value pairs becomes expensive at scale. Compression techniques reduce memory footprint while preserving utility.
class CompressedMemory(nn.Module):
"""
Memory with learned compression for long-term storage.
Recent entries are stored at full resolution.
Older entries are compressed into summary vectors.
"""
def __init__(
self, d_model, short_term_size, long_term_size, compression_ratio=8
):
super().__init__()
self.d_model = d_model
self.short_term_size = short_term_size
self.long_term_size = long_term_size
self.compression_ratio = compression_ratio
# Compression network
self.compressor = nn.Sequential(
nn.Linear(d_model * compression_ratio, d_model * 2),
nn.ReLU(),
nn.Linear(d_model * 2, d_model),
)
# Storage
self.short_term_keys = []
self.short_term_values = []
self.long_term_keys = []
self.long_term_values = []
def write(self, key, value):
"""Write to short-term memory, compressing to long-term when full."""
self.short_term_keys.append(key)
self.short_term_values.append(value)
# Check if compression is needed
if len(self.short_term_keys) >= self.short_term_size:
self._compress_oldest()
def _compress_oldest(self):
"""Compress oldest short-term entries into long-term memory."""
n_to_compress = self.compression_ratio
if len(self.short_term_keys) < n_to_compress:
return
# Take oldest entries
old_keys = self.short_term_keys[:n_to_compress]
old_values = self.short_term_values[:n_to_compress]
# Compress
key_concat = torch.cat(old_keys, dim=-1)
value_concat = torch.cat(old_values, dim=-1)
compressed_key = self.compressor(key_concat)
compressed_value = self.compressor(value_concat)
# Store compressed
self.long_term_keys.append(compressed_key)
self.long_term_values.append(compressed_value)
# Remove from short-term
self.short_term_keys = self.short_term_keys[n_to_compress:]
self.short_term_values = self.short_term_values[n_to_compress:]
# Bound long-term memory
if len(self.long_term_keys) > self.long_term_size:
self.long_term_keys = self.long_term_keys[-self.long_term_size :]
self.long_term_values = self.long_term_values[
-self.long_term_size :
]
# Demonstrate compression
compressed_mem = CompressedMemory(
d_model=128, short_term_size=64, long_term_size=128, compression_ratio=8
)
# Simulate writing many entries
for i in range(200):
key = torch.randn(128)
value = torch.randn(128)
compressed_mem.write(key, value)Compressed Memory Statistics: Entries written: 200 Short-term memory: 56 entries (full resolution) Long-term memory: 18 entries (8:1 compressed) Storage comparison: Without compression: 51,200 parameters With compression: 18,944 parameters Compression ratio: 2.7x
After writing 200 entries, the compressed memory stores only a fraction of what an uncompressed system would require. Short-term memory holds recent entries at full resolution for precise retrieval, while long-term memory contains compressed summaries that capture the essence of older entries. The achieved compression ratio demonstrates significant storage savings, though the actual information retained depends on how well the compression network learns to preserve task-relevant details.
