Retrieval-Augmented Generation (RAG) struggles with complex analytical tasks spanning hundreds of documents, such as financial due diligence or regulatory compliance reviews, because similarity search misses cross-document connections. Task-aware knowledge compression (TAKC) addresses this by pre-compressing entire knowledge bases into task-specific representations using an LLM, with different summaries for different tasks. For example, an annual report compressed for financial analysis retains revenue figures and margins, while the same report compressed for compliance review keeps regulatory citations. TAKC reduces token count by 8x to 64x while targeting task-relevant information.
Multi-rate compression
TAKC maintains four compression tiers per task type: light (8x, ~87.5% context reduction), medium (16x, ~93.8%), high (32x, ~96.9%), and ultra (64x, ~98.4%). A query complexity analyzer routes incoming questions to the appropriate tier based on signals like query length and analytical language. Straightforward factual queries use ultra-compressed cache, while complex analytical queries use lightly compressed cache. This tier-based routing complements existing RAG optimizations.
Architecture on AWS
The implementation runs as two decoupled serverless pipelines: ingestion and query. The ingestion pipeline uses AWS Lambda triggered by Amazon S3 events to chunk documents into 256-token segments with 50-token overlap, then invokes a compression Lambda asynchronously to compress each chunk at all four tiers via Amazon Bedrock. Compressed outputs are stored in Amazon ElastiCache Serverless with keys like takc:financial:medium and backed up to S3. The query pipeline uses Amazon API Gateway, AWS Lambda for query complexity analysis and cache retrieval, and Amazon Bedrock for inference. Amazon Cognito handles JWT authentication, and AWS WAF provides rate limiting. Infrastructure is defined as a single AWS CDK stack.
Cost comparison
For a 100,000-token knowledge base queried 1,000 times per day, TAKC reduces input token costs significantly: full context uses 100 million tokens (100% relative cost), RAG (top-10 chunks) uses 10 million (10%), TAKC light uses 12.5 million (12.5%), medium uses 6.25 million (6.25%), high uses 3.125 million (3.1%), and ultra uses 1.563 million (1.6%). TAKC requires upfront compression cost during ingestion, which amortizes for stable knowledge bases.
When to choose TAKC, RAG, or both
TAKC is better for cross-document reasoning and synthesis, stable knowledge bases, well-defined task types, full corpus coverage, and tight token budgets. RAG is better for narrow factual lookups, frequently changing knowledge bases, unpredictable query patterns, when only a few documents are relevant, and when source attribution is required. In practice, a production system often benefits from both, with a query complexity analyzer routing between them.
Getting started
Prerequisites include an AWS account with Amazon Bedrock model access, AWS CDK CLI, Python 3.12+, and Node.js 18+. Deploy the CDK stack from the aws-samples/sample-bedrock-takc-compression repository, upload a document to the S3 bucket under a task-type prefix, wait 2-3 minutes for ingestion, then query the API endpoint. Clean up by emptying the S3 bucket and destroying the CDK stack.
Images
