Accelerating Retrieval-augmented Language Model Serving with Speculation

We thank the reviewer for their review and address their main concerns below.

Q1. Speculation Stride

When we don’t enable the optimal speculation stride scheduler, the stride of speculation is prefixed to 3. However, if the optimal speculation stride scheduler (OS$^3$) is enabled, the speculation stride will be initialized to 1 and dynamically adapted according to the theoretical optimal value derived from the cost model. Speculation stride is indeed a crucial hyperparameter that controls the tradeoff between speculation gain and overhead. We include an ablation Table 5 and a more fine-grained analysis in Appendix A.6 to show how different speculation stride affects the performance. We also show the table below. The values are the average serving latency per request.

We can observe that larger a stride is better for EDR while smaller a stride is better for ADR/SR. OS$^3$ can capture this dynamism and achieve near-optimal performance. The reason why OS$^3$ performs slightly worse than $s=4, 8$ in the case of EDR is because a warm-up phase is required for OS$^3$ to start its adaptation, i.e., we initialize s=1 when we enable OS$^3$. Thus, the warm-up phase can introduce non-optimal strides and, thus, slightly worse performance.

Besides, we have also added additional experiments on our approach against KNN-LM [1] which is another important RaLM workload. The analysis and results are included in Appendix A.3 of the revised paper. The KNN-LM results also include the ablation study on the speculation stride and our detailed analysis. As a retrieval-intensive workload, our approach can achieve up to 7.59x speed up when the optimal speculation stride scheduler is used.

Q2. OPT2 performance

As the different models can generate contents differently and thus have different retrieval patterns, the main reason is that the OPT2 model tends to generate contents that have a retrieval pattern that is easier to speculate (e.g., more documents have been repeatedly retrieved), which results in a higher speculation accuracy and thus higher speed-up ratio.

Q3. ADR performance

As our approach targets accelerating the retrieval part of the RAG workloads, and the reported speed-up is the end-to-end speed-up ratio, the speed-up ratio is intrinsically upper-bounded by the retrieval ratio with respect to the overall latency. Thus, in the case of the approximate dense retriever, the retrieval latency is much faster than the language model generation latency, which results in a limited acceleration space for our approach, as shown in Figure 1. However, we want to demonstrate that our method can consistently achieve speed-ups even in cases where the acceleration space is limited.

Q4. Related Work

We thank the reviewer for pointing out the drawbacks of our related work presentation, and we have modified the related work section accordingly in the revised paper by reorganizing paragraphs and pruning out unnecessary descriptions to focus more on the position and contribution of our work within the literature.

References:

[1] Khandelwal, Urvashi, et al. "Generalization through memorization: Nearest neighbor language models." International Conference on Learning Representations. 2020.

We thank the reviewer for their review and address their main concerns below.

Q1. The practicality of iterative RAG

We agree with the reviewer that one-shot RAG, where retrieval is conducted once at the start, is a major trend for now in most RAG applications (e.g., LlamaIndex, LangChain). However, the generation quality of such applications largely relies on a powerful language model like GPT-3.5. We note that even in such cases, iterative retrieval can outperform one-shot retrieval (e.g., FLARE [2], IRCoT [4]), especially for longer conversations. A more reasonable scenario for iterative retrieval is when you don’t have access to an LLM as powerful as GPT-3.5 but can use smaller surrogates like GPT-2, OPT-1.3B, or LLaMA-2-7B. In such scenarios, single-shot retrieval performs much worse than iterative retrieval, as shown in IC-RaLM [3] and IRCoT [4]. Additionally, KNN-LM [1] is another retrieval-intensive workload that has been shown to improve the performance of the base language model with retrieval augmentation. To demonstrate our approach also works for KNN-LM, we provide additional experiments in Appendix A.3 of our revised paper to verify the effectiveness of RaLMSpec. More specifically, our approach can achieve a speed-up ratio of up to 7.59x compared with the baseline implementation as shown in Figure 6.

Q2. The difference with the Caching mechanism in Bing-Chat-like systems

The local cache in our system design is used for speculative retrieval instead of query matching. In the scenario where different queries might retrieve the same entry in the cache, Bing-Chat-like systems will have a cache miss due to query mismatch, but our design will still return a retrieved entry in the cache. This design is two-sided. On the one hand, it can benefit from a correct speculative retrieval. On the other hand, we need to add an additional verification mechanism to correct false speculations. The interaction between speculation and verification is crucial (e.g., how many speculation steps will be performed before a verification step, which we call the speculation stride). To this end, we further propose the optimal speculation stride scheduler to automatically adjust to the theoretical optimal strides. To sum up, though our system design also leverages the caching mechanism, the cache is used differently (speculative retrieval rather than common-entry matching); thus, additional efforts are required to guarantee correctness.

Q3. Mature search engines

The reported speed-up ratios are all against end-to-end latencies. Our approach is intrinsically targeting reducing the retrieval latency, not the model generation latency. Thus, if the proportion of the retrieval latency is relatively small compared to the overall latency, the speed-up ratio can be limited, as shown in Figure 4. However, we do believe that there are certain trade-offs between retrieval efficiency and model generation efficiency in different setups, where some prefer a slower but more accurate retriever while others might prefer efficiency over accuracy. However, we want to note that our approach can automatically adjust to the optimal configurations in all different scenarios and achieve speed-ups accordingly, which is enabled by prefetching, the optimal speculation stride scheduler, and the asynchronous verification techniques we proposed in the paper. This can be further verified by our original and updated empirical results. Regarding how well the proposed method will work with mature search engines, it depends on the proportion of retrieval latency over the total serving latency. If the retrieval operation is optimized by approximation and parallelization, the acceleration space of our approach can be limited. However, the proposed algorithm can achieve a more significant speed-up for slow retrieval scenarios due to hardware/software limitations.

Q4. Comparison between one-shot and iterative RAG

We thank the reviewer for pointing out the need for this comparison. As our paper mainly focuses on accelerating iterative RAG while preserving the same model output, we refer the reviewer to [1, 2, 3, 4] for more detailed motivation and experimental results of iterative RAG. For one, Table 1 in [1] has shown that KNN-LM can consistently outperform the base language model with retrieval. Figure 5 in [3] also shows that with a more frequent retrieval step, the generation quality (measured by perplexity) of the RAG model can be improved monotonically.

We thank the reviewer for their review and address their main concerns below.

Q1. Ablation Study

We thank the reviewer for acknowledging our contribution. Indeed, the proposed techniques (cache-based speculative retrieval, prefetching, optimal speculation stride scheduler, and asynchronous verification) are not orthogonal to each other. Thus, it is not as trivial as it seems to directly decompose the speed-up ratio to different techniques. However, we conduct experiments for each technique individually as well as all together, as shown in Table 1, where P stands for prefetching, S stands for optimal speculation stride scheduler and A stands for asynchronous verification. If we ignore the PSA results (the one that enabled all techniques), we can see from Table 1 that the optimal speculation stride scheduler contributes most to the speed-up in most cases compared with prefetching and asynchronous verification. This is due to the fact that the speculation stride is a crucial hyperparameter that controls the tradeoff between speculation gain and overhead, and the proposed OS^3 technique can adaptively find a theoretical optimal value. Besides, we have also included an additional Table 4, Figure 7, and more fine-grained analysis in Appendix A.5 of the revised paper to further show the contribution of each component. We have also included the additional ablation table below. The results are averaged serving latency per request.

We observe that prefetching works pretty well when an exact dense or sparse retriever is used because the overhead of doing prefetching is relatively small compared to when an approximate dense retriever is used. For workloads with a retrieval latency similar to the language model generation latency, asynchronous verification can bring up the most significant speed when optimal speculation stride scheduler is enabled but not on itself. To sum up, the optimal speculation stride scheduler is the most important component while the prefetching and asynchronous verification can be beneficial depending on which retriever or language model is being used.

Additional Experiments

To provide more information to the reviewer, we have also added additional experiments on KNN-LM [1] with our approach in Appendix A.3 of the revised paper, which further verifies that our method can achieve latency savings consistently in different scenarios. As a retrieval-intensive workload, our approach can achieve up to 7.59x speed up when the optimal speculation stride scheduler is used.

References:

[1] Khandelwal, Urvashi, et al. "Generalization through memorization: Nearest neighbor language models." International Conference on Learning Representations. 2020.

We thank the reviewer for their review and address their main concerns below.

Q1. Contribution and Novelty

Though the idea of speculative execution is not new, as far as we know, we are the first work that incorporates the speculation idea into accelerating retrieval augmented language model serving. More specifically, we design a local retrieval cache to enable speculative retrieval that is novel and unseen before. In addition to the design for a local retrieval cache, we introduce prefetching, optimal speculation stride scheduler, and asynchronous verification to boost the performance further. While prefetching and asynchronous verification are relatively trivial to think of, the optimal speculation stride scheduler is a cost model-based dynamic speculation scheduler that can automatically adjust the speculation stride to the optimal setup in various scenarios, which is novel and effective.

Besides technical contributions, we also evaluate our approach against three different types of retrievers and three different sizes of LMs (GPT, OPT, and LLaMA) across four QA datasets. The overall results and ablation study demonstrate that our method with all components enabled (prefetching+OS3+AsynchVerify) can consistently achieve the best speed-up ratio. We further include additional experimental results for the KNN-LM[1] workload in Appendix A.3 as well as the results for LLaMA-2-13B in Appendix A.4. For KNN-LM, our approach can achieve up to 7.59x speed-up ratio. For LLaMA-2-13B, our approach can still achieve a speed-up from 1.02x to 1.85x.

With the aforementioned technical and empirical contributions, we hope the reviewer can reconsider the novelty and contribution of our work.

Q2. Speed-up Variations

The speed-up ratio varies because different retrievers, language models, or datasets are used. More specifically, as our approach is intrinsically optimizing the retrieval latency, but we are reporting the end-to-end latency speed-up ratios, different retrievers and different language models might contribute differently (latency-wise) in the end-to-end latency. This can be demonstrated in Figure 4, where each latency bar constitutes two sub-bars, one for retrieval latency and one for language model generation latency. When an approximate retriever is used, the retrieval latency will not dominate the serving latency, thus the end-to-end speed-up ratio is limited. When a dense retriever is used, the retrieval overhead is much higher, therefore, a higher speed-up ratio can be achieved accordingly. Figure 4 and the analysis paragraph in section 5.2 present the figure and written explanation. In order to make this more clear, we have updated the analysis further in the revised paper.

References:

[1] Khandelwal, Urvashi, et al. "Generalization through memorization: Nearest neighbor language models." International Conference on Learning Representations. 2020.