Query-aware compression for Amazon Bedrock RAG: reduce input tokens, costs, and hallucinations

Reduce Bedrock RAG costs with query-aware compression

If your Amazon Bedrock bill spiked after a knowledge‑base pilot, one small architectural change can shrink the biggest cost driver: input tokens. Instead of sending tens of thousands of retrieved tokens straight to an expensive foundation model, run a cheap model to extract verbatim, query‑relevant spans, then ask the large model to answer from that compressed evidence. On many enterprise workloads this “query‑aware compression” pattern cuts tokens and dollars, often with better citation fidelity, while adding a modest latency cost and a touch more operational complexity.

How it works (simple flow)

User query → retriever returns top‑k chunks (commonly 5-20).
An orchestrator (example: AWS Lambda) receives query + these chunks.
Call a smaller model via the Bedrock Converse API to extract verbatim spans from each chunk (the compression call).
The compressor outputs only those spans and preserves chunk IDs.
Send the query + compressed context to the primary (larger) model via Converse.
The primary model answers, citing the compressed evidence.
Return the grounded answer to the user.

Why it works: retrieval maintains high recall. The compressor reduces token volume sent to the expensive final model. The final model reasons only over verbatim evidence rather than a noisy mass of text.

Quick cost intuition (back‑of‑envelope)

Let R be total retrieved input tokens, c the compression ratio (so compressed tokens = R/c), and A the final answer tokens. Using per‑token prices P_small_in, P_small_out, P_large_in, P_large_out, the per‑query cost formulas are:

Baseline (no compression): R · P_large_in + A · P_large_out
With compression: R · P_small_in + (R/c) · P_small_out + (R/c) · P_large_in + A · P_large_out

Legend: P_small_in is the per‑token price billed for the compressor’s input, P_small_out for its output, and similarly for the primary model. This algebra assumes the compressor emits mainly the compressed evidence tokens and minimal prompt/metadata. Add prompt overhead when you run your own numbers.

For concrete prices use the Amazon Bedrock pricing page. Model families show big spreads, example rows on that page illustrate why shaving input tokens to an expensive model often produces outsized savings. Always plug your workload’s R, c, and current Bedrock prices into the formula above.

A practical example (illustrative arithmetic)

Using the Bedrock pricing table as a reference, Sonnet‑class models and small Gemma‑class models have orders‑of‑magnitude per‑token differences. Plugging typical numbers shows why compression pays off. If R is large (tens of thousands of tokens) and a cheap compressor reduces that by 5-10×, the compressor’s small per‑token cost is more than offset by savings on the expensive model’s input tokens. Tailor the math to your exact model pairing and remember to include prompt tokens, rerank calls, retrieval compute and storage, Lambda invocations, and network egress in full cost estimates.

What a compressor prompt looks like (example)

“You extract evidence from retrieved documents.
You will receive a user QUESTION and a list of CHUNKS.

Your job:
1. For each chunk, identify the spans (verbatim sentences or short paragraphs)
that contain evidence directly relevant to answering the QUESTION.
2. Output ONLY those spans, copied verbatim from the source. Do not paraphrase,
summarize, or rewrite.
3. Preserve chunk identifiers so the downstream system can cite sources.
4. If a chunk contains no relevant evidence, output the chunk identifier
followed by NO_RELEVANT_EVIDENCE.

Output format (strict):
[CHUNK_ID: ]

[CHUNK_ID: ]
NO_RELEVANT_EVIDENCE

Do not add commentary, headings, conclusions, or your own words. Only verbatim
spans from the source chunks, grouped by chunk identifier.”

This prompt text is the compression instruction used in an Amazon Bedrock implementation. It enforces verbatim extraction and preserves chunk IDs so the final answer can cite sources. Pair that with a tight answer prompt such as “Answer using only the evidence provided. Cite sources.”

Benchmarks and measured trade‑offs (what was observed)

On a large enterprise corpus (roughly 500, 000 documents across multiple enterprise source types) and a 500‑question test set, an implementation of query‑aware compression on Amazon Bedrock reported these headline results (attributed to the Amazon Bedrock post by Aakanksha Veesam and Amit Maindola):

Cost: baseline = 100%, compression = 67%, rerank + compression = 64%.
Tokens to the primary model: baseline = 100%, compression ≈ 12%, rerank + compression ≈ 10%.
Latency: compression added ≈ +19% latency. Rerank + compression ≈ +12% (relative to baseline).
Composite quality (4 dimensions scored by an LLM judge): baseline = 100%, compression ≈ 97.5%, rerank + compression ≈ 97.6%.
Hallucination rate (answers with at least one unsupported claim): baseline 51%, compression ≈ 44% (−7 pts), rerank + compression ≈ 38% (−13 pts).

Bottom line from that exercise: roughly 33% cost reduction (64-67% of baseline), an 8-10× reduction in tokens presented to the expensive model, a small drop in composite quality, and a measurable reduction in hallucinations. These numbers describe one corpus and a single evaluation. Expect variance by domain, query mix, chunking, and model choices.

Why rerank+compression sometimes shows lower latency than compression alone. Reranking reorders and often prunes the retrieval set before compression, so the compressor processes fewer or denser chunks. In that benchmark the rerank step reduced compressor work enough to offset its own cost, producing a net latency win relative to the compression‑only pipeline. Measure end‑to‑end p50/p95/p99 for your model pair to confirm behavior under realistic load.

When to prototype query‑aware compression

Prototype when these heuristics apply to your workload:

Average retrieved context per query commonly exceeds ~5, 000 tokens.
Your primary model’s per‑token cost is high enough that a 3-10× reduction changes economics materially.
Most queries are narrow relative to the retrieved volume (support KB lookups, runbooks, regulatory citations, financial transcript Q&A).
Latency budget can tolerate one extra model call (typically a few hundred milliseconds to about 1 second; measure for your region and models).
You can build a representative evaluation set to validate evidence recall and faithfulness.

Quick start checklist (must‑have vs nice‑to‑have)

Must‑have: retriever that returns top‑k chunks, small compressor model, primary model, an orchestrator (e.g., Lambda), prompt for verbatim extraction, evaluation set.
Nice‑to‑have: reranker (pre‑compression), prompt caching, Intelligent Prompt Routing, Bedrock Guardrails for grounding checks, telemetry dashboards.

Pseudocode sequence

retrieve(query) → chunks
compressor = call_small_model(query, chunks) → compressed_spans
final = call_primary_model(query, compressed_spans) → answer_with_citations

Evaluation: what to measure and suggested acceptance gates

Run a frozen‑retrieval A/B so both pipelines see identical retrievals. Measure:

Evidence recall: does compressed context contain the ground‑truth supporting spans?
Faithfulness / hallucination rate: combine an LLM judge with human verification on a representative subset. Hallucination here means any asserted fact not present in the provided evidence.
Citation accuracy: are chunk IDs and verbatim spans intact and mapped correctly to sources?
Latency: p50/p95/p99 under realistic load.
Per‑query cost using current Bedrock pricing (include search and orchestration costs).

Example acceptance gates (illustrative, not universal): evidence recall ≥ 95% on the representative set, composite quality drop ≤ 3 percentage points, p95 latency increase ≤ 25%, cost savings > 15%. Use feature flags and a canary rollout if gates are met.

Operational risks and mitigations

Dropped critical evidence: compressors can miss tiny but vital clauses (negations, exceptions). Mitigation: conservative extraction thresholds and a fallback to full context when coverage is low.
Prompt drift: extraction prompts are brittle to format changes. Mitigation: periodic revalidation, drift detection, and prompt maintenance cadence.
PII and redaction risk: verbatim spans can expose PII. Mitigation: run PII detection and redaction or policy checks prior to compression; document data flow and auditing.
Hidden costs: retrieval compute (OpenSearch/Vector DB), Lambda invocations, reranker fees, network egress, and storage. Mitigation: include these in your per‑query cost model and measure at scale.
Edge cases for multi‑hop reasoning: compression may remove intermediate facts needed for multi‑step answers. Mitigation: include hard/multi‑step queries in your eval set and consider combining smarter retrieval (graph techniques, dynamic pooling) with compression.

Complementary tactics

Better retrieval: techniques like dynamic pooling, reranking, and graph‑based expansion reduce tokens before any model call.
Rerank before compression: reduces noise and can improve compressor effectiveness (the rerank+compression variant in benchmarks showed slightly better cost/quality tradeoffs).
Chunking strategy: smaller chunks increase recall but inflate token counts; tune chunk boundaries alongside the compressor prompt.
Guardrails and validation: use Bedrock Guardrails or rule‑based checks to enforce grounding and prevent unsupported claims.

Monitoring and production metrics

Expose these on a dashboard for ongoing health checks:

Tokens per query (retrieved / compressed / primary)
Per‑query cost and rolling cost savings
Evidence recall and hallucination rate
p50/p95/p99 latency
Fallback rate (how often you fall back to full context)
Prompt change / drift alerts

What the Amazon Bedrock experiment reported (attribution)

The benchmark figures and the compression system prompt quoted above come from an Amazon Bedrock implementation authored by Aakanksha Veesam and Amit Maindola. Their experimental setup used a corpus of roughly 500, 000 enterprise documents and 500 test questions. Answers were scored by an LLM judge across correctness, completeness, citation accuracy, and concision, with faithfulness tracked separately. Treat their headline numbers as an instructive single‑corpus case study rather than a universal guarantee, replicate the evaluation on your queries and data.

Key questions you’re likely asking

Will an extra model call always reduce cost?

No. It depends on retrieved token volume (R), compressor per‑token price, compression ratio (c), and primary model prices. For large R and an expensive primary model, a cheap compressor usually produces net savings; for small R or inexpensive primaries, it may not. Run the algebra with current Bedrock pricing and include orchestration/search costs.
Does compression hurt answer quality?

On the reported enterprise corpus composite quality fell only slightly (≈100% → 97.5%) while hallucinations decreased. That’s encouraging but dataset‑dependent, validate on your narrow, multi‑step, and regulatory queries before rolling out.
How much latency does it add?

The experiment observed extra latency measured as a few hundred milliseconds up to ≈1 second; relative slowdowns were roughly +12-19% depending on pipeline variants. Real numbers depend on models, region, payload size, and load, measure p95/p99 in your environment.
Which workloads benefit most?

Support knowledge bases, runbooks, regulated documents, and financial transcripts, workloads where queries are narrow but retrieval returns lots of broadly relevant text. For multi‑hop research tasks, combine retrieval improvements and conservative compression.

Next steps to prototype

Assemble a representative set of 200-500 queries including “typical” and “hard” cases.
Freeze retrieval results and run baseline vs compressed pipelines on identical retrievals.
Measure evidence recall, hallucination rate (LLM judge + human spot checks), p95 latency, and per‑query cost including retrieval and orchestration.
If evidence recall ≥ 95% and cost savings exceed your gate (example: >15%) with acceptable latency, run a small canary with monitoring and rollback rules.

Query‑aware compression is a practical lever that often reduces a dominant recurring cost in RAG systems: input tokens to expensive models. It’s not a hammer for every job, but when retrieval returns thousands of tokens and questions are narrow, the pattern usually deserves a quick prototype, complete with PII checks, fallback rules, and careful evaluation.

“Answer using only the evidence provided. Cite sources.”