APE: Faster and Longer Context-Augmented Generation via Adaptive Parallel Encoding

Thank you so much for the insightful and valuable comments! They are very helpful for further improving the clarity and quality of our paper. We'll revise our manuscript in the next version to address all of your concerns.

Q1: The methodology might be too similar, with fixed step sparse attention with sink token. When we draw the attention mask of APE, the attention mask itself is extremely similar to the step attention mechanism. However, the difference between APE and APE is that APE reuses the RoPE embeddings rather than just extends them. However, treating the RoPE index is the same as streaming with fixed step sparse attention. Therefore, I think the scientific contribution of methodology is limited.

A1: Our APE is designed specifically for context-augmented generation, where multiple contexts are stored in an external database. The attention mask in APE has two distinctive characteristics that set it apart. First, in the context processing phase, each context can be encoded independently without attending to other contexts. This design enables pre-caching of context KV cache, improving computational efficiency. Second, when inputting the query, it can attend to all contexts, ensuring access to information across the entire context length.

These architectural choices fundamentally differentiate APE from previous methods that rely on fixed-step sparse attention with sink tokens. Such earlier approaches face two critical limitations: they require computationally expensive real-time context re-encoding, and their restricted attention masks prevent queries from effectively accessing the complete context, resulting in suboptimal performance.

The primary novelty of APE's attention mask lies in its ability to enable a new setting for faster and longer context-augmented generation, while maintaining a good performance. APE represents the first method to successfully achieve these in context-augmented generation, marking a significant advancement in the field.

Q2: Lack of comparison with training-free context extension method (Self-Extend: https://github.com/datamllab/LongLM).

A2: First, we need to clarify that APE is primarily designed for context-augmented generation, while the long-context serves as a by-product. Here, we added more baselines for experimental results on the LongBench dataset [1], where we evaluated 11 subsets in total and compared APE with Self-Extend [2]. The results in Table R1 show that APE achieves a 6.61% and 2.15%improvement over the base model and Self-Extend, respectively. Therefore, APE can serve as a SOTA method for length extension.

Table R1: Comparison of APE and Self-Extend on the LongBench dataset.

Q3: Lack of comparing with pre-trained long context LLM. This might not be a big problem, but we need to know what performance is the upper limit if the LLM is already trained in a long context. The LLMs used for experiments are all short context (and maybe a little bit old at this point) models. I am concerned that using a long-context model such as Qwen2 or Llama3.1 that supports 128k tokens with a large-scale GPU cluster might lead to better performance than APE.

A3: First, we need to clarify that APE is primarily designed for context-augmented generation, while the long-context serves as a by-product. Here, we compared the performance of APE and sequential encoding using the LLaMA-3.1-8B-Instruct model as the pre-trained long context LLM on the LongBench dataset. Our evaluation in Table R2 revealed that APE exhibited only a marginal performance decrease of 1.67% compared to Sequential Encoding. This minimal difference does not diminish APE's advantages, as achieving the slight performance gain through sequential encoding would require substantial computational overhead during prefilling. Additionally, starting from a pre-trained long context LLM, APE can always further extend to much longer contexts and lead to better performance.

Q4: Lack of comparison with techniques to speed up the prefill stage.

A4: Thank you for your suggestion. We additionally add MInference [3] as another baseline for prefilling acceleration. As shown in Table R3, APE achieves up to 28x acceleration in prefilling compared to Sequential Encoding with MInference for 128K-token contexts. Furthermore, APE's prefilling overhead scales linearly, consuming less than 10% of the total inference time, while baseline methods require over 50% as context length increases. APE also demonstrates superior versatility across various settings, whereas MInference introduces additional overhead and slower inference speeds when processing short contexts and large batches.

Table R3: Latency on H100 GPU: [prefilling / total time (ms)]. We compared different strategies with various context lengths and batch sizes.

Thank you for insightful and extensively verified rebuttal. The response clearly resolve my all concerns. I think I am misunderstand the concept of the whole paper.

However, I think the author's claim "as existing RAG benchmarks do not fully support 128K context length" is not true, because the LOFT benchmark already supports RAG benchmark up to 1M context document. I think LOFT benchmark is the most important for APE. The LOFT is perfect benchmark to test APE, because (1) the context is concat of thousands of short documents, which is APE targeted. (2) The prompt is shared across queries, therefore partial prefix caching that emphasized by authors. (3) The context length is extremely long (1M) therefore the prefill stage is extremely important to be speedup.

I wonder we can compare the performance in LOFT (but I understand we do not have time now...). As far as, I played with that dataset, 128k evaluation on RAG-NQ dataset only took 10 to 20 minutes if the prompt is cached properly. I hope we can see updated result if possible during rebuttal, or final revision.

I check the provided code, and it is little bit disappointing the implementation is not very efficient. I was expected the codebase is based on PagedAttention framework with prefix caching features such as SGLang. I think the efficiency of implementation is extremely important as the paper proposing efficient mechanism. I want to encourage the author can contribute the open-sources such as vLLM and SGlang with APE in future time after the submission.

After re-considering main manuscript with updated rebuttal, I want to raise my score from 3 (reject) to 6 (borderline accept). Again, thank you for extensively verified rebuttal, and pay respect to authors for trying the best until last hours.

Q7: What do you mean that query and generation lengths were fixed at 256 tokens in line 464? Is that mean, the chunk size of APE is 256? As far as I see the latency, I think 256 means the chunk size, which is NOT used in performance evaluation. Do you have any performance evaluation with 256 chunk size? If my understanding is correct, then the reported latency is might be significantly different with real-world scenario.

A7: Here 256 doesn't mean the chunk size. Instead, it means that our query (w/o context) should be no more than 256 tokens, while the model will generate at most 256 new tokens. With the additional context length from 4K to 512K tokens, this setting simulates the real-world scenarios with long context, short query, and short response, as shown in Table 1 of the original paper. For the context part, it may be a very long context, concatenated by multiple chunks. However, the size of each chunk should approach the model's maximum sequence length (e.g., 8K tokens for the LLaMA-3-8B-Instruct model).

References:

[1] Bai, Yushi, et al. "Longbench: A bilingual, multitask benchmark for long context understanding." arXiv preprint arXiv:2308.14508 (2023).

[2] Jin, Hongye, et al. "Llm maybe longlm: Self-extend llm context window without tuning." arXiv preprint arXiv:2401.01325 (2024).

[3] Jiang, Huiqiang, et al. "Minference 1.0: Accelerating pre-filling for long-context llms via dynamic sparse attention." arXiv preprint arXiv:2407.02490 (2024).

[4] Hsieh, Cheng-Ping, et al. "RULER: What's the Real Context Size of Your Long-Context Language Models?." arXiv preprint arXiv:2404.06654 (2024).

[5] Greg Kamradt. Needle in a haystack - pressure testing llms. 2023.

[6] Tang, Jiaming, et al. "Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference." arXiv preprint arXiv:2406.10774 (2024).

[7] Sun, Hanshi, et al. "ShadowKV: KV Cache in Shadows for High-Throughput Long-Context LLM Inference." arXiv preprint arXiv:2410.21465 (2024).

Q5: In efficiency analysis, the claim 976x faster might be a little bit overclaiming because there is no performance (accuracy) evaluation on 512k. Can you add relevant benchmarks? (InfiniteBench: https://github.com/OpenBMB/InfiniteBench/, LOFT: https://github.com/google-deepmind/loft, RULER: https://github.com/NVIDIA/RULER)

A5: First, we need to clarify that APE is primarily designed for context-augmented generation. However, as existing RAG benchmarks do not fully support 128K context lengths, we demonstrated APE's effectiveness by evaluating it on the RULER [4] and Needle In A Haystack [5] datasets with 128K context lengths, where the original context is divided into 8 chunks of 16K positions. As shown in Table R4, APE demonstrates only a 5% performance decrease compared to sequential encoding. This decline is primarily observed in tasks requiring inter-context dependencies, specifically in VT and FWE subtasks. Compared to other baselines, APE even outperforms Quest [6] and approaches ShadowKV [7] that uses 128K positions with dynamic sparse attention. Therefore, APE should work for longer contexts, especially context-augmented generation without inter-context dependencies. These results, combined with the latency analysis, show the effectiveness and efficiency of APE for extremely long contexts.

Table R4: Experiment results of APE on the RULER dataset.

In Table R5, we further validate that APE can always retrieve the needle from the 128K context on the Needle In A Haystack dataset.

Table R5: Experiment results of APE on the Needle In A Haystack dataset.

Additionally, in the updated efficiency analysis Table R3, we only evaluate from 4K to 128K, which matches the maximum length in the above benchmarks. Therefore, the current performance and efficiency results should be consistent.

Q6: Table 5: what is the average context length of APE and real-world, end-to-end latency, including retrieval latency? As far as I understand the CRAG paper, the context length is quite limited with a smaller size (2k~4k) due to the speed of retrieval.

A6: In our implementation of RAG on the CRAG benchmark, we employ sequential encoding to retrieve the top 10 chunks, each approximately 200 tokens in length. In contrast, when using APE, we retrieve the same number of chunks (top 10), but with significantly larger chunks of approximately 2,000 tokens each. Therefore, this longer context led to the improved performance in Table 5. Notably, despite the larger chunk size, APE demonstrates better latency as shown in Table R6, with improvements in both retrieval and prefilling time. Therefore, the retrieval time is not a bottleneck on this benchmark. This efficiency gain occurs because the increased chunk length reduces the size of the database that needs to be processed, thereby accelerating retrieval.

Table R6: End-to-end latency of APE and sequential encoding on the CRAG benchmark.

While the speedup was not significant on CRAG due to its short context length, we conducted additional measurements using the Ruler benchmark to demonstrate acceleration in long-context scenarios. As shown in Table R7, with a 128K context length, prefilling becomes a significant bottleneck for sequential encoding. In this setting, APE achieved a 15 times speedup in inference time.

Table R7: End-to-end latency of APE and sequential encoding on the Ruler benchmark.

Based on your suggestions, we further evaluated our approach on four text retrieval tasks from the LOFT benchmark using LLaMA-3.1-8B-Instruct as our base model, with performance measured by the F1 score between ground truth and model predictions. To simulate realistic scenarios where RAG systems may retrieve suboptimal content, we randomly sampled passages from the corpus as context, enabling analysis of how performance scales with additional context integration (Otherwise, we only need a short-context LLM to achieve a good performance if the retriever is very accurate, which is not the setting that we target on). For the 128k version of LOFT, we compared APE against Sequential Encoding using varying context lengths in Table R8.

Table R8. Comparison of APE and Sequential Encoding on four text retrieval tasks from the 128k version of LOFT benchmark using the LLaMA-3.1-8B-Instruct model.

Our analysis yielded several key findings:

• APE demonstrates superior performance on average compared to all baselines, validating APE’s scalability in processing hundreds/thousands of retrieved texts in parallel.

• APE achieves performance comparable to the 128K Sequential Encoding baseline while utilizing the same context. Although transferring from sequential encoding to APE would typically be expected to cause performance degradation, we hypothesize that APE recovers this by effectively shortening the distance between the query and all examples. This mechanism helps address the “lost in the middle” problem common in long-context LLMs.

• APE demonstrates remarkable computational efficiency, with its 128K variant requiring inference times comparable to Sequential Encoding at much shorter context lengths (between 16K and 32K), while achieving approximately 10 times faster processing compared to Sequential Encoding at equivalent context length.

We also compared APE with Sequential Encoding using the LLaMA-3.1-8B-Instruct model on the 1m version of LOFT, as shown in Table R9. While APE improves 5.26% on average using 8$\times$ more retrieved texts, it only requires 1/3 inference time by pre-caching these texts.

Table R9. Comparison of APE and Sequential Encoding on four text retrieval tasks from the 1m version of LOFT benchmark using the LLaMA-3.1-8B-Instruct model.

Thank you for your valuable comments and suggestions. We will carefully revise our paper based on your comments. Our responses to your questions are detailed below. We would greatly appreciate your input on whether our revisions address your concern

Q1: The performance evaluation is primarily focused on 8k context length, which feels insufficient given that many open-source LLMs now support context lengths of 128k or more. This restricted evaluation scope makes it difficult to assess the method's scalability to longer contexts.

A1: First, we need to clarify that APE is primarily designed for context-augmented generation. However, as existing RAG benchmarks do not fully support 128K context lengths, we demonstrated APE's effectiveness by evaluating it on the RULER [1] and Needle In A Haystack [2] datasets with 128K context lengths, where the original context is divided into 8 chunks of 16K positions. We also include Quest [3] and ShadowKV [4] as two baselines using dynamic sparse attention and 128K positions. As shown in Table R2, APE demonstrates only a 5% performance decrease compared to Sequential Encoding. This decline is primarily observed in tasks requiring inter-context dependencies, specifically in VT and FWE subtasks. Compared to other baselines, APE even outperforms Quest and approaches ShadowKV despite using only 1/8 positions. Therefore, APE should work for longer contexts, especially context-augmented generation without inter-context dependencies.

Table R1: Experiment results of APE on the RULER dataset.

In Table R2, we further validate that APE can always retrieve the needle from the 128K context on the Needle In A Haystack dataset.

Table R2: Experiment results of APE on the Needle In A Haystack dataset.

Q2: Compared to sequential encoding, APE introduces a non-negligible performance degradation at the 8k length scale, raising concerns about its effectiveness at longer context lengths.

A2: We admit that APE will lead to performance degradation compared to sequential encoding. However, this minimal difference does not diminish APE's advantages, as achieving the slight performance gain through sequential encoding would require substantial computational overhead during prefilling. The results in Table R1 and R2 further validates this performance degradation will not increase when scaling to longer context lengths.

Moreover, the reuse of positions enable APE to leverage more contexts, thereby outperforming sequential encoding in practical applications as shown in Section 6.3.

Q3: How does the computational complexity of APE scale with increasing context length compared to sequential encoding?

A3: We add analysis for the computational complexity of APE and Sequential Encoding. Following Section 2.1 and let the context-augmented input to be

$\mathcal{S} = {\underbrace{s_{C_0, 0}, \ldots, s_{C_0, l_0}}\_{\text{Context}}, \underbrace{s_{C_1, 0}, \ldots, s_{C_1, l_1}}\_{\text{Context 1}}, \ldots, \underbrace{s_{C_{N-1}, 0}, \ldots, s_{C_{N-1}, l_{N-1}}}\_{\text{Context N-1}}, \underbrace{s_0, \ldots, s_l}\_{\text{Query}}},$ where $l_0 + \ldots + l_{N-1} \gg l$, we analyze the prefilling time of Sequential Encoding and APE.

• Sequential Encoding: $O((l_0+...+l_{N-1}+l)^2)$

• APE with pre-caching: $O((l_0+...+l_{N-1}+l)\cdot l)$

• APE without pre-caching: $O(\max(l_0^2, ..., l_{N-1}^2)+((l_0+...+l_{N-1}+l)\cdot l$

Therefore, APE achieves linear complexity with respect to total context length when using pre-caching, resulting in significantly improved efficiency. Even without pre-caching, APE maintains an advantage by allowing independent prefilling of each context. In this case, its complexity scales quadratically with the maximum context length, which is more efficient than Sequential Encoding's quadratic scaling with total context length.

Dear Reviewer onHd,

Based on the suggestions from Reviewer kKET, we further evaluated our approach on four text retrieval tasks from the LOFT benchmark to represent long-context RAG scenarios, which further validate that APE's scalability in processing hundreds/thousands of retrieved texts with the total length of 1M in parallel. We hope these results, combined to the results on the RULER and Needle-In-A-Haystack datasets, can comprehensively show the effectiveness of APE in extreme long-context scenarios.

Here, we use LLaMA-3.1-8B-Instruct as our base model, with performance measured by the F1 score between ground truth and model predictions. To simulate realistic scenarios where RAG systems may retrieve suboptimal content, we randomly sampled passages from the corpus as context, enabling analysis of how performance scales with additional context integration (Otherwise, we only need a short-context LLM to achieve a good performance if the retriever is very accurate, which is not the setting that we target on). For the 128k version of LOFT, we compared APE against Sequential Encoding using varying context lengths in Table R9.

Table R9. Comparison of APE and Sequential Encoding on four text retrieval tasks from the 128k version of LOFT benchmark using the LLaMA-3.1-8B-Instruct model.

Our analysis yielded several key findings:

• APE demonstrates superior performance on average compared to all baselines, validating APE’s scalability in processing hundreds/thousands of retrieved texts in parallel.

• APE achieves performance comparable to the 128K Sequential Encoding baseline while utilizing the same context. Although transferring from sequential encoding to APE would typically be expected to cause performance degradation, we hypothesize that APE recovers this by effectively shortening the distance between the query and all examples. This mechanism helps address the “lost in the middle” problem common in long-context LLMs.

• APE demonstrates remarkable computational efficiency, with its 128K variant requiring inference times comparable to Sequential Encoding at much shorter context lengths (between 16K and 32K), while achieving approximately 10 times faster processing compared to Sequential Encoding at equivalent context length.

We also compared APE with Sequential Encoding using the LLaMA-3.1-8B-Instruct model on the 1m version of LOFT, as shown in Table R10. While APE improves 5.26% on average using 8$\times$ more retrieved texts, it only requires 1/3 inference time by pre-caching these texts.

Table R10. Comparison of APE and Sequential Encoding on four text retrieval tasks from the 1m version of LOFT benchmark using the LLaMA-3.1-8B-Instruct model.

In Table R3, we further validate that APE can always retrieve the needle from the 128K context on the Needle In A Haystack benchmark.

Table R3: Experiment results of APE on the Needle In A Haystack benchmark.

Q5: Are the tasks where APE fails.

A5: APE fails in scenarios that requires the dependency among contexts, ,particularly when long contexts need to be split into multiple chunks. This limitation is especially evident in tasks requiring long-range dependency. Our experimental results from the VT and FWE subsets in Table R2 demonstrate significant performance degradation in such scenarios.

Q6: Does APE still hold when used for Humaneval or RepoBench?

A6: Yes, APE still hold for several code generation tasks as showcased in Table R4. The results on LCC and RepoBench-P show that APE can outperforms other long-context methods on code generation.

Table R4: Comparison of APE and other methods on long-context code generation.

Q7: What would be the impact of finetuning with examples that had APE applied to then. Would it stop degradation?

A7: The model's performance should be improved after fine-tuning for both sequential encoding and APE, as the model encodes the knowledge in examples into its parameters. However, this will lead to significant computational overhead during inference if we fine-tune the model on-the-fly using the retrieved examples, making it not practical.

References:

[1] Dao, Tri, et al. "Flashattention: Fast and memory-efficient exact attention with io-awareness." Advances in Neural Information Processing Systems 35 (2022): 16344-16359.

[2] Hsieh, Cheng-Ping, et al. "RULER: What's the Real Context Size of Your Long-Context Language Models?." arXiv preprint arXiv:2404.06654 (2024).

[3] Greg Kamradt. Needle in a haystack - pressure testing llms. 2023.

[4] Tang, Jiaming, et al. "Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference." arXiv preprint arXiv:2406.10774 (2024).

[5] Sun, Hanshi, et al. "ShadowKV: KV Cache in Shadows for High-Throughput Long-Context LLM Inference." arXiv preprint arXiv:2410.21465 (2024).

Thank you for reviewing our paper and for your valuable feedback. Below, we address your concerns point by point and we will revise our paper according to your suggestions. We would appreciate it if you could let us know whether your concerns are addressed by our response.

Q1: Paper does not provide implementation of the method -- this can make it easier for reproducibility. Would it be possible to provide that?

A1: Our PyTorch implementation of APE is available as a proof of concept at https://anonymous.4open.science/r/APE_Rebuttal-93EB. We are currently developing an efficient implementation based on Flash-attention [1].

Q2: The paper fixes a windows size but shows no empirical or theoretical analysis of choosing different sizes of this window: either from an efficiency or quality perspective. It would be great to have empirical analysis. & Can you provide an analysis of the impact of changing the "window" size you use for parallel decoding.

A2: APE supports dynamic window sizes that adapt to different contexts. The maximum window size is determined by subtracting the combined length of the prefix, query, and maximum generation length from the model's maximum length. We set the window size to its maximum to support longer contexts.

Additionally, we conducted experiments to evaluate the impact of window size on model performance. Using the same long-context inputs, we divided the content into chunks with varying window sizes. The results in Table R1 demonstrate that larger window sizes lead to improved accuracies.

Table R1: Ablation study of window size in APE on the LongBench dataset.

Q3: No complexity analysis written out. This could be particularly relevant if we want to vary window above.

A3: Thank you for your suggestion. We add analysis for the computational complexity of APE and Sequential Encoding. Following Section 2.1 and let the context-augmented input to be

$\mathcal{S} = {\underbrace{s_{C_0, 0}, \ldots, s_{C_0, l_0}}\_{\text{Context}}, \underbrace{s_{C_1, 0}, \ldots, s_{C_1, l_1}}\_{\text{Context 1}}, \ldots, \underbrace{s_{C_{N-1}, 0}, \ldots, s_{C_{N-1}, l_{N-1}}}\_{\text{Context N-1}}, \underbrace{s_0, \ldots, s_l}\_{\text{Query}}},$ where $l_0 + \ldots + l_{N-1} \gg l$, we analyze the prefilling time of Sequential Encoding and APE.

• Sequential Encoding: $O((l_0+...+l_{N-1}+l)^2)$

• APE with pre-caching: $O((l_0+...+l_{N-1}+l)\cdot l)$

• APE without pre-caching: $O(\max(l_0^2, ..., l_{N-1}^2)+((l_0+...+l_{N-1}+l)\cdot l$

Therefore, APE achieves linear complexity with respect to total context length when using pre-caching, resulting in significantly improved efficiency. Even without pre-caching, APE maintains an advantage by allowing independent prefilling of each context. In this case, its complexity scales quadratically with the maximum context length, which is more efficient than Sequential Encoding's quadratic scaling with total context length.

Q4: Would it be possible to include RULER as an evaluation for long-context task. It has become more standard for LC. Showing the impact of APE on needle in a haystack would be interesting as well.

A4: First, we need to clarify that APE is primarily designed for context-augmented generation despite its potential in processing long single contexts as featured in RULER [2] and Needle In A Haystack [3] benchmarks. Here, we evaluated the LLaMA-3.1-8B-Instruct model on these benchmarks using 128K context lengths, with the context divided into 8 chunks of 16K positions each. We also include Quest [4] and ShadowKV [5] as two baselines using dynamic sparse attention and 128K positions. As shown in Table R2, APE demonstrates only a 5% performance decrease compared to Sequential Encoding. This decline is primarily observed in tasks requiring inter-context dependencies, specifically in VT and FWE subtasks. Compared to other baselines, APE even outperforms Quest and approaches ShadowKV despite using only 1/8 positions. These new results further show APE's superiority in long-context understanding.

Table R2: Experiment results of APE on the RULER benchmark.

Q3: And Table 5 lacks crucial baseline comparisons, making it difficult to assess the relative performance improvements.

A3: Thank you for your suggestions. In Table 5, we evaluate RAG with Sequential Encoding and APE for two different external databases, with additional comparison to LLM-only generation. In these settings, APE consistently leads to performance improvement. Therefore, Table 5 only serves as one real-world RAG application for APE, where we only compares it with Sequential Encoding.

To compare with more baseline, our evaluation in Tables R1-R6, which includes various methods for both RAG and long-context understanding, can provide evidence supporting the superiority of APE.

Q4: And this paper lacks sufficient component analysis and validation.

A4: To further validate the effectiveness of each component in APE, we expanded our ablation studies beyond the three ICL tasks presented in Table 3 of the original paper. Our additional analysis included 8 subsets from LongBench and 5 subsets from ChatRAGBench. The results in Table R7 demonstrate incremental improvements with each component: incorporating shared prefix enhanced performance by 1.75%, adding scaling factor yielded a 0.63% improvement, and introducing attention temperature provided an additional 0.5% gain. Similarly, Table R8 shows consistent improvements from each component, with gains of 2.3%, 0.29%, and 0.79% respectively. These findings further validate the value of each component in APE. Additionally, the minor performance gap between APE and sequential encoding in Table R5, combined with these detailed results, shows that the challenges identified in Section 3.2 in the main paper have been addressed by these three components.

Table R7: Ablation study for each component in APE on 8 subsets from LongBench.

Table R8: Ablation study for each component in APE on 5 subsets from ChatRAGBench.

References:

[1] Bai, Yushi, et al. "Longbench: A bilingual, multitask benchmark for long context understanding." arXiv preprint arXiv:2308.14508 (2023).

[2] Izacard, Gautier, et al. "Unsupervised dense information retrieval with contrastive learning." arXiv preprint arXiv:2112.09118 (2021).

[3] Liu, Zihan, et al. "Chatqa: Surpassing gpt-4 on conversational qa and rag." The Thirty-eighth Annual Conference on Neural Information Processing Systems. 2024.

[4] Li, Zehan, et al. "Towards general text embeddings with multi-stage contrastive learning." arXiv preprint arXiv:2308.03281 (2023).

[5] Pan, Zhuoshi, et al. "Llmlingua-2: Data distillation for efficient and faithful task-agnostic prompt compression." arXiv preprint arXiv:2403.12968 (2024).

[6] Xiao, Guangxuan, et al. "Efficient streaming language models with attention sinks." arXiv preprint arXiv:2309.17453 (2023).

[7] Jin, Hongye, et al. "Llm maybe longlm: Self-extend llm context window without tuning." arXiv preprint arXiv:2401.01325 (2024).

Table R4: RAG experiments for APE and other methods on the LongBench dataset for the Gemma-2-9b-it model.

• RAG on the ChatRAGBench dataset [3] with multiple documents. In this setting, we employed three distinct retrievers: Dragon-multiturn [3], Contriever [2], and GTE-base [4] to retrieve the top-5 relevant contexts from five multi-document subsets of the ChatRAGBench dataset. We then evaluated the performance of Llama3-ChatQA-1.5-8B using these retrieved contexts, comparing sequential encoding against APE. Additionally, we tested APE's performance when using all available contexts. The results, presented in Table R5, demonstrate that shifting from sequential encoding to APE results in a minimal performance degradation of no more than 1%, which we consider acceptable. Notably, this performance gap becomes even smaller with lower-quality retrievers. Furthermore, when directly processing all contexts, APE not only outperforms all sequential encoding baselines but also shows substantial improvements compared to scenarios using low-quality retrievers.

Table R5: RAG experiments for APE and sequential encoding on the ChatRAGBench dataset using different retrievers.

Q2: The evaluation of other task types in Sec 6.2 is inadequate, missing essential baseline comparisons needed for meaningful performance assessment.

A2: We added more baselines for experimental results on the LongBench dataset, where we evaluated 11 subsets in total and compared APE with LLMLingua2 [5], StreamingLLM [6], Long-context Fine-tuning, and Self-Extend [7]. The results in Table R6 show that APE achieves a 6.61% improvement over the base model. Among other length extension methods, only Self-Extend shows performance gains, though it still lags behind APE by 2.15%.

Table R6: Comparison of APE and other length extension methods on the LongBench dataset.

Thank you for your constructive comments and suggestions. We respond to your questions below and would appreciate it if you could let us know if our response addresses your concerns.

Q1: For multi-document QA tasks, the evaluation is conducted on a limited subset of LongBench, with additional testing only on one RAG dataset in Sec 6.3. This narrow scope of evaluation fails to provide comprehensive evidence of the method's effectiveness.

A1: Thank you for your comment. We have added more experiments to showcase the effectiveness of our methods in context-augmented generation.

• RAG on the LongBench dataset [1]. In this setting, we implement a chunking-and-retrieval approach across eight LongBench subsets following the original paper. The method involved preprocessing long contexts into chunks and utilizing Contriever [2] to retrieve the top-k most relevant chunks. We conducted experiments using four base models: LLaMA-3-8B-Instruct, LLaMA-3.1-8B-Instruct, Mistral-7B-Instruct-v0.3, and Gemma-2-9b-it. For sequential encoding, we segmented texts into chunks of approximately 200 words. When using APE, we employed larger chunks of 4,000 words, except for Gemma-2-9b-it, which used 2,000-word chunks. For LLaMA-3.1-8B-Instruct, we limited context length to 8K tokens to simulate a typical RAG setting. We also included comparisons with PCW, a baseline parallel encoding method. As shown in Table R1, R2, R3, R4, APE leads to 5.6% and 4.5% improvement on average compared to Sequential Encoding and PCW. These findings suggest that APE enhances RAG performance by enabling the retrieval of more and longer contexts without compromising inference speed, as pre-cached contexts can be loaded directly.

Table R1: RAG experiments for APE and other methods on the LongBench dataset for the LLaMA-3-8B-Instruct model. Here C200*10 means using Contriever to retrieve the top10 chunks, with each chunk containing approximately 200 words.

Table R2: RAG experiments for APE and other methods on the LongBench dataset for the LLaMA-3.1-8B-Instruct model.

Table R3: RAG experiments for APE and other methods on the LongBench dataset for the Mistral-7B-Instruct-v0.3 model.

Thank you so much for the insightful and valuable comments! They are very helpful for further improving the quality of our paper. We will revise our manuscript in the next version to address all of your concerns.

Q1: The performance seems quite sensitive on specific parameters.

A1: Our APE, despite introducing three hyperparameters (shared prefix, scaling factor, and attention temperature), is not very senstive on these hyperparameters. First, the shared prefix is only applied to tasks lacking existing prefixes (e.g., GSM8K and TriviaQA). Our experimental results in Table R1 demonstrate that introducing this shared prefix increases performance by 7.38% on average, while variations in prefix length only result in a 0.4% difference in performance. These findings indicate that while the presence of a shared prefix is beneficial, APE's performance is not significantly sensitive to the prefix length.

Table R1: Ablation study of different prefix lengths on GSM8K and TriviaQA using the LLaMA-3-8B-Instruct model.

For the scaling factor $s$ and attention temperature $T$, we evaluated various combinations across four QA tasks from the LongBench dataset, as shown in Table R2. While these hyperparameters are relatively sensitive compared to the prefix length, their optimal values fall within a narrow range. In our experiments, both $s$ and $T$ were selected from {0.7, 0.8, 0.9, 1.0}, making hyperparameter searching efficient on the validation set. Notably, despite examples of each task on LongBench containing varying numbers of contexts, APE's hyperparameters demonstrated generalization ability across examples with different context lengths, indicating low sensitivity.

Table R2: Ablation study of the scaling factor s and attention temperature T on four subsets from LongBench using the LLaMA-3-8B-Instruct model.

Q2: Will we conclude anything between the task and the choice of hyperparameters?

A2: In different tasks, we may incorporate different number of retrieved texts from external sources. With this number increase, we observe a consistent trend that both the scaling factor s and attention temperature T decrease, aligning with our observations in Section 3.2.

Q3: The latency improvement in Table 2 would be more interesting if other baselines are included.

A3: Thank you for your suggestion. We additionally add MInference [1] as another baseline for prefilling acceleration. As shown in Table R3, APE achieves up to 28$\times$ acceleration in prefilling compared to Sequential Encoding with MInference for 128K-token contexts. Furthermore, APE's prefilling overhead scales linearly, consuming less than 10% of the total inference time, while baseline methods require over 50% as context length increases. APE also demonstrates superior versatility across various settings, whereas MInference introduces additional overhead and slower inference speeds when processing short contexts and large batches.

Table R3: Latency on H100 GPU: [prefilling / total time (ms)]. We compared different strategies with various context lengths and batch sizes.

References:

[1] Jiang, Huiqiang, et al. "Minference 1.0: Accelerating pre-filling for long-context llms via dynamic sparse attention." arXiv preprint arXiv:2407.02490 (2024).

References:

[1] Bai, Yushi, et al. "Longbench: A bilingual, multitask benchmark for long context understanding." arXiv preprint arXiv:2308.14508 (2023).

[2] Izacard, Gautier, et al. "Unsupervised dense information retrieval with contrastive learning." arXiv preprint arXiv:2112.09118 (2021).

[3] Liu, Zihan, et al. "Chatqa: Surpassing gpt-4 on conversational qa and rag." The Thirty-eighth Annual Conference on Neural Information Processing Systems. 2024.

[4] Li, Zehan, et al. "Towards general text embeddings with multi-stage contrastive learning." arXiv preprint arXiv:2308.03281 (2023).

[5] Hsieh, Cheng-Ping, et al. "RULER: What's the Real Context Size of Your Long-Context Language Models?." arXiv preprint arXiv:2404.06654 (2024).

[6] Greg Kamradt. Needle in a haystack - pressure testing llms. 2023.

[7] Tang, Jiaming, et al. "Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference." arXiv preprint arXiv:2406.10774 (2024).

[8] Sun, Hanshi, et al. "ShadowKV: KV Cache in Shadows for High-Throughput Long-Context LLM Inference." arXiv preprint arXiv:2410.21465 (2024).

[9] Pan, Zhuoshi, et al. "Llmlingua-2: Data distillation for efficient and faithful task-agnostic prompt compression." arXiv preprint arXiv:2403.12968 (2024).

[10] Xiao, Guangxuan, et al. "Efficient streaming language models with attention sinks." arXiv preprint arXiv:2309.17453 (2023).

[11] Jin, Hongye, et al. "Llm maybe longlm: Self-extend llm context window without tuning." arXiv preprint arXiv:2401.01325 (2024).