adu2021 GradMem: Test-Time Gradient Descent for Context Memory Compression
Large language models struggle with long contexts due to massive KV-cache overhead. GradMem enables models to compress context into compact memory tokens via test-time gradient descent. Rather than storing full cache for each query, the model learns to write information into a small set of learnable prefix tokens by optimizing a self-supervised reconstruction loss. This approach is particularly effective for context removal scenarios where the model must answer questions without accessing the original long context at inference time.
The key innovation is using gradient updates on memory tokens (with frozen model weights) rather than forward-only writes, enabling iterative error correction and much better context compression.
Core Concept
GradMem operates through three phases:
- Context Ingestion — Process long context once to initialize memory tokens
- Gradient-Based Writes — Iteratively update memory tokens using gradient descent on reconstruction loss
- Query Answering — Use compressed memory (without original context) to answer questions
This contrasts with forward-only approaches which write information in a single forward pass. By optimizing iteratively, GradMem can pack more information into the same number of tokens.
Architecture Overview
- Memory Token Initializer — Create learnable prefix tokens initialized from context
- Reconstruction Loss — Self-supervised objective: can we reconstruct context from memory?
- Gradient-Based Writer — Optimize memory tokens via SGD on reconstruction loss
- Query Encoder — Generate question embeddings in the same space as memory
- Memory-Only Reader — Answer questions using only memory tokens (no original context)
- Scaling Mechanism — Multiple gradient steps improve capacity linearly
Implementation Steps
Start by designing the memory tokens and initializing them from context.
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.optim import SGD
class ContextMemoryCompressor:
"""Compress long context into learnable memory tokens."""
def __init__(self, hidden_dim=768, num_memory_tokens=32, context_seq_len=8192):
self.hidden_dim = hidden_dim
self.num_memory_tokens = num_memory_tokens
self.context_seq_len = context_seq_len
# Learnable memory tokens
self.memory_tokens = nn.Parameter(
torch.randn(num_memory_tokens, hidden_dim) * 0.01
)
# Projections for context encoding
self.context_encoder = nn.TransformerEncoder(
nn.TransformerEncoderLayer(d_model=hidden_dim, nhead=8,
dim_feedforward=2048,
batch_first=True),
num_layers=2
)
# Reconstruction decoder
self.context_decoder = nn.TransformerDecoder(
nn.TransformerDecoderLayer(d_model=hidden_dim, nhead=8,
dim_feedforward=2048,
batch_first=True),
num_layers=2
)
def initialize_memory(self, context_embeddings: torch.Tensor):
"""Initialize memory tokens from context."""
# context_embeddings: [batch_size, context_len, hidden_dim]
batch_size = context_embeddings.size(0)
# Summarize context into memory tokens via attention
memory = self.memory_tokens.unsqueeze(0).expand(batch_size, -1, -1)
# Attention between memory and context
cross_attn = nn.MultiheadAttention(self.hidden_dim, num_heads=8,
batch_first=True)
memory_initialized, _ = cross_attn(memory, context_embeddings,
context_embeddings)
return memory_initialized
def reconstruction_loss(self, memory: torch.Tensor,
context: torch.Tensor) -> torch.Tensor:
"""Loss for reconstructing context from memory."""
# Reconstruct context using memory as queries and context as key/value
reconstructed = self.context_decoder(memory, context)
# MSE loss between original and reconstructed
loss = F.mse_loss(reconstructed, context)
return loss
def write_to_memory_gradient_based(self, context_embeddings: torch.Tensor,
num_gradient_steps=5,
learning_rate=0.1) -> torch.Tensor:
"""Iteratively optimize memory tokens via gradient descent."""
batch_size = context_embeddings.size(0)
# Initialize memory
memory = self.initialize_memory(context_embeddings)
# Create optimizer for memory only (model weights frozen)
memory_param = torch.nn.Parameter(memory.clone())
optimizer = SGD([memory_param], lr=learning_rate)
# Gradient descent steps
for step in range(num_gradient_steps):
optimizer.zero_grad()
# Compute reconstruction loss
loss = self.reconstruction_loss(memory_param, context_embeddings)
# Backward pass
loss.backward()
# Gradient step
optimizer.step()
print(f" Gradient step {step+1}: Loss = {loss.item():.4f}")
return memory_param.detach()
def forward_write(self, context_embeddings: torch.Tensor) -> torch.Tensor:
"""Baseline: single forward pass to write context (for comparison)."""
# Simple pooling of context
memory = context_embeddings.mean(dim=1, keepdim=True).expand(
-1, self.num_memory_tokens, -1
)
return memory
Now implement the query answering mechanism that uses only memory without original context.
class MemoryOnlyReader(nn.Module):
"""Answer questions using only memory tokens."""
def __init__(self, hidden_dim=768, vocab_size=50257):
super().__init__()
self.hidden_dim = hidden_dim
# Transformer decoder (query-conditioned)
self.decoder = nn.TransformerDecoder(
nn.TransformerDecoderLayer(d_model=hidden_dim, nhead=8,
dim_feedforward=2048,
batch_first=True),
num_layers=3
)
# Output projection
self.output_proj = nn.Linear(hidden_dim, vocab_size)
def forward(self, query_embeddings: torch.Tensor,
memory: torch.Tensor) -> torch.Tensor:
"""
Args:
query_embeddings: [batch_size, query_len, hidden_dim]
memory: [batch_size, num_memory_tokens, hidden_dim]
Returns:
logits: [batch_size, query_len, vocab_size]
"""
# Use memory as context (key/value), query as decoder input
output = self.decoder(query_embeddings, memory)
# Project to vocabulary
logits = self.output_proj(output)
return logits
Integrate gradient-based writing into full pipeline with benchmarking.
class GradMemModel(nn.Module):
"""Full GradMem system for context compression and QA."""
def __init__(self, base_model, hidden_dim=768, num_memory_tokens=32):
super().__init__()
self.base_model = base_model
self.compressor = ContextMemoryCompressor(hidden_dim, num_memory_tokens)
self.memory_reader = MemoryOnlyReader(hidden_dim)
def process_context(self, context_ids: torch.Tensor,
num_gradient_steps=5) -> torch.Tensor:
"""Compress context into memory tokens."""
# Encode context
context_embeddings = self.base_model.encoder(context_ids)
# Gradient-based memory writing
memory = self.compressor.write_to_memory_gradient_based(
context_embeddings,
num_gradient_steps=num_gradient_steps
)
return memory
def answer_question(self, query_ids: torch.Tensor,
memory: torch.Tensor) -> torch.Tensor:
"""Answer question using memory (without original context)."""
# Encode query
query_embeddings = self.base_model.encoder(query_ids)
# Generate answer
logits = self.memory_reader(query_embeddings, memory)
return logits
def forward(self, context_ids: torch.Tensor, query_ids: torch.Tensor,
num_gradient_steps=5) -> torch.Tensor:
"""Full pipeline: compress context, then answer query."""
# Compress context into memory
memory = self.process_context(context_ids, num_gradient_steps)
# Answer question using memory
logits = self.answer_question(query_ids, memory)
return logits
def benchmark_memory_compression(model, test_cases, baselines=['forward_only']):
"""Measure compression efficiency vs accuracy."""
results = []
for num_steps in [1, 3, 5, 10]:
accuracies = []
reconstruction_errors = []
for case in test_cases:
context = case['context']
query = case['query']
reference_answer = case['answer']
# Compress context
context_ids = tokenizer.encode(context)
context_embeddings = model.base_model.encoder(
torch.tensor([context_ids]))
# Write to memory (gradient-based)
memory = model.compressor.write_to_memory_gradient_based(
context_embeddings,
num_gradient_steps=num_steps
)
# Answer question
query_ids = tokenizer.encode(query)
logits = model.answer_question(torch.tensor([query_ids]), memory)
# Decode answer
predicted_answer = tokenizer.decode(torch.argmax(logits, dim=-1))
# Evaluate
accuracy = compute_em(predicted_answer, reference_answer)
accuracies.append(accuracy)
# Reconstruction error
reconstructed = model.compressor.context_decoder(memory,
context_embeddings)
recon_error = F.mse_loss(reconstructed,
context_embeddings).item()
reconstruction_errors.append(recon_error)
avg_accuracy = sum(accuracies) / len(accuracies)
avg_recon_error = sum(reconstruction_errors) / len(reconstruction_errors)
results.append({
'num_gradient_steps': num_steps,
'accuracy': avg_accuracy,
'reconstruction_error': avg_recon_error
})
print(f"Gradient steps {num_steps}: Accuracy={avg_accuracy:.1%}, "
f"Recon Error={avg_recon_error:.4f}")
return results
Practical Guidance
Hyperparameters and When to Use:
- Number of memory tokens typically 16-64; more tokens improve capacity but increase overhead
- Gradient steps 3-10 work well; each step scales capacity roughly linearly
- Learning rate 0.01-0.5; lower values converge more slowly, higher values may diverge
- Use when context is fixed across multiple questions (e.g., document QA, retrieval)
- Particularly effective for associative retrieval and structured reasoning tasks
When NOT to use:
- For streaming scenarios where context changes frequently (memory becomes stale)
- When exact context retrieval is needed (compression loses information)
- For very short contexts where compression overhead exceeds benefits
- When latency is critical; gradient descent adds test-time cost
Common Pitfalls:
- Memory tokens becoming misaligned with query encoder; use shared embeddings
- Reconstruction loss not capturing important information; weight loss by query relevance
- Gradient descent diverging; use gradient clipping and careful learning rate selection
- Not accounting for information loss in compression; use reconstruction error as proxy for quality
Reference
Paper: GradMem: Learning to Write Context into Memory with Test-Time Gradient Descent