Title: Mesa-Extrapolation: A Weave Position Encoding Method for Enhanced Extrapolation in LLMs
Venue: NeurIPS 2024
Accepted: yes
Keywords: Large language models

extrapolation

position encoding

transformers

self-attention
Abstract: Large language models (LLMs), although having revolutionized many fields, still suffer from the challenging extrapolation problem, where the inference ability of LLMs sharply declines beyond their max training lengths. In this work, we conduct a theoretical analysis to better understand why No Position Encoding (NoPE) fails outside its effective range, as well as examining the power of Position Encoding (PE) in this context. Our findings reveal that with meticulous weave position, PE can indeed be extended beyond effective range. Our theorems establish that LLMs equipped with weave PE can achieve improved extrapolation performance without additional cost. Furthermore, we introduce a novel weave PE method, Mesa-Extrapolation, which utilizes a chunk-based triangular attention matrix and applies Stair PE to manage the final chunk. This method not only retains competitive performance but also offers substantial benefits such as significantly reduced memory demand and faster inference speed. Extensive experiments validate the effectiveness of Mesa-Extrapolation, demonstrating its potential as a scalable solution to enhancing LLMs’ applicative reach.

Mesa-Extrapolation: A Weave Position Encoding Method for Enhanced Extrapolation in LLMs

Abstract

Large language models (LLMs), although having revolutionized many fields, still suffer from the challenging extrapolation problem, where the inference ability of LLMs sharply declines beyond their max training lengths. In this work, we conduct a theoretical analysis to better understand why No Position Encoding (NoPE) fails outside its effective range, as well as examining the power of Position Encoding (PE) in this context. Our findings reveal that with meticulous weave position, PE can indeed be extended beyond effective range. Our theorems establish that LLMs equipped with weave PE can achieve improved extrapolation performance without additional cost. Furthermore, we introduce a novel weave PE method, MesaExtrapolation , which utilizes a chunk-based triangular attention matrix and applies Stair PE to manage the final chunk. This method not only retains competitive performance but also offers substantial benefits such as significantly reduced memory demand and faster inference speed. Extensive experiments validate the effectiveness of Mesa-Extrapolation, demonstrating its potential as a scalable solution to enhancing LLMs' applicative reach. Our code is available at https: //github.com/soacker/Mesa-Extrapolation .

1 Introduction

Large Language Models (LLMs) with their powerful in-context learning capabilities Brown et al. (2020) offer versatile solutions to a wide-range of intelligent applications. However, one pressing challenge, the extrapolation problem Delétang et al. (2022) Zhang et al. (2022), dictates that the inference ability of LLMs sharply declines beyond their max training lengths, imposing a serious limitation on applications with long inputs. An naive solution is to extend the length of training samples. However, the inherent quadratic complexity of calculations presents practical challenges, demanding more resources, longer training time and higher cost.

Positional encoding (PE) has become a pivotal component of Transformer architecture, to compensate for the overlooking of position information by the attention mechanism. In the realm of extrapolation capability, PE is considered as a key factor influencing the extrapolating ability of LLMs. A few PE approaches, such as the popular RoPE Su et al. (2023) and ALiBi Press et al. (2021), claim to offer improved extrapolation capabilities and have gained widespread usage in industrial applications. Meanwhile, a counter-narrative has emerged. Some works demonstrate that Transformer can achieve better extrapolation capabilities by removing position encoding (NoPE) Kazemnejad et al. (2023), and contend that the mask already plays a significant role in capturing position information.

Inspired by Kazemnejad et al. (2023), we conduct a thorough investigation of the extrapolation problem by designing a specific Transformer model for this purpose and presenting detailed theoretical analysis. To the best of our knowledge, this is the first theoretical endeavor to understand the inner workings of extrapolation. We first elucidate the cause of NoPE's failure when input exceeds the effective window length in Theorem 3.1, and analyse PE's failure in 3.2. Building upon Theorem 3.2, we prove in Theorem 3.3 that through meticulous weave position, (i.e., weave PE ), it is feasible to achieve extrapolation beyond the effective window length.

KeyVectors

denotesself-attentionscoreofcurrent token

1-thchunk

denotesself-attentionscorebetween current token and others

discarded

denotes a chunk.1-th denotes the first, and n-th denotes

2-th chunk

the last.Except the first and the last chunks,other chunks

QueryVectors

haveasamesize.

number masked on self-attention score denotes its relative position

...-thchunk

k3

k4

k5

k6

6b

9

8

7

6

5

3

2

Q

attendtoKeysfrom1-thchunk

PE: linear relative positions used in chunk's self-attention score except n-th

(n-1)-th chunk

k3

k4

6b

6

6

5

5

4

4

3

2

attendtoKeysfromallpreviouschunks

Ln-thchunk

Stair PE:weave relativepositions used in the n-th chunk's attention score

We further introduce a novel weave-PE-based method, Mesa-Extrapolation , which demonstrates substantial extrapolation prowess without the need for additional training. Specifically, we use chunk-based triangular attention matrix, aiming at memory-friendly resource consumption (Figure 1 left), and we employ Stair PE to handle the last chunk (Figure 1 right below), effectively making Mesa-Extrapolation a completely free plug-in for LLMs.

Our contributions are summarized as follows:

Theoretical Analysis We provide a comprehensive theoretical analysis on the challenges encountered by NoPE beyond its effective window length, and investigate the effect of PE in this context. Our theorems prove that through meticulous weave position, PE can be effectively extrapolated beyond the effective window length. This theoretical foundation establishes a clear understanding of the extrapolation problem and potential solutions for LLMs utilizing PE.

Introduction of Mesa-Extrapolation We introduce a novel extrapolation approach called "Mesa-Extrapolation." Based on triangular attention matrix, this method strategically organizes input tokens into chunks, with the final chunk employing a weave PE method (e.g., Stair PE) to integrate all token states. Mesa-Extrapolation provides a practical and effective solution to the extrapolation problem.

Empirical Validation Comprehensive experiments demonstrate that our approach achieves competitive performance, offering the benefits of extremely low memory usage and the fastest inference speed compared to existing methods. Importantly, it is an easy plugin and can enhance extrapolation performance of LLMs without any additional resource consumption.

2 Background

Since the self-attention mechanism itself does not contain position information, PE components have become an integral part of the Transformer architecture. Cmmon choices for PE are either absolute , where each absolute position (e.g., 1 , 2 , 3 , ... ) is directly represented, or relative , where the distance between tokens is used as positional information. Absolute Position Embedding (APE) embeds

each absolute position i into position vector p i and adds word embeddings to their corresponding p i , before feeding them to the LLMs Vaswani et al. (2017). However, the extrapolation ability is limited because it cannot generalize to unseen positions. RoPE Su et al. (2023) rotates the query and key vectors with an angle proportional to their absolute positions, so the attention dot production only depends on relative distance between tokens, providing a relative positional encoding. T5's Relative Bias first maps the relative distance ( i -j ) between tokens at positions i and j to a scalar bias value b = f ( i -j ) , where f is a lookup table. ALiBi is similar to T5's Relative Bias but instead subtracts a scalar bias from the attention score Press et al. (2021)(refer to Appendix D.2 for more details). Recent works such as ReRoPE and Leaky-ReRoPE Su (2023b) achieve effective extrapolation without fine-tuning by meticulously weaving relative positions of RoPE. We refer to this class of methods that achieve extrapolation by weaving the relative positions of PE without fine-tuning as Weave PE .

Kazemnejad et al. (2023) indicates that the decoder-only transformer with position encoding removed (NoPE) demonstrates stronger extrapolation capabilities. Furthermore, it theoretically shows that a specific transformer model can get relative and absolute positional information, even in the absence of PE. Haviv et al. (2022) also demonstrates NoPE achieves comparative performance to standard Transformer models. These new studies pose a key challenge regarding the choice of whether using PE or not in Transformer architecture (refer to Appendix A for more related works).

3 Model Extrapolation: NoPE vs. Weave PE

3.1 Problem Definition

We mainly consider relative PE methods and formally define their self-attention dot product as a function f PE , which takes the query q t located on position t , the key k i located on position i , and their relative positions t -i as input parameters, as follows:

where f PE denotes a relative PE method such as RoPE or ALiBi.

For ALiBi,

where m is head index and C = 2 -2 -log 2 (#heads+3) .

For NoPE,

Based on Equ.1 we formally define weave PE as follows:

where W is a weave function which takes the relative position t -i as input parameter.

For example, ReRoPE Su (2023b) can be considered as an example of weave PE, with its W function defined as follows:

where N is a constant. ReRoPE's dot-product attention is:

where R is a rotation matrix that rotates W ( t -i ) θ radians. This is based on RoPE's dot-product attention:

3.2 Motivation

Chen et al. (2023) explores the evolution of hidden state values and reveals a noticable phenomenon: as the position increases, the hidden state values will explode. This finding appears consistent with observed failures in extrapolation. Through probe experiments (refer to Appendix F), we investigate

the alterations in hidden state values across various positions and layers. Results show a significant shift in the hidden state's value range upon surpassing the effective window. Interestingly, when employing extrapolation techniques, hidden state values exhibit noticeable suppression. This indicates that the effective range of hidden state values likely lies within a specific threshold. When the position exceeds the effective window length, the hidden state values surpass this threshold, resulting in extrapolation failures.

Previous work Kazemnejad et al. (2023) utilizes a constructive approach. By constructing the Transformer's weights, it enables the first and second layers to independently generate position information. Drawing inspiration from this, we endeavor to construct a Transformer model capable of mirroring this observation.

3.3 Theoretical Analysis

In a multi-layer neural network, each layer's outputs, a.k.a hidden state values o , become the inputs for the subsequent layer. To maintain stable network behavior, these values must remain within a reasonable range. We define this observable boundary as the threshold H . This threshold can be either an upper bound or a lower bound. For our analysis, we focus on the lower bound of this threshold. A successful extrapolation occurs when a large model consistently generates accurate next tokens for a long input sequence. Conversely, a failed extrapolation happens when the model produces incorrect or nonsensical next tokens. Based on these definitions, we make the following assumptions:

Assumptions. In LLM, there is a lower bound as threshold H for the hidden state value o in specific dimension and specific layer. Let M be the max window length for LLM. Predefine query W Q , key W K , value W V and output W O matrices, and feed-forward sub-layer W 1 , W 2 matrices. When o > H , LLM extrapolates successfully. Once o < H , LLM extrapolation fails.

These assumptions indicate that by observing whether the hidden state value o in this dimension exceed the threshold H , we can predict whether the large model's extrapolation has failed. Building upon these assumptions, theoretical results for NoPE exceeding the effective window length are as follows:

Theorem 3.1 (NoPE Extrapolation) . Let x = [ < bos >, x 1 , . . . , x T ] be an input sequence of length T +1 to the model. Then, there exists W Q , W K , W V , W O , W 1 , and W 2 matrices, such that when T < M , o T > H ; and when T > M , o T < H .

Full proof is given in Appendix E.1. This theorem reveals the internal mechanism of NoPE extrapolation as the input length changes. The theoretical results for PE are as follows:

Theorem 3.2 (PE Extrapolation) . Let x = [ < bos >, x 1 , . . . , x T ] be an input sequence of length T+1 to the model. Consider a simple relative PE schema where dot product between query q t and key k i at positions t and i ( t ≥ i ) can be expressed as: ⟨ q t , k i ⟩ := q T t k i -( t -i ) . Then, there exists W Q , W K , W V , W O , W 1 , and W 2 matrices, such that when T < M , o T > H ; and when T > M , o T < H .

Full proof is given in Appendix E.2. Theorems 3.1 and 3.2 state the failure of length extrapolation in NoPE and PE, respectively.

Building on Theorem 3.2, we further investigate the case for carefully orchestrated weave PE. The theoretical result is as follows:

Theorem 3.3 (Weave PE Extrapolation) . Let N be a positive constant. Consider a simple weave PE extrapolation schema: when t -i < N , W ( t -i ) = t -i ; and when t -i ≥ N , W ( t -i ) = N . Then, the attention dot product is fixed as below:

, where N ≪ M . Then, applying W Q , W K , W V , W O , W 1 , and W 2 matrices from Theorem 3.2, we have when T > M , o T > H .

Full proof is given in Appendix E.3. This theorem suggests that for existing LLMs relying on PE, simply weaving its relative positional encoding can effectively extend the input window.

Through Theorem 3.1 and Theorem 3.2, we formulate pertinent theoretical models for NoPE and PE, respectively, shedding light on the intricate relationship between extrapolation and the effective window length. Building upon these findings, Theorem 3.3 delves deeper into the realm of explicit PE, revealing that a well-designed weave PE scheme can effectively broaden the original effective window length. The theorems show the existence of an effective length M , which is typically related to the maximum training window length of the LLM. Furthermore, within the proof of Theorem 3.3, our results indicate that N ≪ M , consistent to experimental findings in Su (2023b), where although the extrapolation position can theoretically start from 2048, the best extrapolation starting position is 512. More experimental parameters also validate this setting (refer to Appendix B.2).

3.4 Validating Extrapolation Using Observed Thresholds

Further, we design probe experiments (refer to Appendix F.3 for more results) to validate the observed phenomena and our theorems, as shown in Figure 2. From Figure 2, two observations are noted: Firstly, for hidden state values of the same dimension, the first layer undergoes minimal change, while the second layer exhibits a more pronounced transition. Exceeding the threshold implies extrapolation failure. This observation aligns with our theoretical model construction, where the first layer primarily refines positional information, with significant signal changes occurring in the second layer, as demonstrated in Theorem 3.2. Secondly, based on observational thresholds, when the length of the input sequence is around 12k, the values of hidden state corresponding to Dynamic-NTK Liu et al. (2023) surpass the threshold, implying extrapolation failure. Conversely, for ReRoPE Su (2023b), extrapolation succeeds. These two predictive outcomes corroborate with subsequent experimental results.

Layer:1,Dimension:1

Layer:2,Dimension:1

0.3

0.2-

Value

0.1

State

0.0

-0.1

len

P!H

-0.2-

-0.3

-0.4

2k

5k8k11k14k

2k

5k8k

11k

14k

2k

5k8k11k14k

2k5k8k11k14k

Origin

Dynamic-NTK

Position

4 Mesa-Extrapolation

In this section, we begin by introducing a novel weave position encoding method, termed Stair Position Encoding (Stair PE). Following this, we propose a chunk-based triangular attention matrix. Building on these concepts, we introduce the implementation of Mesa-Extrapolation. Lastly, we discuss the theoretical properties of these innovations.

4.1 Stair PE

Following the concept of weave PE in Equ.2, we define a novel weave PE method, namely Stair PE as follows:

where I = N + ⌈ t -i -N E ⌉ . N denotes the extrapolated position, and E denotes the extrapolated width. Both N and E are positive constants. Stair PE can be applied to existing relative PEs such as RoPE and ALiBi. For example, for RoPE:

Taken a sequence of length 10 as an example, the right subplot of Figure 1 shows that the relative positions generated by PE (Both RoPE and ALiBi) are linear, while those generated by Stair PE (here N = 4 and E = 2 ) are non-linear. Since weave PE changes the linear relative position, it has to calculate the attention matrix more than once Su (2023b). Compared with ReRoPE (detailed on Appendix B.3), Stair PE provides a finer-grained extrapolated positions. Meanwhile, compared with Leaky-ReRoPE, Stair PE reuses known positions, reducing possible generalization errors. We also provide an ablation experiment to compare these weave PE methods (refer to Appendix C.6).

While our work was conducted independently, we note that Jin et al. (2024) have recently explored a similar idea through flooring the original positions and obtaining the relative position matrix with grouped self-attention. Although the two parallel thought processes are different, under certain conditions their formulations are equivalent (refer to Appendix G). Consequently, Self-Extend can be categorized as a Weave PE method. Our proposed Chunk-based Triangular Attention Matrix (detailed in Section 4.2) and its corresponding theoretical properties (Section 4.4) are also applicable to this parallel approach.

4.2 Chunk-based Triangular Attention Matrix

We design a chunk-based triangular attention matrix as shown in the left subplot of Figure 1. To achieve approximate linear memory consumption and computational speed, we further split the triangular attention matrix into several chunks and concatenate these chunks. We segment the input sequence into several sub-sequences according to DynamicSplit function (defined in Appendix 2), which divides a sequence into sub-sequences of equal length, with the exception of the first and last sequences. The length of each sub-sequence is determined by both the input token length and the max training length. Each of the generated sub-sequences then undergoes a self-attention operation to generate a corresponding chunk. That is, a sub-sequence of length l will generate a corresponding chunk with the size l × l .

4.3 Implementation

Mesa-Extrapolation mainly utilizes the chunkbased triangular attention matrix and Stair PE. Notice that regular PE (such as RoPE or ALiBi) is applied to all chunks except for the last chunk, for which Stair PE is applied. For the last chunk, all previous chunks are concatenated, and Stair PE is used to rearrange relative positional encoding to achieve extrapolation beyond the effective window length.

In summary, the process of Mesa-Extrapolation mainly contains four steps (Algorithm 1): The first three steps correspond to the prefill stage, which is used to calculate all input tokens. The last step corresponds to the decoding stage, which is used to generate next-token one by one.

Firstly, DynamicSplit function segments the input sequence, and the first segmented subsequence is fed into LLM to generate the first attention matrix chunk (line 3 in Algo.1). Secondly, subsequential sequences are iteratively processed while simultaneously feeding the key

Algorithm 1 Mesa-Extrapolation Algorithm

Require: DynamicSplit, LLM, StairPE

Input: s [0 : T -1] (input tokens with length T)

Output: s [ T, T +1 , ... ]

1: first _ length, chunk _ width ← DynamicSplit( s )

2: K _ cache, V _ cache ← [] , []

3: first _ K,first _ V ← LLM( s [0 : first _ length ] )

4: Append first _ K to K _ cache , first _ V to V _ cache

5: i ← first _ length

6: while i < T -1 -chunk _ width do

7: K,V ← LLM( s [ i : i + chunk _ width ] , first _ K , first _ V )

8: K _ cache append K , V _ cache append V

9: i ← i + chunk _ width

10: end while

11: apply StairPE to fix positions

12: K,V ← LLM( s [ i : T -1] , K _ cache , V _ cache )

# Decoding Stage

13: apply StairPE to fix positions

14: generate next-token one by one

and value pairs of the first chunk into LLM to generate subsequent chunks (line 6-10 in Algo.1). Thirdly, the last sub-sequence is processed by concatenating the key and value pairs of all previous chunks together and using Stair PE to modify the relative positional encoding. Then, it is fed into

the LLM to produce the last chunk (line 11-12 in Algo.1). Finally, Stair PE is applied to process the current token and cached Key and Value pairs to generate next-token one by one (line 13-14 in Algo.1). We establish the effectiveness of our proposed Mesa-Extrapolation in the next section.

4.4 Theoretical Properties

Theorem 3.2 establishes a theoretical measurement for evaluating the effectiveness of extrapolation. We consistently apply this indicator, along with adjustments in relative positioning using Stair PE, to validate the feasibility of extrapolation. The result is as below:

Corollary 4.1 (Mesa Extrapolation) . Let N be a positive constant. Consider a simple Stair PE extrapolation schema, and the attention dot product is fixed as:

where N ≪ M , I = N + ⌈ t -i -N E ⌉ , and the extrapolated width E is a constant. Then, Apply W Q , W K , W V , W O , W 1 , and W 2 matrices from Theorem 3.2. Although T > M , it still o T > H .

Full proof is provided in Appendix E.4. We prove that Mesa-Extrapolation can effectively extrapolate outside the max window length.

5 Experiments

In this section, we validate the performance of Mesa-Extrapolation through experiments measured over multiple metrics. We choose GovReport Huang et al. (2021), Pile Gao et al. (2020), LongBench Bai et al. (2023), and LongEval Krishna et al. (2023) datasets, and also generate a passkey dataset, which has been integrated in the code warehouse. More experimental details and results are referred to Appendix B and C.

Since our method is completely free plug-in and does not require fine-tuning, we choose methods of this type for comparison, including: model self (Origin), ReRoPE Su (2023b), Leaky-ReRoPE Su (2023b), Dynamic-NTK Liu et al. (2023), LM-Infinite Han et al. (2023), and Streaming-LLM Xiao et al. (2023).

We evaluate Mesa-Extrapolation using three prominent LLM families: LLaMA Touvron et al. (2023a) Touvron et al. (2023b) (including LLaMA-3B (Open-LLaMA-3B), LLaMA2-7B-Chat, and Vicuna13B-V1.3), MPT Team (2023) (including MPT-7B), and PyThia Biderman et al. (2023) (including PyThia-6.9B and PyThia-12B). Notably, LLaMA and PyThia incorporate RoPE Su et al. (2023), whereas MPT employs ALiBi Press et al. (2021) - two of the most influential PE techniques in recent research. Furthermore, we validated our approach using the Phi-3-mini-128k-instruct model Microsoft (2024) (refer to Appendix C.9).

We also conduct ablation experiments (refer to Appendix C.6). We use a 2xA800 80GB NVIDIA GPU server as the experimental environment and adopt the PyTorch framework.

5.1 Evaluation on Passkey Retrieval Tasks

We assess the accuracy of Mesa-Extrapolation using the generated passkey dataset. This dataset comprises samples of varying lengths, each storing a random password at a random position. The sample length initiates at 1024 and increments by 1024 . Simultaneously, 100 samples are randomly generated for each length. The proportion of correct answers found by LLMs is calculated for each input length.

Fig.3 shows the results of 6 LLMs on passkey retrieval task. The LLaMA families can employ various methods, including Origin, ReRoPE, Leaky-ReRoPE, Dynamic-NTK, LM-Infinite, Streaming-LLM, and our Mesa-Extrapolation. Note that these methods are model-specific and may not be universally applicable across all model series. For the MPT model, Origin, Streaming-LLM, ReRoPE, and Mesa-Extrapolation can be utilized. Similarly, for the PyThia model, Origin, Streaming-LLM, and Mesa-Extrapolation can be applied.

LLaMA-3B

LLaMA2-7B

Vicuna-13B

1.0

0.8-

0.6-

0.4 -

0.2

0.0

2k5k8k11k14k17k20k23k26k29k

2k5k8k11k14k17k20k23k26k29k

2k

5k

8k

11k

14k

MPT-7B

PyThia-6.9B

PyThia-12B

1.0

0.8-

0.6

0.4 -

0.2

0.0

2k5k8k 11k14k17k20k23k26k29k

1k

3k

5k

7k

9k

11k

1k

2k

3k

4k

5k

6k

-

Origin

-ReRoPE

11

Dynamic-NTK

--LM-Infinite

Streaming-LLM

**.Mesa-Extrapolation

Token Length

Analyzing the LLaMA model series reveals that weave PE-based methods, including ReRoPE, LeakyReRoPE and Mesa-Extrapolation, achieve superior extrapolation capabilities. Additionally, under our existing hardware constraints, Mesa-Extrapolation demonstrates longer extrapolation capabilities.

In the case of MPT model, after surpassing the maximum training length, the extrapolation capabilities of Mesa-Extrapolation and ReRoPE show a decline. We speculate that this may be attributed to an approximate ALiBi PE applied on MPT model. This approximate ALiBi would lead to significant disruptions when weaving its positions. (refer to Appendix C.2 for more details about MPT).

In the PyThia models, an additional observation is that 100 %accuracy is still not achieved within the training length. We attribute this to the PyThia model being a pre-trained model without instruction alignment, resulting in a weakened intrinsic understanding of the task. Results with enhanced prompts are referred to Appendix C.5.

Dynamic-NTK exhibits partial extrapolation capabilities beyond the maximum training length. Streaming-LLM and LM-Infinite also exhibit poor performance. This is because these methods discard portions of the input token, leading to potential information loss and incorrect answers.

LLaMA-3B

LLaMA2-7B

Vicuna-13B

15.0-

12.5

10.0-

7.5

5.0-

2.5

0.0

2k

5k

8k

11k

14k

2k

5k

8k

11k

14k

2k

5k

8k

11k

14k

Origin

Dynamic-NTK

LM-Infinite

Streaming-LLM

Mesa-Extrapolation

Token Length

5.2 Evaluation on Language Modeling

We further assess the fluency of Mesa-Extrapolation utilizing the perplexity metric. Results evaluated on the Pile dataset are presented in Fig.4. X-axis represents the length of the input token, while the

Y-axis corresponds to NLL (Negative Log-Likelihood) values. It can be observed that the NLL value of Origin consistently increases when the maximum training length is exceeded. Other methods maintain low NLL values. LM-Infinite performs marginally better on Vicuna-13B. Dynamic-NTK method exhibits slightly weaker performance after 11 k, and the performance continues to drop as the input length increases. In summary, our Mesa-Extrapolation demonstrates comparable performance to other methods on PPL metric.

5.3 Evaluation on Summary of Tasks

We conduct a summary task using the GovReport dataset and employ ROUGE ROUGE (2004) (ROUGE-1/2/L) as evaluation metrics. ROUGE assess overlapping N-grams by comparing the generated text with reference answers.

For the GovReport dataset, we segment the range from 3 * 1024 to 11 * 1024 based on sample length, with each interval of 1024 units. A test set is created by randomly selecting 8 samples from each interval. We choose LLaMA2-7B-Chat as the evaluated LLM. The experimental results for ROUGE is presented in Tables 1 below:

In Table 1, we record the average scores of Rouge-1, Rouge-2, and Rouge-L within each interval. It is evident that once the effective input window is exceeded, the performance of Origin and StreamingLLM declines rapidly, rendering it useless. For LM-Infinite, the scores exhibit a slight decrease as the length increases. Dynamic-NTK shows slightly better performance within 11 k. However, combined with the Fluency experiment on Fig.4, it seems that the effective extrapolation range of Dynamic-NTK cannot exceed 12 k, which is consistent with the threshold observed in Fig.2. Weave PE-based methods, including ReRoPE, Leaky-ReRoPE and Mesa-Extrapolation maintain similar generation quality as the length increases. Note our Mesa-Extrapolation shows slight variability in performance within mid-length (8k-11k) in the summary task. We speculate that with a fixed extrapolation width param ( E = 50 ), the 7 k11 k range may spread the model's attention more thinly compared to the 4 k6 k range. We hypothesize that optimizing the extrapolation width could alleviate or improve performance in the 7 k11 k range.

Memory Usage

125

100-

150-

75

Origin

Dynamic-NTK

LM-Infinite

Streaming-LLM

Mesa-Extrapolation

Origin

LM-Infinite

Dynamic-NTK

Streaming-LLM

Mesa-Extrapolation

10

15

Latency

(GB)

125

150

100

75

Origin

LM-Infinite

Dynamic-NTK

Mesa-Extrapolation

Streaming-LLM

Origin

Dynamic-NTK

LM-Infinite

Streaming-LLM

Mesa-Extrapolation

·15

10

20

Latency

50

5

50

5

25

25

0

2k

5k

8k

11k

14k

17k

20k

23k

26k

29k

Q

2k

5k

8k

11k

14k

17k

20k

23k

26k

29k

Token Length

Token Length

(a) Memory Usage & Latency for Open-LLaMA-3B

(b) Memory Usage & Latency for LLaMA2-7B

5.4 Latency & Memory Usage

To compare actual memory consumption and inference speed, we conduct experiments using both the 3B and 7B versions of the LLaMA model. The results are presented in Fig.5. In Fig.5, the X-axis

represents the input token length, the left Y-axis denotes memory usage, and the right Y-axis indicates decoding time. It is noteworthy that decoding time is closely related to memory usage, primarily from the computation of the attention matrix.

Observing Fig.5, both memory usage and decoding time for Origin and Dynamic-NTK exhibit a quadratic trend. Similarly, ReRoPE and Leaky-ReRoPE exhibit the highest memory usage and decoding time, showcasing a quadratic trend, which aligns with our analysis (refer to Appendix C.7). Although LM-Infinite demonstrates a linear trend, its increase is substantial. In contrast, Mesa-Extrapolation method also exhibits a linear trend but significantly outperforms other methods in terms of memory usage and decoding time. Furthermore, as the input length increases, this trend becomes more pronounced.

6 Conclusion

Our study addresses the critical challenge faced by Large Language Models (LLMs) when confronted with longer input lengths, commonly referred to as the extrapolation problem. Through theoretical exploration, we uncover the underlying mechanisms of this challenge, shedding light on the reason of extrapolation failure for both NoPE and PE. Furthermore, we present theoretical evidence demonstrating the potential for effective extrapolation using Weave PE. Based on Weave PE, we introduce a practical solution called Mesa-Extrapolation, which strategically organizes input tokens into chunks to achieve competitive performance with minimal resource usage. Empirical validation demonstrates its effectiveness. Our work not only advances the understanding of the extrapolation problem but also offers a practical solution, with a complete free plug-in for LLMs.

Limitations. Mesa-Extrapolation is a plug-and-play method that does not require additional finetuning. However, previous work, "NTK-aware" ntk (2023) also shows that applying further fine-tuning to plug-in extrapolation is possible. Therefore exploring fine-tuning based on Mesa-Extrapolation can be an interesting next step. Due to limitations of resources, we have not yet validated our method at longer lengths.

Broader Impacts We contend that the significance of completely free plug-in extrapolation method lies in two aspects. Firstly, it enables the expansion of the effective window length of already trained LLMs with no additional cost. Secondly, it allows for training LLMs from scratch with short texts and subsequently expanding their effective window length with a free plug-in extrapolation method, which can greatly help the industry improve the extrapolation capabilities of diverse LLMs.

References

Ntk-aware scaled rope allows llama models to have extended (8k+) context size without any fine-tuning and minimal perplexity degradation, 2023. URL https://www.reddit.com/r/ LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_ have/ .

Ba, J. L., Kiros, J. R., and Hinton, G. E. Layer normalization. arXiv preprint arXiv:1607.06450, 2016.

Bai, Y., Lv, X., Zhang, J., Lyu, H., Tang, J., Huang, Z., Du, Z., Liu, X., Zeng, A., Hou, L., Dong, Y., Tang, J., and Li, J. Longbench: A bilingual, multitask benchmark for long context understanding. arXiv preprint arXiv:2308.14508, 2023.

Bertsch, A., Alon, U., Neubig, G., and Gormley, M. R. Unlimiformer: Long-range transformers with unlimited length input. arXiv preprint arXiv:2305.01625, 2023.

Biderman, S., Schoelkopf, H., Anthony, Q., Bradley, H., O'Brien, K., Hallahan, E., Khan, M. A., Purohit, S., Prashanth, U. S., Raff, E., Skowron, A., Sutawika, L., and van der Wal, O. Pythia: A suite for analyzing large language models across training and scaling. arXiv preprint arXiv:xxxx.xxxxx, 2023.

BlackSamorez. TensorParallel: A Library for Efficient Model Parallelization in PyTorch. https: //github.com/BlackSamorez/tensor_parallel , 2023.

bloc97. Add ntk-aware interpolation "by parts" correction, 2023. URL https://github.com/ jquesnelle/scaled-rope/pull/1 .

Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.

Chen, S., Wong, S., Chen, L., and Tian, Y. Extending context window of large language models via positional interpolation. arXiv preprint arXiv:2306.15595, 2023.

Delétang, G., Ruoss, A., Grau-Moya, J., Genewein, T., Wenliang, L. K., Catt, E., Cundy, C., Hutter, M., Legg, S., Veness, J., et al. Neural networks and the chomsky hierarchy. arXiv preprint arXiv:2207.02098, 2022.

Elhage, N., Nanda, N., Olsson, C., Henighan, T., Joseph, N., Mann, B., Askell, A., Bai, Y., Chen, A., Conerly, T., DasSarma, N., et al. A mathematical framework for transformer circuits, 2021. URL https://transformer-circuits.pub/2021/framework/index.html . Transformer Circuits Thread.

Gao, L., Biderman, S., Black, S., Golding, L., Hoppe, T., Foster, C., Phang, J., He, H., Thite, A., Nabeshima, N., et al. The pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00027, 2020.

Geng, Xinyang and Liu, Hao. Open LLAMA 3B - Large Language Model. https://huggingface. co/openlm-research/open_llama_3b , 2023.

Han, C., Wang, Q., Xiong, W., Chen, Y., Ji, H., and Wang, S. Lm-infinite: Simple on-the-fly length generalization for large language models. arXiv preprint arXiv:2308.16137, 2023.

Haviv, A., Ram, O., Press, O., Izsak, P., and Levy, O. Transformer language models without positional encodings still learn positional information. arXiv preprint arXiv:2203.16634, 2022.

Hendrycks, D. and Gimpel, K. Gaussian error linear units (gelus). arXiv preprint arXiv:1606.08415, 2016.

Hsieh, C.-P., Sun, S., Kriman, S., Acharya, S., Rekesh, D., Jia, F., and Ginsburg, B. Ruler: What's the real context size of your long-context language models? arXiv preprint arXiv:2404.06654, 2024.

Huang, L., Cao, S., Parulian, N., Ji, H., and Wang, L. Efficient attentions for long document summarization. arXiv preprint arXiv:2104.02112, 2021.

Jin, H., Han, X., Yang, J., Jiang, Z., Liu, Z., Chang, C.-Y ., Chen, H., and Hu, X. Llm maybe longlm: Self-extend llm context window without tuning. arXiv preprint arXiv:2401.01325, 2024.

Kazemnejad, A., Padhi, I., Ramamurthy, K. N., Das, P., and Reddy, S. The impact of positional encoding on length generalization in transformers. arXiv preprint arXiv:2305.19466, 2023.

Krishna, K., Bransom, E., Kuehl, B., Iyyer, M., Dasigi, P., Cohan, A., and Lo, K. Longeval: Guidelines for human evaluation of faithfulness in long-form summarization. arXiv preprint arXiv:2301.13298, 2023.

Langley, P. Crafting papers on machine learning. In Langley, P. (ed.), Proceedings of the 17th International Conference on Machine Learning (ICML 2000), pp. 1207-1216, Stanford, CA, 2000. Morgan Kaufmann.

Lindner, D., Kramár, J., Rahtz, M., McGrath, T., and Mikulik, V. Tracr: Compiled transformers as a laboratory for interpretability. arXiv preprint arXiv:2301.05062, 2023.

Liu, X., Yan, H., Zhang, S., An, C., Qiu, X., and Lin, D. Scaling laws of rope-based extrapolation. arXiv preprint arXiv:2310.05209, 2023.

Microsoft. Phi-3-mini-128k-instruct. https://huggingface.co/microsoft/ Phi-3-mini-128k-instruct , 2024.

Mohtashami, A. and Jaggi, M. Random-access infinite context length for transformers. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.

Papineni, K., Roukos, S., Ward, T., and Zhu, W.-J. Bleu: a method for automatic evaluation of machine translation. In Proceedings of the 40th annual meeting of the Association for Computational Linguistics, pp. 311-318, 2002.

Peng, B., Quesnelle, J., Fan, H., and Shippole, E. Yarn: Efficient context window extension of large language models. arXiv preprint arXiv:2309.00071, 2023.

Press, O., Smith, N. A., and Lewis, M. Train short, test long: Attention with linear biases enables input length extrapolation. arXiv preprint arXiv:2108.12409, 2021.

Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, P. J. Exploring the limits of transfer learning with a unified text-to-text transformer. The Journal of Machine Learning Research, 21(1):5485-5551, 2020.

ROUGE, L. C. A package for automatic evaluation of summaries. In Proceedings of Workshop on Text Summarization of ACL, Spain, volume 5, 2004.

Roziere, B., Gehring, J., Gloeckle, F., Sootla, S., Gat, I., Tan, X. E., Adi, Y ., Liu, J., Remez, T., Rapin, J., et al. Code llama: Open foundation models for code. arXiv preprint arXiv:2308.12950, 2023.

Su, J. Understanding the scale operation of attention from the perspective of entropy invariance, Dec 2021. URL https://kexue.fm/archives/8823 .

Su, J. The road to transformer upgrade: 7. length extrapolation and local attention, Jan 2023a. URL https://kexue.fm/archives/9431 .

Su, J. Rectified rotary position embeddings. https://github.com/bojone/rerope , 2023b.

Su, J., Ahmed, M., Lu, Y., Pan, S., Bo, W., and Liu, Y. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, pp. 127063, 2023.

Team, M. N. Introducing mpt-7b: A new standard for open-source, commercially usable llms. www.mosaicml.com/blog/mpt-7b , 2023. Accessed: 2023-05-05.

Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971, 2023a.

Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288, 2023b.

Tworkowski, S., Staniszewski, K., Pacek, M., Wu, Y., Michalewski, H., and Miło´ s, P. Focused transformer: Contrastive training for context scaling. arXiv preprint arXiv:2307.03170, 2023.

Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I. Attention is all you need. Advances in neural information processing systems, 30, 2017.

Von Oswald, J., Niklasson, E., Randazzo, E., Sacramento, J., Mordvintsev, A., Zhmoginov, A., and Vladymyrov, M. Transformers learn in-context by gradient descent. In International Conference on Machine Learning, pp. 35151-35174. PMLR, 2023.

Weiss, G., Goldberg, Y., and Yahav, E. Thinking like transformers. In Meila, M. and Zhang, T. (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 11080-11090. PMLR, 18-24 Jul 2021. URL https://proceedings.mlr.press/v139/weiss21a.html .

Xiao, C., Zhang, P., Han, X., Xiao, G., Lin, Y., Zhang, Z., Liu, Z., Han, S., and Sun, M. Infllm: Unveiling the intrinsic capacity of llms for understanding extremely long sequences with trainingfree memory. arXiv, 2024.

Xiao, G., Tian, Y., Chen, B., Han, S., and Lewis, M. Efficient streaming language models with attention sinks. arXiv preprint arXiv:2309.17453, 2023.

Xie, S. M., Raghunathan, A., Liang, P., and Ma, T. An explanation of in-context learning as implicit bayesian inference. arXiv preprint arXiv:2111.02080, 2021.

Zhang, P., Liu, Z., Xiao, S., Shao, N., Ye, Q., and Dou, Z. Soaring from 4k to 400k: Extending llm's context with activation beacon. arXiv preprint arXiv:2401.03462, 2024.

Zhang, Y., Backurs, A., Bubeck, S., Eldan, R., Gunasekar, S., and Wagner, T. Unveiling transformers with lego: a synthetic reasoning task. arXiv preprint arXiv:2206.04301, 2022.

A Related Work

Extrapolation for PE A distinct line of research is dedicated to refining PE for enhanced extrapolation capabilities. Notable contributions include the introduction of Rotary Position Encoding (RoPE) Su et al. (2023), implementing relative PE through absolute position information. Similarly, Press et al. (2021) proposes a novel PE method ALiBi , fusing position information by directly introducing the relative position distance term in dot multiplication. In addition, there are other PE methods such as absolute position embedding (APE) Vaswani et al. (2017) and T5's Relative PE Raffel et al. (2020).

Another school of research focuses on enhancing the extrapolation performance of existing LLMs, categorized by whether to include further training. The first subcategory involves further fine-tuning, enlarging the effective window length by training LLMs on longer input texts. Chen et al. (2023) demonstrates that Position Interpolation (PI) method has a superior fine-tune effect, resulting in extended extrapolation capabilities with fewer fine-tuning steps. Mohtashami & Jaggi (2023) utilizes a new landmark token to represent individual input blocks, and enables model to select relevant blocks by further training, allowing for longer extrapolation. Zhang et al. (2024) follows similar idea by training a special token. Bertsch et al. (2023) proposes Unlimiformer , a k-nearest-neighbor (kNN) indexed encoder-decoder Transformer, enhancing efficiency by retrieving top-k keys for each decoder layer. Despite claiming support for decoder-only Transformer, differences in retrieval results across decoder layers may lead to potential failure. Tworkowski et al. (2023) introduces Focused Transformer (FOT), utilizing a contrastive learning-inspired training process to enhance the (key, value) space structure and enable effective context extension. The second subcategory explores methods that require no fine-tuning yet offer improved extrapolation through plug-ins. Su (2023b) introduces Rectified Rotary Position Embeddings ( ReRoPE ), a method that rectifies extrapolated relative positions based on RoPE within a specified interval to extend the effective window length. In addition, Leaky-ReRoPE offers an alternative by allowing a controlled leakage of position extrapolation within an interval. Both approaches are well-suited for LLaMA model families Touvron et al. (2023b) Touvron et al. (2023a), eliminating the need for fine-tuning. These methods do suffer from some drawbacks, such as twofold increase in memory consumption and longer inference time. We refer to this class of methods that achieve extrapolation by weaving the relative positions of PE without fine-tuning as Weave PE . We also notice that, recent work, InfLLM Xiao et al. (2024), proposes an additional memory units, which lookup token-relevant units for attention computation. In addition, it uses a modified encoding scheme similar to ReRoPE to achieve longer extrapolation. It is worth noting that the methods of the weave PE class can be applied seamlessly to these new designs. However, there is still no theory to explain why the methods of weave PE class can work . In a different line of inquiry, ntk (2023) proposes the "NTK-aware" method by drawing from the Neural Tangents (NTK) idea, which explores the high-frequency extrapolation and low-frequency interpolation concept for RoPE. Based on "NTK-aware", recent works bloc97 (2023), Peng et al. (2023), Roziere et al. (2023), Liu et al. (2023) perform further fine-tuning for optimal results.

Han et al. (2023) proposes LM-Infinite, which employs a Λ -shaped mask to prevent surpassing the effective window length by discarding central tokens. However, this design choice inevitably results in a loss of information. Likewise, Streaming-LLM Xiao et al. (2023) adopts a similar strategy to avoid surpassing the effective window length by discarding a portion of input tokens.

Extrapolation for NoPE There also exists a counter perspective asserting that the position information of an input sequence can be perceived without utilizing PE. In a comprehensive experimental comparison presented in Kazemnejad et al. (2023), the decoder-only Transformer is shown to exhibit superior extrapolation properties with no position encoding (NoPE). Haviv et al. (2022) conjectures that causal attention enables the Transformer to infer position information without the help of PE, demonstrating its comparable performance with standard Transformer models through probing experiments. These new studies pose a key challenge regarding the choice of whether using PE or not in Transformer architecture.

Theorems for Extrapolation We mainly focus on Transformer architecture. Although the Transformer architecture is considered as a "black box", there are ongoing efforts trying to explain its inner workings. Von Oswald et al. (2023) begins by providing a specific weight construction that elucidates the mechanistic understanding of in-context learning within optimized Transformers. Lindner et al. (2023) introduces interpretability for Transformers through programming, allowing the derivation of a program Transformer architecture tailored for specific tasks. In Han et al. (2023), the authors conclude that long-distance attention logits will explode. However, this does not clarify the

relationship between the effective window length and the existing supremum. It merely indicates an increase in the supremum, without directly implying an increase in the obtained actual attention logits. In Kazemnejad et al. (2023), inspired by programming Transformers Lindner et al. (2023), the authors construct a specific Transformer with predefined matrix parameters, theoretically demonstrating that NoPE can recover both absolute and relative position information even without the use of PE. However, these results fall short in explaining the underlying reasons for the failure of NoPE extrapolation. Built upon the entropy increase theory, Han et al. (2023) proves that with the expansion of input length, the entropy of attention experiences a corresponding increase. Additionally, Su (2021) introduces a scaling factor, denoted as log train -len ( n ) , which serves to mitigate entropy increase. Nevertheless, their theorems do not show that extrapolation fails beyond the effective window length.

We observe that extrapolation failure is intricately linked to the max window length. Nevertheless, existing theories on Transformers do not reveal the internal mechanism about extrapolation and the max window length. In light of this, our endeavor is to establish a correlation between these two aspects, aiming to offer theoretical insights that can guide us towards designing applicable extrapolation methodologies.

B Experiments Details

B.1 Passkey Retrieval Dataset

The data used to perform the passkey task is constructed as follows:

TASK-DESCRIPT = "There is an important info hidden inside a lot of irrelevant text. Find it and memorize it. I will quiz you about the important information there."

DEFAULT-CONTENT = "The grass is green. The sky is blue. The sun is yellow. Here we go. There and back again."

KEY-CONTENT = "The pass key is {KEY}. Remember it. {KEY} is the pass key."

The TASK-DESCRIPT is placed at the beginning of the sample. Subsequently, the DEFAULTCONTENT is repeated multiple times to serve as filler text. A randomly generated password is then created and inserted into the designated KEY-CONTENT section. Finally, a random position is selected, and the KEY-CONTENT, containing the generated password, is seamlessly incorporated into the sample.

The LLM is required to find the correct password from the sample.

B.2 Params Setting

Mesa-Extrapolation We design the DynamicSplit function used in Algorithm 1 to dynamically divide the input sequence into several subsequences according to the input length. The length of each subsequence corresponds to the size of the chunk. Let C denote the length of the each chunk (except the first chunk length F and the last chunk length L ). Define I as input token length and T as max training length.

The DynamicSplit function is detailed as follows:

In general, we set F = 100 , M max = 200 and L = 512 . Based on these criteria, we have devised a method for dynamically segmenting sequence. Leveraging the total length of the input token and the aforementioned requirements, we dynamically determine the length of each individual chunk.

Additionally, Stair PE is primarily employed in the manipulation of the last chunk. We need to concatenate all chunks together. After splicing chunks, the required positons will far exceed the max training length. At this time, we need to rearrange the positions according to Equ.3. In Equ.3, we generally set the extrapolated position N = 512 and set the extrapolated width E = 50 .

Implement Stair PE on RoPE and ALiBi The handling of RoPE Su et al. (2023) involves assigning positions to the query vector and key vector. The length of the query corresponds to the length of the last chunk, while the lengths of the key and value align with the combined lengths of all chunks. We assign position encoding to the query vector according the same sequence positions. Regarding the key's position encoding, we utilize the extrapolated width E and extrapolated point N . By designing

Algorithm 2 DynamicSplit function

Require:

LLM parameters

Input:

s (input token sequence)

Output:

F , C

1: set the values of F and L respectively as constants

2: set T according to the LLM parameters

3: get I according to input sequence s

6: if M < M max then

8: else

10: end if

the position encoding of query and key respectively, the relative position scheme defined by Stair PE can be implemented. And, for computational efficiency, we ensure that the relative PE of the last token in the last chunk completely follows Stair PE. The other tokens of the last chunk are similar to Stair PE.

For ALiBi's Press et al. (2021) position processing, we directly transform the relative position of the last chunk into the aforementioned structure within the mask matrix, refer to RoPE.

It is emphasized that the considerations outlined above primarily address the generation mechanism of LLM. The position encoding of the last chunk must align as closely as possible with the original position setting. This consideration is rooted in the belief that the latent concept Xie et al. (2021) formation heavily relies on there, with subsequent chunks providing essential information support. Experiments affirm that the last chunk significantly influences the accurate output of LLMs. This serves as the experimental foundation for our Mesa-Extrapolation method design.

ReRoPE & Leaky-ReRoPE We follow Su (2023b) for configuring the parameters of ReRoPE and Leaky-ReRoPE. The extrapolated position for ReRoPE is established at 512, mirroring the setting for Leaky-ReRoPE. Simultaneously, the extrapolated increment for Leaky-ReRoPE is calculated in relation to the input length, following the formula below:

where w is the extrpolated position as 512 , T is the maxium training length, and I is the input token length.

The ReRoPE and Leaky-ReRoPE methods are currently only applied to the LLaMA model series. We implemented ReRoPE based on its ideas on the MPT models.

LM-Infinite & Streaming-LLM LM-Infinite Han et al. (2023) introduces two branches for attention masking: a global branch on the left and a local branch on the right. The global branch enables each token to attend to the preceding n global tokens if they appear before the current token. Conversely, the local branch allows each token to attend to preceding tokens within a distance of n local . Tokens outside these two branches are disregarded during the attention operation. Subsequently, following its setting, set n local = I , where I still represents the training length. The choice of n global has a minor impact on model performance within the range [10 , 100] , and thus, we fix n global = 100 as its code setting. The distance limit entails the "effective distance" d within I , and we set d = I .

Streaming-LLM Xiao et al. (2023) adopts a similar design to LM-Infinite. It introduces an initial token known as Attention Sinks, whose length is denoted by x . It also defines a Rolling KV Cache for retaining the most recent tokens, whose length is denoted by y . It set x = 4 and y = T x , where T denotes the maximum training length. The central tokens are referred to Evicted Tokens.

The distinguishing factor between the two approaches lies in the selection of the initial token length, with LM-Infinite opting to splice longer tokens. We speculate that this is also the reason why LM-Infinite performs better than Streaming-LLM in the experiments.

Origin & Dynamic-NTK For Origin, it refers to loading the model directly, without using any methods. Dynamic-NTK differs from Origin by only adjusting the angle value of its RoPE component, according to the input length Liu et al. (2023).

The Dynamic-NTK method is currently applied to the LLaMA model series.

B.3 Weave PE-based Schemes

We list the details of the weave PE-based methods, including ReRoPE and Leaky-ReRoPE, as follows:

where left ReRoPE shows relative position on attention matrix, including (1) Numbers marked on each cell denote its relative position ( t -i ) between q t and k i . (2) It start to extrapolate position from relative position N . And right Leaky-ReRoPE shows relative position on attention matrix, including (1) Numbers marked on each cell denote its relative position ( t -i ) between q t and k i . (2) It start to extrapolate position from relative position N , with an incremental factor 1 k .

With Stair PE defined on Equ.3, the extrapolated position after N exhibits a fixed extrapolated width E , resembling a stair-like increment. Our Mesa-Extrapolation is designed with a finer position granularity compared to ReRoPE. Simultaneously, unlike Leaky-ReRoPE, ours utilizes previously trained positions, ensuring greater stability.

While these weave PE schemes, including ReRoPE, Leaky-ReRoPE, and Stair PE, demonstrate theoretical feasibility, their practical implementation may encounter computational complexities. For instance, both ReRoPE and Leaky-ReRoPE are primarily employed in the LLaMA models, which is based on RoPE Su et al. (2023). Due to the distinctive nature of RoPE, these methods necessitate the calculation of the attention matrix more than once, resulting in double the normal memory consumption. Moreover, given the quadratic complexity of the calculations, as the input length expands, both inference time and memory consumption grow proportionally.

B.4 ALiBi on MPT-7B

It is worth noting that the ALiBi implementation of MPT is only an approximation, as follow:

where showcase an input token length 10 for ALiBi PE mask. Based on this approximate implementation of ALiBi, we still use splitting chunk and Stair PE to implement Mesa-Extrapolation. We speculate that the approximation of ALiBi is the reasons for the instability of the extrapolation on Accuracy.

C More Experimental Results

C.1 BLEU Results on GovReport

We use the GovReport dataset and employ BLEU Papineni et al. (2002) as evaluation metric, with a same setting applied on Table 1. The experimental result for BLEU is presented in Table 2. In Table 2, we present the mean and standard variance for each method across various interval lengths. It is evident that, as the input token length increases, the performance consistently deteriorates for LM-Infinite. Streaming-LLM also experiences extrapolation failure beyond the model's inherent effective window length of 4k.

Dynamic-NTK shows slightly better performance within 11 k. Here, the BLEU scores exhibit almost consistent results with those of ROUGE on Table 1. It is also noteworthy that additional experiments show performance degradation of Dynamic-NTK beyond 11k, likely suggesting a limited effective extrapolation range for Dynamic-NTK.

Weave PE-based methods, including ReRoPE, Leaky-ReRoPE, and Mesa-Extrapolation, consistently demonstrates competitive results, maintaining similar generation quality as the length increases.

Overall, despite the chunking applied to our Mesa-Extrapolation method, it maintains a competitive edge compared to ReRoPE and Leaky-ReRoPE. This is achieved even with the trade-off between losing a certain amount of information and reducing memory consumption and inference time.

The results of Mesa-Extrapolation in the summary task, focusing on the generation of summary text within 1000 tokens across different input lengths, are detailed in Table 5.

C.2 Perplexity (PPL) metrics on MPT-7b

Following the same metrics setting about LLaMA models on Fig.4, we plot the NLL result about MPT-7B on Fig.6 as below:

14

Origin

Streaming-LLM

12

Mesa-Extrapolation

10

NLL

8

6

4

2k

5k

8k

11k

14k

TokenLength

In Fig.6, both Mesa-Extrapolation and Streaming-LLM demonstrate effective negative log-likelihood (NLL) values. For the original MPT model, after the extrapolation is extended 3 . 5 k, it will cause the perplexity (PPL) level to rise rapidly.

ALiBi Press et al. (2021) can still effectively extrapolate to around 3 . 5 k position after surpassing the maximum training length of 2 k. This observation affirms the extrapolation capability of the ALiBi method, especially when compared with RoPE's extrapolation results.

However, it is essential to highlight that while ALiBi achieves extrapolation to 3 . 5 k in above NLL task, its performance in the Passkey and LongEval tasks (refer to C.3) reveals limitations. This discrepancy is attributed to an inherent characteristic of ALiBi. ALiBi's expression (refer to D.2) indicates a constant decrease in the value of its position encoding as the relative distance increases. Considering the SoftMax operation within the attention mechanism, under identical conditions, the attention score obtained becomes smaller. This implies that tokens situated farther away receive less attention, possibly leading to the neglect of these distant tokens. Therefore, despite ALiBi's apparent success in effectively extrapolating the NLL task to 3 . 5 k, it falls short in attending to tokens beyond the maximum training length in practice. A similar analysis process is as Su (2023a).

It's noted that, NoPE Kazemnejad et al. (2023) claims to possess extrapolation capabilities akin to ALiBi. We hypothesize that the underlying reasons for NoPE's extrapolation performance align with the analysis presented here.

C.3 Evaluation on LongEval

We conduct additional testing on LongEval Krishna et al. (2023) lines task, a recently prominent evaluation task for long texts.

1.0

Origin

Origin

0.08 -

LM-Infinite

0.5

Streaming-LLM

0.8 -

Origin

Streaming-LLM

Mesa-Extrapolation

0.6

0.06 -

Mesa-Extrapolation

0.4

Dynamic-NTK

0.3

0.4

LM-Infinite

Streaming-LLM

0.04 -

0.2 -

0.2

Mesa-Extrapolation

0.02

0.1 -

0.0

-00°0

0.0 -

5k

7k

9k

1ik

13k

15k

4k

6k

8k

10k

12k

4k

6k

8k

10k

12k

Token Length

Token Length

Token Length

Figure 7: LongEval Lines Task on LLMs using various methods. Some observations: (1) Origin and StreamingLLM consistently exhibit an inability to extrapolate beyond the effective window length. (2) LM-Infinite shows weak extrapolation ability on the LLaMA2-7B-Chat model and fails to extrapolate on the MPT-7B model. (3) Weave PE-based methods, including ReRoPE, Leaky-ReRoPE, and Mesa-Extrapolation, all exhibit commendable extrapolation capabilities on the LLaMA2-7B-Chat model. (4) For the MPT-7B model, all methods display weak extrapolation performance, with Mesa-Extrapolation slightly outperforming other methods. (5) Mesa-Extrapolation also shows a certain extrapolation ability on the PyThia-6.9B model, with a decrease in extrapolation performance as the input length increases.

Fig.7 show the accuracy of LongEval Lines Task. Origin and Streaming-LLM consistently exhibit an inability to extrapolate beyond the effective window length. LM-Infinite shows weak extrapolation ability on the LLaMA2-7B-Chat model and fails to extrapolate on the MPT-7B model. Dynamic-NTK method, as shown earlier, almost fails after the input token length exceeds 11 k. Weave PE-based methods, including ReRoPE, Leaky-ReRoPE, and Mesa-Extrapolation, all exhibit commendable extrapolation capabilities on the LLaMA2-7B-Chat model. For the MPT-7B model, all methods display weak extrapolation performance, with Mesa-Extrapolation slightly outperforming other methods. We analyze that the reason for the failure of extrapolation in the MPT-7B model is attributed to the approximated implementation of ALiBi PE. This approximation makes it susceptible to interference when use weave PE (refer to C.2).

Mesa-Extrapolation also shows a certain extrapolation ability on the PyThia-6.9B model, with a decrease in extrapolation performance as the input length increases.

C.4 Evaluation on LongBench

We select LongBench Bai et al. (2023) dataset and use 5 major categories of tasks, including SingleDocument QA, Multi-Document QA, Few-shot Learning, Synthesis Tasks and Code Completion. Among them, each task selects a dataset, namely qasper, hotpotqa, samsum, passage-retrieval-en, and repobench-p. Taking into account the varying memory consumption of different methods, we opt to filter out samples with an input length exceeding 10 k to prevent discrepancies caused by out-of-memory (OOM) issues. We use the abbreviations in Table 3: S-Document (Single Document) QA, M-Document (Multi Document) QA, F-Learning (Few-shot Learning), S-Tasks (Synthesis Tasks), C-Completion (Code Completion).

In Table 3, Origin and Streaming-LLM barely work after exceeding the maximum training length of 4 k. LM-Infinite shows slightly weaker performance. Dynamic-NTK shows good performance, especially for Code Completion. Considering the effective scope of Dynamic-NTK, we speculate that this is related to the setting within 10 k length. Mesa-Extrapolation shows better performance on most tasks.

C.5 Enhanced Prompts for Mesa-Extrapolation

Considering that Mesa-Extrapolation involves positional approximation, as the input length increases, attention becomes more dispersed, resulting in increased entropy. One possible approach to mitigate these entropy increases is to utilize instruction-aligned prompts. These prompts may represent the strongest latent concepts that the model can follow, while clearly articulating the goals to be achieved. Similarly, it may be the reason why PyThia cannot achieve 100% accuracy within its training length on Fig.3, due to a lack of instruction alignment.

Based on these considerations, we employ the model's corresponding prompt Geng, Xinyang and Liu, Hao (2023) and augment it with explicit goal statements, referred to as an enhanced prompt, like below:

or

Additionally, we leverage model parallel inference technology BlackSamorez (2023) and perform verification concurrently on 2 A800 GPUs.

Fig.8 show the accuracy on 3 LLaMA models with enhanced prompts using Mesa-Extrapolation. We declare that the longest input length that these 3 models can currently reach is the limit of our existing hardware resources. For LLaMA-3B, the enhanced prompts improve the Accuracy and is extrapolated to 60 k. (2) Both of LLaMA2-7B-Chat and Vicuna-13B models, show good improvements by the enhanced prompts.

Through this experiment, we hypothesize that by writing prompts more carefully and accurately, we can maximize the extrapolation performance of free plug-in.

\n A:

1.0 -

1.0 -

1.0

0.8 -

0.8-

-80

0.6

0.6

0.6 -

0.4 -

0.4-

0.4 -

0.2

0.2

0.2

Mesa-Extrapolation

Mesa-Extrapolation

Mesa-Extrapolation

0.0-

Mesa-Extrapolation (enhanced-prompt)

0.0

Mesa-Extrapolation(enhanced-prompt)

0.0-

Mesa-Extrapolation (enhanced-prompt)

Token Length

Ok

4k

8k

12k

Token Length

ok

3k

6k

9k

TokenLength

12k

15k

18k

21k24k27k

(a) Accuracy on LLaMA-3B

(b) Accuracy on LLaMA2-7B-Chat

(c) Accuracy on Vicuna-13B

Figure 8: Accuracy about Passkey Retrieval Tasks using enhanced prompts for Mesa-Extrapolation. The gray line represents the max training length at 2 k, 4 k, and 2 k, respectively. Some observations: (1) On the LLaMA-3B model, the enhanced prompts significantly improves the Accuracy and is extrapolated to 60 k. (2) Both of LLaMA2-7B-Chat and Vicuna-13B models, show good improvements by the enhanced prompts.

C.6 Ablation Experiments

We design ablation experiment about weave PE based methods. We apply the encoding schemes of ReRoPE and Leaky-ReRoPE to chunk-based triangular attention matrices, aligning them with our implemented Mesa-Extrapolation. The entire design of the ablation experiment involved controlling only one variable, the weaving PE scheme, applying different weaving PE methods to the processing of the last chunk, while using the same dataset and testing environment. The experimental results 9 are as follows:

1.0 -

1.0

1.0 -

0.8 -

0.8-

0.8-

0.6

0.6

0.4

0.4

0.4

0.2

Mesa-Extrapolation

0.2

..d..

Mesa-Extrapolation

0.2 -

Mesa-Extrapolation

0.0-

0.0-

0.0-

Ok

5k

10k

15k

20k

25k

30k

ok

5k

10k

15k

20k

25k

30k

ok

2k

4k

6k

8k

10k

12k

14k

16k

Token Length

Token Length

Token Length

(a) Accuracy on LLaMA-3B

(b) Accuracy on LLaMA2-7B-Chat

(c) Accuracy on Vicuna-13B

Figure 9 demonstrate that both ReRoPE and Leaky-ReRoPE experience a certain loss in accuracy as the input sequence grows. We speculate that this may be due to ReRoPE repeatedly using the same extrapolation positions, while Leaky-ReRoPE, employing some fractions, may not be as precise as Stair PE in comparison to normal relative positions, resulting in a slight decrease in effectiveness.

C.7 Theoritical Speed & Memory

We assess the computational memory usage and inference time about decoding speed on various methods.

Table 4 shows the theoretical result about the memory usage. For Origin, the attention matrix calculation requires the current token to compute the attention score with each previous token, resulting in a memory footprint of O ( n 2 ) . Dynamic-NTK only alters the angle base, making it different from the Origin but still quadratic. For ReRoPE and Leaky-ReRoPE, as showed in Su (2023b), since the attention matrix needs to be calculated twice, their memory footprint is 2 × O ( n 2 ) . For Mesa-Extrapolation, its strategy involves splitting chunks while avoiding the

quadratic term. Accounting for splicing in the last chunk, its total memory usage scales proportional to O ((2 + √ 2) n ) . LM-Infinite and Streaming-LLM adopt a similar Λ -shaped mask, that scales proportional to O ((1 + √ 2) n ) .

C.8 Generated Summary using Mesa-Extrapolation

We present here the results of the summary generation task using GovReport Huang et al. (2021). In this experiment, task is to generate a summary for texts of varying lengths, limited to 1000 tokens. The outcomes yielded by Mesa-Extrapolation are outlined below on Table 5:

Table 5: Mesa-Extrapolation undertakes summarization tasks and generates text based on varying input lengths from the GovReport dataset. Some observations: (1) Mesa-Extrapolation adeptly and systematically summarizes key points one by one across various input lengths. It showcases logical coherence and demonstrates commendable summarization capabilities. (2) Additionally, Mesa-Extrapolation accurately incorporates punctuation and exhibits an automatic stopping mechanism for proper stopping.

Continued on next page

Continued on next page

Continued on next page

Continued on next page

Continued on next page

Continued on next page

Table 5 shows that Mesa-Extrapolation adeptly and systematically summarizes key points one by one across various input lengths. It showcases logical coherence and demonstrates commendable summarization capabilities. It's worth noting that the initial symbols used for sub-items beneath each summary vary, including symbols such as "*", "first, second,...", or the number "1, 2, ...", etc. We attribute this variability to the inherent content of the input text. LLM follows the expression format of the text itself and generates corresponding output. Additionally, Mesa-Extrapolation accurately incorporates punctuation and exhibits an automatic stopping mechanism for proper stopping.

C.9 Evaluation of Phi-3-mini-128k-instruct Model on Ruler Datasets

We further conducted experimental validation on the Ruler datasets Hsieh et al. (2024), focusing on the single-keys NIAH task. The needle-in-a-haystack (NIAH) test assesses the ability to retrieve a specific piece of information (the 'needle') from long distractor texts (the 'haystack').

In this experiment, we employed the microsoft/Phi-3-mini-128k-instruct model, which has practical value in real-world applications and is officially claimed to support extrapolation lengths of up to 128k tokens. To evaluate performance, we gradually increased the input length by 32k increments, comparing the original model's results with those achieved using our Mesa-Extrapolation method.

As shown in Figure 10, the original Phi-3-mini-128k-instruct model can only extrapolate up to its official limit of 128k tokens, beyond which it fails. In contrast, when utilizing the Mesa-extrapolation method, the model successfully extrapolates to at least 192k tokens. It's important to note that this 192k limit is due to our current hardware resource constraints-attempting to extend beyond this length results in an out-of-memory error.

1.0

0.8

Accurac

2

0.6

0.4

0.2

Origin

0.0

·.★..

Mesa-Extrapolation

ok

32k

64k

96k

128k

160k

TokenLength

D Background

D.1 Preliminary

In this section, we lay the groundwork and introduce the notation for below theorem analysis. Bold letters typically denote vectors or matrices, while normal letters represent scalars. We mainly follow Kazemnejad et al. (2023) to define these notations.

Let f θ be a decoder-only Transformer model, where θ denotes the model parameters. f θ processes the input sequence x = [ < bos >, x 1 , . . . , x T ] by applying a series of layers, where < bos > represents the first token Tokenizer transformation.

Each layer l , consisting of self-attention heads and a feed-forward sub-layer, produces the hidden state H l at layer l , after reading the previous hidden state H ( l -1) .

Each head is parameterized by a query W m Q , key W m K , value W m V , and output W m O matrices, where m represents the attention head index, and W m Q , W m K , W m V ∈ R h × d and W m O ∈ R d × h . d is the model's hidden state size and h is the attention dimension ( h = d # heads ). Note that we drop the attention head index m where it is clear from the context.

The Transformer layer TLayer ( l ) ( H ( l -1) ; θ l ) consist of self-attention heads and a feed-forward sub-layer, and input the previous hidden state H ( l -1) , and generate the hidden state H ( l ) at layer l . l is the layer index, and θ l is the set of parameters of the l -th layer. Each hidden state H ( l ) ∈ R d × ( T +1) is a matrix, and h ( l ) t denotes its hidden state at column t , i.e. at position t .

Each feed-forward sub-layer FF is parameterized by W 1 , W 2 ∈ R d × k.d matrices, where k denotes a multiplier of the hidden state size in this sub-layer, and is usually set to 4 in common implementations of the Transformer.

The Transformer layer TLayer ( l ) processes each column of H ( l -1) independently and in parallel to generate the output. The computation of the t -th column of H ( l ) is as follows:

where λ is layer normalization, and a t ∈ R d is the output of the multi-head self-attention sub-layer at position t .

a t is computed as:

where Attn ( m ) denotes the m -th attention head. Let o t ∈ R d denote the output of an attention head at position t . Then, o t is computed as:

where ˆ α = softmax( α ) ∈ R ( t +1) , and α is the attention weight vector such that:

where q t = W Q h ( l -1) t ∈ R h , k i = W K h ( l -1) i ∈ R h , and ⟨· , ·⟩ denotes the dot product operation.

The feed-forward sub-layer FF( · ) ∈ R d is a two-layer MLP as below:

where σ is a non-linear activation function(usually ReLU or GeLU Hendrycks & Gimpel (2016)), and σ ( · ) ∈ R d is layer normalization Ba et al. (2016). It's worthy noted that the additive view of attention heads Elhage et al. (2021) in Eq.(5) is mathematically equivalent with the view of concatenate and multiple Vaswani et al. (2017). The additive view is easier to understand and analyze.

The hidden state H (0) is initialized with a learned embedding of the input sequence X , i.e. H (0) = W E X , where W E ∈ R d × V is the embedding matrix and X ∈ R V × ( T +1) is the one-hot encoded input sequence. V is the vocabulary size.

D.2 NoPE & PE

NoPE We emphasize here that NoPE generally refers to these architectures that do not use position encoding components or remove the position encoding schemas, which is reflected in the attention dot product operation and is formally expressed as:

PE For PE, this generally refers to position encoding components, including Alibi, Rope, APE, etc. Their main difference from NoPE can be reflected in the attention dot product operation.

For ALiBi Press et al. (2021), it takes the form:

where m is head index and C is a constant defined as:

For example, if the number of heads is 8, we have 1 2 , 1 2 2 , . . . , 1 2 8 .

For RoPE Su et al. (2023), it's a relative PE that applies a rotation to the query and key representations based on their absolute positions before dot product attention. We formulate RoPE for model dimension d = 2 , and its dot product as below:

where R is a rotation matrix that rotates ( i -t ) θ radians:

For d > 2 , RoPE applies the same approach on every two consecutive dimensions of q t and k i , but with different θ angles.

For other PE methods, we recommend readers to consult the corresponding papers.

E Proofs

We provide proof about Theorems 3.1, 3.2, 3.3, and Corollary 4.1. Our proofs are inspired by Kazemnejad et al. (2023) Weiss et al. (2021) Lindner et al. (2023) and relies on the causal attention mask in the decoder-only Transformer and the SoftMax function.

Please refer to Appendix D for some necessary background knowledge.

Theorem E.1 (NoPE Extrapolation) . Let x = [ < bos >, x 1 , . . . , x T ] be an input sequence of length T +1 to the model. Then, there exists W Q , W K , W V , W O , W 1 , and W 2 matrices, such that when T < M , o T > H ; and when T > M , o T < H .

Proof. Our proof only specifies the weights of a single attention head in the first layer. In this parameterization, we only require the first three dimensions of the hidden states and regard the third dimension as the specific dimension. As for the remaining heads, they can be arbitrary as long as they don't override the first three dimensions. This doesn't pose any challenges, considering that Transformers used in practice usually have a very large model dimension d .

First, we construct the word embedding matrix W E ∈ R d × V , where each column is the embedding of a token in the vocabulary. We construct W E such that it always sets the first dimension of every embedding vector to be 1 , and sets the second dimension to 0 except the token <bos> . We always assume <bos> is the first token in the vocabulary, i.e. the first column. Then, we have:

where e d,i ∈ R .

Then, we use the word embedding matrix W E to compute the embedding H (0) :

Second, for head dimensions h ≥ 1 , we construct the weights W K , W V , W O of the first attention head in the first layer. W Q can be any arbitrary matrix. Specifically,

W K reads from the first dimension of the hidden state. Since all word embeddings have 1 in their first dimension, this parameterization will result all key vectors to be the same. Considering k i = W K h (0) i , then:

We use W Q to compute the query vector q t by applying q t = W Q h (0) t :

where q j ∈ R can take any arbitrary value.

Next, we compute the attention weight vectors α :

where α ∗ = q 1 + q 2 + · · · + q h .

Then, we apply SoftMax to compute the attention weight score. Since all α ∗ are the same, we have:

Now, we compute the value vectors by applying v i = W V h (0) i :

Then, we compute the attention value:

Finally, we compute the output of the attention head by applying W O :

From Equ.24, it can be observed that when t < M , the hidden state value on the third dimension is greater than H , indicating a successful extrapolation. Conversely, when t > M , the hidden state value on the third dimension is less than H , indicating a failed extrapolation.

Next, we provide the proof for Theorem 3.2 as below:

Theorem E.2 (PE Extrapolation) . Let x = [ < bos >, x 1 , . . . , x T ] be an input sequence of length T+1 to the model. Consider an simple relative PE schema where dot product between query q t and key k i at positions t and i ( t ≥ i ) can be expressed as: ⟨ q t , k i ⟩ := q T t k i -( t -i ) . Then, there exists W Q , W K , W V , W O , W 1 , and W 2 matrices, such that when T < M , o T > H ; and when T > M , o T < H .

Proof. We regard the third dimension as a specific dimension. And there is a threshold in this dimension in the second layer. Our proof construct the matrices on the first two layers. The first layer is used to extract position information. The second layer will generate a significant hidden state value in its third dimension. If the input length exceeds the max training length, the hidden state value will exceed the threshold.

(1)Consider the first layer operation.

First, following Theorem E.1, we construct the same word embedding matrix W E .

For head dimensions h ≥ 3 , we construct the weights W K and W V for the first attention head in the first layer. W Q still can be any arbitrary matrix. Specifically,

Since the elements in W K is all 0 , by applying k i = W K h (0) i , we have:

We use W Q to compute the query vector q t by applying q t = W Q h (0) t :

where q j ∈ R can take any arbitrary value.

Next, we compute the attention weight vectors α . Consider q T t k i = 0 and position information t and i , we have:

Then, apply SoftMax to α , we have:

where define S ( t ) = ∑ t -1 j =0 e -j . We adopt its abbreviation S .

Now, we compute the value vectors by applying v i = W V h (0) i :

Then, we compute the attention value:

After that, we compute the output of the attention head by applying W O (using the same construction as Equ.16 in Theorem E.1):

Then, e -( t -1) S can be regarded as a monotonically decreasing function of t . Then through the MLP feedforward layer, we can always let the MLP layer recover the value of t .

Therefore, we can get that:

where position information is embedded in the third dimension of hidden state.

(2)Consider the second layer operation.

For head dimension h ≥ 3 , we construct the weights W Q and W K for the first attention head in the second layer. W V follows the Theorem 3.1 setting. Specifically,

Then, by applying q i = W Q h (0) i , we have:

where q j can be arbitrary value for j ≥ 4 .

By applying k i = W K h (0) i , we have:

Then, we have:

Next, according to the PE definition, apply the position information t and i to the dot production, we have:

̸

, which means ⟨ q t , k i ⟩ = ⟨ q t , k l ⟩ for any l = i .

Then, we know that, where α ∗ = ∑ 4 ≤ j ≤ t q j .

Then, apply SoftMax to α , we have:

Next, follow same process as Equ.(21) (22) (23) (24) in Theorem E.1, we can get same result in the second layer. It shows the limitation of extraordinary for position encoding.

Next, we provide the proof for Theorem 3.3 as below:

Theorem E.3 (Weave PE Extrapolation) . Let N be a positive constant. Consider a simple weave PE extrapolation schema: when t -i < N , W ( t -i ) = t -i ; and when t -i ≥ N , W ( t -i ) = N . Then, the attention dot product is fixed as below:

, where N ≪ M . Then, applying W Q , W K , W V , W O , W 1 , and W 2 matrices from Theorem 3.2, we have when T > M , o T > H .

Proof. Our proof is closely related to Theorem 3.2. We completely adopt the LLM setting by Theorem 3.2, and only modify the representation of the relative position. The purpose is to illustrate that effective extrapolation beyond the maximum window length can be achieved only by rearranging the relative position.

Following Theorem E.2, we can get H (0) .

Then, compute the attention weight vectors α . Consider q T t k i = 0 and position information t and i . By applying this Extrapolation schema, we have:

Follow Eqs.28 and 29, and apply SoftMax to α , we have:

where t > N and define S # ( t ) = ∑ N -1 j =0 e -j +( t -N ) e -N . We adopt the abbreviation S # where there is no confusion.

Then, compute the value vector by applying v i = W V h (0) i , still have:

Then, we compute the attention value:

After that, we compute the output of the attention head by applying W O :

Notice that we set the MLP feed-forward process simulate a a piecewise function about e -( t -1) S ( t ) , which produce a Integer t to recover the position.

Considering t > N , we know e -( t -1) S ( t ) < e -( t -( t -N )) S # ( t ) < e -( t -( t -N )) S ( N +1) . And the e -( t -( t -N )) S # ( t ) function decreases extremely slowly as t increases, almost equal to e -( t -( t -N )) S ( N +1) .

At the same time, the e -( t -1) S ( t ) function decreases very quickly when t is relatively small; As t increases, it also decreases extremely slowly.

Combining the characteristics of these two functions, we can always choose a relatively small N to keep it away from M (that is N ≪ M ), so that the value recovered after e -( t -( t -N )) S # ( t ) passing through MLP process is around N +1 , because e -( t -( t -N )) S ( N +1) can recover N +1 and e -( t -( t -N )) S # ( t ) is almost equal to e -( t -( t -N )) S ( N +1) . For the sake of simplicity, considering the rounding characteristics of piecewise functions, we might as well set e -( t -( t -N )) S # ( t ) to recover N +1 .

Therefore, we can get the hidden state like that:

Continue the second layer operations.

By applying q i = W Q h (0) i , we have:

where q j can be arbitrary value for j ≥ 4 .

By applying k i = W K h (0) i , we have:

Then, consider t > N , and apply Extrapolation .

when t -i < N ,

when t -i ≥ N and i > N ,

when t -i ≥ N and i < = N ,

⟨

q

t

,

k

i

⟩ -

f

rel

(

t

-

i

) =

4

∑

j

h

Let τ = ∑ 4 ≤ j ≤ h q j . We can get the attention score:

where third -part represents Equ.49, second -part represents Equ.50, and first -part represents Equ.51.

Apply SoftMax to α , we can get that:

Since the first element of α is the maximum value, ˆ α 1 must be greater than 1 t .

Next, similar process, we compute the value vectors by applying v i = W V h (0) i :

Then, we compute the attention value:

q

j

+(

N

+1

-

i

)

-

N

(51)

Then, we compute the output of the attention head by applying W O :

When t > M , due to ˆ α 1 > 1 t , it's possible to satisfy ˆ α 1 · M -1 > 0 condition, such that the hidden state value o t in the third dimension is greater than H .

We conjecture that while the attention mechanism in RoPE Su et al. (2023) involves multiplying by the rotation matrix of the relative position, it's essential to acknowledge that neural networks possess the capability to model any function. This implies the potential to decompose the relative position information into an additive form, as exemplified in our extrapolation scheme detailed in Theorem 3.3.

Next, we provide the proof for Corollary 4.1 as below:

Corollary E.4 (Mesa Extrapolation) . Let N be a positive constant. Consider a simple Stair PE extrapolation schema, and the attention dot product is fixed as:

where N ≪ M , I = N + ⌈ t -i -N E ⌉ , and the extrapolated width E is a constant. Then, Apply W Q , W K , W V , W O , W 1 , and W 2 matrices from Theorem 3.2. Although T > M , it still o T > H .

Proof. Due to the striking similarity between Stair PE and the positional arrangement scheme proposed in Theorem 3.3, the proof process largely mirrors Theorem 3.3. Following the proof structure of Theorem 3.3, in the first layer, according to Equ.45 and .46, we obtain the hidden state matrix as follows:

where I is defined on Stair PE extrapolation schema.

In the proof of the second layer, we still get the similar result like Equ.52, as below:

where C is a constant, and each element within repeat-part is smaller than τ -0 .

Apply SoftMax to α , we can get that:

which also shows that the first element of α is the maximum value and ˆ α 1 must be greater than 1 t . Consequently, following Equ.56, we reach the same conclusion.

F Probe Experiment Visualization

We hypothesize that when the input length surpasses the effective window length of the model, some dimensions' values in the exceeded positions will experience a jump as the position changes.

To investigate this jump phenomenon's correlation with extrapolation failure, we design the following experiment: we adopt a standardized input by repeating the word "hello" 8000 times, resulting in 8001 tokens (automatically fill in initial token) after tokenization. For a more accurate explanation, we take the LLaMA2-7B-Chat model and list the sequences converted by the tokenizer, as follows:

where 1 denotes the initial token < bos > , and 22172 denotes the word "hello". It's noted that the initial token < bos > is filled in automatically.

F.1 Normal Case

We input this token sequence into the model and observe the hidden state. We focus on the hidden states produced by the first 11 layers, specifically selecting the position intervals from 4000 to 5000 and the last 1000 positions. We then concatenated the hidden states from these two intervals. This selection was deliberate, considering that the LLaMA2-7B-Chat model's training length is 4096 , implying the effective input window is in proximity to this location.

Following this, we created a matrix graph, yielding the following results Fig.11:

0-

layer: 0

layer: 1

layer: 2

layer: 3

layer: 4

layer: 5

layer: 6

layer:7

layer:8

layer: 9

layer: 10

layer: 11

100-

200-

Dimension

300

500

0

9000

500

400

1000

500

009

0

1000

0

1000

0

1000

0

1000

0

1000

0

1000

0

1000

Q

1000

0

1000

Q

1000

0

1000

0

1000

Token Position

In Fig.11, the X-axis denotes the position of the input token, ranging from 0 to 1000 , representing tokens at positions 4000 to 5000 . The scale of 1000 -2000 represents tokens at positions 7001 to 8001. The Y-axis signifies the token dimension, ranging from small to large. The LLaMA2-7B-Chat model has a total of 4096 dimensions, and we display the first 640 dimensions in Y-axis. Red color means greater than 0 , and blue color means less than 0 . Fig.11 shows some noticeable jump from 0 -th to 640 -th dimensions, especially at the 1000 scale for the original model.

F.2 Extrapolation Case

Given the effectiveness of the extrapolation method ReRoPE in extending to lengths of up to 8 k, we apply it to the LLaMA2-7B-Chat model. Employing the same conditions as those in the Normal Case settings, the results are as follows:

In Fig.12, the obvious jumping phenomenon is successfully suppressed. It can be seen that each dimension at different positions still maintains consistent values. This illustrates that by suppressing sudden changes in the values of these dimensions, extrapolation can be made successful even outside the effective window length.

0 -

layer: 0

layer: 1

layer: 2

layer: 3

layer: 4

layer: 5

layer: 6

layer: 7

layer: 8

layer: 9

layer: 10

layer: 11

100 -

200-

0000

Dimension

300

500

500

400

1000

500

009

0

1000

0

1000

0

1000

0

1000

0

1000

1000

0

1000

1000

1000

口

1000

0

1000

0

1000

Token Position

Overall, comparing Normal Case and Extrapolation Case again, it shows that using the extrapolation method can suppress mutations in dimensions, thereby making the extrapolation successful.

Additionally, for layers 0 and 1 , the hidden state values of different positions in each dimension are nearly identical. According to our Theorem 3.2, layer 1 may extract position information, while subsequent layers determine distances beyond the effective window, generating a positive or negative jump signal. We speculate that this signal change may not be strictly positive or negative, but there is a threshold. When the signal exceeds the threshold, the extrapolation will fail; when the signal does not exceed the threshold, the extrapolation will be successful.

We also use the LLaMA-3B model and the Vicuna-13B model to compare and demonstrate the effects of using Origin and ReRoPE, as follows on below Fig.15 and .16.

F.3 Validating extrapolation using observed thresholds

We continue our probe experiments using the "hello", and input sequences is up to 16000 tokens. These were validated on the LLaMA2-7B-Chat model. Given that the hidden state in LLaMA2 has a total of 4096 dimensions, we selected the first 20 dimensions to search observable thresholds. In the main text, we present the prediction results for the first and 6 -th dimensions. The results for the 7 -th and 9 -th dimensions are shown in Figure 13.

Layer:1, Dimension:7

Layer:2, Dimension:7

Layer:1, Dimension:9

Layer:2,Dimension:9

0.2-

0.1

0.0-

S

0.1-

len

P!H

-0.2

0.3 -

2k

5k

8k

11k14k

2k

5k

8k

11k

14k

2k

5k

8k

1ik14k

2k

5k

8k

11k

14k

Origin

Dynamic-NTK

Position

Figure 13 also shows the same results as Figure 2 in the main text. When the length of the input sequence is around 12k, the hidden state values of Dynamic-NTK in the 7-th and 9-th dimensions surpass the thresholds, implying extrapolation failure. Due to not outside the thresholds, ReRoPE shows successful extrapolation.

Our experiments also reveal that not every dimension exhibits a clear threshold. Considering that the Transformer model is a type of neural network, we hypothesize that these dimensions with threshold undergo significant signal changes due to the characteristics of activation functions. These changes are amplified through successive layers, eventually leading to model output failures. We believe

this internal mechanism explains why observable thresholds can predict the success or failure of extrapolation.

We also validate our findings using the Vicuna-13B model, and the results are shown in Figure 14:

Layer:1, Dimension:4

Layer:4, Dimension:4

Layer:1, Dimension:10

Layer:4, Dimension:10

0.01

0.8 -

1

-60°0

0.20

0.00-

0.7 -

0.08 -

0.15 -

StateValue

0.6 -

0.10

0.01

0.5

0.07 -

0.05

0.4

0.3

0.06

0.00-

0.2

0.05 -

0.05 -

0.03-

0.1

0.10

2k

5k

8k

11k

14k

0.0

2k

5k

8k

11k

14k

0.04

2k

5k

8k

11k

14k

2k

5k

8k

11k

14k

Origin

Dynamic-NTK

Token Position

Figure 14 shows the hidden state value on the specific 4 -th and 10 -th dimensions. Similar to the observations with LLaMA2-7B-Chat model, we find no significant differences for the hidden state values at the first layer. However, substantial differences are observed at the fourth layer. Based on these observed thresholds, we can predict that Dynamic-NTK will fail to extrapolate when the input length approaches around 9 k. Conversely, ReRoPE can extrapolate further. These predictions align well with our subsequent experiments, further validating the feasibility of using observed thresholds to determine extrapolation success.

In contrast to LLaMA2-7B-Chat model, we observed threshold phenomena at the fourth layer in the Vicuna-13B model. We hypothesize that although our theoretical model predicts the threshold to occur at the second layer, considering the integration of positional information from preceding layers, this still aligns with our theoretical framework.

G Stair PE and Self-Extend

The formulation of Self-Extend is as follows:

Let t denote the position of query, i denote the position of key, i denote the position of key, W denote the neighbor size and G denote group size.

According to its equation 3 in Self-Extend (Jin et al. (2024)), P = P // G , and the shifted relative position ( W -W // G ) , we can get the position of query as:

By RoPE, their relative position between t and i (consider t > i ), is remapped to:

It's worthy noting that only under the necessary and sufficient condition ( t mod G > = i mod G ), we get that

Next, when W mod G == 0 , the equation of Self-Extend becomes:

and the position of key as:

If we adjust this flooring operation to ceiling operation, and replace N with W and E with G, then Stair-PE is equivalent to Self-Extend.

In summary, under conditions t mod G > = i mod G and W mod G == 0 , as well as changing the flooring operation to ceiling operartion in Self-Extend, these two formulas are equivalent. Except for this condition, they yield slightly different results.

Note because W and G are constants, the condition W mod G == 0 can be met easily. But i and t change with positions, therefore there will always be positions where the condition t mod G > = i mod G is not satisfied .

̸

0

layer: 0

layer: 1

layer: 2

layer: 3

layer: 4

layer: 5

layer: 6

layer: 7

layer: 8

layer: 9

layer: 10

layer: 11

100-

200

0000

Dimension

300

500

400

500

500-

009

0

1000

0

1000

0

1000

0

1000

0

1000

0

1000

0

1000

0

1000

0

1000

0

1000

0

1000

0

1000

Token Position

0

layer: 0

layer: 1

layer: 2

layer: 3

layer: 4

layer: 5

layer: 6

layer: 7

layer: 8

layer: 9

layer: 10

layer: 11

100 -

200

Dimension

300

500

400 -

500

500-

600-

0

1000

0

1000

0

1000

0

1000

0

1000

0

1000

0

1000

0

1000

0

1000

0

1000

0

1000

0

1000

Token Position

(b) Visualization of Hidden states with ReRoPE for the first 11 layers and the first 640 dimensions. Observation: Each dimension stays consistently at each position

0 -

layer: 0

layer: 1

layer: 2

layer: 3

layer: 4

layer: 5

layer: 6

layer: 7

layer: 8

layer: 9

layer: 10

layer: 11

100-

200-

uo(suawjg

300

1000

2000

0

1000

400-

2000

500

- 009

0

1000

0

1000

0

1000

0

1000

0

1000

0

1000

0

1000

1000

0

1000

1000

0

1000

0

1000

Token Position

0

layer: 0

layer: 1

layer: 2

layer: 3

layer: 4

layer: 5

layer: 6

layer: 7

layer: 8

layer: 9

layer: 10

layer: 11

100-

200

Dimension

300

1000

2000

0

1000

400 -

2000

500-

600-

0

1000

1000

0

1000

0

1000

1000

0

1000

1000

1000

0

1000

0

1000

0

1000

0

1000

TokenPosition

(b) Visualization of Hidden states with ReRoPE for the first 11 layers and the first 640 dimensions. Observation: Each dimension stays nearly same without color jump as position changes.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The abstract and introduction of our paper accurately reflect its contribution and scope.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We discuss the limitations on Section Discussion.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: We give a full proof on Appendix.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We disclose all the needed information.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We provide code and data.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We specify all the training and test details.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We provide detailed explanations about report.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g., negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We provide the hardware and software information about experiments.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We follow the NIPS Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: We discuss the impacts on section Discussion.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: Our paper poses no such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We follow the license and terms.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: We don't release new assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: Our paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: Our paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper conducts a theoretical analysis to help understand the No Position Encoding. Also, the paper proposes weave position encoding to achieve improved extrapolation performance without additional cost. Also, the paper introduces the weave PE method, Mesa-Extrapotion, which recalculates the position ID to reduce the gap between training and inference. Finally, the paper conducts experiments to prove the effectiveness of Mesa-Extropoation.\n\nSTRENGTHS:\n* The presentation of wave PE, Star PE, and Mesa-Extrapolation is clear. The author also provides the details of wave PE, Star PE and Mesa-Extrapolation to help understand the concepts\n* The author conducts experiments to prove the effectiveness of the proposed Mesa-Extrapolation.\n* The author also further analyzes the Latency & Memory Usage of the proposed Mesa-Extropoaltion.\n* The paper discusses the limitations for further discussion.\n\nWEAKNESSES:\n**Major Concerns**: It seems that the proposed method Star PE is the same as Self-Extend LLM [1]. If possible, I sincerely hope that the author can address the following concerns:\n* **Concern 1**: The Figure 1 Star PE implementation result does not match the Equation proposed in Section 4.1 Page 5. In Figure 5, when t-i is 5, the implementation result of Star PE is 4. However, according to the Equation proposed in Section 4.1 Page 5, the implementation result should be N+ $\\lceil (t-i-N)/E \\rceil$=4+$\\lceil (5-4)/2 \\rceil$=5. Hence, to match the implementation result of Figure 1, the Star PE calculation equation should be N+ $\\lfloor (t-i-N)/2 \\rfloor$.\n\n* **Concern 2**: The Equation of Star PE is almost the same as Self-Extend LLM. When t-i is small than N, both Star PE and Self-Extend LLM employ normal relative distance. When t-i is larger or equal to N, we discuss it below.\n   * The Equation of Star PE is N+ $\\lfloor (t-i-N)/E \\rfloor$ (as shown in Figure 1), while N is called the extrapolated position and E is called the extrapolated width, and t-i is the relative distance. \n   * The Equation of Self-Extend LLM is $(t-i)//W + (W- W//G)$=$W+ (t-i)//G - W//G$, while W is called neighbor window size and G is called group size. Apparently, when W%G==0, the Equation of Self-Exntend LLM becomes $W+ (t-i)//G - W//G$=$W+ (t-i-W)//G$= $W+\\lfloor (t-i-W)/G \\rfloor$. Then, we change the notation W to N and the notation G to E, we have N+ $\\lfloor (t-i-N)/E \\rfloor$, which is the same as Star PE.\n* **Concern 3**: If possible, could the author compare the performance between Mesa-Extropolation and Self-Extend LLM?\n* **Concern 4**: when the output sequence length $L_{generate} \\gg L_{input}$, will the time cost also becomes O($L_{generate}^2$)?\n\nBased on the above concerns, the paper may need to rethink the major contribution. The proposed Mesa-Extrapolation seems to make sense and may benefit society, while the paper should clarify its original contribution.\n\n\nReference:\n\n[1] Jin, H., Han, X., Yang, J., Jiang, Z., Liu, Z., Chang, C. Y., ... & Hu, X. (2024). Llm maybe longlm: Self-extend llm context window without tuning. arXiv preprint arXiv:2401.01325.\n\nQUESTIONS:\nPlease see the above question and concerns in weakness.\n\nLIMITATIONS:\nYes, the authors have addressed the limitations. We may further discuss and analyze the Mesa-Extropolation in other areas.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-24"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the length extrapolation of LLMs.\n1. It provides a theoretical analysis of why NoPE and PE fail to extrapolate beyond a certain length. Previous work has shown that this failure is related to the explosion of hidden states as positions increase. This paper demonstrates that both NoPE and PE suffer from this hidden state explosion, using a constructive approach to illustrate the existence of Transformer weights.\n2. It proposes weave PE, a simple adaptation of PE that theoretically addresses the extrapolation issue. It also provides a simple implementation of weave PE, using a chunk-based triangular attention matrix. Then, it demonstrates that the proposed extrapolation scheme matches the performance of prior length extrapolation methods, such as Dynamic-NTK.\n\nSTRENGTHS:\n- Great theory explains the failure of NoPE and PE in length extrapolation.\n- Proposes weave PE, derived from the theoretical analysis, which also works well in practice.\n- Shows good empirical results in passkey retrieval, language modeling, and summarization.\n\nWEAKNESSES:\n1. Methodological comparison with $\\Lambda$-Attention\n\nThe proposed Stair PE resembles the $\\Lambda$-attention of LM-Infinite & Streaming-LLM, yet with differences in 1) the additional attention at the bottom, and 2) a different length extrapolation scheme, Meta-Extrapolation. In the experiments, Meta-Extrapolation significantly outperforms LM-Infinite & Streaming-LLM. Could the authors provide the intuition behind these empirical gains?\n\n---\n2. Empirical comparison with Dynamic-NTK\n\nDynamic-NTK outperforms Meta-Extrapolation on the summarization task for mid-lengths of 7-11k, while Meta-Extrapolation shows better performance on summarization for shorter lengths of 4-6k and better language modeling fluency for lengths greater than 11k. Could the authors provide the intuition behind these results?\n\n---\n3. Relation between input sequence length $T$ and effective length $M$\n\nThe theorems only show the existence of an effective length $M$, but do not provide intuition on the scale of $M$, such as the ratio over the input length $M / T$. Could the authors provide some intuition on this? If I understand correctly, $M$ is set from the construction of the Transformer weights, so can it be controlled to an arbitrarily large number?\n\n---\nEditorial comments\n\n- The fonts of the figures and tables are too small. Please make them more readable.\n- Some parts of the writing are mechanical. For example, lines 116-120 do not provide meaningful information. It would be great to discuss the implications of the theorems in natural language. For instance, both theorems state the failure of length extrapolation in NoPE and PE, rather than just \"revealing the internal mechanism of extrapolation.\"\n\nQUESTIONS:\nSee Weaknesses.\n\nLIMITATIONS:\nWell discussed. Extending this approach to fine-tuning would be an interesting next step.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-05"}

comments: SUMMARY:
The paper proposes a positional embedding scheme to address the extrapolation issue: train on short sequences, evaluate on longer sequences. Authors propose a theoretical framing of the positional embeddings contribution to attention. They apply their analysis to NoPE (No Positional Embedding) and to standard PE, and RoPE. They propose the Mesa-Extrapolation idea where input tokens are organized so that attention is paid to nearby tokens and those at other key positions. Authors validate their findings with empirical evidence on several benchmarks and applications.

STRENGTHS:
The paper is about a very relevant topic which has attracted a lot of attention lately. The paper proposes a simple approach to solve the problem which seems to be easy to adapt to different positional embedding models. Some of the numerical experiments are encouraging.

WEAKNESSES:
The theory part of the paper is hard to read and I am not sure about its usefulness. Result appear hand-wave-y and vaguely stated. For example the definition of the threshold H in the Assumption is surprising (see questions). Numerically, experiments on language modeling and Summary of Tasks do not seem to show the method's claims.

QUESTIONS:
1. Can authors explain the threshold definition: "When o > H, LLM extrapolates successfully. Once o < H, LLM extrapolation
fails." Is there a typo and inequalities are reversed?

2. In Fig 2, why dim 1 & 6 are of interest?

LIMITATIONS:
--
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2024-07-10

comments: SUMMARY:
This paper introduces a new LLM length extrapolation method, called Mesa-extrapolation, which utilizes a chunk-based triangular attention matrix and applies stair PE. The proposed method is based on theoretical analysis. The paper conducts extensive experiments on passkey, PPL, summarization to demonstrate the effectiveness.

STRENGTHS:
1. The paper provides a theoretical analysis to prove the effectiveness of meticulous weave position with PE for length extrapolation.
2. The proposed method is efficient and is proved to be effective through extensive experiments.

WEAKNESSES:
1. The passkey retrieval experiment is simple, good performance on the passkey is far from a real usable context window. Please consider to add evaluations on Ruler[1] and RepoQA[2]

2. The achieved context window is limited. 



[1] https://arxiv.org/abs/2404.06654

[2] https://evalplus.github.io/repoqa.html

QUESTIONS:
1. Length extrapolation is a necessary technique, but the current extrapolation length is very limited. Considering that there are already many models that have undergone long context window extension, such as phi3-mini-128k, can your proposed method continue to perform length extrapolation on these long context LLMs? If so, it will significantly enhance the impact of your method

2. If I understand correctly, the proposed method is mainly for those with PE. Why is there a need to prove NoPE? Is NoPE your baseline?

3. The proposed Mesa-extrapolation is somehow similar as a variant of "sliding window attention + attention sinks". Could the author explain why mesa-extrapolation is theoretically superior compared to sliding window attention and attention sinks?

LIMITATIONS:
See the weakness and question sections.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-12

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose a weave position encoding method to enhance LLMs’ inference performance when the input context window exceeds the training context window. This method can be integrated into existing pretrained LLMs without additional finetuning. To support their findings, the authors conducted theoretical analyses on the failure reasons of various position encoding methods, including those without position encodings. They demonstrate that the significant shift in the hidden state’s value range, when input token positions exceed the maximum context length, is the cause of this phenomenon.\n\nSTRENGTHS:\nOne of the strengths of the proposed method is that it can be integrated into existing pretrained LLMs without requiring any additional finetuning. This makes the method highly practical and easy to implement, saving both time and computational resources.\n\nThe method has demonstrated excellent performance in pass key retrieval tasks, showcasing its effectiveness in real-world applications. This indicates that the proposed approach not only works in theory but also delivers tangible improvements in practical scenarios.\n\nThe authors have conducted comprehensive theoretical analyses to understand the failure reasons of various position encoding methods, including those without position encodings. This thorough investigation provides a solid foundation for the proposed method and enhances its credibility\n\nWEAKNESSES:\nThe proposed position encoding method, while promising, does not consistently improve performance across different tasks. This inconsistency suggests that the method may not be universally applicable or reliable in every context, potentially limiting its overall utility. \n\nAdditionally, the main narrative of the paper emphasizes the method’s ability to handle extrapolation beyond the training context window. However, given the observed variability in improvements, it would be more accurate to adjust the claims to better reflect the method’s performance, providing a more balanced and realistic presentation of the work.\n\nQUESTIONS:\nThe caption for Figure 1 is not sufficiently informative.\n\nAdditionally, it is unclear how the failure of an LLM is measured in Section 3.4 and Figure 2. \n\nThe experiments visualizing hidden state values in Figure 2 would have been more effective if conducted on the same task and with the same setup as Figure 3. This alignment would allow for a clearer connection between the findings in Figures 2 and 3.\n\nminor typos:\n\nTheorem 3.2: an simple -> a simple\n\nLine 161: defer to -> refer to\n\nLine 242-243: a significant disruptions\n\nLIMITATIONS:\nThe authors have included a limitations section; however, it reads more like a discussion of future work rather than addressing the actual limitations of the current study.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-15"}