Originally published on tamiz.pro.
The paradox of Large Language Model (LLM) agents is increasingly glaring: as we grant them more context, their reliability often decreases. This phenomenon, frequently observed in Retrieval-Augmented Generation (RAG) pipelines and long-context agent workflows, stems from the "Lost in the Middle" effect. When an agent is fed thousands of tokens of retrieved documents, chat history, and tool outputs, the probability of it attending to irrelevant noise or conflating distinct facts rises exponentially, leading to hallucinations.
Context compression is not merely a cost-saving measure; it is a critical engineering control for reliability. By intelligently reducing the semantic density of the input context without discarding high-fidelity information, we can guide the model’s attention mechanism toward the most relevant signals. This article explores the architectural patterns, algorithmic techniques, and engineering practices for implementing robust context compression in production AI agents.
The Mechanics of Context Collapse
To understand how to compress context, we must first understand why uncompressed context fails. Transformer-based models rely on the attention mechanism, which calculates a weighted sum of value vectors for every token in the sequence. In long sequences, the attention weights are distributed across many tokens. If the relevant information is buried in a sea of noisy retrieval results, the signal-to-noise ratio drops.
Furthermore, standard RAG implementations often retrieve top-k documents based on vector similarity. However, semantic similarity does not always equate to relevance for the specific task at hand. A document might be topically similar but factually contradictory or outdated. Without a mechanism to prune or summarize this retrieved data, the agent is forced to reason over conflicting premises, a primary driver of hallucination.
Strategic Compression Techniques
Context compression can be approached at three distinct layers: the retrieval layer, the processing layer, and the prompt engineering layer. Each layer offers different trade-offs between latency, cost, and fidelity.
1. Vector Pruning and Re-ranking
The first line of defense is ensuring that only the most relevant chunks are included in the context window. Simple cosine similarity on dense embeddings is often insufficient for complex queries.
Cross-Encoder Re-ranking
Cross-encoders are transformer models that process the query and the document together, allowing them to capture fine-grained interactions between the two. While computationally more expensive than bi-encoders (used for initial retrieval), they provide a much more accurate relevance score.
Implementation Strategy:
- Use a bi-encoder (e.g.,
text-embedding-3-small) for fast initial retrieval of top 50-100 chunks. - Pass these chunks and the original query to a cross-encoder model (e.g.,
BGE-Reranker-v2-m3or Cohere’s Rerank API). - Re-rank and select only the top N results (e.g., top 5-10).
from sentence_transformers import CrossEncoder
import numpy as np
# Initial retrieval would happen here, yielding 'retrieved_chunks'
# Let's assume we have a list of retrieved document texts
query = "What are the specific constraints on data retention in the new compliance policy?"
retrieved_docs = [doc.text for doc in initial_retrieval_results]
# Prepare pairs for cross-encoder
pairs = [[query, doc] for doc in retrieved_docs]
# Initialize cross-encoder
cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')
# Get relevance scores
scores = cross_encoder.predict(pairs)
# Sort by score descending and pick top 5
indexed_scores = list(zip(scores, retrieved_docs))
indexed_scores.sort(key=lambda x: x[0], reverse=True)
top_context = [doc for _, doc in indexed_scores[:5]]
This approach ensures that the context window is populated with chunks that have a high semantic intersection with the user's specific intent, drastically reducing noise.
2. Summarization Chaining
When dealing with massive knowledge bases or long conversation histories, simple retrieval is often insufficient. Summarization chaining allows us to compress historical or retrieved data into a concise narrative.
Recursive Summarization
This technique involves breaking a long document into chunks, summarizing each chunk, and then summarizing those summaries recursively until a manageable size is reached. This preserves the global structure and key facts while discarding verbose details.
Architectural Pattern:
- Split document into chunks.
- Summarize each chunk using an LLM.
- Combine summaries.
- If the combined summary is still too large, repeat the process.
def summarize_chunk(chunk, previous_summary=""):
prompt = f"""
You are an expert analyst. Summarize the following text,
retaining only facts relevant to the previous summary if provided.
Previous Summary: {previous_summary}
Current Text:
{chunk}
New Summary (concise): """
return llm.generate(prompt)
def recursive_summarize(documents):
summaries = []
prev_summary = ""
for doc in documents:
s = summarize_chunk(doc, prev_summary)
summaries.append(s)
prev_summary = f"{prev_summary} {s}".strip()
return summaries
By maintaining previous_summary as context, we ensure that the new summary is not just isolated but connected to the broader narrative, preserving nuance across chunks.
3. Query Expansion and Reformulation
Often, the source of hallucination is not too much context, but too little relevant context. If the user’s query is vague, the retriever may fetch irrelevant documents. A compression technique here is to expand the query before retrieval, or to reformulate the query based on retrieved context.
Step-back Prompting
This involves asking the model to generate a more general, abstract version of the query. This step-back query can retrieve broader conceptual contexts that might contain the answer, which are then used to guide the retrieval of specific facts.
# Step 1: Generate a step-back query
step_back_prompt = f"""
Given the following question, generate a more general question to guide reasoning.
Question: {user_query}
General Question: """
step_back_query = llm.generate(step_back_prompt)
# Step 2: Retrieve using the step-back query
general_docs = vector_store.search(step_back_query)
# Step 3: Retrieve using the original query
specific_docs = vector_store.search(user_query)
# Step 4: Combine and compress
combined_context = general_docs + specific_docs
# Apply re-ranking here
This technique helps the model understand the conceptual space of the query, allowing it to filter out specific but irrelevant details in the subsequent generation step.
Advanced Compression: Semantic Pruning and Memory Management
For agent architectures that maintain long-term memory or process multi-step reasoning, static compression is not enough. We need dynamic, semantic-aware compression.
1. Importance-Aware Memory
Instead of treating all retrieved chunks equally, we can assign importance scores based on:
- Recency: More recent interactions are often more relevant.
- Frequency: How often a concept has been referenced.
- Surprise: How much the new information deviates from existing context (high surprise might indicate a critical update).
An agent can use a sliding window of high-importance memories, discarding low-importance ones when the context window approaches its limit. This is particularly useful for long-running conversations.
2. Structured Output Compression
LLMs perform better when context is structured. Instead of pasting raw text, we can compress context into structured formats like JSON or XML. This reduces token count by removing conversational filler and focuses the model on data.
Example:
Instead of:
"The user said their name is John. They mentioned they are a software engineer. They are looking for a job in New York."
Compress to:
{
"user": {
"name": "John",
"role": "Software Engineer",
"location_preference": "New York"
}
}
This not only saves tokens but also provides a clear schema for the LLM to follow, reducing the chance of misinterpreting relational data.
Mitigating Hallucinations During Generation
Even with perfect compression, hallucinations can occur if the generation process is not constrained. Context compression must be paired with generation controls.
1. Citation-Backed Generation
Require the LLM to cite the specific source chunk for every claim it makes. If it cannot cite a source, it should state that it does not know. This forces the model to ground its output in the compressed context.
Prompt Template:
Answer the question based ONLY on the provided context.
For each fact, cite the source ID [ID:xxx].
If the answer is not in the context, say "I don't know."
Context:
{compressed_context}
Question: {question}
2. Self-Consistency and Verification
For critical tasks, use self-consistency. Generate multiple answers using different random seeds or slightly varied prompts, and then aggregate the results. If the answers diverge significantly, it indicates low confidence, potentially due to ambiguous or missing context.
3. Guardrail Models
Deploy a secondary, smaller LLM as a guardrail. This model checks the generated response against the compressed context for factual consistency. If inconsistencies are found, the response is rejected or flagged for review.
Engineering Trade-offs and Best Practices
Implementing context compression is not a one-size-fits-all solution. Engineers must balance several trade-offs:
| Technique | Latency Impact | Cost Impact | Fidelity Retention | Best Use Case |
|---|---|---|---|---|
| Vector Pruning | Low | Low | High | General RAG, high-noise environments |
| Summarization | High | High | Medium | Long documents, long histories |
| Query Expansion | Medium | Medium | High | Complex, ambiguous user queries |
| Structured Output | Low | Low | High | Data extraction, API interactions |
| Guardrails | High | High | Very High | Critical decision-making, compliance |
Best Practices
- Start with Retrieval Optimization: Before implementing complex compression, ensure your embedding model and retrieval metrics are optimized. Often, better retrieval reduces the need for aggressive compression.
- Benchmark Reliability: Use a gold-standard dataset of questions and answers. Measure hallucination rates (factuality) and relevance scores before and after implementing compression techniques.
- Monitor Context Window Usage: Track the token count of your context windows. If you are consistently near the limit, it’s a signal that your compression strategy needs tightening.
- Human-in-the-Loop for Critical Paths: For high-stakes applications, allow human reviewers to validate the compressed context and the final output. Use this feedback to fine-tune your compression algorithms.
Conclusion
Context compression is a foundational technique for building reliable AI agents. By moving beyond naive retrieval and embracing strategic pruning, summarization, and structuring, engineers can significantly reduce hallucinations and improve the nuance of agent responses. The key is to view context not as a static block of text, but as a dynamic signal that must be filtered, refined, and focused to guide the model toward accurate, grounded reasoning.
As LLMs continue to evolve, the ability to manage and compress context efficiently will become a differentiator between fragile prototypes and robust, production-grade AI systems. The future of AI engineering lies not just in larger models, but in smarter context management.
Top comments (0)