Research map / Architectures and efficiency / Attention and transformers
APE: Faster and Longer Context-Augmented Generation via Adaptive Parallel Encoding
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes APE, a training-free inference-time modification of parallel KV encoding (shared prefix, attention temperature, and scaling factor) that closes most of the accuracy gap to sequential encoding for RAG and ICL while enabling cached, faster long-context generation.
Problem
Context-augmented generation (RAG and ICL) is slow because long concatenated contexts must be re-encoded per request, while naive parallel encoding with cached KV states degrades accuracy substantially, especially for ICL and reasoning tasks.
Approach
Align the attention distribution of parallel encoding with sequential encoding through three training-free steps: prepending a shared prefix to avoid duplicated attention-sink states, lowering attention temperature for contexts, and applying a scaling factor to the context LogSumExp term.
Key results
- On ChatRAG-Bench APE keeps about 98% of sequential encoding performance (e.g., 36.78 average with all texts vs. 36.38 for sequential Dragon-multiturn retrieval).
- APE outperforms parallel encoding (PCW) by 15.4% on GSM8K, 4.7% on TriviaQA and 3.5% on MMLU, and preserves 93% of full-shot sequential ICL accuracy.
- APE achieves up to 4.5x speedup over sequential encoding and 28x prefill-time reduction vs. MInference at 128K context on H100.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add a shared prefix (ablation) | Parallel encoding · GSM8K, TriviaQA, and MMLU | Accuracy: 38.25% -> 50.42%; 67.99% -> 70.76%; 63.09% -> 63.70% | Helped |
| Add attention temperature to the shared-prefix variant (ablation) | APE · GSM8K, TriviaQA, and MMLU | Accuracy: 50.42% -> 51.15%; 70.76% -> 71.03%; 63.70% -> 64.49% | Helped |
| Add a scaling factor to the shared-prefix and temperature variant (ablation) | APE · GSM8K, TriviaQA, and MMLU | Accuracy: 51.15% -> 53.62%; 71.03% -> 72.64%; 64.49% -> 66.62% | Helped |
| Use APE rather than sequential encoding with the same retrieved chunks | ChatRAG-Bench with Contriever · ChatRAG-Bench (Contriever) | Average: 29.53 -> 29.02; -0.51 | Hurt |
| Use APE rather than sequential encoding with the same retrieved chunks | ChatRAG-Bench with GTE-base · ChatRAG-Bench (GTE-base) | Average: 32.91 -> 31.99; -0.92 | Hurt |
| Use APE rather than sequential encoding with the same retrieved chunks | ChatRAG-Bench with Dragon-multiturn · ChatRAG-Bench (Dragon-multiturn) | Average: 36.38 -> 35.24; -1.14 | Hurt |
| Retrieve more and longer contexts using APE | LongBench RAG on LLaMA-3-8B-Instruct · LongBench RAG | Average: 37.60 -> 39.62 | Helped |
| Retrieve more and longer contexts using APE | LongBench RAG on Mistral-7B-Instruct-v0.3 · LongBench RAG | Average: 25.06 -> 32.55 | Helped |
| Retrieve more and longer contexts using APE | LongBench RAG on Gemma-2-9b-it · LongBench RAG | Average: 41.28 -> 44.18 | Helped |
| Add APE and process 20 parallel segments of 4K tokens | Llama-3-8B-Instruct on CRAG · CRAG Task 1 | Accuracy (%): 23.28 -> 25.53 | Helped |
| Add APE and process 20 parallel segments of 4K tokens | Llama-3-8B-Instruct on CRAG · CRAG Task 2 | Accuracy (%): 24.46 -> 27.04 | Helped |
| Segment long-context inputs and apply APE | LLaMA-3-8B-Instruct · LongBench long-context tasks | Average: 31.50 -> 38.11 | Helped |
Surprising findings
- Trainable parallel-encoding method CEPED underperforms parallel encoding on the MultiNews summarization task, suggesting fine-tuning on simple data hurts complex-task generalization.
- Rotating the initial token's key state reveals that separate per-context rotation axes degrade parallel encoding earlier (from 90 degrees) than a shared axis.
Limitations
- Sensitive to hyperparameters, particularly attention temperature T and scaling factor S, requiring a greedy search on a validation set.
- Difficulty automatically aligning parallel and sequential encoding when contexts vary in length, quantity, and content.
- Limited ability to model relationships between contexts, so it is not emphasized for long-context tasks.
Experimental setup
Datasets: ChatRAG-Bench LongBench GSM8K TriviaQA MMLU LOFT CRAG
Benchmarks: ChatRAG-Bench LongBench GSM8K TriviaQA MMLU LOFT CRAG
Metrics: accuracy F1 Rouge-L EM latency prefill time CRAG score
Baselines: Sequential encoding Parallel Context Windows (PCW) CEPE-Distilled (CEPED) Block-Attention MInference LLMLingua2 StreamingLLM Self-Extend
Builds on: Parallel Context Windows StreamingLLM Flash Attention CEPE
Closest papers by meaning
- RAPID: Long-Context Inference with Retrieval-Augmented Speculative Decoding (ICML 2025 · similarity 0.88)
- Inference Scaling for Long-Context Retrieval Augmented Generation (ICLR 2025 · similarity 0.87)
- SCBench: A KV Cache-Centric Analysis of Long-Context Methods (ICLR 2025 · similarity 0.86)
- Accelerating Inference of Retrieval-Augmented Generation via Sparse Context Selection (ICLR 2025 · similarity 0.85)
- ArkVale: Efficient Generative LLM Inference with Recallable Key-Value Eviction (NeurIPS 2024 · similarity 0.85)
- CateKV: On Sequential Consistency for Long-Context LLM Inference Acceleration (ICML 2025 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.