RAPID: Long-Context Inference with Retrieval-Augmented Speculative Decoding

📄 arXiv: 2502.20330v2 📥 PDF

作者: Guanzheng Chen, Qilong Feng, Jinjie Ni, Xin Li, Michael Qizhe Shieh

分类: cs.CL

发布日期: 2025-02-27 (更新: 2025-06-23)

备注: ICML 2025 Spotlight

💡 一句话要点

RAPID：检索增强推测解码加速长文本LLM推理并提升生成质量

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 长文本推理 检索增强生成 推测解码 大型语言模型 知识转移

📋 核心要点

  1. 长文本LLM推理计算开销大，传统推测解码在长文本场景下效率降低，主要受限于KV缓存的内存瓶颈。
  2. RAPID利用RAG，通过在缩短的检索上下文中运行的RAG drafter来推测长文本目标LLM的生成，实现加速和质量提升。
  3. 实验表明，RAPID有效整合了RAG和长文本LLM的优势，在长文本推理中实现了显著的性能提升和加速效果。

📝 摘要（中文）

长文本大型语言模型（LLMs）的出现为处理大量文档提供了一种有希望的替代方案，以取代传统的检索增强生成（RAG）。然而，长文本推理的计算开销带来了巨大的效率挑战。虽然推测解码（SD）传统上使用较小的草稿模型来加速推理，但由于受限于内存的KV缓存操作，其有效性在长文本场景中大大降低。我们引入了检索增强推测解码（RAPID），它利用RAG来加速和增强长文本推理中的生成质量。RAPID引入了RAG drafter——一个在缩短的检索上下文中运行的草稿LLM——来推测长文本目标LLM的生成。我们的方法实现了一种新的范例，即相同规模甚至更大的LLM可以充当RAG drafter，同时保持计算效率。为了充分利用更强大的RAG drafter的潜在优势，我们开发了一种推理时知识转移方法，通过RAG来丰富目标分布。在LLaMA-3.1和Qwen2.5主干上的大量实验表明，RAPID有效地整合了RAG和长文本LLM的优势，实现了显著的性能提升（例如，在LLaMA-3.1-8B的InfiniteBench上从39.33提升到42.83），并且长文本推理的速度提高了2倍以上。我们的分析还揭示了RAPID在各种上下文长度和检索质量上的鲁棒性。

🔬 方法详解

问题定义：论文旨在解决长文本LLM推理效率低下的问题。现有方法，如传统的推测解码，在长文本场景中由于KV缓存的内存限制，加速效果显著降低。此外，如何有效利用RAG增强长文本LLM的生成质量也是一个挑战。

核心思路：论文的核心思路是利用RAG来辅助推测解码过程。具体来说，使用一个RAG drafter（一个在缩短的检索上下文中运行的LLM）来预测长文本目标LLM的输出。通过这种方式，可以在保证计算效率的同时，利用RAG的知识增强能力。

技术框架：RAPID框架包含以下主要模块：1) RAG Drafter：负责在缩短的检索上下文中生成草稿。2) Target LLM：长文本目标LLM，用于验证草稿并生成最终输出。3) 检索模块：负责从外部知识库中检索相关信息，提供给RAG Drafter。4) 推理时知识转移：利用RAG增强目标LLM的生成分布。整体流程是，首先检索相关文档，然后RAG Drafter基于检索结果生成草稿，最后Target LLM验证草稿并生成最终结果。

关键创新：RAPID的关键创新在于将RAG与推测解码相结合，提出了一种新的推理范式。与传统的推测解码不同，RAPID使用RAG drafter，可以在保证效率的同时，利用RAG的知识增强能力。此外，推理时知识转移机制进一步提升了生成质量。

关键设计：RAG Drafter可以使用与Target LLM相同规模甚至更大的模型，这与传统推测解码中使用较小模型作为drafter不同。检索模块的选择和检索质量对RAPID的性能有重要影响。推理时知识转移的具体实现方式（例如，如何融合RAG drafter和Target LLM的输出分布）也是一个关键设计。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，RAPID在LLaMA-3.1和Qwen2.5等模型上取得了显著的性能提升。例如，在LLaMA-3.1-8B的InfiniteBench测试中，RAPID将性能从39.33提升到42.83。同时，RAPID实现了超过2倍的推理速度提升。实验还验证了RAPID在不同上下文长度和检索质量下的鲁棒性。

🎯 应用场景

RAPID可应用于需要处理大量文本信息的各种场景，例如：法律文档分析、金融报告解读、医学文献检索、长篇小说创作等。该研究可以显著提升长文本LLM在这些领域的应用效率和生成质量，降低计算成本，并促进相关技术的进一步发展。

📄 摘要（原文）

The emergence of long-context large language models (LLMs) offers a promising alternative to traditional retrieval-augmented generation (RAG) for processing extensive documents. However, the computational overhead of long-context inference presents significant efficiency challenges. While Speculative Decoding (SD) traditionally accelerates inference using smaller draft models, its effectiveness diminishes substantially in long-context scenarios due to memory-bound KV cache operations. We introduce Retrieval-Augmented Speculative Decoding (RAPID), which leverages RAG for both accelerating and enhancing generation quality in long-context inference. RAPID introduces the RAG drafter-a draft LLM operating on shortened retrieval contexts-to speculate on the generation of long-context target LLMs. Our approach enables a new paradigm where same-scale or even larger LLMs can serve as RAG drafters while maintaining computational efficiency. To fully leverage the potentially superior capabilities from stronger RAG drafters, we develop an inference-time knowledge transfer that enriches the target distribution by RAG. Extensive experiments on the LLaMA-3.1 and Qwen2.5 backbones demonstrate that RAPID effectively integrates the strengths of both RAG and long-context LLMs, achieving significant performance improvements (e.g., from 39.33 to 42.83 on InfiniteBench for LLaMA-3.1-8B) with more than 2x speedups for long-context inference. Our analyses also reveal the robustness of RAPID across various context lengths and retrieval quality.