APE: Faster and Longer Context-Augmented Generation via Adaptive Parallel Encoding
Authors: Xinyu Yang
Tianqi Chen
Beidi Chen
Venue: ICLR 2025 Poster (2025)
Keywords: Parallel Encoding; Context-Augmented LLM; Efficient Inference; Length Extrapolation
Primary area: generative models

Abstract: Context-augmented generation (CAG) techniques, including RAG and ICL, require the efficient combination of multiple contexts to generate responses to user queries. Directly inputting these contexts as a sequence introduces a considerable computational burden by re-encoding the combined selection of contexts for every request. To address this, we explore the promising potential of parallel encoding to independently pre-compute and cache each context's KV states. This approach enables the direct loading of cached states during inference while accommodating more contexts through position reuse across contexts. However, due to misalignments in attention distribution, directly applying parallel encoding results in a significant performance drop. To enable effective and efficient CAG, we propose Adaptive Parallel Encoding (**APE**), which brings shared prefix, attention temperature, and scaling factor to align the distribution of parallel encoding with sequential encoding. Results on RAG and ICL tasks demonstrate that APE can preserve 98\% and 93\% sequential encoding performance using the same inputs while outperforming parallel encoding by 3.6\% and 7.9\%, respectively. It also scales to many-shot CAG, effectively encoding hundreds of contexts in parallel. Efficiency evaluation shows that APE can achieve an end-to-end 4.5$\times$ speedup by reducing 28$\times$ prefilling time for a 128K-length context. The code is available at 
https://github.com/Infini-AI-Lab/APE.

Published as a conference paper at ICLR 2025
APE: FASTER AND LONGER CONTEXI-AUGMENTE GENERATION VIA ADAPTIVE PARALLEL ENCODING
Xinyu Yang
Tianqi Chen
Beidi Chen CMU
CMU
CMU, NVIDIA
xinyuagi i @ cmu. edu
tgchen@cmu.com edu
be i di c/
ABSTRACT
Context-augmented generation (CAG) techniques, including RAG and ICI., require the efficient combination of multiple contexts to generate responses to user queries. Directly inputting these contexts as a sequence introduces a considerable compu- tational burden by re-encoding the combined selection of contexts for every request. To address this, we explore the promising potential of parallel encoding to indepen- dently pre-compute and cache each context's KV states. This approach enables the direct loading of cached states during inference while accommodating more contexts through position reuse across contexts. However, due misalignments in attention distribution, directly applying parallel encoding results in significant performance drop. To enable effective and efficient CAG, we propose Adaptive Parallel Encoding (APE), which brings shared prefix, attention temperature, and scaling factor to align the distribution of parallel encoding with sequential encoding. Results on RAG and ICL tasks demonstrate that APE can preserve 98% and 93% sequential encoding performance using the same inputs while outperforming parallel encoding by 3.6% and 7. 7.9%, respectively. It also scales to many-shot CAG, effectively encoding hun- dreds of contexts in parallel. Efficiency evaluation shows that APE can achieve an end-to-end 4.5 speedup by reducing 28x prefilling time for 128K-length context. The code is available at httts::: htth//ggttttbttt com/ Infini AI Lab/APE APE
INTRODUCTION
Recent advances in context-augmented generation (CAG), particularly
generation
(RAG) (Gupta et al., 2024 Gao et al., 2023) and in-context learning (ICL) (Dong et al. [2022] Weil ct al. 2022), have been widely adopted in large language models (LLMs) (Dubey et al. [024 Achiam ct al. 2023). improving their ability to generalize to unseen tasks with contextual information, as demonstrated in Figure (top). These techniques employ sequential encoding process ground LLM inputs with knowledge from external sources: concatenating the retrieved texts into one sequence, and encoding the sequence into key-value (KV) states as the context for subsequent queries. While this new, significantly longer input improves performance, the increased latency in context prefilling becomes bottleneck in tasks that require long inputs but generate short outputs (Bai et al. 2023 Agarwal et al. 2024 Jiang et al.) 2024b). For example, prefilling a 128K context can take 17 seconds, whereas generating 256 tokens requires only seconds. This discrepancy leaves sigmificant room improve the practical efficiency of CAG systems when deployed in real-world applications (Liu) 2022 Chase
( 2022).
Since these contexts are typically predetermined and stored independently in external databases (Za) yarni et al. 2024 Douze et al., 2024). naively/bruteforeely one might accelerate CAG systems by pre- caching them for direct loading during inference. However, for autoregressive LLMs, the KV states inherently context- context-dependent, which means that the KV states for the same text vary based on the preced- ing context. This dependency necessitates caching all possible context permutations, leading to factorial growth in memory requirements as the database size increases. For instance, caching all permutations of just ten 256-token texts for the LLAMA-3-80 model would demand an impractical 22 PB of memory.
To address this issue, parallel encoding (Ratner et al. 2022 Yen et al. 2024 Li et al. 2024 Sun et al., 2024) encodes each context into KV states separately, ensuring that tokens from different contexts can- not attended in the encoding phase. Next, the on-the-fly fly generation is started by prefilling user queries, which can attend to the cached KV states from all contexts without re-encoding, offering two benefits:
Pre-caching Contexts for Fast Inference: Contexts from the external sources can be pre-computed and cached into KV states for direct loading during inference. Additionally, this approach allows for cost-free manipulation of contexts, including insertion, deletion, deletion. replacement, and swapping operations.
Re-using Positions for Long Context: Contexts can be inserted into the same range of positions in an LLM's context window, allowing for more and longer context chunks. This also avoids the "positional bias" in context ordering (Liu al. 2024a), each context is equally "close" to the generated tokens.

---

Published as a conference paper at ICLR 2025
Problem Setup: Context-Augmented Generation
Passage I: Pugitive Lady may refer to: Fugitive Lady (1934 film)
Who was bom first? A:
Marino Girolami.
Additional Contexts C: numerous and lengthy
Query Q
Response R
Sequential Encoding
Latency: 80s
Accuracy: 37.60%
Position 05200
Position 2500-5000 5000
Position 5000-75000
Inference start: compute the KV states of contexts on-the-fly
Parallel Encoding
Latency: 19s
Accuracy: 36.73%
Position 0- 7500: More and longer contexts prepended to the Query
Pre-compute and store the KV states of contexts
Inference start: load pre-oached contexts
Our Approach: Adaptive Parallel Encoding
Latency: 19s
Accuracy: 39.62%
Answer the question.
Position 20- 7500
Shared Prefix
Low Temperature
Scaling Factor
Our adaptive alignments recover the accuracy
Figure 1: Overview of our approach. Context-augmented generation leverages additional contexts to improve LLM response quality to user queries. Sequential encoding prefills selected context chunks as a long sequence during inference, leading to high latency from on-the-fly re-encoding and low accuracy due to context window limitations. Parallel encoding offers an alternative method pre-compute more and longer contexts within the same positional range but results in worse performance. To address these challenges, we propose Adaptive Parallel Encoding (APE) to re-align the attention weight distribution of parallel encoding with sequential encoding via three maining-free steps: shared prefix, scaling factor, and adaptive temperature, leading to fast, long, and accurate CAG systems in real-world applications.
Despite these advantages. parallel encoding leads to significant performance degradation across multi- ple scenarios as shown in Figure 2 with average declines of 4. 9%% (despite using 2-10x more contexts) and 49.0%. respectively. While prior works (Sun al. [2024] Yen et al. 2024) have attempted to correct this with fine-tuning, these methods continue to exhibit reduced accuracy in reasoning tasks. This performance drop can be attributed to the limited generalization ability from simple to complex tasks.
However, our results Figure reveal that pncooing holds promise, as LLMs can still gener- ate reasonable responses due to their inherent alignments with sequential encoding. Based on this obser- vation, we aim to strengthen these alignments while addressing the remaining discrepancies to achieve more accurate parallel encoding. Our insight from Figure Jand Figure is that KV states from indepen- dent contexts can be naturally merged into one sequence due to their similarity in direction and magni- tude, altributed to the presence of an attention sink (Xiao et al. [023)) This initial alignment influences the direction of following KV states, which consistently keep a similar large angle with the sink tokens (see Figure 5.. Based on this observation, we can reduce our challenges addressing residual misalign- ments, which manifest as anomalous distributions at the initial and recent positions within each context.
Motivated by these insights, we propose Adaptive Parallel Encoding (APE), which aligns the distri- bution of attention weights between sequential encoding and parallel encoding, APE enables accurate and fast CAG, as demonstrated in Figure (bottomm. The main contribution of this paper is as follows:
We systematically analyze the distribution properties of attention weights in parallel encoding, focus- ing the direction and magnitude of KV states across various samples and positions. Our analyses identify alignments and residual misalignments between parallel encoding and sequential encoding.
We propose APE to recover the accuracy of parallel encoding with three alignment steps: (i) Prepend a shared prefix to avoid the duplication of abnormal initial token distributions. (ii) Adjust a lower atten- tion temperature to sharpen the distribution, focusing on important tokens. (iti) Apply scaling factor offset the increase in the magnitude of the LogSumExp value of attention scores from the context.
We empirically show that (i) APE keeps 98% and 93% of the sequential encoding performance in RAG and ICL tasks, respectively. (ii) APE outperforms parallel encoding in RAG and ICL, yielding improvements of 3.6% and 7. 7.9%, respectively. (iii) APE effectively scales to handle hundreds of con- texts, approaching or even surpassing sequential encoding in many-shot scenarios. (iv) APE acceler- ates long-context generation, achieving up to a 4.5x speedup through 28 x reduction in prefilling time.
2

---

Published as a conference paper at ICLR 2025
2
BACKGROUND AND RELATED WORK
2.1
GENERATION
This work explores CAG problems using LLMs, where user queries are enhanced with additional con- texts from external databases. CAG involves two scenarios: RAG (Asai et al.) 2024; Gupta et al. 2024
Gao et al.
2023), which focuses directly retrieving relevant information, and ICL (Dong et al. 2022 Wei et al.
2022 which emphasizes further acquiring emergent capabilities from in-context examples.
2.2 PARALLEL ENCODING
Next, we present the formulation of using parallel encoding in LLMs for CAG. Let S represent the input sequence including N contexts CN and one query Q. Formally, this can be denoted as:
S = (SC,,1 15. SCN SC dy SQ.11.....)
(I)
Context
Context
Context N
Query
For simplicity, we can express this as: S= (Sc, ,Sc, p....... per.,,,r Sr.,SSy SQ). Given two models Onc and DDe (which may be identical), a response R generated to input S using parallel encoding through two steps:
Pre-caching Contexts. The initial step involves encoding and caching the KV states for each context in- dependently using O For a given context Sc., we compute the KV states as (Kc,,Vc, Vc..) ) Gme ( (Sc.) and store them for subsequent direct use, where Kc, = I..........Kk.... } and Vc, (fc,,1,...NC...
Generating Response. Next, we input the query and load all relevant KV states to generate response: R = Dec (Sq KK. Vc), where Kc, Ve are subsets of {Kc, Kew w and
}. respectively.
Parallel encoding significantly improves efficiency compared to sequential encoding by reducing (i.e., quadratic) to linear with respect to context length. With pre-caching, the cost becomes ((4 +I. (lQ)) In the absence of pre-caching, the complexity is ((..x(((......3)) ((4 +.. +10) lQ)) which is efficient for multiple contexts of similar length.
complexity from ( +/y +
+l2)²) +
Prior parallel encoding approaches vary in their design of O and θ Parallel Context Windows (PCW) (Ratner et al. 2022) directly employs pre-trained LLMs as both, resulting in significant performance drops. Block-Attention (Sun et al., (024) further fine-tunes the model, successfully recovering performance in RAG tasks. Alternatively, CEPE (Yen et al.) (024) and FocusLLM (Li encoders using encoder-only and decoder-only architectures, respectively. These methods also differ in Dee: CEPE trains additional cross-attention layers for processing contexts. whereas other methods directly input the context into original self-attention layers. While these trainable methods show promising results in RAG tasks, challenges remain regarding their training overheads and generalization abilities to more complex ICL scenarios. Moreover, applying parallel encoding in CAG can be viewed as a kind of memory- nemoor-aaugmetee neural networks (Burtsev et al.) 2020
al. 2024 train new Transformer-basee
De Jong
2021
Févry al. 2020), where external memory is directly stored into KV states.
2.3
ATTENTION MECHINISM
In standard Softmax Attention, we attend the query to all KV states using the following formula:
QKT Vd ) QeRaxd K.VER
O=Softmax(
(2)
where Qs the query state, and K and V denote the key and value states, respectively. Previous research has revealed several significant insights into the distribution of attention weights (i.e., Softmax( OK? )). These findings serve the foundation for parallel encoding in LLMs, motivating our design of APE.
Attention Sink. StreamingLLMM (Xiao et al. 2023 identifies the presence of an "attention sink" in LLMs a token that receives significantly higher attention score than other tokens but provides limited semantic information. It observes that the attention sink only exists in the initial token.
Position Embedding. To effective process sequential input, L.LMs necessitate the explicit encoding of positional information. Common techniques include absolute positional embedding Vaswani (2017); Devlin (018)) and relative positional encoding (Su et al. 2024, Press et al. 2021). l). However, the introduction of position embedding not only limits the context window to the training length (Chen) ct al. 2023), but also results in the "lost the middle" phenomenon (Liu et al. 2024a). where LLMs struggle to produce correct answers when relevant information is located in the middle of the context.
3

---

Published as a conference paper at ICLR 2025
Tedill
CaPAD
Soe Arcond
32
Aearey
26-
336
HotpoqQ
IWRIMUQA
MuSiQue
MultNNew
GSMKK ull-shet
TrivaAAA (hul-aattt))
300
MMLL (fulk-bett
(a) Retrieval-augrmentee Generation
(b) In-context Learning
Figure 2: Comparison of sequential encoding, parallel encoding, and CEPED in RAG ICL scenarios.
Parallel encoding and CEPED dramatically degrade performance, especially for math reasoning.
In this section, we evaluate sequential encoding, parallel encoding. and CEPE-Distilled (CEPED) Yen et et. 2024) using the LLA MA-2-AA--BBBB model I on various RAG and ICL tasks in Figure 2 Our analysis yields several key observations of parallel encoding. In Section 3.1
we explain the limitation of trainable approaches in their generalization abilities to complex tasks. Then, we explore why parallel encoding does not completely break down, despite LLMs being trained sequentially in Section 3.2 This analysis illuminates both alignments and misalignments between parallel and sequential encoding.
3.1 TRAINABLE APPROACHES ARE ONLY EFFECTIVE FOR EASY TASKS.
In Figure 2 we compare the performance of different context encoding methods on RAG and ICL tasks, with detailed setups described in Appendix x B Our analysis of the long-context RAG ability on LongBench (Bai et al. 2023) showcased in Fi Figure 2a Despite accessing larger context windows, CEPED only surpasses the sequential baseline in 2/3 QA tasks, and it even notably underperforms parallel encoding in the Summarization task (MultiNews), which requires synthesizing information from the entire context. We hypothesize that CEPED cannot process complex tasks since the encoder and decoder are only trained unlabeled pre-training corpus, without instruction-tune tunnng on nigh-quality QA samples. This is further supported by the results of ICL tasks (see Figure 2b). where CEPED performs on par with 1-shot sequential encoding baseline on TriviaQA but falls short of it on GSM8K and MMLU, despite using much more examples, respectively. The latter involves reasoning steps hard for the ill-trained model to understand. In conclusion, fine-tuning models to improve parallel encoding requires (i) more diverse and labeled data, and (ii) resource-intent
instruction-tuning (c.g.. SFT or RLHF (Ouyang et al. 2202))). Given this unfavorable trade-off' between training costs and model capabilities, we propose developing training-free solution for better parallel encoding performance.
3.2 COMPARING PARALLEL ENCODING AND SEQUANTIAL ENCODING.
In Figure we surprisingly observe that parallel encoding, despite reducing performance, can generate reasonable responses without further modifications. This is non-trivial as contexts are encoded into KV states separately, with no guarantee that these states can be compared or combined. However, our analysis reveals that the attention mechanism in L.Mss naturally builds an inherent alignment across independent contexts similar to sequential encoding, despite the change in context and position. To clarify this, Figure 3 focuses the impact of the attention sink (Xiao et al. (023), where visualize the direction (cosine similarity) of KV states for different differentssmpless and positions. In Figure 4 we further visualize the distribution of different components in the Softmax attention, resulting in several findings.
Key states from different contexts are similar. In Figure 3l 3p. we measure the cosine similarity of the key states for different tokens for the L.LAMA-3, --3--8B--- INSTRUCT and MISTRAL--BB SEN.--BBB-INNTUCTT----T V0.3 models, which consistently yields a value close to This indicates that the direction of the initial key state remains largely invariant across different inputs. Figure 3 and 3d further analyze the similarity be- tween the initial key states and their subsequent states, where observe comparable negative values from different positions. Therefore, the angles between the initial key states and their subsequent states similar and significantly larger than the angles between different initial key states, visualized in Figure This suggests that the direction of key states remains relatively consistent across contexts, as they are primarily decided by the initial key states, which exhibit similar directions across examples. These findings, combined with the small variance in magnitude across examples in Figure[4b] show that key states from different contexts share similar directions and magnitudes, making them comparable.
kuaker
Kinital
Figure 5: Geometry Key States.
We use I.I.MA A-2 for CEPED, it is the only supported model. For other analyses, we employ I.LAMA-3.
4

---

Published as a conference paper at ICLR 2025
Layer
Layer
Layer 8
Layer 12
Layer 16
Layer 20
Layer 24
Layer 25
1.0
1.0
1,00
Cosine Similariyy
0.9
0.75
0.8
0.50
0.0
0.8
0.4
-0.5
-0.25
10"
10"
1.0
1.00
1.00
0.8
Similarit
0.5
0.75
0.8
0.8
0.50
0.4
0.25
0.4
0.25
0.2
0.2
10
10
10
10"
10"
10"
10"
Similarity betwoen samples different positions
Similarity between the Initial token and tokens In differrnnt positions
Figure 3: Top Left: Both LLAMA-3 --BBB---B--- I-STUTTRUCT (a) and MISTRAL- - states exhibit a cosine similarity larger than 0.9 for the key states from distinct initial tokens. Top Right: Initial token's 's key states show similar negative values to those from other positions for LLAMA-3-88B INSTRUCT (c) and VISTRAL--B-INSTRUCE 7B-INSTRUCT--O (d) models. Bottom: Value states exhibit patterns
(b)
similar to those observed in key states. The X-axis shows positions of key states logarithmic scale.
Costra simiturtt
Magnkude
Magnitu
(a) Query-Key Similarity
(b) Key Magnitude
(c) Value Magnitude
(d) Query-Keerr Key Product
Figure 4: (a) The cosine similarity between query and key states increases as the distance between their positions decreases. (b) The magnitudes of key states show a slowly upward trend as position increases. (c) The magnitude of value states remain constant across positions. (d) Query-key key dot products keep consistently low values except at initial and recent positions. red dashed line marks the anomalous region for the first two tokens in all subfigures. The X-axis shows positions of key and value states log scale. Results are measured on the Hotpot QA dataset using the LLAMA-3 ----8-3- 3-8---8-BB-- SB-INNSRUCT model.
To further understand this key observation, we con- ducted small experiment using the LLAMMA---BB INSTRUCT model on the HotPotQA dataset. Our analysis involved applying rotations of varying de- grees around random axes to the key states of initial tokens. For parallel encoding, we explored two rota- tion modes: one using the same rotation axis across different contexts, and another employing random ro- tation axis for each context. Figure 6 presents our findings, revealing that sequential encoding maintains performance across various rotation degrees. In con- trast, both modes in parallel encoding deteriorate when rotations exceed 150 degrees. This effect arises from the duplication of initial tokens' key states, which intensifies the impact of our rotations. Notably, using separate axes for each context leads to an earlier breakdown beginning at 90 degrees. This disrupts the directional similarity of key states with different initial tokens (i.e., kinitial) in Figure enlarges the angle between key states from different contexts.
F1 Score
Sequertial Uncoding
Parilla Encreiing: Some Acis
Pandlid Encuding; Difficcsi
36
6
20
130
150
164
Rotation Degree
Figure 6: Rotation Analysis
Values states from different contexts can be combined, In equation value states are combined through weighted summation in the attention calculation, where the Softmax operator would normalize the weights (i.e., attention scores) of all value states to sum to 1. This normalization indicates that the magnitude of current value states is determined solely by those from previous positions,
5

---

Published as a conference paper at ICLR 2025
104
300
(a) Sequential
(b) Parallel (T= 1.0)
Parallel (T (T=0.2) 0.2)
(d) Parallel vs. Sequential
Figure 7: (a) Sequential encoding allocates high attention scores to neighboring tokens. (b) Parallel encoding distributes attention scores more uniform across neighboring tokens from all contexts. (c) Adjusting the temperature T sparsifies the distribution. (d) After adjustment, the distribution in parallel encoding becomes similar to sequential encoding. The X-axis represents token positions.
resulting a similar L norm across positions, as shown in Figure 4c, Additionally, the small variance demonstrates that the magnitudes are comparable amang among samples. This, coupled with a similar direction (Bottom). indicates that the value states can be combined. Opportunities for improvement. Previous analyzes demonstrate that key and value states exhibit natural alignment across contexts both in direction and magnitude for most positions. However, residual misalignments in Figure still severely reduce performance. We summarize them as follows:
across samples and positions in Figure 3
In 4 observe notable discrepancy in direction and magnitude for the first few positions, leading to large QK dot products in Figure[4d] This is designated as an abnormal region in the context. Figure 4d shows dot products between the final query state and all key states, revealing notable increase when states are close to each other in position due to the larger cosine similarity Figure 4al
4 ADAPTIVE PARALLEL ENCODING
With all the lessons learned in Section B] we will design our Adaptive Parallel Encoding (APE) to ad- dress the observed misalignments. APE enables seamless shift to parallel encoding without requiring retraining, while maintaining the majority of the model's is capabilities. Our approach adaptively aligns the distribution of attention weights between sequential encoding and parallel encoding through a three-step procedure as showcased in Figure (Bottom), thereby boosting both efficiency and accuracy.
4.1 PREPENDING SHARED PREFIX.
Fllutrees that the distribution of various components for the first few tokens differ significantly from those of subsequent tokens. This discrepancy poses a challenge when encoding contexts in parallel from the beginning, as would result in the duplication of these abnormal KV states. To address this issue, we propose simple yet effective solution: prepending shared prefix to all contexts. This approach ensures that these KV states appear only once in each generation step. In practice, the choice of prefix varies with the model and task. We use existing system prompts and instructions as the shared prefix when available. Otherwise, we will insert few newline characters (i.e., "In") before all contexts.
4.2 ADJUSTING ATTENTION TEMPERATURE.
In Figure the value of QK dol products increases as the relative distance decreases, with notably sharper rise when the distance approaches zero. To show its impact on parallel encoding, we 50-token prefix and query, encoding the remaining 900 tokens sequentially or in five parallel chunks, with attention distributions shown in Figure 7 Comparing Figure 7b wbiht 7a| duplicating neighboring KV states in parallel encoding will disperse the query's 's attention to multiple contexts, resulting in more uniform attention distribution. We adjust the attention temperature T a value less than to refocus on the most relevant tokens. The comparison between different T is shown in Figure 7c
4.3 ADDING SCALING FACTOR.
While adjusting the temperature sharpens the attention distribution among context tokens but also alters the overall attention allocated to them. as shown by the LogSumExp value with different in Figure 8 Specifically, when the sum of the original QK dot product values significantly greater than 0, reducing T amplifies these positive values, resulting in an increased, positive LogSumExp value. Conversely, when the sum closer O, lowering T has stronger effect on the negative QK dot products, leading to decreased, negative LogSumEXxp value. These effects generally increase the absolute value of LogSumExp. To mitigate this, add scaling factor S 1 to directly reduce the absolute value.
LegSumEk0
Layer
Layo Layer
Layer
Layer
Figure 8: Parallel w/ Diverse T.

---

Published as a conference paper at ICLR 2025
4.4 FORMULATION.
Given these three steps, we can formulate the modified attention in APE. We begin with the standard Softmax attention, where Q, K, and V are the query, key, and value states from the input or output:
Q
KC.,...... KI.K,....KK
0 OSSoftmax
)x VV.,,...,vo
Vd
[AC, 1 Acy
x xIV
N
Li= i=
2j=1
QkG..1
Q and aC,,
Qk vd
where Ac, = oxx
vd rpp
Similar for A and aj.
vd
After incorporating the proposed changes, the formula for our refined attention calculation becomes:
AA,,Ac ..... A'...A E ap
O'
x
E 1e, OP. (E )
Qho..1 where Acc -lexp TVd exp TVd
Qu
le
Qke
J.
dc..) ( 8-1
and d'., = exp Tvd
Ap represents the attention weights for the shared prefix while A denotes that for query and generated tokens. The attention temperature T and the scaling factor S for the context are less than 1. Appendix[D provides detailed deduction of this formula. All these modifications are compatible with fast attention implementations such as flash attention (Dao et al. 2022) by computing the context and non-context KV states separately and merging them the attention output, which only incur negligible overhead.
We tune hyperparameters on a validation set with greedy search. If no prefix is provided, we begin by adding two "'\n" and increase the prefix length by 10, 20, and 40. S and T are searched in the ranges 10.1, 1.0] using 0. step sizes. We use S-T instead of S as the scaling factor to simplify our search.
5
EXPERIMENTS
this section, we present effectiveness and efficiency of APE, focusing on CAG scenarios such as RAG and ICL. Therefore, we do not include comparisons with long-context LLMs here. Specifically,
AG-Bench compared to sequential encoding when retrieving the same text input. Furthermore, it improves performance by 3.3% on average for RA on LongBench through its ability to retrieve and process more and longer contexts. parallel encoding by 7.9% on average in ICL. tasks. Moreover, APE keeps 93% of accuracy achieved by sequential encoding when using the same number of examples.
S.1 APE can maintain 98% of the accuracy on
In Section 5.2 APF
Section 5.3
APE scales many-shot CAG tasks, effectively encoding hundreds of texts in parallel.
Section 5.4 APE can achieve up to 4.5x faster inference through 28 x reduction in prefilling time.
5.1
RETRIEVAL-AUUGMENTER
GENERATION.
For RAG tasks, we validate that APE retains most of the sequential encoding capability while accom- modating more and longer contexts, mitigating retrieval errors and outperforming encoding baselines.
5.1.1
RETRIEVAL FOR MULTI-TUL MURTI-TURN QUESTION ANSWERING.
Setup. APE is evaluated five conversational QA tasks using ChatR AGBench (Liu al. 2024b) For each query, 1-000 text chunks are prepared, with retrievers employed to select the top-n most relevant chunks as input. We utilize three different retrievers of varying quality: Contriever (Izacard et al. 202Tb, GTE-base Li et al. (2023), and Dragon-multiour/ Liu et al. (2024b). b) and retrieve up to the top-5 chunks for evaluation. The base model used is LLAMA3 LLAMA3-CHATQA, A 1.5-8B. To measure performance drop after our alignment steps, the same retrieved texts are used for both APE and sequential encoding.
Results. Table shows that switching from sequential encoding APE results in performance drops of %, 0.92%, %, and 1. 14% across different retrievers, respectively. Notably, while this drop increases with retriever quality, APE still keeps 97% of the sequential encoding performance for the best retriever. Moreover, APE can directly input all texts without retrieval process, achieving superior performance.
5.1.2 RETRIEVAL FOR LONG-CONTEN. UNDERSTANDING.
Setup. Our evaluation involves eight tasks on LongBench (Bai et al. 2023). Given the long context, split it into chunks with default size of M words, employ Contriever (Izacard al. 2021 to compute the embeddings of the text chunks and the query, and retrieve the top-N chunks according to the cosine similarity of their embeddings to the query embedding. M and N vary across different
7

---

Published as conference paper at ICLR 2025
Table I: Comparison between APE and sequential encoding using three retrievers on ChaiRAG-Bench.
Method
INSCIT
Doc2Dial
TopicCQA
Qrecc
QuAC
Average
Contriever, Sequential
19.97
23.85
30.49
46.75
26.57
29.53
Contriever, APE
19.88
23.28
28.84
46.28
26.80
29.02
-0.09
-0.57
-1.65
-0.47
+0.23
-0.51
GTE-base, Sequential GTE-base, APE
21.58
32.35
33.41
46.54
30.69
32.91
20.85
30,99
31.92
45.83
30.35
31,99
A
-0.73
-1.36
-1.49
-0.71
-0.34
-0.92
Dragon-multiuune multiturn, Sequential
25.42
36,27
36.10
49.01
35.12
36.38
Dragon-multinne APE
23.84
34.93
33.80
48.70
34.92
35,24
-1.58
-1.34
-2.30
-0.31
-0.20
-1.14
All texts, APE
27.22
36.13
35.72
49.15
35.70
36.78
Table Comparison between APE and baselines on LongBench across different models using RAG.
denooess Contriever, and M N indicates retrieval of the top-N chunks, each containing M words.
Model
MuSiQue Qasper 2WikiMQA DuRead HotpotQA NarratQA MFQA MFQA_co en Avg.
LLAMA-3-8B.
20.70
41.05
30.02
9.55
45.90
20.98
58.54
45.04
33.97
C200x20, Sequential
27.93
42.71
38.35
12.65
49.60
22.78
57.82
48.94
37.60
C4000x20, PCW
18.82
42.59
40.99
21.57
47.09
$4.40
45.05
36.73
C4000x20, APE
26.19
42.32
44.43
23.13
49.71
30.71
55.03
45.41
39.62
MISTRAL-77B BUSTRAI-7B-INSTEIN 1EA-AA--7--NSTUCC----O INSTRUCT-VO.3)
10.05
31.08
22.12
17.68
32.09
19.68
32.03
40.38
25.64
C200x20, Sequential
11.58
21.98
24.44
20.80
32.79
16.06
34.43
38.40
25.06
C4000x20, PCW
17.58
35.57
32.97
18.70
37.05
14.10
34,69
40.14
28.85
C4000x20, APE
20.30
36.81
34.37
21.89
42.33
20.49
40,20
44.03
32.55
GEMMA-2-9B-1
22.57
39.99
48.06
27.40
47.49
23.11
50.81
45.35
38.10
C200x10, Sequential
30.69
42.86
53.55
28.04
$2.05
24.45
50.25
48.34
41.28
C2000x20, PCW
26.27
46.69
47.59
23.43
48.95
27.11
56,69
49.81
40.82
C2000x20, APE
33,38
47.72
49.49
28.43
56.62
30.41
56.52
50.84
44.18
LLAMA-3.1 1-A--3.--B--ISSTRUCC
22.18
46.81
40.58
34.61
43.97
23.08
61.60
51.89
38,98
128K, Sequential
28.35
47.20
40.81
33.34
53.46
30.57
61.97
53.25
42.24
C200x20, Sequential
30.62
42.33
44.39
33.51
49.97
23.87
56.87
55.14
40.22
C4000x20, PCW
21.23
41.52
44.87
31.11
49.47
19.98
60.90
51.19
38.44
C4000x20, APE
26.88
43.03
50.11
32.10
55.41
30.50
62.02
52.51
42.86
encoding methods. We compare with sequential encoding with and without RAG. and PCW, using LLAMA-3-8B. SEL-LA----BINNSTRUA (Dubey et al. (2024). MISTRAL-7B-- ---IIITRUTTTT--- v0.3 (Jiang et al. 2023).
REMMA-2-98--2 (Team at. 2024)) and LLAMA-3. --8-- -NSTRUUT as base models.
Results. In Table 2 APE consistently improves performance across all models, achieving a 5.6% average gain over base models and outperforming sequential RAG baselines by 3.3% through retrieval of more and longer contexts. The superior performance over PCW further showcases the effectiveness of our alignments. Notably, APE surpasses the 128K-context variant of LLAMA-3. 1--1-----INNRUCE by placing retrieved texts within the 8K context window, mitigating the "lost the middle" phenomenon.
5.2 IN-CONTEXT LEARNING
Setup. We evaluate APE on three ICL tasks using the LM Evaluation Harness (Gao al. 2024 GSM8K (8-shot) (Cobbe et al.) 2021a), TriviaQA (5-shot) (Joshi et al., (017). and MMLU (5- shot) (Hendrycks al. 2020a). Experiments use the same base models as in our LongBench evaluations.
Baselines. We compare parallel encoding (PCW) to show the improvement from APE's alignment steps. Sequential encoding with varying numbers of shots (i.e., 1-shot, half-shots, and full-shots) are also employed to measure the gap from the ideal scenarios. All methods can access all input examples.
Results. In Figure 9 APE significantly surpasses parallel encoding with average improvements of 15.4% on GSM8K, 4.7% on TriviaQA, and 3.5% MMLU. When compared with the 1-shot sequential baseline with comparable context length, our method consistently yields superior results. Moreover, APE achieves better performance than half-shot sequential encoding in 8/12 settings and preserve 93% accuracy comparing to the full-shot sequential encoding. Additionally, our results suggest that the LLAMA family exhibits enhanced compatibility with parallel encoding, potentially due to the stronger directional alignment of initial tokens from different contexts (see Figure Across different tasks, the performance gap between APE and full-shot sequential encoding the largest on GSMSK, suggesting that while APE maintains most capabilities, its effectiveness may decrease task complexity increases.
5.3 MANY-SHOT CONTEXT-AUGMENTT
GENERATION
Setup. To demonstrate the scalability APE, we employ four RAG and ICL tasks from the LOFT benchmark (Lee et al. 2024), each involving hundreds of texts that provide additional information. We
8

---

Published as a conference paper at ICLR 2025
150
Agcuracy 10
85
11
69
79
(a) LLAMA- IN-INNTRUCT
(b) LLAMA-3. LLLLLLA-.-----INSTTEEEE -INSSRUTTT
70
124
34
MA
00
12
GSMKK
TrA404
MMUU
GEMAK
KNLE
(c) MISTRAL-7B DINIRALL-7--SSTUCC--VOOO INSTRUCT-VO.3
(d) GEMMA 12AMMA-2-992 2-9B-IT
Figure 9: Performance comparison of APE, parallel encoding, and sequential encoding on ICL tasks.
Table 3: Comparison between APE and sequential encoding in various many-shot RAG and ICL tasks.
Retrieval-augmennte Generation
In-contexi Learning
Method
ArguAna
FEVER
NQ
SciFact
Date
Salient
Tracking?
Web
Sequential, Zero-shot
11.15
7.78
17.78
7.74
20.00
8.89
1.12
8.89
Sequential, Few-shot
11.20
9.78
17.81
9.49
36.64
38.89
6.67
38.89
Sequential, Half-shot
15.34
13.12
19.64
16.12
45.55
42.22
8.89
55.56
Sequential, Full-shot
12.84
14.19
24.54
16.88
46.67
46.67
8.89
58.89
APE, Full-shot
16.32
14.70
21.91
15.72
43.33
45.55
8.89
58.89
employ LLAMA-3. 1-B. --NSSSTUUT as our base model compare APE with sequential encoding, both applied to the same many-shot long-contert inputs. The total context lengths for RAG and ICL tasks are 128K and 32K, respectively. We also include the zero-shot. few-shot (< 5), and half-shot variants of sequential encoding as baselines. For metrics, FI score and EM are used in RAG and ICL.
Results. As shown in Table APE achieves performance comparable to sequential encoding when processing the same many-shot long-context inputs, showcasing its ability to efficiently encode hundreds of texts in parallel. Notably, for RAG tasks, it outperforms sequential encoding ArguAna and FEVER. While APE is expected to reduce performance, it recovers this drop by positioning all texts close to the query. This mechanism addresses the "lost in the middle" problem in long-context LLMs. Additionally, for ICL tasks, APE can learn from examples as effective as sequential encoding.
5.4 EFFICIENCY EVALUATION
Setup. We measure the latency for sequential encoding, MInference (Jiang et al. 2024a). and APE on the Llama-3.1 --BB-BB..-BBBBB- Instructt (Dubey et al. (024) model. Our evaluation is conducter ucted on an H100 GPU with batch sizes of 1 and 4. The query and generation lengths are fixed at 256 tokens, while the input context lengths ranged from 2K to 128K tokens. We employ VLLM (Kwon et al.) (023) as our inference engine and measure both prefilling time (Time to first token) and total inference time.
Results. Comparing to sequential encoding and MInference. APE can accelerate inference up to 4.5x Notably, for 228K-toker contexts, APE reduces prefilling time by 28x compared MInference. APE's prefilling overhead exhibits linear scaling and consumes less than 10% of inference time, while baselines require over 50% as
and 2.2x respectively for long- contecttext scenarios in Figure TO
Suguestii
Suque
16
60
41
12
40
30
70
10
14%
1355
105
(a) Prefill Time (bsz=1)
(b) Prefill Time (bsz=4)
(c) Total Time (bsz=1) 1)
(d) Total Time (bsz=4)
Figure 10: Latency on H100 GPU: [prefill and total time (s)]. The gray text in brackets is batch size.
9

---

Published as a conference paper at ICLR 2025
Table 5: Performance and latency comparison using the CAA-MMA------INSSRURE model CRAG.
Task
Model
Latency (ms)
Accuracy (%)
Hallucination Missing
Score,
LLM only
LLAMA-3-8DD -LAMA-3-8B-INSTRUCT
682
22.14
48.97
28.90
-26833
LLAMA-3-8B-INSTRUUCE
1140
23.28
29.49
47.22
-6.21
Task
+APE
1054
25.53
21.30
37.93
-0.41
1830
24.46
28.38
47.15
-3.92
Task 2
+APE
1672
27.04
18.74
37.32
2.16
Table 6: Performance comparison across different long-context tasks on LongBench (Bai et al.
2023
Method
NarratQA
Qasper
MultiFQA
GovReport
QMMum
LCC
LLLAMA3--BB-ISTTTE
19.32 26.87
32.83
43.38 59.12
27.89
22.40
53.22
+APE
39.14
29.10
23.08
66.09
Method
RepoBench-P
HotpotQA
2WikiMQA
MuSiQue
MultiNews
Average
LLAMA-3 -1-MM-3-8B--A SALLA------INSTRUUE INSTRUCT
38.15
44.24
21.01
20.47
23.63
31.50
+APE
49.43
50.11
28.06
25.79
22.40
38.11
context length increases. APE also demonstrates superior versatility across various settings, whereas MInference slows inference with additional overhead when processing short contexts and large batches.
6 ANALYSIS
In this section, we present analyses to answer the following research questions: RQ1: How does each component in APE contribute to the performance? RQ2: Can APE improve performance for real- world RAG tasks? RQ3: Can APE extend LLM context window size in long-context scenarios without RAG?
6.1 How DOES EACH COMPONENT IN APE CONTRIBUTE TO THE PERFORMANCE?
In Table we conduct an ablation study to examine each alignment process in APE, including the shared prefix (P), attention temperature (T), and scaling factor (S). We present results averaged across the four models evaluated in Figure 9 Our findings indicate that incorporating each of these components can consistently improve performance across all tasks, with average improvements of 5.19%, 0.59%, and 2.07%, respectively. Among them, adjusting the attention temperature yields minimal performance gains without the complementary effect of the scaling factor.
Table 4: Ablation study of APE align- ments on ICL tasks. P: shared prefix, T: attention temperature, S: scaling factor.
P T S GSM8K TriviaQA MMLU
38.25%
67.99%
63.09%
/
50.42%
70.76%
63.70%
/
51.15%
71.03%
64.49%
53.62%
72.64 %
66.62% %
6.2 CAN APE IMPROVE PERFORMANCE FOR REAL-WORLD RAG APPLICATIONS?
In 5 we evaluate APE's performance in real-world RAG scenarios using the CRAG benchmark (Yang al. 2024). Task I augments the model with several webpages, while Task provides an additional knowledge graph as another retrieval source. In our experiments, the sequential encoding baseline limited to retrieving 4K tokens of context, whereas APE can process 20 parallel segments of 4K tokens each. By incorporating significantly more external data during generation, APE consistently outperforms sequential encoding that has limited context sizes while reducing latency. Moreover, the improvement in Task 2 further shows the effectiveness APE in merging text from multiple sources.
6.3 CAN APE EXTEND CONTEXT LENGTHS IN LONG-CONTEXT SCENARIOS WITHOUT RAG?
Table examines the effectiveness APE when processing a single long context input for the LLAMA- B------NSRRCE model on LongBench (Bai et al. 2023). To accommodate the long context within APE, we split it into multiple segments with of less than 7,500 tokens, appending the final 500 tokens to the query. Our results indicate that APE enhances performance across 10/11 tasks, yielding an average improvement of 6.6% compared to the sequential encoding baseline with limited context window size.
CONCLUSION
This work explores parallel encoding in context-augmented generation, which can pre-cache KV states for fast inference and re-use positions for long context but lead to worse performance. To address this, we propose APE, a training- -free method to enable accurate and fast CAG systems. It achieves this by aligning the attention distribution of parallel encoding with sequential encoding via three steps: shared prefix, adaptive temperature, and scaling factor. We show that APE improves accuracy and efficiency in various RAG and ICI scenarios while successfully scaling process hundreds of chunks in parallel.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENT
This work is supported in part by NSF award CNS-2211882 and gift from Qualcomm. We gratefully acknowledge additional support from Amazon, Intel, Li Auto, and Moffett AI. Additionally, we extend our gratitude the authors of ChaiQA (Liu et al., 2024b), Longbench (Bai et al. (023). CRAG Yangg et al. 2024). LM Evaluation Harness (Gao et al. 2024), VLLM (Kwon et al. 2023). and MInference (Jang al. 2024a) for their valuable codebases, benchmarks, and models. We also sincerely thank Yixin Dong, Hanshi Sun, Zhuoming Chen for their insightful discussions and feedback.
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, llge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et. Gpt-4 technical report. arXiv preprint (wr.iv:2333307770 2023.
Rishabh Agarwal, Avi Singh, Lei M Zhang, Bernd Bohnet, Stephanie Chan, Ankesh Anand, Zaheer Abbas. Azade Nova, John D Co- Co-Reyes, Eric Chu, al. Many-shot in-context learning. arXiv preprint arXiv22404.11018, 2024.
Akari Asai, Zexuan Zhong, Danqi Chen, Pang Wei Koh, Luke Zettlemoyer, Hannaneh Hajishirzi, and Wen-tau Yih. Reliable, adaptable, and attributable language models with retrieval. arXiv preprint
arXiv:2403.03187, 2024.
Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, Yuxiao Dong. Jie Tang, and Juanzi Li. Longbench: A bilingual, multitask benchmark for long context understanding. arXiv preprint arXiv::3008 /4508, 2023.
Mikhail S Burtsev, Yuri Kuratov, Anton Peganov, and Grigory V Sapunov. Memory transformer. arXiv preprint arXiv:2006.11527, 2020.
Harrison Chase.
Longchain, 2022.
URL bttps://gttuu..co//aannggh
Langchain
Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian. Extending context window of large language models via positional interpolation. arXiv preprint arXiv:2306./5555 2023.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. Training verifiers to solve math word problems. arXiv preprint arXiv:2 10. /4/68, 2021a.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. Training verifiers to solve math word problems. arXiv preprint arXiv:2 2021b.
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré, Flashattention: Fast and memory- efficient exact attention with io-awareness. Advances in Neural Information Processing Systems, 35: 16344-16359, 2022.
Michiel De Jong, Yury Zemlyanskiy, Nicholas FitzGerald, Fei Sha, and William Cohen. Mention memory: incorporating textual knowledge into transformers through entity mention attention. arXiv preprint arXiv:2/10.066 2021.
Jacob Devlin. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint Xiv: /810004805.. 2018.
Qingxiu Dong, Lei Li, Damai Dai, Ce Zheng, Zhiyong Wu, Baobao Chang, Xu Sun, Jingjing Xu, and Zhifang Sui. A survey on in-context learning, arXiv preprint arXiv:2301.00234, 2022.
Matthijs Douze, Alexandr Guzhva, Chengqi Deng, Jeff Johnson, Gergely Szilvasy, Femran-m Mazaré, Maria Lomeli, Lucas Hosseini, and Hervé Jégou. The faiss library. 2024.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The llama herd of models. arXiv preprint arXiv: 2024.
11

---

Published as a conference paper at ICLR 2025
Alexander R Fabbri, Irene Li, Tianwei She, Suyi Li, and Dragomir R Radev. Multi-news. A large- scale multi-document summarization dataset and abstractive hierarchical model. arXiv preprint arXiv:1906.01749 2011
Thibault Févry. Livio Baldini Soares, Nicholas FitzGerald, Eunsol Choi, and Tom Kwiatkowski. Entities as experts: Sparse memory access with entity supervision. arXiv preprint arv::2200 07202. 2020.
Lco Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac h. Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A framework for few-shot language model evaluation, 07 2024. URLhttts://zenoot http://roooooorrrrrecor/o/////// 602
Yunfan Gao, Yun Xiong, Xinyu Gao, Kangxiang Jia, Jinliu Pan, Yuxi Bi, Yi Dai, Jiawei Sun, and Haofen Wang. Retrievalauugmente generation for large language models: A survey, arXiv preprint wwXiv:2312.0099,, 2023.
AI Gradient. Llama-3-4 3-mna------itttutt2625- instruct-262k.
2024.
Aman Gupta, Anup Shirgaonkar, Angels de Luis Balaguer, Bruno Silva, Daniel Holstein, Dawei Li, Jennifer Marsman, Leonardo O Nunes, Mahsa Rouzbahman, Morris Sharp. et al. Rag vs fine-tuning: Pipelines. tradeoffs, and a case study agriculture. arXiv preprint arXiv:2401.08406, 2024.
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. arXiv preprint arXi:200990033000 2020a.
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. arXiv preprint arXiv:2009.03300 2020b.
Xanh Ho, Anh-Khoa Duong Nguyen, Saku Sugawara, and Akiko Aizawa. Constructing a multi-hop qa dataset for comprehensive evaluation of reasoning steps. arXiv preprint arXiv:201 2020.
Gautier Izacard, Mathilde Caron, Lucas Hosseini, Sebastian Riedel, Piotr Bojanowski, Armand Joulin, and Edouard Grave. Unsupervised dense information retrieval with contrastive learning. arXiv preprint arXiv:21 18, 2021.
Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al. Mistral 7b. arXiv preprint arXiv:2310.06825, 2023.
Huiqiang Jiang, Yucheng Li, Chengruidong Zhang, Qianhui Wu, Xufang Luo, Surin Ahn, Zhenhua Han, Amir H Abdi, Dongsheng Li, Chin-Yew Lin, et al. Minference 1.0: Accelerating pre-filling for long-context Ilm via dynamic sparse attention. arXiv preprint arXiv:2407.022490, 2024a.
Ziyan Jiang, Xueguang Ma, and Wenhu Chen. Longrag: Enhancing retrievallauggmette generation with long-context llms. arXiv preprint arXiv:2406. /53/9, 2024b.
Hongye Jin, Xiaotian Han, Jingfeng Yang, Zhimeng Jiang, Zirui Liu, Chia- Yuan Chang, Huiyuan Chen, and Xia Hu. Llm maybe longlm: Self-extend Ilm context window without tuning. arXiv preprint arXiv:2401.01325 2024.
Mandar Joshi, Eunsol Choi, Daniel S Weld, and Luke Zettlemoyer. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension. arXiv preprint arXiv: 2017.
Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th Symposium on Operating Systems Principles, pages 61 16-266, 2023.
12

---

Published as a conference paper at ICLR 2025
Jinhyuk Lee, Anthony Chen, Zhuyun Dai, Dheeru Dua, Devendra Singh Sachan, Michael Boratko, Yi Luan, Sébastien MR Arnold, Vincent Perot, Siddharth Dalmia, et Can long-conten language models subsume retrieval, rag, sql, and more? arXiv preprint mrXiv:2406.13121, 2024.
Zehan Li, Xin Zhang, Yanzhao Zhang, Dingkun Long, Pengjun Xie, and Meishan Zhang. Towards general text embeddings with multi-stage contrastive learning. arXiv preprint arXtv:2308.002281 2023.
Zhenyu Li, Yike Zhang, Tengyu Pan, Yutao Sun, Zhichao Duan, Junjie Fang, Rong Han, Zixuan Wang, and Jianyong Wang. Focusllm: Scaling I1m's 's context by parallel decoding. arXiv preprint arXiv:2408. 11745, 2024.
Jerry Liu. Llamaindex, 11 2022.
index
Nelson Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. Lost in the middle: How language models use long contexts. Transactions of the Association for Computational Linguistics, 12:157-173, 2024a.
Zihan Liu, Wei Ping, Rajarshi Roy, Peng Xu, Chankyu Lee, Mohammad Shoeybi, and Bryan Catanzaro. Chatqa: Surpassing gpt-4 on conversational qa and rag. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024b.
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in neural information processing systems, 35:27730- 27744, 2022.
Zhuoshi Pan, Qianhui Wu, Huiqiang Jiang, Menglin Xia, Xufang Luo, Jue Zhang, Qingwei Lin, Victor Rühle, Yuqing Yang. Chin Yew Lin, el. LImlingua-22 2: Data distillation for efficient and faithful task-agn prompt compression. arXiv preprint 2r0v 12968, 2024.
Ofir Press, Noah A Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. arXiv preprint arXiv:2108. 12409, 2021.
Nir Ratner, Yoav Levine, Yonatan Belinkov, Ori Ram, Inbal Magar, Omri Abend, Ehud Karpas, Amnon Shashua, Kevin Leyton-Brown, and Yoav Shoham. Parallel context windows for large language models. arXiv preprint arXiv:2212.10940 2022.
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568: 127063, 2024.
East Sun, Yan Wang, and Lan Tian. Block-attentioo for efficient rag. 2024. https://:/ semanticscholar: htt.://cc.....oooo.oooooooaaaaaa
Gemma Team, Morgane Riviere, Shreya Pathak, Pier Giuseppe Sessa, Cassidy Hardin, Surya Bhu- patiraju, Léonard Hussenot, Thomas Mesnard, Bobak Shahriari, Alexandre Ramé, al. Gemma Improving open language models at practical size. arXiv preprint arXiv:2408.000 18, 2024.
AI Together.
2023.
Harsh Trivedi, Niranjan Balasubramanian. Tushar Khot, and Ashish Sabharwal. Musique Multihop questions via single-hop question composition. Trunsuctions the Linguistics, 10:539-554 554, 2022.
the Association for Computational
Vaswani. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Zoph. Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, et al. Emergent abilities large language models. arXiv preprint arXiv:2206.0744 2022.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. arXiv preprini arXiv:2309. 17453, 2023.
13

---

Published as a conference paper at ICLR 2025
Xiao Yang, Kai Sun, Hao Xin, Yushi Sun, Nikita Bhalla, Xiangsen Chen, Sajal Choudhary, Rongze Daniel Gui, Ziran Will Jiang, Ziyu Jiang, et al. compreenpeeen rag benchmark. arXiv preprint arXiv:2406.0444 04744, 2024.
Zhilin Yang, Peng Qi, Saizheng Zhang, Yoshua Bengio, William W Cohen, Ruslan Salakhutdinov, and Christopher D Manning. Hotpotqa: A dataset for diverse, explainable multi-hop question answering. arXiv preprint 18X.::800000900 2018.
Howard Yen, Tianyu Gao, and Danqi Chen. Long-context language modeling with parallel context encoding. arXiv preprint arXiv:2402. 16617, 2024.
André Zayarni, Andrey Vasnetsov, et al. Qdrant, 2024. URL https: /ttplp::://///at..../// ceirant.tecll
14

---

Published as a conference paper at ICLR 2025
A LIMITATIONS
While APE preserves the efficiency of parallel encoding and the effectiveness of sequential encoding through inference-time modifications to the attention distribution, its performance remains highly sensitive to hyperparameter selection, particularly the attention temperature T and scaling factor S. In real-world applications, where contexts vary in length, quantity, and content, automatically aligning the distribution between sequential and parallel encoding during runtime presents significant challenge.
B DETAILED EXPERIMENTAL SETUPS FOR SECTION 3.1
We evaluate four tasks from the LongBench dataset BBai et al. 2023) that require processing multiple independent documents: HotpotQA (Yang et al. 2018)) 2 2WikiMulthhoQAA (Ho et al. | (2020) MuSiQue (Trivedi et al. 2022). and MultiNews (Fabbri et al. 20199) For the three QA tasks, we use the Fl score as the evaluation metric, while Rouge-L is used for the summarization task. Both parallel encoding and CEPED process each document independently using E If document exceeds the window size of θ Enc- these methods split it into multiple chunks before encoding. In contrast, sequential encoding handles lengthy inputs by truncating them from the middle to fit within the context window.
RAG.
ICL, We select three few-shot learning tasks from LM Evaluation Harness (Gao al. 2024) to evaluate the ICL ability of different encoding methods of different encoding methods: GSM8K (Cobbe et al. 2021b), TriviaQA Joshi et al. (2017), and MMLU (Hendrycks et al. 2020b). For parallel encoding and CEPED, each example is encoded independently. with all resulting KV states stored as input to Dee- To thoroughly analyze these methods, we evaluate variants of sequential encoding with different numbers of chots-O-shot 1-shot, half-shot, and full-shot, and compare with the full-shot parallel ones.
C MORE VISUALIZATIONS FOR SECTION 3.2
C.I SIMILARITY BETWEEN KEY STATES FROM DIFFERENT SAMPLES IN EACH POSITION.
In Figure II we show that key states across different layers maintain consistently high cosine similarity values for various initial tokens, with only the first layer exhibiting slightly reduced similarities. Our analysis reveals that I.LAMA A-AM--3--- A-3-8BB-- 1-----IISTUUCE and A.L.MA-... -INIRRTTUCT maintain nearly identical directions (cosine similarity approximately 1.0) for different tokens beyond the first layer, while MISTRAL- BEETRAL--B-INSRRUCCE 1E--INSTUCCT-V v0.3 and GEMMA-2-99 2-9---99- IT display substantial but comparatively lower similarities ranging from 0.8 to 0.9. These findings indicate inherent alignments across contexts while highlighting the potential for further improvements through the shared prefix in Section 4.1
Layer
Layer
Layer
Layer
Layer
Layer 20
Layer
Layer
1.0
1.0
0.9
Cosine Similarity
0.9
0.9
0.8-
0.8
0.8
0.7-
0.7
0.7
0.8
0.5
10
13
10
10"
10
10
10"
10
10"
10"
10"
10'
(a) LLaMA--34B-Anstruee
(b) LLaMA-3.1 1...1-1-----stuuctt
(c) Mistrrll
(d) Gemma-2-9b-it
Figure 11: For all base models, key states from distinct initial tokens exhibit a cosine similarity larger than 0.8 for most layers, where the LLaMA family even approaches 1. In contrast, later tokens show markedly lower cosine similarity. The X-axis shows the positions of key states logarithmic scale.
C.2 SIMILARITY BETWEEN VALUE STATES FROM DIFFERENT SAMPLES IN EACH POSITION.
[ that value states keep high cosine similarity across layers for various initial tokens, with two exceptions: the first layer in all models and some layers in GEMMA-2 -2-...IT. This observation aligns with the requirement that GEMMA-2-9B- necessitates a system prompt for proper functioning.
C.3 SIMILARITY BETWEEN THE INITIAL KEY STATES AND FOLLOWING KEY STATES.
Figure 13 illustratess how the cosine similarity between the initial and subsequent key states stabilizes as position increases. This similarity converges to near-constant value for all models after 10 tokens.
15

---

Published as a conference paper at ICLR 2025
Layer
Layer 4
Layer
Layer 12
Layer
Layer 20
Layer
Layer 20
0.8
Cosine Similarity
0.8
0.0
0.8-
0.6
0.6
0.4
0.4-
0.4
0.4
0.2
0.2
10
10
10"
10
10'
10"
10'
10'
10
10"
10"
10"
10
10"
10"
(4) LBaMM-34444nstruut
(b) -LLLMMA-11-BB--nstrret
(c) Misttral7B-4nstuuc--v00
(4) Gemma-2-9b-it
Figure 12: Across all four models, value states from distinct initial tokens exhibit higher cosine similarity than other positions, which is larger than 0.8, except for the first layer in all models and some layers in GEMMA-2-98 --T-- The X-axis displays the positions of value states on logarithmic scale.
Layer
Layer 12
20
Layer 24
1.0
1.0
1.0
1,0
Cosine Similarty
0.8
Awiil
0.5
0.5
0.5
0.0
0.2
0.0
0.0
-0.5
(a) LLaMMA3-6B-hhstruct
(b) LLa.A-M. Instrrtt
Mistral-7B -.n
Figure 13: For all base models, the similarity between the initial key state and subsequent key states stabilizes as the position increases. The X-axis shows the positions of key states logarithmic scale.
C.4 SIMILARITY BETWEEN THE INITIAL VALUE STATES AND FOLLOWING VALUE STATES.
Similar key states, the value states exhibit stable similarity between the initial token and subsequent tokens in Figure [4] with all models convergent to a hear-constant value after approximately 10 tokens.
Layer
Layer 12
20
Layer 24
28
1.0
1.0
1.0
1,0
Cosine Similarity
0.8
0.8
0.8
0.8
0.2
0.2
0.0
0.0
0.0
-0.2
4.2
.0.2
LLaMA-3-8E-Inetruct
(b) L...MAA-.--8---eeeee
2-9b-it
Figure 14: For all base models, the similarity between the initial value state and subsequent value states stabilizes the position increases. The X-axis shows value states' positions on a logarithmic scale.
C.5 SIMILARITY BETWEEN THE QUERY STATE AND PAST KEY STATES.
In Figure the query states across layers, and base models exhibit higher cosine similarity with the initial tokens. Additionally, neighboring positions tend to receive higher cosine similarity.
C.6 MAGNITUDE OF KEY STATES FROM DIFFERENT POSITIONS.
Figure (iillsstillsssttess that the magnitude key states slowly increases with position, except for the first few tokens, which exhibit significantly smaller magnitudes.
C.7 MAGNITUDE OF VALUE STATES FROM DIFFERENT POSITIONS.
In Ii the value states across all positions exhibit similar magnitude, except for the first few positions, which show noticeable deviation. We indicate this region with red dashed line.
16

---

Published as a conference paper at ICLR 2025
Layer
Layer
Layer 8
Layer 12
Layar
Layer
Layer
Layer
0.6
0.5
Cosine Similarity
0.4
0.2
0.2
0.0
0.0
0.0
-0.2
-0.2
-04
-0.4
10
10"
10
10"
10"
10
10"
10"
10
10
10
10"
10
10"
10
10
(a) LLaMA-3-60R 4B -nnstruc-B-lnstructt Instruct
(b) LLaMA-3.1 Instruct
(6) -nnssruct...
(d) Gemma-2-9b-it
Figure 15: The cosine similarity between the query and key states stabilizes for most positions, except for the initial and recent positions. The X-axis shows the positions of key states on logarithmic scale.
Layer 0
Layer 4
Layer
Layer 12
Lay er 16
Layer 20
Layer 24
Layer 28
150.0
120.0
100.0
125.9
128.0
Mognitude
103.0
100.0
100.0
80.0
80.0
78.0
60.0
60.3
50.0
50.0
40.0
40.0
25.0
25.0
20.0
10
10
10
10
10
10
10
10
10"
10
10"
10'
10"
10"
(6) LLaMA-3-0B-Instreee
(b) LLaMMA-.1-08--matte
(c) Mistral-7B-ttattrut---..3
(d) Gemma-2-0b-it
Figure 16: The magnitude of key states increases slowly with position, with red dashed line marking the anomaly in initial tokens. The X-axis displays the positions of key states on a logarithmic scale.
Layer
Layer
Layer
Layer 12
Layer 16
Layer 20
Layer 24
Layer 28
60.0
30.0
Magnitude
40.0
40.0
10.0
10.0
20.0
20.0
0.0
0.0
0.0
10'
10
10
10'
LLaMMt--BB-isstruet
Gemman2 --mmma-2-9b-it
Figure 17: The magnitude of value states remains consistent, except for the first few positions high- lighted by red dashed line. The X-axis represents the positions of value states on logarithmic scale.
C.8 DOT PRODUCT BETWEEN THE QUERY STATE AND PAST KEY STATES.
In Figure 18 the query states across all layers, and base models exhibit larger dot product values with the initial tokens. Additionally, tokens at neighboring positions tend to receive higher values as well.
D FORMAL DERIVATION OF APE
D. HIERARCHICAL FORMULA FOR SOFTMAX ATTENTION.
We begin with the standard Softmax attention, where Q, K, and V represent query, key, and value states. To distinguish between components from different sources, we use the subscript C, for elements from the context, while those without a subscript correspond to user queries or generated responses.
QIKT ..... QIKT....,KK
-T..,.KKK K Cy
0 = Softmax
(6)
Vd
[AC, ACN:
Ver.
(7)
> i= L/ ac.. + E/
17

---

Published as a conference paper at ICLR 2025
Layar
Layer 4
Layer
Layar
Layer
Layer
Layer
Layer
5.0
Product
10.0
0.0
0.0
3,3
6.0
-2.5
-2.5-
0.0
-7.5
-10.0
-10.0
10
10
10
10
10
10
10
10"
(2) LLaMAA-3-8B-Instree
(b) LLaMA-3.1 --8tttrruuut
(c) PeatrallBBlnnstuutt-v.....
(d) Gemma-2-9b-4t
Figure 18: The dot product between the query state and key states stabilizes for most positions, except for the initial and recent positions. The X-axis shows the positions of key states on logarithmic scale.
Quc vi
Qke va
where Kc,
lexp
exp
J.A
Q4 [exp exp Vi bac,,
Qh J exp
Qkc,
QK vi and = exx vá are the denotations of each component.
We can rewrite this formula in hierarchical way, first computing Vt and A for each context C₁:
Q (
Ko.se. Ve........V...ee. A. = =ogSSumExp
Va = Softmax
Vd
Vd
(8)
Similarly, for those non-context tokens, we have:
(QIN) ,....K]] Vd
Qik p....k]]]
yh Softmax
x [V1...., Vi], A =oogSumxpp
Vd
After computing the values of all components, we can combine these intermediate value states while renormalizing with the aggregated attention scores, the formula is shown below:
O O=Softmax(A$y re..... Aex Va. ...., Vē V4)
(10)
D.2 HIERARCHICAL FORMULA FOR APE.
After incorporating all components in APE, we have new V/l and A for each context C::
@lbd,ar-thrveeee T.Vd
VA Softmax
(VCc, DAA =
- LogSSuxxx
T.Vd
(11)
For the non-content tokens, including our shared prefix, the formulas of V and A remain unchanged. Here, we introduce separate terms Vbd and Ali' for the shared prefix. Combining them, we have:
= Softmax
AA..AK....A AA) IVB
(12)
18

---

Published as a conference paper at ICLR 2025
D.3 RELATION WITH EQUATION 5
At last, we show that it can be rewritten as Equation with the only difference being that all contexts are treated as whole. For an token from the position j in context C, the final attention score "C,.. is
a
exp(Qk /TV
decd
E Els .TTdd
exp
S L-LoSSumExp
T.Va
)) 2 m. caiiin p
texp TeeppLooSSuBBxx
Q1kT monk?]
exp
LLogSuunExp
T.vd
T.V/
(13)
Exp(Qko.g//VV)
(Ene 1o exp(Q1 / /TV)))
exp(xx( [V E Ee /TV))s+ E /Va)
E C0
(14)
exp(Qk (VV) /M EE E
dc
+E /Va) EE
5
1=
(15)
This formula is equivalent to Equation 5 (i.e., query and generated tokens) for simplicity. For an token from position j. we can derive a" as
except it combines the prefix and other non-context tokens
exp(Qk] /Va)
Eis +1 --= - EN Ex ex /Va) EEN
aj
ΣN
(16)
Combining these two components, we obtain the final formula presented Equation
5
D.4 EFFICIENT IMPLEMENTATION.
To combine the computation for context and non-context tokens, employ flash attention twice- once for each part-and and then merge these results. This only introduces a marginal computational overhead.
def ape_ attention (query, key, value, temperature, scale):
split key and value states into context and non context parts key_ key_context, key_other= key value calue_context, value value_other other = value
logits_context, Ise_context flash_ (query, key, value, temperature) logits_other. lse_other flash_attn (query, key, value)
lse_context = ise_context scale
attn weights = [Ise_ cote_ccotetttttt Ise _eteorrhr]
attn_weights Softmax >
value_states
ingits_other
attn_output attn_ sitn_weights @ value_states
D.5 FUTURE DIRECTIONS.
The hierarchical formulation of APE naturally extend more complex tree structures, as illustrated in Figure 19 This flexibility allows each user query to be enriched with external knowledge organized in such structures, demonstrating the capability of APE in handling structured external data effectively.
E COMPARING APE WITH LONG-CONTEXT LLMS.
In 7 we further compare APE with Long-context LLM, including: (i) Prompt Compression: Truncation, L.MMiinua2 (Pan et al. 2024). (ii) KV Cache Eviction: StreamingLLM (Xiao et al. 2023), (iii) Long-conten FT: Llama-3 3-BB Instruct-2621 (Gradient) (024)). Hlama-2-7B-Instteee
32K (Together) 2023), (iv) length extrapolation: Self-Exte Self-Extend (Jin et al. 2024). Experimental results
showcase that APE consistently outperforms all existing long-context I.LM methods. We hypothesize
19

---

Published as a conference paper at ICLR 2025
Shared Prefix
Context 1
Context 2
Context 3
(a) Parallel Cache Structure
(b) Tree Cache Structure
Beyond the parallel cache structure in our paper, APE handle complex cache structures, where each context forms a tree-like hierarchy. In this setup, computations can be performed hierarchically along each branch, progressively merging intermediate results into the final value state.
Figure 19:
Table 7: Comparison between APE and long-conted LLMs on the LongBench (Bai et al.
2023
Method
NarratQA
Qasper
MultiFQA GovReport
QMSum
LCC
BLAMA-3-8BBIINTRRUUS
19.32
32.83
43.38
27.89
22.40
53.22
LLMLingua2
21.00
25.78
48.92
27.09
22.34
16.41
Streamingl.LM LLM
16.99
28.94
11.99
25.65
19.91
40.02
Long-context FT
14.88
21.70
47.79
32.65
24.76
55.12
Self-Extend
24.82
37.94
50.99
30.48
23.36
58.01
+APE
26.87
39.14
59.12
29.10
23.08
66.09
Method
RepoBench-P HotpotQA
2WikiMQA
MuSiQue
MultiNews
Average
38.15
44.24
21.01
20.47
23.63
31.50
LLMLingua2
20.56
40.16
24.72
20.85
21.34
26.29
StreamingLLMM
26.16
32.76
20.12
17.32
21.49
23.76
Long-context FT
43.05
15.89
10.49
8.74
24.28
27.21
Self-Extend +APE
41.83
51.09
24.17
28.73
24.11
35.96
49.43
50.11
28.06
25.79
22.40
38.11
that this improvement stems from APE enabling queries to access all past contexts, thereby enhancing LLM's retrieval ability. However, due to the limitations APE identifying relationships between contexts, do not emphasize its performance current long congecontext benchmarks in main paper.
20

---

Published as a conference paper at ICLR 2025
F
APE CACHE VERSUS PREFIX CACHE
Context
Context 2
Context 3
Context 4
Hit!
42% Hit Rate
Miss!
(a) Prefix Cache
Hit!
100% Hit Rate
Dispatch Cache
Hit!
(b) APE Cache
Figure 20: APE cache can keep a 100% hit rate while the Prefix cache only has a 42% hit rate.
Finally, we compare our APE cache with the prefix cache to demonstrate its advantages when serving multiple queries in the CAG setting. Figure 20 illlsttsatsss scenario with four contexts where both caching strategies are allocated identical budgets. Each query retrieves three contexts. Under these conditions, the prefix cache can only match limited number of combinations, resulting in average hit rate of 41.7%, whereas the APE cache consistently achieves a 100% hit rate. This performance gap becomes more pronounced as the number of contexts increases, resulting in many more combinations.
21


Reviews:
[{"review_id": "FBM0Djy1X4", "reviewer": "Reviewer_YT3n", "summary": "The paper presents Adaptive Parallel Encoding (APE) in order to improve efficiency and performance in language models handling RAG and ICL tasks. APE overcomes limitations in sequential encoding, which requires re-encoding context data and suffers from context window restrictions, by employing parallel encoding with adaptive adjustments. These changes include a shared prefix, scaling factors, and a modified attention temperature, aligning parallel encoding closely with sequential encoding. Experiments show APE achieves a high speedup in long contexts and outperforms slightly other methods in accuracy for long contexts.", "questions": "- Will we conclude anything between the task and the choice of hyperparameters?\n- The latency improvement in Table 2 would be more interesting if other baselines are included.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The method sounds intuitive. APE extends the usable context length for language models, overcoming limitations in the context window and enabling efficient handling of much larger inputs without additional training.\n- The performance is great esp. in latency\n- The analysis in Section 3 looks quite novel and interesting", "weaknesses": "- The performance seems quite sensitive on specific parameters", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ald5mFUlo0", "reviewer": "Reviewer_BcRN", "summary": "The papers introduce a simple and effective modification to parallel decoding which can be used a training-free drop-in method to substantially increase efficiency of inference while retaining most of its accuracy.", "questions": "* 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.\n* Are the tasks where APE fails. For example, Does APE still hold when used for Humaneval or RepoBench?\n* What would be the impact of finetuning with examples that had APE applied to then. Would it stop degradation?\n* Can you provide an analysis of the impact of changing the \"window\" size you use for parallel decoding.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "* The paper evaluates across many different models of different architectures. \n* The method is training-free which enables it be easily tested across many model -- assuming a working implementation.\n* Shows significant gains on efficiency and would have an real-world impact on deployed models. For example, enabling Pre-caching contexts.", "weaknesses": "* Paper does not provide implementation of the method -- this can make it easier for reproducibility. Would it be possible to provide that?\n* 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.\n* No complexity analysis written out. This could be particularly relevant if we want to vary window above.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "UrpIZzs7X9", "reviewer": "Reviewer_Yenp", "summary": "This paper investigates the performance degradation issue in parallel encoding and analyzes the distribution patterns of key and query vectors. This paper analysis reveals that key vector distributions remain relatively similar and query vector distributions can be combined. Based on these findings, they propose APE (Adaptive Parallel Encoding) to mitigate performance losses in LLMs during parallel encoding. Specifically, APE incorporates three key components: shared prompt prefix, position-aware dynamic scaling, and attention temperature adjustment. We evaluated APE on three LLMs (Llama-3-8B, Llama-2-7B, Gemma-2-9B) across multi-document QA and few-shot ICL tasks. Results demonstrate performance improvements compared to full attention and baseline methods, with latency benchmarks showing significant speedup over full attention.", "questions": "1. Formatting issue: In Table 1 #413 Gemma-2-9B line has miss alignment.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. The research focuses on a practical and significant challenge, particularly relevant to real-world LLM applications like RAG.\n2. The paper provides insightful observations and analysis regarding the feasibility of parallel encoding.", "weaknesses": "1. This paper demonstrates limited evidence for performance superiority over baselines across various tasks. Specifically:\n  - The evaluation is restricted to only two task categories from LongBench (multi-document QA and few-shot ICL), which is insufficient to demonstrate the method's effectiveness across diverse scenarios.\n  - 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. And Table 5 lacks crucial baseline comparisons, making it difficult to assess the relative performance improvements.\n  - The evaluation of other task types in Sec 6.2 is inadequate, missing essential baseline comparisons needed for meaningful performance assessment.\n\n2. And this paper lacks sufficient component analysis and validation.\n  - The paper lacks comprehensive ablation studies to isolate and validate the contribution of each proposed component. This makes it impossible to determine whether performance improvements stem from specific modules (e.g., attention temperature adjustment) or their combination.\n  - There is insufficient analytical evidence demonstrating how the three proposed components effectively address the challenges identified in the earlier sections of the paper. The causal relationship between the proposed solutions and the observed improvements needs stronger empirical support.\n  - The absence of detailed component-wise analysis makes it difficult to justify the necessity and effectiveness of each module in the proposed architecture.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "h4IT5htvAf", "reviewer": "Reviewer_onHd", "summary": "This paper introduces Adaptive Parallel Encoding, a method to improve the efficiency of large language models when processing multiple external contexts. APE pre-caches key-value states of contexts separately and enables position reuse during inference. The method consists of three key components: a shared prefix to align initial token distributions, a scaling factor to offset increased attention weights, and lower attention temperature to focus on semantically important tokens. The authors demonstrate that APE achieves a 976× speedup for long context generation while maintaining 93% accuracy compared to sequential encoding.", "questions": "1. How does the computational complexity of APE scale with increasing context length compared to sequential encoding?\n2. Why does the method perform particularly poorly on code completion tasks with continuous long contexts?\n3. How sensitive is the method to the choice of hyperparameters (scaling factor and attention temperature) across different models and tasks?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The authors provide thorough empirical analysis to understand the behavior of attention mechanisms and KV states. \n2. The method is practical as it requires no additional training and can be implemented with minimal modifications to existing architectures. \n3. Comprehensive evaluation across multiple tasks (ICL, RAG, long-context understanding)", "weaknesses": "1. 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. \n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WhULHPiMis", "reviewer": "Reviewer_kKET", "summary": "APE proposed the adaptive parallel encoding of LLM prompt, to speedup the prefill stage without limited by pretrained context length. APE add sink tokens (shared prefix) and split the context into chunks and perform attention inside of it. During decoding stage, it attends to every previous KV without adjusting the RoPE, so it will not exceed the pretrained context length.", "questions": "1. 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)\n\n2. 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.\n\n3. 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.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 2, "contribution": 3, "strengths": "- Nice empirical analysis to show value state can be merged & key states are similar for each context chunk.\n- Speeding up prefill stage drastically.", "weaknesses": "1. 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. \n\n2. The performance of the method is mostly proved by empirical results. I do not think these empirical results are weaknesses, but I have concerns about the lack of important comparisons with previous technologies to extend the context window and speed up prefill.\n\n- Lack of comparison with training-free context extension method (Self-Extend: https://github.com/datamllab/LongLM).\n\n- Lack of comparing with pretrained 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.\n\n- Lack of comparing with techniques to speed up the prefill stage.\n\n--- Some minor improvements ---\n\n- Figure 3 and 4 are quite hard to understand. Can you increase the font size? And also I do not think we need visualized every layers. Can you randomly sample some layers and show them only? (e.g., range(0, 32, 4)).\n- In Table 1, there is a typo in Gemma2. Maybe you shifted the columns to the left.\n- Tables 2, 4, and 5 might have a too large a size.\n- Table 2 should improve the formatting.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)