Research map / Large language models / Retrieval and knowledge
RAPID: Long-Context Inference with Retrieval-Augmented Speculative Decoding
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 RAPID, a speculative decoding method where a RAG-based draft model operating on retrieved short contexts speculates for a long-context target LLM, with a retrieval-augmented target distribution that transfers knowledge from the drafter during inference.
Problem
Speculative decoding loses its speed benefit in long-context inference because KV cache operations are memory-bound, and strict rejection sampling discards good candidates from stronger drafters.
Approach
Use a RAG drafter (the same-scale or even larger LLM given a retrieved short context) as the draft model for a long-context target, and modify the target distribution with an inference-time logit shift toward the drafter, derived as a distillation gradient step with the drafter as teacher.
Key results
- On InfiniteBench, RAPID with LLaMA-3.1-8B self-speculation reaches 42.83 average vs 39.33 for the long-context LLaMA-3.1-8B, with 2.10x speedup on LongBench v2.
- Upward speculation with LLaMA-3.1-70B as RAG drafter lifts LLaMA-3.1-8B to 49.98 on InfiniteBench.
- On multi-turn dialogue (~122K tokens), RAPID scores 4.21 quality vs 2.82 for the target and 2.94 for naive SD, with 76.94% acceptance rate.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use a same-scale RAG drafter with retrieval-augmented target distribution | LLaMA-3.1-8B · ∞Bench AVG | score: 39.33 -> 42.83 | Helped |
| Use a larger LLaMA-3.1-70B RAG drafter instead of the same-scale drafter | LLaMA-3.1-8B · ∞Bench AVG | score: 42.83 to 49.98 | Helped |
| Use a same-scale RAG drafter with retrieval-augmented target distribution | LLaMA-3.1-70B · ∞Bench AVG | score: 45.07 -> 50.62 | Helped |
| Use a same-scale RAG drafter with retrieval-augmented target distribution | Qwen2.5-7B · ∞Bench AVG | score: 38.12 -> 42.48 | Helped |
| Use a larger Qwen2.5-72B RAG drafter instead of the same-scale drafter | Qwen2.5-7B · ∞Bench AVG | score: 48.72 vs 42.48 | Helped |
| Use retrieval-augmented speculative decoding instead of TriForce | LWM-Text-Chat-128K · LongBench v2 (Long, CoT) | accuracy: 18.0 -> 21.6 | Helped |
| Use retrieval-augmented speculative decoding instead of TriForce | LWM-Text-Chat-128K · LongBench v2 (Long, CoT) | speedup: 1.27 -> 2.56 | Helped |
| Use retrieval-augmented speculative decoding instead of MInference | LLaMA-3.1-8B · LongBench v2 (Long, CoT) | accuracy: 30.9 -> 34.2 | Helped |
| Increase η from 20 to 30 with an LLaMA-3.1-8B drafter under unrelated retrieval (ablation) | RAPID on LLaMA-3.1-8B · LongBench v2 (Long, CoT), unrelated retrieval | ΔAccuracy: 1.20 -> -2.40 | Hurt |
| Increase η from 20 to 30 with an LLaMA-3.1-70B drafter under unrelated retrieval (ablation) | RAPID on LLaMA-3.1-8B · LongBench v2 (Long, CoT), unrelated retrieval | ΔAccuracy: 4.40 -> 6.60 | Helped |
| Increase η from 40 to 50 with an LLaMA-3.1-70B drafter under unrelated retrieval (ablation) | RAPID on LLaMA-3.1-8B · LongBench v2 (Long, CoT), unrelated retrieval | ΔAccuracy: 6.60 -> 6.00 | Hurt |
| Increase η from 0 to 5 with an LLaMA-3.1-8B drafter under unrelated retrieval (ablation) | RAPID on LLaMA-3.1-8B · LongBench v2 (Long, CoT), unrelated retrieval | ΔAccuracy: 1.20 -> 2.80 | Helped |
Surprising findings
- Upward speculation with a 70B RAG drafter yields 0.93x-1.14x throughput, so gains are mainly in quality rather than speed in that setting.
- Naive RAG alone sometimes underperforms the long-context model (e.g., Qwen2.5-72B En.QA drops from 39.21 to 30.72), yet RAPID still improves on both.
Limitations
- Upward speculation needs extra GPUs to host the larger drafter and yields little or no speedup (0.93x-1.14x).
- Sensitive to the transfer strength eta: large values degrade performance on unrelated retrieval contexts for the 8B drafter.
- Evaluation of multi-turn dialogue relies on a single LLM judge (GPT-4-Turbo) on a custom-constructed dataset.
Experimental setup
Datasets: InfiniteBench LongBench v2 MT-Bench-101
Benchmarks: InfiniteBench LongBench v2 MT-Bench-101 (adapted multi-turn dialogue)
Metrics: F1 accuracy ROUGE-L-Sum speedup prefill time acceptance rate quality score (1-10)
Baselines: LC (long-context target) RAG naive speculative decoding MagicDec TriForce MInference
Builds on: speculative decoding MagicDec StreamingLLM BGE-M3
Closest papers by meaning
- Inference Scaling for Long-Context Retrieval Augmented Generation (ICLR 2025 · similarity 0.94)
- Accelerating Inference of Retrieval-Augmented Generation via Sparse Context Selection (ICLR 2025 · similarity 0.92)
- MagicDec: Breaking the Latency-Throughput Tradeoff for Long Context Generation with Speculative Decoding (ICLR 2025 · similarity 0.92)
- Long-Context LLMs Meet RAG: Overcoming Challenges for Long Inputs in RAG (ICLR 2025 · similarity 0.92)
- Scaling Up, Speeding Up: A Benchmark of Speculative Decoding for Efficient LLM Test-Time Scaling (ICLR 2026 · similarity 0.92)
- Accelerating LLM Inference with Lossless Speculative Decoding Algorithms for Heterogeneous Vocabularies (ICML 2025 · similarity 0.92)
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.