RAPID: Long-Context Inference with Retrieval-Augmented Speculative Decoding
一、论文摘要
长上下文大型语言模型（LLM）的出现为处理大量文档提供了传统检索增强生成（RAG）的有力替代方案。然而，长上下文推理的计算开销带来了重大效率挑战。推测解码（SD）传统上使用较小的草稿模型加速推理，但其效果在长上下文场景中因内存受限的 KV 缓存操作而大幅下降。本文介绍 RAPID（Retrieval-Augmented SPeculatIve Decoding，检索增强推测解码），利用 RAG 同时加速和提升长上下文推理的生成质量。RAPID 引入了 RAG 草稿器——在缩短的检索上下文上运行的草稿 LLM——来推测长上下文目标 LLM 的生成。我们的方法开创了一种新范式，使同等规模甚至更大的 LLM 可以作为 RAG 草稿器，同时保持计算效率。为充分利用更强 RAG 草稿器的潜在优势，我们开发了一种推理时知识转移机制，通过 RAG 丰富目标分布。在 LLaMA-3.1 和 Qwen2.5 骨干网络上的大量实验表明，RAPID 有效整合了 RAG 和长上下文 LLM 的优势，实现了显著的性能提升（例如，LLaMA-3.1-8B 在 InfiniteBench 上从 39.33 提升至 42.83），同时在长上下文推理上实现超过 2 倍的加速。分析还揭示了 RAPID 在各种上下文长度和检索质量下的鲁棒性。
二、基本信息
论文 ID: 2502.20330
标题: RAPID: Long-Context Inference with Retrieval-Augmented Speculative Decoding
作者: Guanzheng Chen, Qilong Feng, Jinjie Ni, Xin Li, Michael Qizhe Shieh
（* 表示同等贡献）
单位:
- National University of Singapore（新加坡国立大学）
- DAMO Academy, Alibaba Group（阿里巴巴达摩院）
- Hupan Lab, Hangzhou, China（湖畔实验室）
会议/期刊: ICML 2025
代码: https://github.com/NUS-TRAIL/RAPID
原文保存位置: ~/.openclaw/workspace/papers/20260305_RAPID/source/
报告生成日期: 2026-03-05
三、论文主体分析
1. 引言
大型语言模型（LLM）传统上依赖检索增强生成（RAG）来处理大量文档，通过选择性检索相关文本片段完成任务。虽然 RAG 有效，但其性能本质上受限于检索器从不同查询中提取相关信息的能力。近期出现的长上下文 LLM 可直接处理百万级 token 的文档，提供了复杂 RAG 流程的有力替代方案。然而，这一突破受制于长上下文推理的计算效率——处理大量键值（KV）缓存变得内存受限，带来显著的延迟。
推测解码（SD）是一种加速 LLM 推理而不损失生成质量的主流方法。通过利用较小的草稿模型提出多个候选 token，由目标模型在单次前向传播中验证，SD 在候选被接受时实现显著加速。SD 的收益取决于两个关键因素：草稿模型生成候选的计算效率，以及生成高质量、可接受候选的能力。然而，SD 在长上下文场景中效果减弱，因为内存受限的 KV 缓存操作使较小 LLM 无法保持对较大模型的显著速度优势。
图 1：LLaMA-3.1-8B（单张 A800）和 LLaMA-3.1-70B（8 张 A800）在 LongBench v2（长文本）上跨不同检索上下文长度的性能（准确率，左轴）和吞吐量（tokens/秒，右轴）对比。随着上下文从 1K 增加到 128K，8B 模型的吞吐量增益从 23.6 骤降至 9.4。
:::
为此，本文提出 RAPID，通过引入 RAG 草稿器——在 RAG 缩短上下文上运行的草稿 LLM——来弥补 SD 在长上下文推理中的不足，同时提升生成质量。RAPID 提出 RAG 草稿器可以作为长上下文目标 LLM 的理想草稿模型，因为它在接近长上下文 LLM 能力的同时提供了更高的计算效率。
RAPID 支持两种设置：
- 自推测（Self-speculation）：长上下文目标 LLM 与 RAG 草稿器规模相同
- 上向推测（Upward-speculation）：RAG 草稿器参数规模大于目标 LLM
2. RAPID：检索增强推测解码
2.1 背景：推测解码
自回归生成传统上需要顺序前向传播，每个 token $x_i$ 从分布 $p_\phi(x_i | x_{<i})$ 中采样。SD 使用较小的草稿模型 $q_\psi$ 生成 $\gamma$ 个候选 token，由目标模型 $p_\phi$ 在单次前向传播中通过拒绝采样验证。
对于每个推测 token $x'_i \sim q_\psi(x_i | x_{<i})$，接受准则为：
$$r \leq \min\left(1, \frac{p_\phi(x'_i| x_{<i})}{q_\psi(x'_i| x_{<i})}\right)$$其中 $r \sim U(0,1)$。拒绝时，从残差分布重新采样：
$$x_i \sim \texttt{norm}({\max(p_\phi(x_i|x_{<i}) - q_\psi(x_i|x_{<i}), 0)})$$该过程保证重采样 token 遵循与直接从目标模型 $p_\phi$ 采样完全相同的分布。
2.2 方法概述
RAPID 由两个关键组件组成：
RAG 草稿器：RAPID 利用 RAG 草稿器为长上下文 LLM 生成候选，RAG 草稿器在选择性检索的上下文片段上运行，在保持访问相关信息的同时实现显著加速。
检索增强目标分布：SD 的严格接受准则可能拒绝高质量候选，因为它要求候选严格匹配目标 LLM 分布。RAPID 引导检索增强目标分布，实现推理时从 RAG 草稿器到目标模型的知识转移。
2.3 RAG 草稿器
处理大量上下文 $\mathcal{C}$ 的查询时，朴素 SD 的目标分布为：
$$p(x_i) = p_\phi(x_i \vert [\mathcal{C};x_{<i}])$$即使使用较小的草稿模型，由于对完整上下文 $\mathcal{C}$ 的内存受限 KV 缓存操作，计算收益也会大幅减少。
我们提出利用 RAG 作为草稿模型的基础。RAG 草稿器在压缩上下文 $\mathcal{C^S}$ 上运行，通过余弦相似度选择性检索相关文本片段。形式化地，草稿分布定义为：
$$q(x_i) = q_\psi(x_i \vert [\mathcal{C^S};x_{<i}])$$其中严格控制压缩比，使 $|\mathcal{C^S}| \le |\mathcal{C}|/\lambda$，$\lambda \gg 1$。
对候选 token 的验证采用修改后的接受准则：
$$r \leq \min\left(1, \frac{p(x_i)}{q(x_i)}\right) = \min\left(1, \frac{p_\phi(x'_i| [\mathcal{C};x_{<i}])}{q_\psi(x'_i| [\mathcal{C^S};x_{<i}])}\right) \quad \text{（公式 3）}$$RAG 草稿机制的两大优势：
- 通过压缩上下文操作（$|\mathcal{C^S}| \ll |\mathcal{C}|$）显著降低内存开销和计算成本
- 通过选择性检索相关信息，相比处理稀释的完整上下文，可能增强推测质量
2.4 检索增强目标分布
LLM 有效利用上下文的能力常随无关信息的引入而下降。经验分析表明，LLM 专注于检索的相关块时，有时在生成质量上可超过完整上下文利用。然而，传统 SD 的严格接受准则可能不必要地拒绝这些卓越的生成结果，导致质量下降和计算效率降低。
RAPID 引入检索增强目标分布，正式定义为：
$$\hat{p}(x_i) = \text{softmax}(z(x_i) / T + \eta \cdot (q(x_i) - p(x_i))) \quad \text{（公式 4）}$$其中 $\eta$ 是控制知识转移强度的超参数，$z(x_i)$ 是目标 LLM 的未归一化对数分数（logits），$p(x_i) = \text{softmax}(z(x_i)/T)$，$T$ 是温度。
命题 2.1（知识蒸馏梯度）：设 $p(x) = \text{softmax}(z(x)/T)$ 为以 logits $z(x)$ 和温度 $T$ 参数化的学生模型分布，$q(x)$ 为教师模型分布。知识蒸馏损失 $\mathcal{L} = T^2 \cdot \text{KL}(q(x) \| p(x))$ 关于 $z(x)$ 的梯度为：
$$\frac{\partial \mathcal{L}}{\partial z(x)} = T \cdot (p(x) - q(x))$$基于此，将 RAG 草稿器视为教师、目标模型视为学生，蒸馏 logits 偏移为：
$$\frac{\partial \mathcal{L}}{\partial z(x_i)} = T \cdot (p(x_i) - q(x_i))$$通过梯度上升步骤可以推导出由 RAG 草稿器增强的"蒸馏" $z(x_i)$：
$$\hat{z}(x_i) = z(x_i) - \eta \frac{\partial \mathcal{L}}{\partial z(x_i)} = z(x_i) + \eta T(q(x_i) - p(x_i)) \quad \text{（公式 5）}$$因此，检索增强目标分布等价于归一化后的 $\hat{z}(x_i)$，即 $\hat{p}(x_i) = \text{softmax}(\hat{z}(x_i)/T)$。
完整算法（RAPID 算法 1）如下：
- 输入：目标 LLM $p_\phi$，RAG 草稿器 $q_\psi$，上下文 $\mathcal{C}$，检索上下文 $\mathcal{C^S}$，推测 token 数 $\gamma$，温度 $T$，转移强度 $\eta$
- 循环生成：
- RAG 草稿器生成 $\gamma$ 个推测 token
- 顺序验证每个推测 token：计算目标和草稿分布，计算检索增强目标分布，拒绝采样
- 若被拒绝，从调整后的残差分布采样
拒绝时，从调整后的残差分布采样：
$$x_i \sim \texttt{norm}({\max(p(x_i)-\hat{p}(x_i), p(x_i)-q(x_i))})$$可证明（见原文附录）该采样策略保证了理论保障，即结果 token 遵循与直接从原始目标模型 $p(x_i)$ 采样相同的分布。
3. 实验设置
3.1 实现细节
目标与草稿 LLM：在 LLaMA-3.1（8B、70B）和 Qwen2.5（7B、72B）上评估 RAPID。实现两种推测设置：(1) 自推测，RAG 草稿器与目标 LLM 规模相同；(2) 上向推测，较大的 RAG 草稿器辅助较小的目标 LLM。每步生成 $\gamma=10$ 个 token 供目标 LLM 验证。$\eta$ 在自推测时从 $\{5, 10, 20\}$ 中搜索，上向推测时从 $\{40, 50\}$ 中搜索。
RAG 设置：将长上下文分割为 512-token 块，使用 BGE-M3 编码，基于余弦相似度检索 top-k 片段，过滤相似度低于 0.3 的片段。检索上下文长度限制在 4096 tokens 到输入长度的 1/24 之间。
3.2 评估协议
基线方法：
- 长上下文目标 LLM（LC）：直接在长上下文上生成
- RAG：在草稿 LLM 的检索上下文上生成
- 朴素推测解码（SD）：使用朴素长上下文目标分布
- MagicDec：利用 StreamingLLM 压缩草稿模型的 KV 缓存
评测基准：
- $\infty$Bench：三项任务——长书籍问答（En.QA，指标：F1）、多选问答（En.MC，指标：准确率）、摘要（En.Sum，指标：ROUGE-L-Sum），上下文长度超过 100K
- LongBench v2：跨 8K 到 2M 词多种上下文长度的多选任务
4. 结果与分析
4.1 主要结果
| 目标模型 | 方法 | 草稿模型 | En. QA | En. MC | En. Sum | AVG. | LB v2 整体 | LB v2（CoT） | 预填充时间(s) | 加速比 |
|---|---|---|---|---|---|---|---|---|---|---|
| LLaMA-3.1-8B | LC | - | 34.58 | 53.28 | 30.14 | 39.33 | 28.0 | 30.4 | 25.89 | 1.00× |
| RAG | - | 31.91 | 62.01 | 27.27 | 40.40 | 29.2 | 33.4 | 0.36 | 3.35× | |
| SD | - | 32.90 | 55.90 | 30.11 | 39.64 | 29.4 | 31.0 | 26.37 | 1.63× | |
| MagicDec | - | 29.83 | 52.03 | 30.18 | 37.35 | 29.2 | 30.6 | 26.05 | 0.71× | |
| RAPID | LLaMA-3.1-8B (RAG) | 34.90 | 63.32 | 30.27 | 42.83 | 32.4 | 34.2 | 26.37 | 2.10× | |
| RAPID | LLaMA-3.1-70B (RAG) | 40.94 | 79.04 | 29.96 | 49.98 | 38.8 | 40.2 | 28.04 | 1.14× | |
| LLaMA-3.1-70B | LC | - | 36.48 | 68.56 | 30.18 | 45.07 | 31.6 | 36.2 | 160.54 | 1.00× |
| RAG | - | 38.66 | 76.86 | 27.17 | 47.56 | 38.0 | 39.4 | 2.81 | 4.44× | |
| RAPID | LLaMA-3.1-70B (RAG) | 40.56 | 81.66 | 29.64 | 50.62 | 40.2 | 40.2 | 163.43 | 2.69× | |
| Qwen2.5-7B | LC | - | 16.93 | 66.81 | 30.62 | 38.12 | 30.2 | 33.2 | 20.32 | 1.00× |
| RAG | - | 20.28 | 75.11 | 25.60 | 40.33 | 31.2 | 33.8 | 0.34 | 6.47× | |
| RAPID | Qwen2.5-7B (RAG) | 19.81 | 75.98 | 31.64 | 42.48 | 32.0 | 35.4 | 21.62 | 2.65× | |
| RAPID | Qwen2.5-72B (RAG) | 30.10 | 83.84 | 32.21 | 48.72 | 35.6 | 41.2 | 23.45 | 0.93× | |
| Qwen2.5-72B | LC | - | 39.21 | 81.66 | 32.45 | 51.11 | 40.0 | 43.9 | 162.42 | 1.00× |
| RAG | - | 30.72 | 80.22 | 28.63 | 46.52 | 38.8 | 39.8 | 3.09 | 3.60× | |
| RAPID | Qwen2.5-72B (RAG) | 40.52 | 85.59 | 32.94 | 53.02 | 42.9 | 44.1 | 164.80 | 1.98× |
表 1：RAPID 与基线方法的综合评估。LC 和 RAG 分别表示在长上下文和检索上下文上评估目标模型。
主要发现：
- 自推测整合 RAG 和长上下文 LLM 的优势：在自推测设置中，RAPID 在所有模型系列上均取得一致改进。LLaMA-3.1-8B 在 $\infty$Bench（42.83 vs 39.33 LC, 40.40 RAG）和 LongBench v2（34.2% vs 30.4% LC, 33.4% RAG）上均表现优异。
- 更大的 RAG 草稿器通过知识转移进一步提升性能：LLaMA-3.1-8B 使用 70B RAG 草稿器在 $\infty$Bench 上达到 49.98，整体准确率达 40.2%，甚至超过 LLaMA-3.1-70B 的 LC 性能（36.2%）。
- RAPID 实现 >2× 加速：自推测设置中，RAPID 实现 2.10× 加速（LLaMA-3.1-8B）和 2.69×（LLaMA-3.1-70B）。
4.2 收益整合分析
图 2：LLaMA-3.1 系列不同目标-草稿模型配置在 LongBench v2 上相对于目标 LLM 的性能对比。RAPID 整合了目标和草稿 LLM 两者的优势，在不增加失败率的情况下实现更高的相对成功率。
:::
RAPID 成功整合了 RAG 草稿器和目标 LLM 的优势。相比于直接使用 RAG 应用，RAPID 实现了更高的准确率提升，同时维持显著更低的准确率下降。此外，与朴素 SD 相比，RAPID 在准确率提升方面大幅胜出，尽管共享相同的目标 LLM 和 RAG 草稿器。
最值得注意的是，RAPID 表现出一种"涌现现象"——成功处理了目标 LLM 和 RAG 草稿器单独失败的案例（图中 "RAPID Only" 部分）。该涌现准确率随 RAG 草稿器增强而增加，在上向推测设置中尤为明显。
4.3 上下文与检索长度的影响
图 3：上下文和检索长度对 RAPID（自推测）性能和效率的影响（LLaMA-3.1-8B）。RAPID 持续优于朴素 SD，在上下文长度超过 32K、检索长度适中（≤16K）时实现加速。上图：相对于目标 LLM 的 ΔAccuracy；中图：接受率（被接受草稿 token 的比例）；下图：相对于目标 LLM 推理的加速比。
:::
关键发现：
- RAPID 在所有配置中均实现 2-8% 的 $\Delta$Accuracy 改善，而朴素 SD 仅有边际甚至负的增益（-5 到 2%）
- RAPID 在上下文长度超过 32K 时实现加速（speedup > 1.0×），而 SD 需要超过 64K 的上下文才能展示加速
- 检索长度 4K 和 8K 实现几乎相同的加速比；超过 16K 时，较长草稿上下文的开销变得明显
4.4 生成质量分析
| 方法 | 质量评分（1-10） | 接受率（%） | 吞吐量（tokens/s） |
|---|---|---|---|
| 目标 LLM | 2.82 | - | 10.64 ± 0.98 |
| RAG 草稿器 | 3.95 | - | 40.49 ± 0.47 |
| SD | 2.94 | 56.34 ± 0.13 | 14.07 ± 3.08 |
| RAPID | 4.21 | 76.94 ± 0.13 | 18.18 ± 3.23 |
表 2：多轮对话生成（~122K token 聊天历史）的评估结果。质量评分由 GPT-4-Turbo-1106 使用 LLM-as-a-Judge 协议评定。
RAPID 生成质量评分（4.21）显著优于目标 LLM（2.82）、RAG 草稿器（3.95）和朴素 SD（2.94），同时接受率高达 76.94%（vs SD 的 56.34%），吞吐量为 18.18 tokens/s（目标 LLM 的 1.7× 加速）。
4.5 检索质量鲁棒性
| η | LLaMA-3.1-8B（草稿） | LLaMA-3.1-70B（草稿） | ||
|---|---|---|---|---|
| ΔAccuracy | Speedup | ΔAccuracy | Speedup | |
| 0 | 1.20 | 1.62× | -1.30 | 0.67× |
| 5 | 2.80 | 1.75× | 0.40 | 0.69× |
| 10 | 1.60 | 1.77× | 1.20 | 0.72× |
| 20 | 1.20 | 1.78× | 4.40 | 0.75× |
| 30 | -2.40 | 2.07× | 6.60 | 0.80× |
| 40 | -2.60 | 2.08× | 6.60 | 0.84× |
| 50 | -6.30 | 2.10× | 6.00 | 0.87× |
表 3：RAPID 在不同知识转移参数 $\eta$ 下的鲁棒性研究（使用无关检索上下文的压力测试）。
关键发现：
- 自推测（8B 草稿器）：$\eta \leq 20$ 时维持正的性能增益和 1.62×-1.78× 加速
- 上向推测（70B 草稿器）：在所有 $\eta$ 值下均保持正向性能增益（最高 6.60%），展示更强的鲁棒性
5. 相关工作
5.1 推测解码
推测解码通过小草稿模型提出多个候选 token 供单次验证来加速 LLM 推理。TriForce 和 MagicDec 尝试通过 KV 缓存压缩技术（StreamingLLM）将 SD 扩展到长上下文场景，但压缩方法往往产生弱化的草稿模型，在复杂应用中加速有限。RAPID 采用 RAG 草稿器，在各种应用中保持高质量推测和显著加速。
5.2 长上下文推理加速
研究主要关注两个方向：通过选择性保留或量化优化 KV 缓存操作，以及探索提示压缩方法。这些方法在提高效率的同时常常在没有质量保证的情况下压缩上下文信息。RAPID 利用 SD 通过长上下文 LLM 的明确验证保持生成质量，提供更可靠的效率与性能平衡。
5.3 RAG 与长上下文 LLM
近期研究揭示了 RAG 与长上下文 LLM 的互补优势，两者尽管表现特征不同，但有大量的预测重叠。RAPID 通过直接将 RAG 优势集成到解码过程中，提供了更原则性的解决方案，实现了动态适应，同时保留了两种范式的优势。
6. 结论
本文介绍了 RAPID，一种新颖的解码方法，通过检索增强推测弥补了推测解码在长上下文推理中的效率差距，同时提升了生成质量。RAPID 的核心在于利用 RAG 草稿器为长上下文目标 LLM 实现高效推测，并配以检索增强目标分布，有效整合可能更强的草稿器的知识。大量实验证明，RAPID 在不同模型规模和任务中均成功实现了计算效率和提升的生成质量。
四、论文简评
创新点
新颖的 RAG-SD 范式融合：首次将 RAG 作为推测解码的草稿机制，优雅地解决了长上下文 SD 效率下降的核心问题，思路清晰且具有实用价值。
上向推测的突破：打破了传统 SD"草稿模型必须小于目标模型"的范式，提出了用更大模型作为草稿器加速小模型推理的新思路，在保持合理效率的前提下实现大幅性能提升。
推理时知识转移的理论基础：将检索增强目标分布设计巧妙地与知识蒸馏梯度理论联系起来，有严格的数学推导支撑，不仅直觉清晰而且有理论保证。
涌现能力的发现：观察到 RAPID 能够解决目标 LLM 和 RAG 草稿器单独均失败的案例，这一"涌现现象"启示了模型协作的新可能性。
局限性
系统复杂性：RAPID 需要同时维护长上下文 LLM 和 RAG 系统，增加了部署复杂度。上向推测还需额外 GPU 资源运行更大的草稿器，限制了在资源受限环境中的适用性。
超参数敏感性：知识转移强度 $\eta$ 的选取需要仔细调优，在自推测中 $\eta > 20$ 时性能会下降，需要针对不同场景调整。
有限的加速场景：RAPID 主要在上下文长度超过 32K 时才能实现加速，在较短上下文场景中收益有限甚至为负。
检索质量依赖：尽管具有一定鲁棒性，但检索质量仍会影响最终性能，尤其在自推测设置中，当 $\eta$ 较大且检索质量差时性能可能退化。
应用场景
- 长文档问答系统：法律文档分析、科研文献综述、企业知识库问询
- 长对话系统：客服系统、多轮对话 AI 助手，需要理解冗长对话历史的场景
- 大规模文档摘要：从数百页报告中提取关键信息的企业场景
- 推理加速部署：对推理延迟敏感、文档超长的在线服务场景
可改进方向
动态检索策略：探索自适应调整检索长度的机制，根据问题复杂度动态决定检索范围，而非固定比例。
多轮迭代蒸馏：研究多轮知识转移的可能性，在多步推测中逐步精炼目标分布。
端到端联合优化：探索联合微调 RAG 草稿器和目标 LLM 的方法，以提升在特定领域的协作能力。
扩展至更多架构：探索 RAPID 在 Mamba、混合 Transformer 等新型架构上的适用性，以及与 KV 缓存压缩方法的结合。
检索质量自动估计：开发自动评估检索相关性的机制，动态调整 $\eta$ 参数，在检索质量差时降低知识转移强度，提升整体鲁棒性。