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

Dear Reviewer s5Em,

We sincerely appreciate your thorough review of our paper. Your constructive feedback will help us strengthen this work. Below are our responses to your concerns:

1. The paper does not introduce entirely new theoretical frameworks but rather combines existing ideas in a way.

A1: We believe big ideas are always composed of many small and existing ideas. While individual components (speculative decoding, RAG, long-context LLMs) may be familiar, our contribution lies in the novel integration of RAG drafter for long-context target models using speculative decoding, with theoretically guaranteed inference-time transfer to combine both benefits from long-context LLMs and RAG. Our method is not an extension of any previous works, but provides a new perspective and solution for a long-time debate regarding “long-context LLMs or RAG”, which is meaningful for the application of long-context LLMs.

2. The upward-speculation setting requires additional computational resources (extra GPUs) to serve the larger RAG drafter.

A2: We agree and especially point out that a larger RAG drafter will introduce more latency and require extra GPUs to serve. However, our RAPID can operate on two modes: (1) self-speculation for low-resource scenarios and (2) upward-speculation for high-resource scenarios. For a low-resource scenario, the self-speculation of RAPID can work well with both improved performance and efficiency and without additional resource requirements, which demonstrates great potential in real applications.

While we believe the upward speculation of RAPID can serve as a turbo mode, which maximizes performance with comparable generation speed, but consumes more resources. For example, our RAPID enables LLaMA-3.1-8B (LC) with LLaMA-3.1-70B drafter to achieve 10 points accuracy gains with a comparable generation throughput to naive LLaMA-3.1-8B (LC), which allows us to explore the higher-bound performance wth a similar level of speed but more GPUs.

3. Although the paper demonstrates robustness to suboptimal retrieval contexts, the method's performance can still be influenced by the quality of retrieval.

A3: We agree that RAPID’s performance can still be influenced by the quality of retrieval. We are not only demonstrating the robustness of RAPID to retrieval quality in Section 4.5, but also seek to highlight that our inference-time knowledge transfer will not hinder the target LLM from utilizing its long-context capabilities to reject low-quality candidates from RAG drafter as long as $\eta$ is properly set. Our adjusted resampling distribution in Eq. (10) for rejected candidates guarantees that the resampled tokens follow the exact distribution as direct sampling from the target model. In other words, our RAPID can guarantee the lower-bound performance at the level of the target LLM as long as $\eta$ is properly set, while the retrieval quality only affects the performance “gains” and will not introduce too many drops.

4. Missed Reference.

A4: Thanks for pointing it out. REST[1] proposed directly selecting possible continuation by retrieving from a built datastore rather than generating with a draft LLM. This work should be conceptually related to our work, but the method is distinct and orthogonal. We will discuss the mentioned works (and further) in a revised version.

[1] REST: Retrieval-Based Speculative Decoding (He et al., NAACL 2024)

We hope the responses above can address your concerns and contribute to a reconsideration of review score. We also appreciate your careful review to point out missed references, which will be fixed in a revised version. Looking forward to discussing more with you.

Best,

Authors

Dear Reviewer SMJY,

We sincerely appreciate your thorough review of our paper. Your constructive feedback will help us strengthen this work. Below are our responses to your concerns:

1. Missed References.

A1: Thanks for pointing it out. REST[1] proposed selecting possible continuation by retrieving from a built datastore rather than generating with a draft LLM. TRIFORCE[2] introduced KV Cache compression for draft LLM based on chunk-wise similarity, which we have discussed in the related work and compared with in the response to Reviewer KbZ7. Ouroboros[3] sought to produce longer and more acceptable candidates from draft LLM per step based on draft phrases. Speculative RAG [4] did not utilize speculative decoding, but proposed a parallel draft-then-verify mechanism to improve RAG quality. These works may be conceptually related to our work, but our method is quite distinct and orthogonal to them. We will discuss the mentioned works (and further) in a revised version.

[1] REST

[2] TRIFORCE

[3] Ouroboros

[4] Speculative RAG

2. Overhead of each component.

A2: Unlike regular RAG pipeline, which builds indexes for a large external corpus (hundreds of millions of documents), we only index/retrieve the chunks for the input long context (<128K) on-the-fly during inference. Therefore, the RAG component latency in our method will become marginal compared to the inference latency over long context. We now list the latency of each component of our RAPID on LongBench v2 (Long, CoT) with LLaMA-3.1-8B below:

3. Potential Bias of the RAG Drafter.

A3: In section 4.5 of our paper, we have discussed the robustness of our RAPID to the retrieval quality. The results indicate that our RAPID can maintain a stable (or better) performance even when the retrieval text is irrelevant to the long context. This demonstrates the robustness of RAPID to retrieval quality, which is because the long-context target model in target distribution (Eq. 6) will preserve the long-context ability to verify candidates effectively, while RAG drafter provides benefits rather than reconstruction.

4. Explanation regarding upward speculation (using larger LLMs as drafter).

A4: For the rationality, as the RAG based on compressed context introduces bias, even the superior LLMs may still lose some crucial information. As shown in Fig 2 of our paper, LLaMA-3.1-70B (RAG) on LongBench v2 achieves 23.66 points gain with another 13.72 points drop compared to LLaMA-3.1-8B (LC). Despite the final gains of LLaMA-3.1-70B (RAG) over LLaMA-3.1-8B (LC) approach 10 points, there is still a large proportion of samples that cannot be solved by the superior LLaMA-3.1-70B (RAG). However, our RAPID can integrate benefits from both long-context target LLM and stronger RAG drafter, achieving performance improvements by incorporating the gains from drafter with minimum extra drops.

For the efficiency, larger RAG drafter will introduce more latency and require extra GPUs to serve. However, our RAPID allows LLaMA-3.1-8B (LC) with LLaMA-3.1-70B drafter to achieve a comparable generation throughput to naive LLaMA-3.1-8B (LC). Empirically, the upward speculation of RAPID can serve as a turbo mode, which maximizes performance with comparable generation speed but consumes more resources. For a low-resource scenario, the self-speculation of RAPID can work well with both improved performance and efficiency and without additional resource requirments.

5. Use models from another family

A5: Thanks for the insightful question. We now conduct a pilot experiment that uses LLaMA-3.1-8B-Instruct (LC) target LLM and Qwen2.5-7B-Instruct RAG drafter for RAPID. We convert the draft logit to target logit space and cut off the probability mass which is mismatched. The implementation is inspired by https://github.com/huggingface/transformers/blob/786d9c5ed920a099573ea7b6dbf265f1aeb32fc0/src/transformers/generation/candidate_generator.py#L783

Surprisingly, using Qwen2.5-7B as RAG drafter can further improve the performance (though lower speedup due to logit alignment overhead). The gains may be due to the better short-context abilities of Qwen2.5-7B compared to LLaMA-3.1-8B upon RAG context. This indicates the essentiality of LLM capability in RAPID and opens more potential exploration of LLM combinations. Thanks for the constructive discussion.

We hope the responses above can address your concerns and contribute to a reconsideration of review score.

Best,

Authors

Dear Reviewer KbZ7,

We sincerely appreciate your thorough review of our paper. Your constructive feedback will help us strengthen this work. Below are our responses to your concerns:

1. Comparison with LC models like Qwen-1M

We've evaluated RAPID on Qwen2.5-7B-1M (released post-submission). Results show RAPID on Qwen2.5-7B-Instruct is comparable to Qwen2.5-7B-1M (35.4 vs 35.6), while RAPID on Qwen2.5-7B-1M further improves both efficiency and performance:

2. Comparison with Triforce

Triforce wasn't included because it can't be directly applied to new LLMs with GQA. We've now conducted comparisons on LWM-Text-Chat-128K (based on LLaMA2-7B), setting retrieval budget at 4096, chunk size at 8, and draft cache budget at 256 for Triforce:

While Triforce achieves efficiency gains, RAPID demonstrates more significant speedup and performance improvements. Triforce recalls information based on chunk-wise attention scores, but higher attention scores don't necessarily indicate greater semantic relevance (e.g., initial tokens often attract high attention as "sinks" despite lacking semantic importance [1]). Our RAD drafter better recalls semantically relevant information, hence achieving higher acceptance rate and speedup for more challenging tasks.

[1] Efficient Streaming Language Models with Attention Sinks

3. Additional Results for MagicDec and SD

We haven’t included results of MagicDec and SD on more models since they achieve quite similar efficiency gains (significantly below our method) across models, while they aren’t designed for performance improvement. We now include more results below.

4. Comparison with Sparse Attention (MInference)

Per your suggestion, we've compared with MInference on LLaMA-3.1-8B. Results show MInference achieves impressive prefill speedup, while RAPID demonstrates significant performance and decoding throughput advantages:

We believe sparse attention is orthogonal to our work and combining it with RAPID holds potential for future research.

5. Retrieval Hyperparameter Settings

We haven't extensively tuned hyperparameters except for $\eta$, as our goal is to propose a generally effective method rather than overfitting benchmarks. The RAG hyperparameters only affect retrieval quality, and RAPID has demonstrated such robustness (Section 4.5). To clarify our selection criteria:

• Cosine similarity threshold (0.3) was selected from {0.1, 0.2, 0.3, 0.4, 0.5} for RAG on LLaMA-3.1-8B, with overall scores of {28.8, 29.2, 29.2, 29.0, 29.0}.
• Compression ratio of 24 (120K/5K) was chosen as retrieval length beyond 5K showed no significant benefits.

For draft tokens per step, we set 10 without tuning. Our ablation study below shows RAPID maintains stable performance with <15 draft tokens, though excessive tokens may reduce performance gains despite increasing throughput.

6. Robustness Analysis Intuition

The robustness analysis in Section 4.5 was conducted on LongBench v2 (Long, CoT) subset, which involves generating reasoning paths before providing answers. Given some preliminary chains, we believe a strong drafter is capable of generating continual CoTs with higher quality to introduce the performance gains. In our LLaMA-3.1-8B target / 70B draft analysis, acceptance rates significantly increased after generating 32 tokens despite irrelevant retrieval context, supporting this intuition.

7. Generation Quality Analysis

Our generation quality analysis (Section 4.4) is a pilot experiment demonstrating RAPID's potential effectiveness in real-world applications. We acknowledge that real-world long-context conversations typically contain similar themes/topics, and hope for robust benchmarks with better evaluation metrics beyond LLM-as-judge in the future. (Human evaluation is both expensive and logistically complex at scale.)

We hope the responses above can address your concerns and contribute to a reconsideration of review score. We also appreciate your careful review to point out some typos/missed references, which will be fixed in a revised version.

Best,

Authors

Dear Reviewer bhTe,

We sincerely appreciate your thorough review of our paper. Your constructive feedback will help us strengthen this work. Below are our responses to your concerns:

1. Possible explanation for the sudden drop in Figure 1.

A1: This is an observed issue in RAG that performance may drop at certain lengths when more retrieval chunks are included [1]. The explanation from [1] is that the presence of certain hard negatives can mislead the LLMs and hinder their ability to generate accurate answers (even if relevant information is included). This issue cannot be directly addressed by using a stronger retrieval model. We believe the retrieval chunks between 16k and 64k may contain many hard negatives. Note that the drop in Fig 1 of our paper is not as “sudden” as displayed, as the x-axis is log-scaled.

[1]Long-Context LLMs Meet RAG: Overcoming Challenges for Long Inputs in RAG

2. FLOPs comparison

A2: Thanks for your suggestion. We now list the FLOPs of our RAPID and baselines per step (generating $\gamma$ tokens) below. As our RAPID consistently demonstrates higher acceptance rates than naive speculative decoding (SD) (shown in Fig. 3 of our paper), the FLOPs order will be Long Context >> SD > RAPID > RAG Drafter.

Target Model:

• $T$ : Number of parameters in the target model.
• $L$: Long context length.

Draft Model:

• $D$: Number of parameters in the draft model
• $L^{R}$: The retrieval length of draft LLM input.

Speculation:

• $\gamma$: the generation length of draft model per step.
• $\beta^{\text{SD}}$: the expectation of speculative decoding acceptance rate.
• $\beta^{\text{RAPID}}$: the expectation of our RAPID acceptance rate.

3. The prefill gets delayed compared to the baseline

A3: Yes, the speculative decoding mechanism does introduce extra prefill latency for the draft model. In our settings, however, the extra latency will be marginal (e.g, 26.37s vs 25.89s) as the draft input length is far shorter than the target one. Moreover, the two prefill stages are possible to be further optimized to overlap at the infrastructure level, which has the potential to totally erase the extra latency.

4. The generation throughput benefits should be clearly highlighted.

A4: We primarily list the generation throughput speedup in Table 1 and discuss it in Section 4.1. The efficiency benefits are also concluded in the abstract and introduction sections. We apologize for any confusion and are glad to strengthening the presentation of these efficiency benefits in revised version.

5. The prefill delay overhead should be discussed associated to each component.

A5: The prefill delay overhead in Table 1 of our paper is only related to the prefill of the draft model upon the retrieval input. In addition, the RAG pipeline will introduce avg 1.43s latency per data sample. We will state the latency more clearly in a revised version.

6. The decode throughput: when it starts to become beneficial compared to the baseline SD, how it is a function of seq length.

A6: In section 4.3 of our paper, we have discussed the impact of context and retrieval length for both performance and efficiency. The conclusions are (1) Our RAPID consistently improve the performance when the seq length > 8K. (When seq len < 8K, it is not really a long context and doesn’t need to apply RAPID) (2) Our RAPID starts to improve efficiency when the seq length is beyond 32k (with retrieval length < 16k). The longer the seq length, the more significant the speedup.

7. The memory and compute overhead to support the inference-time knowledge transfer.

A7: As all required variables in Eq. (6) are necessary to compute for naive SD, our RAPID would not introduce any extra memory overhead for the inference-time knowledge. While the computation is just a tensor add operation, the extra latency is also very marginal and can be ignored.

We hope the responses above can address your concerns and contribute to a reconsideration of review score. Looking forward to discussing more with you.

Best,

Authors