arXiv:2502.20330cs.CL2025-02ICML被引 14

用检索增强的猜测解码,让长文本推理更快更准

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

  • 用检索后的短上下文做小模型猜测,加速长文本生成
  • 在LLaMA-3.1上推理速度提升2倍以上,得分从39.33升至42.83
  • 适合需要高效处理长文档的场景,如法律、科研摘要

长上下文大语言模型为处理长文档提供了新思路,但其推理计算开销大。传统猜测解码(SD)依赖小模型加速,但在长上下文场景中因键值缓存操作受限而效果下降。本文提出检索增强的猜测解码(RAPID),利用RAG构建基于缩短检索上下文的草案模型,用于推测长上下文目标模型的生成。该方法允许同规模甚至更大模型作为草案模型,同时保持计算效率。为进一步挖掘强草案模型的优势,引入推理时知识迁移,丰富目标分布。在LLaMA-3.1和Qwen2.5模型上的实验表明,RAPID显著提升性能(如LLaMA-3.1-8B在InfiniteBench上从39.33升至42.83),长上下文推理速度提升超2倍,且对不同上下文长度与检索质量均表现出鲁棒性。

原文摘要 · Abstract (English)

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.

长文本推理推测解码RAG增强高效生成

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。