# Gradmem Context Memory Compression > Compress long context into compact memory tokens via iterative gradient descent. Learn to write information into prefix memory without storing full KV-caches, enabling efficient long-context reasoning and retrieval. - Skill: `adu2021/gradmem-context-memory-compression` (Agent Skill) - Install (CLI): `npx skillmds@latest add adu2021/gradmem-context-memory-compression` - Raw SKILL.md: https://api.skillmd.com/api/skills/adu2021/gradmem-context-memory-compression/raw - Safety review: pending - Works with: Claude Code, Claude.ai, OpenAI Codex - Category: Coding & Dev Tools - License: MIT - Author: adu2021 (https://skillmd.com/u/adu2021) - Updated: 2026-09-17 - Page: https://skillmd.com/skills/adu2021/gradmem-context-memory-compression --- # 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: 1. **Context Ingestion** — Process long context once to initialize memory tokens 2. **Gradient-Based Writes** — Iteratively update memory tokens using gradient descent on reconstruction loss 3. **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. ```python 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. ```python 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. ```python 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](https://arxiv.org/abs/2603.13875)