ReAttention: Training-Free Infinite Context with Finite Attention Scope
Authors: Xiaoran Liu
Ruixiao Li
Zhigeng Liu
Qipeng Guo
Yuerong Song
Kai Lv
Hang Yan
Linlin Li
Qun Liu
Xipeng Qiu
Venue: ICLR 2025 Poster (2025)
Keywords: long context
length extrapolation
large language model
natural language processing
Primary area: foundation or frontier models, including LLMs
TL;DR: We propose ReAttention, a training-free approach that enables LLM to support an infinite context in length extrapolation with finite attention scope.

Abstract: The long-context capability of the Large Language Models (LLM) has made significant breakthroughs, but \textit{the maximum supported context length in length extrapolation} remains a critical bottleneck limiting their practical applications. The constraint of context length in LLMs arises from the self-attention mechanism, which cannot effectively and efficiently capture the semantic relationships within infinitely long contexts via the limited pre-trained positional information and attention scope. In this work, we propose \textbf{ReAttention}, a training-free approach enabling LLM based on the self-attention mechanism to support an infinite context with a finite attention scope under sufficient memory resources. ReAttention performs the position-agnostic top-$k$ attention before the ordinary position-aware self-attention, freeing LLMs from the length extrapolation issue. We validate the performance of ReAttention on the LongBench, L-Eval, and InfiniteBench and demonstrate that it is on par with traditional methods. Furthermore, we also apply ReAttention on mainstream LLMs, including LLaMA3.1-8B and Mistral-v0.3-7B, enabling them to support context lengths of at least 1M and even expanding the context length of LLaMA3.2-3B-chat by 128$\times$ to 4M without any further training in Needle-In-A-Haystack tests. We also improve the efficiency of ReAttention with Triton and achieve an efficient extrapolation without additional overhead. The code is available at \url{https://github.com/OpenMOSS/ReAttention}.

Published as a conference paper at ICLR 2025
REATTENTION: TRAINING-FREE INFINITE CONTEXT WITH FINITE ATTENTION SCOPE
Xiaoran Liu 134: Ruixiao ;i Zhigeng Liu'; Qipeng Guo Yuerong Song Kai Ly Hang Yan', Linlin Li², Qun Liu", Xipeng Qiu "School of Computer Science, Fudan University, HHuawei Noah's Ark Lab, "Shanghai AI Lab, "Shanghai Innovation Institute xr rrliuu4 ém fudan. edu. cn, xpqi fudan. .eutaaneeddccm
ABSTRACT
The long-context capability of the Large Language Models (LLM) has made significant breakthroughs, but the maximum supported context length in length extrapolation remains critical bottleneck limiting their practical applications. The constraint of context length in LLMs arises from the self-attention mechanism, which cannot effectively and efficiently capture the semantic relationships within infinitely long contexts via the limited pre-trained positional information and attention scope. In this work, we propose ReAttention a training-free approach enabling LLM based on the self-attention mechanism to support infinite context with a finite attention scope under sufficient memory resources. ReAttention performs the position-agnostic top-k attention before the ordinary position-aware self-attention, freeing LLMs from the length extrapolation issue. We validate the performance of ReAttention on the LongBench, L-Eval, and InfiniteBench and demonstrate that it is on par with traditional methods. Furthermore, we also apply ReAttention on mainstream LLMs, including LLaMA3.1 1--M.-1-BBBBB and Mistral-v0.3-7BB enabling them to support context lengths of at least 1M and even expanding the context length of 1...MAA..------..
by 128x to 4M without any further training in Needle- An-A-Haystack tests. We also improve the efficiency of ReAttention with Triton and achieve efficient extrapolation without additional overhead. The code is available at https:: http:aawwb...ooooeeooeee//eeeeee
INTRODUCTION
Large Language Models (LLM) based on Transformer (Vaswani et al., 2017 OpenAl 2023 Reid et al. 2024 Sun et al.) 2024) has made great progress in the applications of Natural Language Processing (NLP). Particularly. in long context modeling, a considerable amount of research has been dedicated to extending the length capabilities of LLMs (Chen et al.) [0233] Peng al. 2023 Liu et al. 2023 Xiong et al. 2023). increasing the maximum supported context length from the initial 2K Touvron et al. (023) 2M (Ding al. 2024) in open-source LLMs. However, the maximum context length in length extrapolation remains a bottleneck that limits the practical applications of
LLMs (Press et al., 2022 Chen et al. 2023). To achieve infinite context length under sufficient
memory for Transformer-basee LLMs, the following three conditions must be satisfied:
a The position information in the inference phase should not be
(OOD)
compared to the training phase; (Han et al. 2023 Liu et al. 2023)
The self-attention entropy in the inference phase should not increase with the length of the input: (Han et al. 2023 Xiao et al., 2023)
с LLM should keep effectively aware of critical contextual information at each inference step. (Dong et al. 2024 Zhang et al. 2024a)
Early works in extrapolation focus on the first point, primarily by interpolating the position em- Xiong et al. 2023) or limiting the relative position
bedding (Chen et al. [003] bloc97 2023b
Equal contribution. This work is done during Xiaoran Liu's internship in Shanghai AT Lab

---

Published as a conference paper at ICLR 2025
within the pre-training context length (Su 2023 Jin et al. 2024 An al. 2024). However, later
studies have shown that even if positional information is not OOD, the attention entropy tends to increase logarithmically with the length of the self-attention window (Han et al. 2023). Specifically, the self-attention distribution becomes increasingly diffuse as inference length expands, making difficult to effectively gather information from the context, which leads to unstable model out-
puts (Peng et al.) [023 Han et al., (023). In response, methods like LM-Infinite (Han et al., (2023) and StreamingLLM, (Xiao et al. (023) have been proposed, which retain the starting and ending parts of the input, allowing LLMs to maintain stable outputs as input length increases. However, they compromise the global awareness of contextual information, harming the downstream performance.
For human beings, although we have long-term memory, we only need a small amount of information to think and act in real-time. Similarly, while an L.L. requires the full context to complete the inference process, it only needs limited contextual information at each inference step. This fact has inspired research like InfLLM (Xiao al. 2024a) and LongHeads (Lu et al., 2024) to perceive and extract critical information from the context before performing self-attention. However, LongHeads still face upper limits on positional information due to the use of NTK or PI for prefill (Lu et al. 2024), while InfLLM suffers from bias in representing contextual information, raised from the chunk representation and ReRoPE-style position embedding when extracting key information.
In this work, we introduce ReAttention, training free method that extends LLMs with finite context lengths to process infinite contexts. ReAttention regards extracting critical contextual information additional attention process preceding traditional self-attention, akin to "thinking before you act." At each reasoning step, ReAtteetion selects the most critical finite segments from the KV cache without position information, concatenates them, applies positional embedding, and performs self-attention. By controlling the length of the selected KV cache segments, ReAttention can achieve infinite context length with a finite attention scope, while ensuring the position information and atten- tion entropy are not out-of- of-distribution (Han et al., (023). Furthermore, inspired by the optimization techniques used in self-attentio Triton (Tillet et al. (019), GPU programming language, to minimize read and write overheads in top-k attention. With our custom Triton kernel, ReAttention avoids the extra computational overhead and reduces memory usage for long contexts. Our contributions are summarized as follows:
particularly Attentnn (Dao al. 2022 Dao 2023), we use
We outline three requirements for infinite context in Transformer- LLMs, position embedding not OOD, stable attention entropy, and effective contextual awareness. We also find that the last one can be satisfied via the attention score without positional embedding.
Based this observation, we propose training-free method, ReAttention, satisfying the aforementioned three conditions, thus extrapolating the context length of L.LMs to infinity with a finite attention scope, and making LLMs free from the issue of length extrapolation.
We validate that ReAttention matches the performance of traditional self-attention long contexts with computational overhead and less memory usage. Specifically, ReAttention extends the context length of leading LLMs, such as LLaMA3.1 1-MMA...--B-B1A 128K, to at least I million tokens. For smaller models like LLaMA3.2-3B--00 increased by 128x, reaching up to 4 million tokens, without additional training.
the context length can be
2
METHOD
The overall structure of ReAttention is illustrated Figure which consists of a A attention responsible for full-context cache selection and traditional self-attention transform with position embedding. ReAttention achieves a training-free integration between the two components.
top-
2.1 FULL-CONTEXT CACHE SELECTION
While LLMs require complete long context perform the entire inference process, only a limited context segment is needed at each inference step (Xiao et al.) 2023 2024a Lu et al. 2024). Consider- ing the beginning and end of the input context correspond to globally important prompts and local information for inference (iioo et al. 2023), ReAttention retains both the global and local segments of the KV cache for inclusion in the self- process.
Keache
1C Agloball Kmidile Kiocall
V = [V globaa V Vmiddle, V Plocat]
(1)
2

---

Published as a conference paper at ICLR 2025
Then, ReAttention uses the query vector of the current step to perform top-k selection on the middle part of the KV cache (Ribar et al., 2023), to identify the most important cache segments for the current step as shown in Figure
Output
Short-Term Memory
Indices top-k (q:
Self Attention
L-Layer
Kuelect miidll [Indices].
t
Vselect
V middle [Indices].
TopK-Attention without PE
ReAttention performs full-context selection on the KV cache in each layer, allowing differ- ent layers to choose different KV caches for the calculation. Furthermore, since each attention layer has multiple attention heads and ReAtten- tion adopts chunked streaming imput during the prefilling stage, multiple query vectors may ex- ist simultaneously. In this case, ReAttention
Long-Tern Term Memory
Cache)
Embedding
Figure I: Overview of ReAltention.
votes based on the top-k selections from different heads and query vectors to identify the top- KV caches. Additionally, ensure semantic coherence, ReAttention not only retains the top-k' elements themselves but also extracts m neighboring entries as a whole. Overlapping parts are deduplicated.
Unlike chunk-based selection in previous works (Lu al. 2024 Xiao et 2024a), ReAttention extracts critical contextual information with the dot-pree dot-product between q: and Keachee rather than chunk-wise representation of Kacche- This method offers strong adaptability and avoids the semantic fragmentation caused by fixed chunking (Luo et al., (024). Moreover, since the intermediate result of Kinddae is too large for long-con mong-cotext scenarios, and would incur huge read and write overhead, we follow the approach of FlashAttention (Dao al., [2022] Dao (202). fusing the whole process into one kernel using Triton (Tillet et al. (011)) as detailed in Section 3.4
2.2 TRAINING-FREE INTEGRATION
ReAttention concatenates the selected KV cache segments between the global and local parts, applies positional embedding sequentially, and preserves the relative order while ignoring the absolute distance between the selected segments, as shown in Figure | Self-attention can then be applied to the concatenated KV cache. The pseudocode of the whole process detailed in Appendix
Keache' [Keloohal, Keyllet, Kiocal).
V eache"
(Volobul, V. Vselect V
(3)
Ct, Kenche' PE (qt, Keacche),
SelfAttn
(qr. Kennhee' V Venche')
It's important to note that, unlike the conventional implementation in Huggingface Transformers (Wolf al. 2020), where the position embedding applied before KV caching, in ReAttention, the position embedding is separated from the KV cache and performed after the KV cache selection. That is to say, the cached KV does not include the positional information. This design offers several advantages. the attention score without positional embedding is more conducive locating the key information the context. On the other hand, since the concatenated cache length remains within the pre-training context length or the extrapolation upper bound, the position embedding will never be OOD
On the one hand, as mentioned in Section 4
(Han al. 2023).
Furthermore, the unselected KV cache segments are unimportant in the current inference step, as their self- scores are minimal (Zhang al.) 2024d). Therefore, this modification does not harm the self-attention output and can eliminates interference from irrelevant information (Zhang et al. 2024d, Ge et al. 2023). Compared with InfLLM (Xiao et al., 2024a). ReAttention's position embedding and attention transform within the pre-training window do not introduce the untrained attention pattern, thereby preventing bias accumulation in the KV cache during the prefilling stage. Thus, ReAttention disregard position information, achieve unlimited attention context with a limited attention scope without any training, and remain compatible with the existing attention acceleration methods (Dao et al. 2022 Dao 2023
3

---

Published as a conference paper at ICLR 2025
S-Doc
M-Doc
Sum
ICL
Syn.
Code
Avg.
LLaMA3-8B-8K
18.08
13.17
21.42
63.95
10.45
69.89
35.48
LL
15.14
8.90
18.36
58.19
7.28
70.27
32.49
InfLLM
12.43
8,93
13.20
66.08
7.55
70.03
32.51
ReAttention (Ours)
19.08
11.33
20.69
63.40
8,95
70.38
35.03
LLaMA3. 1-8B- 1-883118---8 328K
23.61
15.29
16.15
64.61
31.81
70.63
38.79
+ StreamingLLM
21.67
11,46
14.01
61.05
33.43
70.50
37.06
InfLLM
17.11
9.70
15.12
68.46
34.30
68.37
37.14
ReAttention (Ours)
23.97
15.22
15.66
64.54
31.20
70.48
38.63
-LaMA3.2-3B-1255.
34.45
26.24
24.46
59.29
40.11
55.73
40.76
StreamingLLM
28.05
19.41
22.67
55.23
22.63
54.16
35,08
InfLLM
17.47
6.77
17.44
49.11
23.63
48.09
28.21
RcAttention (Ours)
32.53
24.39
24.13
58.13
37.39
55.17
39.43
Table 1: Results of LLaMA Meries(Meta 2024a et al. 2023). ReAttention achieves consistent superiority over Streaming LaLLLLM(Xaa et al. 2023)
Dubey al. 2024
Meta
2024b)
LongBench Bai
and InfLLM Xiao al. 2024a) and shows comparable performance with LLMs with full attention.
32K
64K
128K
Avg.
MC
QA
Sum
MC
QA
Sum
MC
QA
Sum
LLaMA3-8B- 90-MA33---Instruct---
Instruct-8K
50.66
7.44
21.76
StreamingLLM
27.95
4.40
18.41
28.38
4.30
17.61
34.93
4.36
17.64
17.55
InfLLM
40.61
4.15
20.30
35.37
5.12
18.26
33.19
5.78
19.70
20.28
ReAttention (Ours)
46.29
5.33
21.76
44.54
5.25
20.15
49.78
5.98
20.36
24.38
LLaMA3.1 1-8B
BEEIAAA---S---stuut------ ---ntrucct-22R
37.12
18.85
1.60
31.88
23.32
1.39
21.83
25.43
1.40
18.09
StreamingLLM
34.50
6.82
19.61
34.93
6.83
20.82
35.81
6.71
20.10
20.68
InfLI.M
32.31
5.48
21.00
34.50
7.56
19.82
39.30
8.92
19.27
20.91
+ ReAAtention (Ours)
36.68
12.71
19.68
35.37
12.14
18.84
40.61
12.63
18.14
22.98
LLaMA3.2-33 SuuMMA3.-3B-Insttutt-222222 1---nxtrrut-128K
14.41
15.65
1.38
11.79
18.16
1.27
13.10
17.63
1.30
10.52
StreamingLLM
19.21
8.41
18.81
19.65
8.91
18.73
20.09
8.74
18.80
15.71
InfLLM
20.09
4.35
13.66
23.14
9.38
12.52
22.27
12.08
14.25
14.64
ReAAtention (Ours)
19.21
11.80
16.75
17.47
11.82
15.06
20.09
11.91
14.13
15.36
Table 2: Results of LLaMA Series on InfiniteBench (Zhang et al.
2024b) in different context lengths.
means LLM could not provide stable output in certain context length. ReAttention achieves superiority over StreamingLLM (Xiao et al. 2023), InfLLM (Xiao et al. 2024a) and full attention.
EXPERIMENT
3.1
SETUP
conduct experiments on LLaMA3-8B-SK Meta 2024a), LLaMA3.1-BB-123K (Dubey et al.)
2024). 1--...--..-BBB.... 128K- (Dbeey et al. 2024).
LaMA3.2-3B. 128K (Dubey et al. 2024). Mistral- v0.3-7B-333 32K (mistralai) 2024). Intern M2.5- B IM (nternLM 2024). Qwen2- 128K (Yang et al.) 2024a), Qwen2-72B T28K YYngg al. 2024a). Qwen2- 1B-32K Yang al. 2024a). For all models, we set the length of KCglobal to 32, the length of Kiacal to 4096, and selected span size to 32. Moreover, set 4, K'' 127 in top-k attention. Importantly, the attention scope in each step remains within the maximum attention window. For example, for the LLaMA3-888 3-B--8K with ReAttention, the maximum attention scope size 32 4096 + 127x32, which exactly matches the maximum supported attention window of 8192. We use OpenCompass (Contributors 2023b) for validation. All experiments are performed with FP16 precision and accelerated with FlashAttention) (Dao 2023).
3.2 LONG-CONTEN. BENCHMARK EVALUATION
We first evaluate all 9 LLMs on the commonly used long-context benchmark LongBench (Bai et al. 2023) and L--Eval (An et al. 2023), with default context length of 32K and middle truncation.
4

---

Published as a conference paper at ICLR 2025
(a) LLaMA3.1-8B-11
(b) LLaMA3.2-BB-1284
(c) Mistral-vO.3-7B--33
(d) M2.57-LM...-77
(e) Owen2-7B-1282
(f) Qwen2-1B-32K
Figure 2: Results of 8K and v0.-a-v-330---32 on Boedle-In-A-Haystaae (Contributors 2023a implemented
existing mainstream LLMs, including LLaMA3-8B-
OpenCompass (Contributors) 2023b).
For LLAMA3-8B-8K which has context length of less than 32K, we report its performance with Dynamic NTK (bloc97. 2023a). For the Dynamic NTK implementation, we use the default settings from the Huggingface Transformers (Wolf et al. 2020), setting the scaling factor to 4. Additionally, we compare the performance all LLMs with StreamingLLM (Xiao et al. 2023) as an ablation study, using the same global and local segment lengths as the ReAttention settings.
As shown Table and supplemented in Table 3 in Appendix ReAttention outperforms Streaming LLM across all 9 models, indicating that the full-context selection acquires useful in- formation for long-conten inference. Furthermore, ReAttention performs on par with full attention and even surpasses it in some cases, such LLaMA33. 1-7.--1---1111 128K (Meta 2024a) and Qwen2-1 1B- 32K (Yang et al. 2024a). This demonstrates that Atteention can be applied to LLMS various sizes and achieve performance close full attention on downstream tasks (Jiang et al. 2024).
To further demonstrate the superiority and extrapolation capability of ReAttention, we validate our method InfiniteBench (Zhang et al. 2024c), more challenging benchmark with longer context length. We choose 3 commonly tested subtasks, En.MC, En.QA and En.Sum, evaluate models with varying context lengths, and compare ReAttention with DynamicNTK (bloc97 (023a) and InfLLM (Xiao et al. 2024a) with the same selection config. The results are shown in Table 2 Remarkably, ReAttention consistently outperforms full-attention and InfL.LM at the 128K context length and in the average score. While DynamicNTK performs well at 32k, it encounters a clear upper bound on extrapolation, beyond which the model fails to produce stable outputs. Moreover, while InfLLM can extend context length indefinitely (Xiao et al. (024a). it still lags behind ReAttention in downstream tasks due to inaccurate extraction of critical information and differences in position embedding format compared to the pre- pre-training phase.
3.3 NEEDLE-IN-AL AAI-II----H-ASAAAA EVALUATION
Building upon the long- long-context benchmark evaluation, we use models with strong retrieval capabilities within their training context lengths and conduct the Needle- Boodll-In-A-Hasssee H-yystackk (NIAH) evaluation (Contributors. 2023a b). We perform experiments 8 A100 GPUs and extend the context lengths of L.LMS with ReAttention to least 1M tokens. As shown Figure 2 L.MMs with ReAttention maintain remarkably high retrieval accuracy across the entire range of context lengths they could support, regardless of their original attention windows. Importantly, we also extend the context of mainstream LLMs like LLaMA3-SB-88 (Metal 2024a) and Mistral-vO.3-7B -0.--7B-32KK (mistralai) 2024) to at least 1M, providing the community with an effective solution for deploying long-context LLMs.
Due to the extensive evaluation for all 9 models, and considering that InfLLM has limited support for larger models such 1.M...--70B-ch. and other series like Intern M2.5 and Qwen2, we only report results of the LLaMA series models the 8B and 3B scales.
5

---

Published as a conference paper at ICLR 2025
(a) Multi NIAH of LLaMA3. 1-8BB Instruct1228K
(b) Single NIAH of LLaMA3.2-3B
(c) Single NIAH of Qwen2-18
Instruct-1288 in 2M Context
Instruct-32K in 4M Context
Figure 3: Results of Caalti-Needle-In--1 --aasstack (Reid et al. 2024) and Single NIAH in a longer
context length implemented in OpenCompass (Contributors) 2023b
Building on this, we increase the difficulty of the evaluation by conducting experiments on multi- NIAH and single-NIAH in longer contexts, with results shown in Figure Given the increased diffi- culty, we employ the corresponding instruct versions of the models. On the mainstream LLaMA3.1- SB-Instruct-122 128K (Dubey et al., 2024). we achieve a context length of 1M in the multi-NIAH task. For smaller models, such as LLLMA3.2--3B--ll 128K (Meta 2024b), we extend context length to at least 2M. Notably, for Qwen2- IB-nstrucc--22. (Yang al. 2024a) to 4M on 4 A100 GPUs, achieving a training-free context expansion of 128x. To the best of our knowledge, it is the largest amplification of context length for LLMs achieved without additional training. This demonstrates that ReAttention can extend the context length while disregarding the position information, using finite attention window achieve infinite attention context.
extend the context length
3.4 EFFICIENCY ANALYSIS
In the Py PyTorch framework, operators execute independently, requiring frequent I/O the GPU memory, which introduces unnecessary overhead and latency. Inspired by FlashAttent tion (Dao et al. 2022 Dao 2023), develop GPU kernel function for the top-k attention described in Section 2.1 using Triton. As shown in Figure our kernel fuses operators for attention score computation and top-k calculation, enabling the entire process run within the GPU cache. Thus GPU mem- ory I/O is significantly reduced, improving both GPU memory usage and runtime. We added top-k attention that reduces the time for the self-attention, and our kernel keeps top-A: atten- tion's 's overhead minimal, leaving overall latency unchanged.
60-
TopK
Time (ms)
MatMul
20
Tensor
Fused Kernel
We analyze kemel''s efficiency in terms of GPU execution time and memory usage, showing that using Triton rather than Py PyTorch for the top-k attention greatly enhances GPU perfor- mance while reducing GPU memory consumption. Addition- ally, compare ReAttention with the standard implementation in the HuggingFacc Transformers (Dai et al.) 2019), measuring time first token (TTFT) and memory cost context lengths 32k to 256k. The TTFT and throughput comparison results are detailed in Appendix C system with 48-core CPU, 256GB RAM, and an A800-80GB GPU.
PyTorch
Custom Triten Kernel
Figure 4: Overview of the kernel fusion in our customized top-k at- tention kernel. The performance measurements reflect the execution time of the corresponding kernel functions, with the input length 8K for Llama3.1 inference tasks.
All experiments were conducted
Triton Operator Efficiency We tested our Triton operator using real inference inputs (c.g.. Q and matrices extracted from real-world inference tasks) to measure exccution time and device memory usage. As shown in Figure 5a| the standard PyTorcc implementation exceeds 80GB of memory for sequences over 64k, while our Triton operator minimizes memory usage, limited to input and output matrices. It also achieves hundreds of times faster performance 64k sequence lengths and scales efficiently as length increases (Figure 5b)
Memory Usage in Prefilling Stage
Memory overhead limits the prefilling for longer sequences. Figure shows that our method surpasses the standard implementation HuggingFace Transformers

---

Published as a conference paper at ICLR 2025
PyTonhh
GPU Memery Cost (GB)
60
1000
Trisn
150
800
100
40
Time
400
200
0
4k
32k
64k
128k
256k
4k
32k
64k
128k
256k
Context Length
Context Length
(a) GPU memory cost different context lengths.
(b) Latency across different context lengths.
Figure 5: The efficiency cooperation between our Triton kernel and its Py PyTorch version. Our fused top-k attention kernel performs better than PyTorch implementation across all context lengths.
Time To First Token (s)
FullAthh
GPU Memory Cost (GB)
Fullttn
100
ours
ours
75
50
40
25
0 32k 64k
20
128k
256k
32k
64k
128k
256k
Context Length
Context Length
Figure 6: First token (FFe) and GPU memory consumption in the prefilling stage. FullAttn refers to the official HuggingFace Transformers implementation of Llama3, 1-8B-B.1-BBBB 1-8--Basc, which runs out of memory after 192k.
as context length grows, continuing to work when the standard implementation runs out of memory. We also record first-token latency, which remains comparable to the standard im- plementation in Huggingface Transformers.
En.
Zh.
Code
Avg.
StreamingLLM
29.48
13.22
70.27
32.49
ReAttention
33.42
18.74
70.38
35.03
chunk size 1024
25.50
18.07
70.32
34.75
chunk size 2048
25.39
18.07
70.30
34.72
span size 8
13.02
8.74
58.54
22.96
DISCUSSION
span size 16
14.13
9.43
60.92
24.29
span size 64
25.27
17.98
70.83
34.69
4.1 ANALYSIS ON HYPER-PARAMETER
+ span size 128
25.23
16.97
70.54
34.57
+ top-k=
25.04
70.75
34.45
We first discuss the selection of hyper- parameters. We evaluate the performance of the LLaMA3-8B--K (Meta) 2024a) on the LongBench benchmark (Bai et al. 20023, com- paring
+ top-k=8
25.40
17.51
70.64
34.70
size 2048
24.51
16.77
70.39
34.08
local size = 1024
22.31
13.64
69.48
32.23
including chunk size, span size, local size, and top-k values. Since the prefilling chunk is always contained within the local part self-attention, the chunk size must be smaller than the local size. Addition- ally, to ensure faimess in our comparisons, we maintain consistent maximum attention window size across different settings, specifically keeping the summation of global size, span size times k' and local size equal to 8192. To achieve this, we set global size equal span size, retaining the earliest portion the first selection segment. As
Table 3: Analysis of ReAttention hyper-paramede on LLaMA3-8B-8 chunk size 512, span size 32 as well global size, local size 4096, and top-k Selection.
and LongBench, by default
7

---

Published as a conference paper at ICLR 2025
(a) Attention distribution (b) Attention distribution (c) Attention distribution (d) Attention distribution of a correct case with posi- of correct case without of wrong case with po- of a wrong case without tion embedding.
position embedding.
sition embedding.
position embedding.
Figure 7: Visualization of the attention distributions for LnterLLM2-77-2000 (Cai et al. 2024b) eval- uated on Beeell-In-A-Haystaee each heatmap represents the attention distribution of single layer, with the y-axis corresponding to the 32 attention heads and the x-axis corresponding to the 32K context length. The first 32 tokens, the instruction part, and the last 6 tokens, the generated part, are not excluded. The values represent the cumulative self-attention distributions for each token generation. The brighter the color, the higher the cumulative attention. A pooling operation with kemnel size of 100 is applied sequentially to show the pattern more clearly. The red box represents the position of the "needle" in the context.
(Contributors) 2023a b) in 32K context length. In each sub-figure,
span size increased, we correspondingly reduced the k'' to limit the upper bound of cache selection. Additionally, when local size decreases, k is increased. The final results are shown Table 3
We find that as the chunk size decreases from 2048 to 512, the performance gradually improves. The prefilling stage becomes more precise in extracting critical information during each top-k attention step. However, since prefilling takes more time chunk size decreases, we select 5 512 as our default value without further reduction. Next, we examine the impact different span sizes. Among these prperpparamttes, span size has the most significant effect downstream performance. When span size is relatively small, fragmented segments weaken coherence and may mislead the LLM's predictions, resulting in poorer performance compared to StreamingI.LMM When the span size reaches 32, the performance on LongBench peaks. Beyond this point, due to the constraint of pre-trained position embedding, the number of extracted segments decreases, hindering the effective capture of critical contextual information. We also compare the influence of the top-k number and find that top-4 yields the best results. A smaller top-k number makes it difficult for LLM to identify critical information, while a larger top-k number may introduce irrelevant content that interferes with judgmeen. Finally, we analyze the effect of local size. Smaller local size compromises LLM's ability to maintain semantic coherence, and excessive selection beyond the local context may skew the KV cache during the prefilling stage. Therefore, we set the local size half of the pre-training context length for LLLMA3-SB--- 8K and apply this hyper-pppaaae parameter configuration to other LLMs.
4.2 POSTION-AGNOSTICC CACHE SELECTION
In ReAttention, the full-context cache selection is based on the dot product between the query and key vectors without position embedding. This may create gap between the selection of the KV cache and the self-attention distribution. Therefore, here raises the question of whether using semantic vectors without position information can effectively locate the critical information in the context. To analyze this, we select a correct case and a wrong case from the results of Needle-In- In-A-Haystack evaluation (Contributors) 2023a b) on InternI.M2-1 M2-7B-200K (Cai al. 2024b) with full attention in
8

---

Published as a conference paper at ICLR 2025
32K context length. For each case, we calculate the attention distribution with position embedding, Le., the real self-attention distribution, and the attention distribution without position embedding.
The result is shown in Figure 7 For the correct case, the attention distributions, both with and without position embedding, can locate the position of the "needle" in the "haystack". However, the attention distribution with position embedding appears more dispersed, and this effect becomes more pronounced in the higher layers of the model. For the wrong case, the attention distribution with position embedding is also very dispersed and highly influenced by a particular attention head misled by irrelevant information from the lower layers. This ultimately results in the model's 's failure.
Interestingly, when observing the attention distribution without position embedding, we find that the model can identify the position of the "needle" using the inner product of semantic vectors and effectively filter out large amount of irrelevant noise. Practically, the ReAAtention-ennaanee L2ternL.M2-7B-2000 can successfully locate the "needle" in this case. Therefore, the inner product of semantic vectors in the cache selection is reasonable, and it proves to be more effective in finding the relevant context compared with position-aware inner products. Additionally, due to the absence of position embedding in the KV cache, ReAttention can achieve an effective context extension.
4.3 OBSERVATIONS ON SYNTHETIC TASKS
The above analysis has demonstrated ReAttention's good performance in various long-context scenar- ios, suggesting that it can extrapolate to infinite context lengths. However, these evaluations primarily focus on common natural contexts and lack currently widely-discussed benchmarks (Kuratov et al.) 2024 Li et al., (024)), such RULER (Hsieh et al.) 2024). Unfortunately, ReAttention performs poorly on the RULER benchmark. In fact, the extrapolation methods not based on full attention, including ReAttention and InfLLM, fail to pass the RULER benchmark. For instance. using the LLaMA3-8B-SK (Meta 2024a), we compare Dynamic NTK extrapolation with InfLLM (Xiao et al.)
2024 and ReAttention 8K and 16K context lengths. The results are shown in Table Only Dynamic NTK achieves effective extrapolation, while both ReAttention and InfL_MM exhibit more pronounced performance degradation the context length increases. Specifically, We focus on reporting the results of two subtasks from RULER: NIAH-Single3 and NIAH-Multikyy (Hsich) al. 2024). Both tasks involve extracting key information from a context that contains misleading items. The difference is that the context for Single3 consists of natural text, where the key information a string of random alphanumeric characters, whereas MultiKey3 features a stack of key-value pairs composed of those strings. We find that while all three methods can succeed in the Single3 test, only full-attention-based Dynamic NTK can effectively handle MultiKey3.
To uncover the mechanisms behind the observation, we perform a t-SNE visualization (Van der Maaten & Hinton 2008) of the K cache in LLaMA3-8B-8K using Dynamic NTK for both Single3 and MultiKey3, using the K cache in the last layer as example. The results are shown in Figure g For Single3, which primarily consists natural text, the reduced cache exhibits a single manifold
S3
MK3
All
MalSkey Lager31 CacherAtter RaFG) E-ME Rausht
Singhel Lager21 K Caslechfter RuPS) SNS Reah
8K Context Length
DynamieNTK
100.00
98.00
91.47
InfLLM
45.00
38.00
44.91
STRRE SURAA
ReAttention
95.00
78.00
90.10
16K Context Length
DynamieNTK 91.00
95.00
88.13
InfLLM 18.00
17.00
29.61
ReAttention 94.00
15.00
66.57
Figure 8: Visualization of the K Cache from the final layer of LLaMA3-8B-SK with Dynamic NTK extrapo- lation after inputs randomly sampled from the Single3 and MultiKey3 subsets in RULER. The visualization uses 2D t-SNE projection, with each token repre- sented as a point in the image and the input index shown via the color changing.
Table 4: The performance of Dynamic NTK, InfLLM, and ReAttention on RULER bench- mark in 8K and 16K context length. S3 and MK3 are the short forms of NIAH-Single3 and NIAH-Multikey respectively.
9

---

Published as a conference paper at ICLR 2025
with a local cluster and an overall coherent feature as the input length increases. In contrast, for MultiKey3, which is dominated by chaotic text (Lv et al. 2024). the reduced K cache reveals multiple overlapping manifolds rather than an extension along single curve. This results in the fact that during the K cache selection in Mul regardless of whether we use the same index encoding or direct sequential encoding for the selected segments, the subsequent self-attention fails to identify which manifold branch the segments belong to. Only by encoding all K caches according to their original position indices can LLM implicitly recognize (Chi et al. 2023 Kazemnejad et al. 2224) which manifold branch it resides on. This is why only full attention can effectively handle such tasks.
Nevertheless, the synthetic chaotic long texts are quite rare in practical scenarios (Lv et al. 2024). In the common natural long text benchmarks we have tested, ReAttention consistently demonstrates strong performance. Therefore, ReAttention still remains competitive for application in extending the context length of LLMs to infinity. We provide an in-depth discussion of the characteristics reflected in the QK sequence extension patterns between synthetic and natural texts in Appendix [E]
5 RELATED WORK
Length extrapolation is important issue for LLMs (Press et al. 2022), namely training in short context, and maintaining good performance in a longer context. The mainstream extrapolation research mainly focuses on adjusting the Rotary Position Embedding (RoPE) (Su et al. 2021). For example, Linear PI (Chen et al. 2023) first achieves the length extrapolation in LLMs by scaling the position indices to the pre pre-training range with little tine-tuning. The NTK method (bloc97 2023b a Peng et al.] (203)) then adjusts the rotary base in RoPE (Su et al.] [021] to achieve plug-and-pln length extrapolation. Subsequently, amplifying the rotary base and training on longer lengths has become the dominant approach for length extrapolation (Rozière et al. 2023 Xiong et al., 2023 Liu ct al. 2023 Ding et al. 2024), but these methods all have an explicit extrapolation upper bound. In addition, ReRoPE (Su 2023). Self-Extend (Jin ct al. 2024), and ChunkLLaMA (An et al.) 2024) also achieve plug-and anyy extrapolation by limiting the relative position. However, all of the above methods are based on full attention, facing the problem of attention entropy soaring with input length, and thus fail to achieve infinite context length (Han et al. 2023 Wang et al. (002)).
In contrast, another line of research has tried extend the context length of models through sparse attention. Considering the self-attention distribution tends to focus on the global and local context, StreamingI.LMM and IM- proposed the A-shaped attention window to achieve almost unlimited input length (Xiao et al.) 20233 Han et al. 2023). However, since the input is accompanied by discarding the previous context, still fails extend the context length (Dong et al. 2024). Based on StreamingLLM, InfLLM (Xiao et al. 2024a) and LongHeads (Lu et al. (022)) attempt to extend context through chunkwise retrieval from the middle cache, but the semantic fragmentation and block representation issues also affect downstream task performance (Luo et al.) (024). Recently, et al. (2024a) utilize a dynamic cache
MInference (Jiang et al.) 2024) and RetrievalAttentiorr selection to achieve significant speedup, but not attempt context extrapolation. In contrast, we propose ReAttention, which regards the cache selection as the preceding attention before the ordinary self-attention and achieves an infinite context with finite attention scope, making LLMs free from the challenge of length extrapolation. We also minimize the top-k attention kernel-wise like FlashAttention in self-attention (Dao et al. 2022 Dao 2023). Additionally, since ReAttention is conducted based on cache selection, we also expand our discussion on cache optimization, especially on token eviction methods like H2O (Zhang et al. 2024d) and SnapKV, in Appendix D
6 CONCLUSION
In this paper, we introduce ReAttention, using a finite attention window to realize an infinite context length in each inference step. We evaluate ReAttention is on par with full attention in performance with LongBench, L-Eval, and InfiniteBench. Furthermore, ReAttention has been shown to success- fully extend the context of mainstream LLMs, including the LLaMA series and Mistral, to up to 1M tokens and even expand the context length of LLaMA3.2-3B-che training in Needle- Mnnelle-I-A--BEe A-Haasstakck tests. We also improve the efficiency of ReAttention with Triton and achieve an efficient extrapolation without additional overhead.
by 128x to 4M without any further
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENT
This work is supported by the National Key Research and Development Program of China (No. U24B20181). We also appreciate the constructive comments from reviewers in the rebuttal and add discussion on methodology, effectiveness, efficiency, and related work in Appendix A to Appendix
REFERENCES
Chenxin An, Shansan Gong, Ming Zhong, Mukai Li, Jun Zhang, Lingpeng Kong, and Xipeng Qiu. L- eval: Instituting standardized evaluation for long context language models. CoRR, abs/2307. 1088. 2023. doi: 10-48550ARXV.2230..1111
0ttrg/10..855/rx/// URL httppss//oo...rrg/04445555/////0
2307 110110080 1088
Chenxin An, Fei Huang, Jun Zhang, Shansan Gong, Xipeng Qiu, Chang Zhou, and Lingpeng Kong. Training-frre long-context scaling of large language models. arXiv preprint arXiv:2402.17400 2024.
Sotiris Anagnostidis, Dario Pavllo, Luca Biggio, Lorenzo Noci, Aurelien Lucehi, and Thomas Hofmann. Dynamic context pruning for efficient and interpretable autoregressive transformers. Advances in Neural Information Processing Systems, 36, 2024.
Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, et al. Longbench: A bilingual, multitask benchmark for long context understanding. arXiv preprint arXiv:2308. 14508, 2023.
bloc97. Dynamically scaled rope further increases performance of long context llama with zero fine-tuning, July 2023a. htt.htts::/ //wp::::://///w/// wwddit http://wwweeattccoommc/////ooaao hawpammiiallllyallllleeee htg.://ccaalaaaeeoeeeoeeeeeeeeer
bloc97. Ntk-aware scaled rope allows Ilama models have extended (8k+) context size without any fine-tuning and minimal perplexity degradation. June 2023b. URL
https: /wwpw//wwww....mo/
www.c/oo/////eeeee LLaMA/
sca Laaaled_roe red_rope_aat alows____www___ ows_ amma_mmoose s_tt___hht_ haveeave/eee
Zefan Cai, Yichi Zhang, Bofei Gao, Yuliang Liu, Tianyu Liu, Keming Lu, Wayne Xiong, Yue Dong, Baobao Chang, Junjie Hu, et al. Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling. arXiv preprint arXiv:2406.0200 2024a.
Zheng Cai, Maosong Cao, Haojiong Chen, Kai Chen, Keyu Chen, Xin Chen, Xun Chen, Zehui Chen, Zhi Chen, Pei Chu, Xiaoyi Dong, Haodong Duan, Qi Fan, Zhaoye Fei, Yang Gao, Jiaye Ge. Chenya Gu, Yuzhe Gu, Tao Gui, Aijia Guo. Qipeng Guo, Conghui He, Yingfan Hu, Ting Huang. Tao Jiang, Penglong Jiao, Zhenjiang Jin, Zhikai Lei, Jiaxing Li, Jingwen Li, Linyang Li, Shuaibin Li, Wei Li, Yining Li, Hongwei Liu, Jiangning Liu, Jiawei Hong, Kaiwen Liu, Kuikun Liu, Xiaoran Liu, Chengqi Lv, Haijun Lv, Kai Lv, Li Ma, Runyuan Ma. Zerun Ma, Wenchang Ning, Linke Ouyang, Jiantao Qiu, Yuan Qu. Fukai Shang, Yunfan Shao, Demin Song, Zifan Song, Zhihao Sui, Peng Sun, Yu Sun, Huanze Tang, Bin Wang, Guoteng Wang, Jiaqi Wang, Jiayu Wang, Rui Wang, Yudong Wang, Ziyi Wang, Xingjian Wei, Qizhen Weng, Fan Wu, Yingtong Xiong, Chao Xu, Ruiliang Xu, Hang Yan, Yirong Yan, Xiaogui Yang, Haochen Ye, Huaiyuan Ying, Jia Yu, Jing Yu. Yuhang Zang, Chuyu Zhang, Li Zhang, Pan Zhang, Peng Zhang, Ruijie Zhang, Shuo Zhang, Songyang Zhang, Wenjian Zhang, Wenwei Zhang, Xingcheng Zhang, Xinyue Zhang, Hui Zhao, Qian Zhao, Xiaomeng Zhao, Fengzhe Zhou, Zaida Zhou, Jingming Zhuo, Yicheng Zou, Xipeng Qiu, Yu Qiao, and Dahua Lin. Internim2 technical report, 2024b.
Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian. Extending context window of large language models via positional interpolation. CoRR, abs/2306. 15595, 2023, doi: 10.48550/ ARXIV2230.115595 URL https::///// http::///dd....01....50.5//0/
Ta- Ta-Chung Chi, Ting- Han Fan, Li-Wei Chen, Alexander Rudnicky, and Peter J Ramadge. Latent positional information is in the self-attention variance of transformer language models without positional embeddings. arXiv preprint 2rXiv:2305.33577 2023.

---

Published as a conference paper at ICLR 2025
Contributors. Needle in haystack pressure testing LIMTest_ Mecleenltt 2023a.
OpenCompass Contributooss Opencompass: A universal evaluation platform for foundation models. wwtp/opppcccccssppepe/ 2023b.
https:://ithhue http://:hh.aaaacooooooomoaaoaaaa http://pen.cccmaaaaaas/ecccppapp
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime G Carbonell, Quoc Le, and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond fixed-length context. In Proceedings of the 57th Anmual Meeting of the Association for Computational Linguistics, pp. 2978-2988 2019.
Tri Dao. F Faster attention with better parallelism and work partitioning. CoRR, abs/2307.086911, 2023. doi: 10.4855/ARXIV.300.0006
URL https:://ci..o///////
48550/arXiv 08691
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. In Sanmi Koyejo, S. Mo- hamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems 35: Anmual Conference on Neural Information Process- ing Systems 2022, NeurlPS 2022, New Orleans, LA, USA, November 28 December poappppeeecppppeefeeeee/eppprrre 2022, 2022. URL ott.:///eeee..aeeeeeoooooaoroor- Conferencee html 6 67657c22600F00ar3322200....8.380
/
Yiran Ding, Li Lyna Zhang, Chengruidong Zhang, Yuanyuan Xu, Ning Shang, Jiahang Xu, Fan Yang, and Mao Yang. Longrope: Extending Ilm context window beyond million tokens. arXiv preprint 20Xiv:2402.13753, 2024.
Harry Dong, Xinyu Yang, Zhenyu Zhang, Zhangyang Wang, Yuejie Chi, and Beidi Chen. Get more with less: Synthesizing recurrence with kv cache compression for efficient Ilm inference. arXiv preprint 200i:::2002003998. 2024.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, al. The llama 3 herd of models. arXiv preprint arXiv:2407.21783 2024.
Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, and Jianfeng Gao. Model tells you what to discard: Adaptive kv cache compression for llms. arXiv preprint arXiv:2310.01801, 2023.
Suyu Ge, Xihui Lin, Yunan Zhang, Jiawei Han, and Hao Peng. A little goes long way: Efficient long context training and inference with partial contexts. arXiv preprint arXiv:2410.01445, 2024.
Zhiyu Guo, Hidetaka Kamigaito, and Taro Watanabe. Attention score is not all you need for token importance indicator in kv cache reduction: Value also matters. arXiv preprint arXiv: 12335, 2024.
Chi Han, Qifan Wang, Wenhan Xiong, Yu Chen, Heng Ji, and Sinong Wang. Lm-infinite: Simple on-the-fly length generalization for large language models. CoRR, abs/2308.16137, 2023. doi: 10. 48550/ARXIV.23008 1613377788.117. URL httpthtp: http::///ooo.o///..o.oo wtt.://..//....////r//....
16137
Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, and Boris Ginsburg. Ruler: What's the real context size of your long-context language models? arXiv preprint arXiv:2404.0660 2024.
InternLM. Internlm2.5-7b, July 2024. URL hctpa:///aogginofcccccoo//
internlm2_5-77
Huiqiang Jiang, Yucheng Li, Chengruidong Zhang, Qianhui Wu, Xufang Luo, Surin Ahn, Zhenhua Han, Amir H Abdi, Dongsheng Li, Chin-Yer Lin, et al. Minference 1.0: Accelerating pre-filling for long-context Ilms via dynamic sparse attention. arXiv preprint arXiv:2407.02499 2024.
Hongye Jin, Xiaotian Han, Jingfeng Yang, Zhimeng Jiang, Zirui Liu, Chia- Yuan Chang, Huiyuan Chen, and Xia Hu. Lim maybe longlm: Self-extend Ilm context window without tuning. arXiv preprint ar 2x11:2401.013223 2024.
12

---

Published as a conference paper at ICLR 2025
Amirhossein Kazemnejad, Inkit Padhi, Karthikeyan Natesan Ramamurthy, Payel Das, and Siva Reddy. The impact positional encoding on length generalization in transformers. Advances in Neural Information Processing Systems, 36, 2024.
Minsoo Kim, Kyuhong Shim, Jungwook Choi, and Simyung Chang. Infinipot: Infinite context processing on memory--constrained llms. arXiv preprint arXiv:2410001518 2024.
Yuri Kuratov, Aydar Bulatov, Petr Anokhin, Ivan Rodkin, Dmitry Sorokin, Artyom Sorokin, and
Mikhail Burtsev. Babilong: Testing the limits of Ilms with long context reasoning in-a-haystack. arXiv preprint arXiv:2406. 10149, 2024.
Tianle Li, Ge Zhang, Quy Duc Do, Xiang Yuc, and Wenhu Chen. Long-context Ilms struggle with long in-context learning. arXiv preprint arXiv:2404.02000 020600 2024.
Yuhong Li, Yinghing Huang, Bowen Yang. Bharat Venkitesh, Acyr Locatelli, Hanchen Ye, Tianle Cai, Patrick Lewis, and Deming Chen. Snapky: Llm knows what you looking for before generation. Advances in Neural Information Processing Systems, 37:22947-22970, 2025.
Di Liu, Meng Chen, Baotong Lu, Huiqiang Jiang, Zhenhua Han, Qianxi Zhang, Qi Chen, Chen- gruidong Zhang. Bailu Ding, Kai Zhang. al. Retrievalattention: Accelerating long-context llm inference via vector retrieval. arXiv preprint arXiv:2409. /6, 2024a.
Xiaoran Liu, Hang Yan, Shuo Zhang, Chenxin An, Xipeng Qiu, and Dahua Lin. Scaling laws of rope-based extrapolation. CoRR, abs/2310.05209,, 2023. doi: 1044550/ARXIV2330000000 URL
https: org/10 10200.485555 ht1.55//xxi...3130
Zichang Liu, Aditya Desai, Fangshuo Liao, Weitao Wang, Victor Xie, Zhaozhuo Xu, Anastasios Kyrillidis, and Anshumali Shrivastava. Scissorhands: Exploiting the persistence of importance hypothesis for Ilm kv cache compression test time. Advances in Neural Information Processing Systems, 36, 2024b.
Yi Lu, Xin Zhou, Wei He, Jun Zhao, Tao Ji, Tao Gui, Qi Zhang, and Xuanjing Huang. Longheads: Multi-head attention is secretly a long context processor. arXiv preprint vrXiv:2402.1066 2024.
Kun Luo, Zheng Liu, Shitao Xiao, and Kang Liu. Bge landmark embedding: A chunking- freee embedding method for retrieval augmented long-context large language models. arXiv preprint 2XXi:22022.1573. 2024.
Kai Lv, Xiaoran Liu, Qipeng Guo, Hang Yan, Conghui He, Xipeng Qiu, and Dahua Lin. Longwanjuan: Towards systematic measurement for long text quality. arXiv preprint 133533,402.13583, 2024.
Al Meta. Introducing meta Ilama 3: The most capable openly available Ilm date. Meta AI., 2024a.
AI Meta. Llama 3.2: Revolutionizing edge and vision with open, customizable models. Meta Al., 2024b.
mistralai. Mistral-7b-v0.3, May 2024. URL http:://bugghnffcc.cc//mteeete internim2_5 5-
OpenAl. Gpt-4 technical report. Technical report, 2023.
Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. Yarn: Efficient context window
extension of large language models. CoRR, abs/2309 000001,0000001,, 2023. doi: (0.48550/ARXV..2309. 00071. URL https: /tt.:p:://oi..oooo wtg///............ /t..z///...../////// 48550//5220 www.vv.iiii....//.via/.aaaaaaai. 2309. 00071
Ofir Press, Noah A. Smith, and Mike Lewis. Train short, test long: Attention with linear bi- ases enables input length extrapolation. In The Tenth International Conference on Learning Representutions, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022. URL httpss://ppenevv ew. net wwr.orrriideeee (oorummid--ss@ppGcC
Machel Reid, Nikolay Savinov, Denis Teplyashin, Dmitry Lepikhin, Timothy Lillicrap. baptisae Alayrac. Radu Soricut, Angeliki Lazaridou, Orhan Firat, Julian Schrittwieser, et al. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context. arXiv preprint arX5v:2403.05530 2024.
13

---

Published as a conference paper at ICLR 2025
Luka Ribar, Ivan Chelombiev, Luke Hudlass Hudlass-Galley, Charlie Blake, Carlo Luschi, and Douglas Orr. Sparq attention: Bandwidth-efficien IIm inference. arXiv preprint arXiv:2312.04985, 2023.
Baptiste Rozière, Jonas Gehring, Fabian Gloeckle, Sten Sootla, Itai Gat, Xiaoqing Ellen Tan, Yossi Adi, Jingyu Liu, Tal Remez, Jérémy Rapin, Artyom Kozhevnikov, Ivan Evtimov, Joanna Bitton, Manish Bhatt, Cristian Canton-Ferrer, Aaron Grattafiori, Wenhan Xiong, Alexandre Défossez, Jade Copet, Faisal Azhar, Hugo Touvron, Louis Martin, Nicolas Usunier, Thomas Scialom, and Gabriel Synnaeve. Code Ilama: Open foundation models for code. CoRR, abs/2308. 1 12950, 2023. doi: 10. 48550/ARXIV..308. 12950.
htttps://dee org/10. 48550 rXiv. 2308 12950
Zhenmei Shi, Yifei Ming, Xuan-Phi Nguyen, Yingyu Liang, and Shafiq Joty. Discovering the gems in early layers: Accelerating long-context Ilms with 1000x input token reduction. arXiv preprint arXiv:2409. 17422, 2024.
Jianlin Su. Rerope: Rectified rotary position embeddings, July 2023. URL https://gtgttttt
com/bo joojjon/reeog
Jianlin Su, Yu Lu, Shengfeng Pan, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with
rotary position embedding. CoRR, abs/2104.0985 098644 2021. URL /RLLitttp///aaax..orr///// 2104. 09864
Tianxiang Sun, Xiaotian Zhang, Zhengfu He, Peng Li, Qinyuan Cheng, Xiangyang Liu, Hang Yan, Yunfan Shao, Qiong Tang. Shiduo Zhang, Xingjian Zhao, Ke Chen, Yining Zheng, Zhejian Zhou, Ruixiao Li, Jun Zhan, Yunhua Zhou, Linyang Li, Xiaogui Yang, Lingling Wu, Zhangyue Yin, Xuanjing Huang, Yu-Gang Jiang, and Xipeng Qiu. Moss: An open conversational large language model. Machine Intelligence Research, 2024. ISSN 2731-5398, doi: 10.10007/111633-22--133 8.
URL https: //g1 thub. com/ ctw.://boccc///eem//////// /ttp///p//MMMS///////
Philippe Tillet, H. T. Kung, and David Cox. Triton: an intermediate language and compiler for tiled neural network computations. In Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages, MAPI. 2019, pp. 10-19, New York, NY, USA, 2019. Association for Computing Machinery. ISBN 9781450367196 196. doi: URL httpp://:// 0tt.://...111//33311 1we.11111331511 ..0.0083339997 3329973
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurélien Rodriguez, Armand
Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language models. CoRR, htts/200..13771 2023. doi: 1044550/ARXIV.2302.155 13971. URL https
7/doi org/10, 48 48550 2302 13971
Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine learning research, 9111), 2008
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need, In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4.9, 2017, Long Beach, CA, USA, pp. 5998 2017. URL htttpps://prcceeedngge neurips cc http//papr//21///aaal/// AAtaaaooSSscceeeeeaaaaatttaaaa.a Aast html
Jie Wang, Tao Ji, Yuanbin Wu, Hang Yan, Tao Gui, Qi Zhang, Xuanjing Huang, and Xiaoling Wang. Length generalization of causal transformers without position encoding. arXiv preprint arXiv:2404. 12224, 2024.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma. Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Transformers: State-of-the-an art natural language processing. In Proceedings of the 2020 Conference Empirical Methods in Natural Language Pro- cessing: System Demonstrations, pp. 38 45, Online, October 2020. Association for Computational Linguisties. _ wttptt:t////.... www.w//wwwwww///////////.../. wew. wew..orgg//nccooolgg/////0 emn em pep-dmmss... 6
14

---

Published as a conference paper at ICLR 2025
Chaojun Xiao, Pengle Zhang. Xu Han, Guangxuan Xiao, Yankai Lin, Zhengyan Zhang, Zhiyuan Liu, Song Han, and Maosong Sun. Inflm: Unveiling the intrinsic capacity of llms for understanding extremely long sequences with training free memory. arXiv preprint 1rX11:24022.4611,, 2024a.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. arXiv preprint arXi:2339.174553 2023.
Guangxuan Xiao, Jiaming Tang, Jingwei Zuo, Junxian Guo, Shang Yang, Haotian Tang, Yao Fu, and Song Han. Duoattention: Efficient long-context Ilm inference with retrieval and streaming heads. arXiv preprint ar wx0.::241..08880 2024b.
Wenhan Xiong, Jingyu Liu, Igor Molybog, Hejia Zhang. Prajjwal Bhargava, Rui Hou, Louis Martin, Rashi Rungta, Karthik Abinav Sankararaman, Barlas Oguz, et al. Effective long-context scaling of foundation models. arXiv preprint arXiv:2309. 2tt::2309.63999 2023.
An Yang, Baosong Yang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Zhou, Chengpeng Li, Chengyuan Li, Dayiheng Liu, Fei Huang, Guanting Dong, Haoran Wei, Huan Lin, Jialong Tang, Jialin Wang, Jian Yang. Jianhong Tu, Jianwei Zhang, Jianxin Ma, Jin Xu, Jingren Zhou, Jinze Bai, Jinzheng He, Jumyang Lin, Kai Dang, Keming Lu. Keqin Chen, Kexin Yang, Mei Li, Mingfeng Xue, Na Ni, Pei Zhang, Peng Wang, Ru Peng, Rui Men, Ruize Gao, Runji Lin, Shijie Wang. Shuai Bai, Sinan Tan, Tianhang Zhu, Tianhao Li, Tianyu Liu, Wenbin Ge, Xiaodong Deng, Xiaohuan Zhou, Xingzhang Ren, Xinyu Zhang, Xipin Wei, Xuancheng Ren, Yang Fan, Yang Yao, Yichang Zhang, Yu Wan, Yunfei Chu, Yuqiong Liu, Zeyu Cui, Zhenru Zhang, and Zhihao Fan. Qwen2 technical report. arXiv preprint arXiv:2407. 240.::2407..0671, 2024a.
Dongjie Yang, XiaoDong Han, Yan Gao, Yao Hu, Shilin Zhang, and Hai Zhao. Pyramidinfer: Pyramid kv cache compression for high-throughput 11m inference. arXiv preprint arXiv:2405. 12532, 2024b,
Amir Zandieh, Insu Han, Vahab Mirrokni, and Amin Karbasi. Subgen: Token generation in sublinear time and memory. arXiv preprint arXiv:2402.06082 2024.
Peitian Zhang, Zheng Liu, Shitao Xiao, Ninglu Shao, Qiwei Ye, and Zhicheng Dou. Soaring from 4k to 400k: Extending Ilm's context with activation beacon. arXiv preprint arXiv:2401. 0.3462. 2024a.
Xinrong Zhang, Yingfa Chen, Shengding Hu, Zihang Xu, Junhao Chen, Moo Hao, Xu Han, Zhen Thai, Shuo Wang, Zhiyuan Liu, et al. Infinitebench: Extending long context evaluation beyond 100k tokens. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 15262-15277, 2024b..
Xinrong Zhang, Yingfa Chen, Shengding Hu, Zihang Xu, Junhao Chen, Moo Khai Hao, Xu Han, Zhen Leng Thai, Shuo Wang, Zhiyuan Liu, et al. Inifinitehench: Extending long context evaluation beyond 100k tokens. arXiv preprint arXiv:2402.13718, 2024c.
Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, et al. H2o: Heavy-hitter oracle for efficient generative inference of large language models. Advances in Neural Information Processing Systems, 36, 2024d.
A PSEUDOCODE OF REATTENTION
In this section, present the pseudocode for ReAttention's 's prefilling and decoding phases. Our approach uses aosition-agnostic cache selection at each generation step, as shown in Section 2.1
B MORE VALIDATION ON EFFECTIVENESS
Table includes detailed comparison of various LLMs including LLaMA (Meta
2024a Dubey
et al. 2024; Meta 2024b). Mistral(mistrall) 2024), InternLM (Cai et al. 2024b InternL.) 2024).
and Qwen Yang et al. 2024a) series in Lench (Bai et al. 2023 and L-Eval (An et al. 2023
15

---

Published as a conference paper at ICLR 2025
Algorithm 1: Prefilling Phase
The input token indices imput ids, number tokens L, the number of starting tokens Iglohat, the number of ending tokens houal, size of prefilling chunk lchukk Initialize: 0, gloobal + hocals past_key_valuee None
Input:
while < do
> Process next chunk
-ey-luee None
input_chunk = inpputidss:: s:min(e, L)]
hidden states embedding higbbeddinnnnuuuu (input chunk)
& Forward through each transformer layer
for < to Mayer do > Project to Q. K. V query q-pro ((hiddn key = k. pro jhiddeenn states) value = v p states) > Update KV cache for this layer
10
11
12
14
_values then
key www.tcttttttttteeccccccccco jil key)
16
value = hetp:ccattppstt.e..yllllllll [illues(i].vauuee value)
17
past.key.valuesij] (key, value)
> Position-agmostic cache selection
19
score = query @ keyl/giobai--hecal
20
idx i topkssoee,,,,,, dim=0, k=k)
key = concat(
vest_key-values(1 (rryy/allesss].keyllll. keyyllgtou
22
past key (allues(j].kyylle fidx recall], past key (eeuess[j.kky[[-baaa
24
26
value value_recall =
27
mistkkeyvvllus[2 valuue past key key.kkuuslls]]]]]] v.luussi]..valuef value[:-/ocaal
28
recall],
29
30
31
& Compute self-attention with position embedding attention_outge rope_attent i (quueryy, key_recall, recul_) hidden_states = f En(attention.ouuput
33
34
Move the next chunk
35
e
36
e Lchunk
return hidden states, past key values
C MORE COMPARISON ON EFFICIENCY
Regarding the comparison between full attention and ReAttention in efficiency, the throughput and TTFT (Time To First Token) are two different measuring dimensions. TTFT measures performance under identical conditions (same batch size and context length). As shown Table 6 our TTFT results demonstrate that ReAttention achieves latency par with full attention equipped with FlashAttention2 under these controlled conditions. Using a single A800 GPU, we test our ReAttention approach and the standard implementation of Hugging Face's Llama3-8B (Meta) 2024b) modeling with FlashAttetioo22 (Dao 2023). We keep all the settings unchanged and set the batch size to
Throughput, the other hand, evaluates maximum processing capacity when fully utilizing available GPU memory. As shown in Table 7 ReAttention demonstrates higher efficiency by supporting larger batch sizes while maintaining comparable better token processing rates. Besides, beyond context length 64K, full attention equipped with FlashAttention2 (Dao 2023) encounters memory limitations, while ReAttention continues working. Using single A800 GPU, we test our ReAttention approach
16

---

Published as a conference paper at ICLR 2025
Algorithm 2:
Decoding Phase
Input: KV cache past key values, the last token index last token, number of tokens L, the number of starting tokens Iglobal, the number of ending tokens Hocal- hidden_states embedd iinidiiggllsst ng((lastttoken)
Forward through each transformer layer
for i < 0 to Mayer
do
> Project to Q. K. V
query q-pro f (hidden_states)
key k_pro (hidden states)
value = v pro proj(hidden) states)
Update KV cache for this layer
key = www.onnccatttttkeeeallllltaeaaae value woncat (pastt_tpppatkkkkk key -aalues[]l]..veee value, value)
key)
10
11
key values (key, value)
12
b Position-agnooticc cache selection
13
score = query @ keyl/ghoaa--hooaal idx_recall (score, dim=0, k=k)
15
key_recall = concat(
past key values[jj.key|l Ulgobai:
17
past salyy/vllusss]__e_ key fidx recall],
19
20
value_recall concat(
21
22
.esslkeevalles]]]] val.ll..llll.lllleeeeeee past key values vaabues -alae[:-hhall
23
24
25
& Compute self-attention with position embedding
26
attention attention_output = rope_attens ion(query, key_ recall, recall)
hidden_stane states f (fn(aatentionnoutpll
28 > Get the next token prediction
29 logits hm.head(hdddenn states)
30 next next_tokee token = ar gmax(l))
31 return nexe,token, past.key_vel values
and full attention with FlashAttention2 (Dao 2023 with LLaMA3-8B (Meta 2024b). We keep all the settings unchanged. The generation length is 1024. This performance pattern indicates not trade-off where we sacrifice one capability for another. Rather, it represents extending context length while maintaining competitive throughput and enabling larger batch sizes.
D MORE RELATED WORK
The evolution of token eviction techniques begins with DCP (Anagnostidis et al. 2024). Scis- sorhands (Liu et al. 2024b), and H2O (Zhang et al. (024d). These approaches demonstrate that re- taining only a small subset of tokens with high attention mance while significantly reducing memory and computational overhead. Then, Streaming LLMM (Xiao) al. 2023) advances this concept by identifying that critical tokens always occur the beginning and end part of the attention scope. By preserving only corresponding tokens, StreamingLLM enables stable model outputs after infinite input lengths, dramatically reducing computational and memory costs. However, this approach substantially compromised long-context performance (Li et al. 2025 Xiao et al. 2024b). Subsequent researches explore mitigation strategies, such as head-adaptive meth- ods like FastGen (Ge et al. 2023) and DuoAttention (Xiao et al.) 2024b). Specifically, SnapKV (Li et al. 2025) introduced more nuanced approach to pruning previous KV caches by querying the above context with later input segments, substantially enhancing long-contexs performance for token eviction techniques. Building upon these foundations, researchers such as PyramidKV (Cai et al.
could maintain near-equivalent perfor-
2024a), PyramidInfer (Yang et al. 2024b), Gemfilter (Shi al. 2024), and LongGen (Ge al.
17

---

Published as a conference paper at ICLR 2025
LongBench
L-Eval
Close- Ended
Open- Ended
S-Doe
M-Doc
Sum
ICL
Syn.
Code
Avg.
Avg.
LLaMA3-8B-88
18.08
13.17
21.42
63.95
10.45
69.89
35.48
46.11
2.79
17.23
StreamingLLM
15.14
8.90
18.36
58.19
7.28
70.27
32.49
36.52
1.36
13.08
+ ReAttention
19.08
11.33
20.69
63.40
8,95
70.38
35.03
44.94
2.43
16.60
LLaMA3.1
23.61
15.29
16.15
64.61
31.81
70.63
38.79
45.30
0.50
15.44
+ StreamingLLM
21.67
11.46
14.01
61.05
33.43
70.50
37.06
42.36
0.38
14.37
ReAttention
23.97
15.22
15.66
64.54
31.20
70.48
38.63
45.25
0.50
15.42
LLaMA3.2-33
34.45
26.24
24.46
59.29
40.11
55.73
40.76
55.75
20.63
32.33
StreamingLLM
28.05
19.41
22.67
55.23
22.63
54.16
35.08
43.27
18.66
26.87
ReAttention
32.53
24.39
24.13
58.13
37.39
55.17
39.43
53.09
20.12
31.11
LLMMA3.2-3B-1288
26.46
15.60
21.88
67.39
66.62
74.89
46.02
60.78
0.15
20.36
StreamingLLLM
23.72
12.14
20.66
65.35
56.36
75.06
43.20
55.32
0.16
18.54
ReAttention
27.23
15.38
21.67
67.58
66.35
75.26
46.15
61.11
0.18
20.49
(istrulls0.3-7B-5222
19.93
11.22
22.32
62.05
20.97
65.13
35.61
41.97
4.33
16.88
StreamingL.LM
17.60
9.79
20.68
59.32
19.42
64.85
34.01
36.49
4.45
15.13
+ ReAttention
20,20
10.64
22,41
61.57
20.58
65.32
35.48
41.95
4.09
16.71
Internl 52.5--R--.-.-.... IM
44.18
41.06
23.17
61.85
63.17
52.62
47.27
62.96
24.99
37.64
+ StreamingI.IMM
41.26
34.46
22.26
59.58
53.67
52.97
44.03
57.97
23.30
34.85
ReAttention
43.85
41.10
23.14
62.15
63.67
52.72
47.35
63.47
24.36
37.40
Qwen2-7B- 7B-K28K
37.75
14.82
26.16
62.81
29.00
67.65
41.34
53.43
15.41
28.08
StreamingI.IMM
35.46
11.22
25.11
61.46
23.00
67.39
39.16
50.07
14.30
26.22
ReAttention
37.64
14.39
26.24
62.90
27.00
67.56
41.01
53.18
15.35
27.96
Qwen2-72B-122
44.61
54.04
26.33
65.50
64.33
71.15
54.71
68.17
25.44
39.68
StrcamingLLM
40.11
40.24
25.29
64.48
54.83
70.91
50.08
60.43
22.38
35.06
ReAttention
44.96
53,20
26.18
66.00
62.67
71.39
54.53
67.31
25.34
39.33
Qwen2-1B-322
24.07
12.96
23.45
58.14
10.17
57.98
33.21
41.28
15.59
24.16
StreamingI.LMM
23.87
10.13
22.87
56.62
10.00
58.32
32.40
37.46
15.31
22.69
+ ReAttention 24.55
13.38
23.44
58.27
10.17
58.58
33.51
41.00
16.10
24.40
Table 5: Results of existing mainstream LLMs, including LLaMA Series and Mistral on LongBench and L-Eval.. ReAttention achieves a consistent superiority over StreamingLJ and shows comparable performance with LLMs with full attention.
ReAttention (ms)
Baseline (ms)
ReAttention
Baseline
token's
batch
token's
batch
32k
3984.66
3592.32
43.380
40.595
64k
10696.22
9823.82
64k
24.383
5
20.399
96k
20139.42
18942.57
96k
17.005
4
OOM
128k
32328.20
31953.82
128k
11.683
0OM
192k
60
64475.56
192k
7.2813
OOM
256k
108240.09
0OM
256k
5.9680
00M
Table 6: Comparison TTFT.
Table 7: Comparison on Throughput.
2024) take layer-wise optimization, maintaining inference effectiveness while improving memory efficiency further. Besides, works like VATP (Guo et al.) 2024) and InfiniPot (Kim et al., 2024) expand token pruning metrics beyond traditional QK dot product scoring. Our ReAttention draws inspiration from these token eviction methods, specifically applying the principle of selecting KV cache to context extrapolation, enabling comprehensive context perception throughout the inference process while computationally addressing each attention calculation with a finite context, thereby achieving algorithmically infinite context length.
18

---

Published as a conference paper at ICLR 2025
E ANALYSIS OF CACHE FEATURE
E. THE METHOO FOR OBTAINING T-SNE RESULTS
We begin by extracting a lengthy sample (approximately 16K tokens) of text and feeding it into the LLaMA3-8B-88 (Meta 2024a). Subsequently, we capture the Q K. states at each layer, both before and after position embedding, allowing for clearer understanding of the impact of position embedding on data representation. Next, the captured matrices are reshaped to (seqlen, nhead headdim), and we apply 1-SNE (Van der Maaten & Hinton 2008) for dimension reduction, mapping the nigh-dimensional data into two dimensions with a final shape of (seqlen, 2).
E.2 CACHE FEATURE OF LONG NATURAL LANGUAGE TEXTS
(a) t-SNE results Q, from data sample from the PG19 training set in the final layer of LLMMA3-8B-8K
(b) t-SNE results of Q, FC from a data sample from PG19 training set in the 15th layer of LLaMA3-8B-8K
Figure 9: The 2-D reduction results of Q, K. from the final/15th layer, both before and after position embedding. Each figure is divided into four sections, representing four dimensionality reduction results. The left side corresponds to Q. while the right side corresponds to K. The top displays results before position embedding, and the bottom shows results after position embeddingg. Each point represents token. The colorbar indicates token indices(positions) with purple representing the beginning of the sentence and yellow representing the end.
BO
(a) L-SNE results of Q. K. (b) I-SNE results of Q K. (c) I-SNE results of Q K. (d) i-SNE results of Q, AC from data sample from from a data sample from from a data sample from from data sample from LongBench GovReport the 5th layer of LLaMA3- the 15th layer of LLaMA3- 8B-8 8K with NTK
LongBench GovReport
LongBench GovReport in LongBench GovReport in the 25th layer of LLaMA3- the final layer of LLaMA3- 8B-8K with NTK
8B-8K with NTK
8B-8K with NTK
Figure 10: Given the smaller number of tokens of this sample, the clusters are distinctly separated by the algorithm. Despite some dispersion in the intermediate layers, the data representation maintains alignment along a linear manifold in the higher layers, indicating the robustness of the feature of Global Coherence.
19

---

Published as a conference paper at ICLR 2025
The Q, K. states of long natural language text preserves coherence, which can be visualized curve extending through high-dimensional space, encompassing the entire text. When the model extracts contextual information. it effectively accesses specific points along this line. The visualization in Figure and Figure
T0 reveals several key features:
Local Semantic Similarity: Locally, the semantic meaning of individual text segments is preserved, forming small clusters, (Zandich et 2024)
Global Coherence: These local clusters, however, not isolated but part of larger linear manifold representing the text's global coherence. This coherence becomes more apparent after position embedding is applied. (Zandich al. (024)
Out-of-Distributin Position Embedding: When the position embedding exceeds 8K, Llama3-8B struggles to extract semantic information, as seen in the diminished clus- when the position embedding is out-of- distributionnOODD) the text representation continues to align with a linear manifold.
tering. (Liu al. 2023) Nonetheless,
E.3 FROM NATURAL LANGUAGE TEXTS TO CHAATIC TEXTS
In contrast, Ruler's MultiKey3 subset contains unnaturally long texts that deviate from typical linguistic patterns, lacking the structure and semantic coherence of natural language. Here in this subset, the text length is artificially extended to 16k tokens, not by model generation but by repeating specific patterns
One of the special magic UUIDs for A is: B.
like
One of the special magic UUIDs for elB8d5ie-b22-00ffaa422200101133- is: 1044c0-411--a5ed-7777
Moreover, large portions of these patterns consist of non- non-linguistic, non-logical elements like UUIDs.
When applying t-SNE reduce the dimension of MultiKey3 data(Figure) TTp the linear manifold observed in natural language text vanishes, highlighting the impact of unnatural patterns on data representation. Key observations include:
Local Semantie Similarity: By examining Q after position embedding, we can clearly see that the UUIDs and natural language components are completely separated, even in the deep purple section at the start. IC after position embedding, adjacent tokens (i.c., data points of similar colors) are dispersed across multiple lines, thus making the term 'local similarity" questionable.
Global Coherence: The assumption that the entire data representation lies on a linear manifold no longer holds. Unnaturally extended texts, unlike typical natural language, fail to maintain global structural dependencies.
The effectiveness dynamic sparse methods in replacing the full attention mechanism hinges on the collaboration between global and local information. However, when text coherence is lacking, even adjacent tokens in the Q, IC cache become dispersed, disrupting the required linear manifold structure that LLM relies on. This suggests that L.LM focuses its top-k attention natural language segments and successfully extracts critical information. Yet, when the top-k values are spread across multiple manifolds. LLM fails utilize this information efficiently. This contrasts sharply with the typical feature of the Q, K. cache in natural language, where the representation forms continuous, intertwined structure across the text in high dimensional space. It is precisely the absence of this characteristic feature that causes the cache generated during prefilling with ReAttention to be biased when processing RULER Multikey3, leading to the output non-existent UUIDs.
20

---

Published as a conference paper at ICLR 2025
(a) t-SNE results of O, K. from sample in RULER Multikey3 in the 15th layer of LLaMA3-8B-SK with NTK
(b) t-SNE results of O K. from sample in RULER Multikey3 in the final layer of LLaMA3-8B-8KK with NTK
(c) I-SNE results Q, IC from a sample in RULER Multikey3 in the th layer of LLaMA3-8B-8K
(d) I-SNE results Q, from a sample RULER Multikey3 in the final layer of LLaMA3-8B-8K
(c) I-SNE results of Q, IC from a sample in RULER Multikey3 in the 15th layer of Qwen2-7B-128K
( I-SNE results of Q, K. from sample in RULER Multikey3 in the final layer of Qwen2-7B-128K
Figure 11: t-SNE visualization results different models under RULER Multikey3,
21


Reviews:
[{"review_id": "baZXAAKhRE", "reviewer": "Reviewer_xrHG", "summary": "This paper proposes ReAttention, a a training-free approach enabling LLM based on the self-attention mechanism to support an infinite context with a finite attention scope under sufficient memory resources. ReAttention splits the tokens into three parts: initial tokens, middle tokens, and local tokens. It maintains the initial and local tokens (their KV caches), select a subset of middle tokens, and re-concatenates them with positional index calibration, and then conduct attention operation.", "questions": "1.\tThe idea appears similar to token eviction methods like H2O and FastGen. How does ReAttention compare with these methods in terms of accuracy and efficiency?", "limitations": "", "rating": 5, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 1, "strengths": "1.\tThe method is intuitive. ReAttention maintains KV caches of tokens that are more frequently attended by the latest tokens, which reduces the memory overheads and remove redundant contextual segments from being attended. \n2.\tThe evaluation is comprehensive. Authors provide evaluation for many models and benchmarks. \n3.\tThe performance improvement is notable, compared to InfLLM and full attention.", "weaknesses": "1.\tThe two-round attention operation introduces non-trivial overheads, which may result in latency similar to full attention. \n2.\tThe efficiency in terms of latency and throughput can be the major drawbacks of this method. \n3.\tThe method requires sufficient memory resources where full attention is also feasible. The advantage over full attention lies in evicting irrelevant tokens, similar to token-dropping methods, which improves performance. However, efficiency in long-sequence tasks remains a concern. When full attention isn't feasible, the two-round attention calculation might be costly, as the first round still requires full-attention calculations. \n4.\tIn long-sequence generation, the query tokens can change dynamically. How does ReAttention handle middle-cache selection in this setting? Does it re-select middle caches periodically?", "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": "J2NNSbjcXt", "reviewer": "Reviewer_iiPf", "summary": "This paper proposes ReAttention, a training-free approach to extend the context window of Large Language Models (LLMs) through a position-agnostic top-k attention mechanism. The key idea is to perform cache selection before regular self-attention, treating the extraction of critical context info as an additional attention process. The method enables models to handle theoretically infinite contexts with a finite attention scope. The authors validate their approach on multiple benchmarks and models, showing strong performance and efficiency gains through custom Triton kernel optimization.", "questions": "How does ReAttention perform on more structured tasks that require precise positional understanding (e.g., code completion, mathematical reasoning)?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The problem is highly relevant and timely, as context length remains a critical bottleneck for LLMs. The paper clearly articulates three key conditions for infinite context extension (lines 040-046).\n- The proposed method is simple and training-free. The two-stage attention approach (position-agnostic selection followed by regular attention) is intuitive and well-motivated. \n- The empirical results are impressive. They demonstrate: (1) Comparable or better performance vs full attention across multiple benchmarks (Fig 2); (2) Successful extension of LLaMA3.2-3B context by 128x to 4M tokens (Section 3.3); (3) Strong performance on challenging tasks like LongBench, L-Eval, and InfiniteBench.\n- The experimental evaluation is comprehensive, including multiple model architectures (LLaMA, Mistral, Qwen series) and various benchmarks (LongBench, L-Eval, InfiniteBench). It also includes detailed ablation studies on hyperparameters (Table 2) and thorough efficiency analysis with Triton optimization in Section 3.4.", "weaknesses": "- The paper has several writing issues and typos that should be addressed: e.g., Llama vs LLaMA and the wrong quotation in line 067.\n- The related work section omits several relevant recent papers on efficient attention and KV-cache optimization, such as:\nSnapKV [1] for efficient cache management, PyramidKV [2,3] for hierarchical cache structures, GemFilter [4] for attention-filtering, and many so on. These works tackle similar challenges and a comparison would strengthen the paper. It would be good to discuss them. \n\n[1] Li, Yuhong, Yingbing Huang, Bowen Yang, Bharat Venkitesh, Acyr Locatelli, Hanchen Ye, Tianle Cai, Patrick Lewis, and Deming Chen. \"Snapkv: Llm knows what you are looking for before generation.\" arXiv preprint arXiv:2404.14469 (2024).\n\n[2] Zhang, Yichi, Bofei Gao, Tianyu Liu, Keming Lu, Wayne Xiong, Yue Dong, Baobao Chang, Junjie Hu, and Wen Xiao. \"PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling.\" arXiv preprint arXiv:2406.02069 (2024).\n\n[3] Yang, Dongjie, XiaoDong Han, Yan Gao, Yao Hu, Shilin Zhang, and Hai Zhao. \"PyramidInfer: Pyramid KV Cache Compression for High-throughput LLM Inference.\" arXiv preprint arXiv:2405.12532 (2024).\n\n[4] Shi, Zhenmei, Yifei Ming, Xuan-Phi Nguyen, Yingyu Liang, and Shafiq Joty. \"Discovering the gems in early layers: Accelerating long-context llms with 1000x input token reduction.\" arXiv preprint arXiv:2409.17422 (2024).", "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": "a1IvpycDBd", "reviewer": "Reviewer_fRQF", "summary": "This paper proposes a training-free method (ReAttention) for improving the ability of pre-trained language models to process long context lengths. The method is motivated by the insight that attention scores without positional embeddings are still useful for identifying salient tokens.  \n\nThe main claims of the paper are that:\n\n1. ReAttention matches or outperforms the performance of traditional self-attention in long-context scenarios. \n2. ReAttention matches or outperforms the performance of other training-free methods long-context methods (*e.g.* InfLLM). \n3. ReAttention extends the context length of existing LLMs (*e.g.* 8B and 3B llama) “to at least 1 million tokens”. \n4. ReAttention adds no computational overhead and less memory usage than regular self-attention.", "questions": "- *What is the performance of the baselines in the needle-in-a-haystack experiments.*\n- *In Figure 7, is the FullAttn* implementation FlashAttention?\n- *How exactly does ReAttention differ from InfLLM? Is the only difference that position embeddings are not applied during the lookup step?* (I recommend the authors update the paper to make this comparison more clear.)\n- *Minor improvements:*\n    - In Figures 3 and 4, the axes and color bars are too small to be legible.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "*The method is motivated by an interesting observation.* “*We also find that the third condition can be satisfied through the attention score without positional embedding.” This insight that attention scores without positional embeddings can effectively identify salient tokens is quite interesting and, to my knowledge, novel.  \n\n*Method works on pre-trained models.* Unlike approaches like Mamba which require training models from scratch with new architectures, ReAttention can be applied directly to existing pre-trained models. This makes it much more practical for immediate deployment and use with current state-of-the-art models.\n\n*Claim 1 seems to be well-supported by the results.*   The empirical evaluation in Figure 2 and Table 1  across multiple benchmarks (LongBench and InfiniteBench) and model architectures shows that ReAttention matches or outperforms regular attention.\n- Note: in the weaknesses, I highlight some missing baselines.", "weaknesses": "*Baselines in the long context evals are missing (Claim 2). W*hy is StreamingLLM included in the LongBench results but not the InfiniteBench results? And why is InfLLM included in the InffiniteBench results but not the LongBench results?\n\n*Baselines in million-scale context seem to be missing. (Claim 3)* For the needle in a haystack experiments supporting Claim 3, the paper should compare against other methods capable of handling very long contexts like InfLLM and other recent approaches. Without these comparisons, it's hard to assess the relative advantages of ReAttention at extreme context lengths.\n\n*Missing token-throughput comparisons (Claim 3).* The paper only compares time to first token (*i.e.* the cost of prefill).  How does it fare against regular attention in terms of efficiency?\n\n*Apparent tradeoff in memory consumption (Claim 4).* The authors state that ReAttention uses less memory than self-attention. However, based on Figure 7, it appears that this is only true after ~100k in context length. The authors should highlight this tradeoff. \n\n*Method description lacks details.* It’s unclear exactly how the method works from the description in Section 2. For example, what does it mean that ReAttention “adopts chunked streaming input during the prefilling stage.” I recommend that the authors provide a detailed description of how prefill with ReAttention differs from decoding. It would be great if the authors included a pseudo-code description of both the prefill and decoding stages (either in the main paper or the appendix). \n\n*The \"infinite\" in the title is misleading.* It’s not clear from the experiments that ReAttention actually supports infinite context. Further, there are still practical limitations from memory constraints that would prevent the model from really achieving infinite context length. The paper should be more precise about what \"infinite\" means in this context and acknowledge the real-world constraints.", "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": "VlweW3bQ6X", "reviewer": "Reviewer_bDJC", "summary": "The paper introduces ReAttention, a training-free approach for enabling large language models (LLMs) to handle infinite context lengths with a finite attention scope. The method involves a position-agnostic top-k attention mechanism, which precedes traditional self-attention, to select the most relevant context segments. This approach addresses the limitations of context length in LLMs without the need for additional training. The authors validate ReAttention across multiple benchmarks, demonstrating its effectiveness and efficiency compared to existing methods.", "questions": "see weaknesses", "limitations": "", "rating": 5, "confidence": 5, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The use of Triton to optimize the attention mechanism reduces computational overhead and memory usage.\n2. The method supports significant context length extensions, showcasing its potential for real-world applications.", "weaknesses": "1. The paper uses the top-K method to select KV, but this has been studied in H2O. It feels like the paper is a combination of StreamingLLM+H2O, and then rewrites Triton to make up for the calculation time of QK. I don't see any novelty. It would be helpful to provide more specific details on how ReAttention differs from or combines elements of StreamingLLM and H2O. \n\n2. The rewrite of the Triton operator is indeed very effective, but it is more of an engineering contribution.\n\n3. The paper is mainly an improvement on KV. There are so many optimizations on KV, but there are so few baselines added for comparison. It is recommended to select some baselines from SnapKV, PyramidKV, LongGen, InfiniPot, DuoAttention and other methods for comparison.\n\n4. Many experiments in the paper show good results, but there is no intuitive or formal explanation for why this method can achieve such good results. For example, theoretical analysis, or visualizations that might help explain the performance gains.", "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)