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

Thanks for your work.

major concern

Thank you very much for your suggestions. The differences between our method and Self-Extend can be categorized into three aspects:

Firstly, from a methodological perspective, our designed Stair PE is not only applicable to RoPE but also to other positional encoding methods, such as ALiBi. Self-Extend is only applied to RoPE. Additionally, Stair PE differs from the self-extend design approach, which involves implementing group attention and standard attention and then merging them through masking. This implies that it requires calculating the attention matrix twice, resulting in higher memory and time overhead. In contrast, our approach uses a chunk-based triangular attention matrix that significantly reduces computational and memory costs. We utilize Stair PE at the base of this chunk-based triangular attention matrix, and for practical applications, approximate methods can be employed to further save memory and reduce time costs.

Secondly, theoretically, we provide a theoretical proof for the Weave PE approach, demonstrating that it can achieve extrapolation. In contrast, the self-extend approach only provides experimental evidence without a theoretical explanation.

Finally, from an experimental standpoint, we further conducted comparisons of the two approaches and demonstrated that our method shows comparable advantages with lower memory consumption and inference latency (see attached PDF in Fig 5).

concern 1

We apologize for not checking it carefully. However, there is no need to modify the formula. We can simply change it to N=3.

concern 2

When we set the Stair PE formula as your derivation on # concern 1, your derivation is correct and aligns with Self-Extend LLM formulation. However, there is a crucial mismatch: self-extend derives the positional encoding according this formulation just for its query vectors (please refer to its pseudocode on algorithm 1 on its page 12).

we briefly list below:

g_pos = pos // g_size # the floor operation
shift = w_size - w_size // g_size
s_g_pos = g_pos + shift
g_q = apply_pos_emcode (q , s_g_pos )
g_k = apply_pos_emcode (k , g_pos )

In its code, s_g_pos is correspond to your formula, but it’s not relative position encoding.

Note that Self-Extend is only applied to RoPE positional encoding. Thus, with RoPE, the relative positional encoding is obtained by subtracting the key's positional encoding from the query's positional encoding. In contrast, our defined Stair PE directly targets the relative positional encoding. This is also why our Stair PE is more broadly applicable, not only to RoPE but also to other methods like ALiBi, whereas Self-Extend is only applicable to RoPE.

According to these formulations, the relative positional encoding implemented by self-extend is different from that implemented by Stair PE. We have provided a clear comparison in the uploaded PDF (see Fig 6 ).

concern 3

The comparative results are provided in the uploaded PDF(see Fig 1, 2, 3, and 5). In both the NIAH and Passkey tasks (corresponding to Figure 3 in our paper), our method demonstrates comparable advantages. Additionally, our method performs well in terms of memory consumption and inference latency.

concern 4

The output of large models is divided into the prefill stage and the decoding stage. The time savings from mesa-extrapolation primarily occur in the prefill stage, which involves processing the N^2 matrix corresponding to the input sequence. During the subsequent token-by-token output process, the time complexity is O(L_generate + L_input ). This mainly concerns the last token and its computation with all preceding tokens, and this complexity is linear with respect to the sequence length.

Dear Authors,

Thank you very much for your rebuttal. After reading the author's rebuttal, I still have questions about the contribution of Self-Extend and Star PE.

• After the author's rebuttal, I confirm that the Star PE is equivalent to Self-Extend PE.

  • First, their formulation is equivement, after adjusting floor and ceiling operations. The core part is the same: increase 1 position number after $E$ steps and begin the operation at the $N^{th}$ step.
  • Then, the author claims that the Self-Extend PE cannot work on relative positional encoding. However, this is not correct
    • The author may check Figure 2 in the Self-Extend LLM arxiv paper. This paper Figure 2 presents how they process the relative distances between queries and keys. Self-Extend can be directly applied to Alibi, as Alibi only needs to know the relative distance between queries and keys.

• I appreciate the comparison between Mesa-Extropolation between Self-Extend. The result suggests the effectiveness of the Chunk-based Triangular Attention Matrix.

To increase the score, I need the author to

• Rethink the relation between Star PE and Self-Extend
• Rewrite the part of Star PE.
• Correctly claim the contribution of the paper, while I now have relatively high confidence to say that the Star PE and Self-Extend are equivalent. Therefore, the paper cannot claim that the Star PE is its contribution, while the theoretical analysis can be kept as part of contribution.

Overall, I have to keep the current score and increase my confidence, as the paper contribution Star PE is equivalent to Self-Extend.

Thank you very much for recognizing our work and providing your suggestions.

weakness 1

For decoder-only architectures, the model's output is based on the next-token prediction. The last token in the input is crucial, as it generates the next token. For the last token to output the correct next token, it must attend to all previous tokens (principle of self-attention mechanism), especially those with critical information.

The attention matrices used by Streaming-LLM and LM-Infinite both discard middle tokens, inevitably resulting in loss of information from these tokens. This problem becomes severe when the middle tokens carry important information.

Our approach appends all tokens at the bottom, ensuring that information is not lost. For overly long inputs, we use Stair PE to reorganize relative position encoding, maintaining the integrity of the information.

weakness 2

We believe that the inherent capabilities of LLMs set the upper limit for these methods. As the input length increases, especially beyond the maximum training window, errors are bound to accumulate regardless of the method used, which will affect the model's performance.

The NTK method scales the rotation angle of RoPE positional encoding. We hypothesize that within an 11k input length, the errors introduced by adjusting the rotation angle are acceptable for large models.

For mesa-extrapolation, we reuse previous trained positions to minimize errors as much as possible. We speculate that with a fixed extrapolation width param (E = 50), the 7k-11k range may spread the model's attention more thinly compared to the 4k-6k range. For the summary task, the model needs to pay closer attention to the dispersed information within the context, which requires a finer extrapolation width to minimize errors. We hypothesize that optimizing the extrapolation width could alleviate or improve performance in the 7k-11k range.

weakness 3

Yes, M is set within the Transformer weights we constructed. This M is typically related to the maximum training window length of the model. For example, if the maximum window length for Llama is 2k, then M can be considered as 2k.

Therefore, we think that M can't be controlled to an arbitrarily large number.

Editorial comments

Thank you for the modification suggestions. We made corresponding adjustments according to the comments.

Dear Reviewer bJoG,

We sincerely appreciate your positive feedback, and it is incredibly gratifying to have our work highly recognized by you.

As the discussion progresses, if you have additional questions, please let us know, and we will be happy to provide further clarification.

We will express our deep gratitude to you in the Acknowledgments section of our paper.

Once again, thank you very much.

Thanks for your work.

Weakness 1 & question 1: “Theory is hard to read and unclear definition of the threshold H”

Definition of the extrapolation success or failure: When a large model continuously produces valid next-tokens for a given long input sequence, we define this as successful extrapolation. Conversely, if it outputs invalid next-tokens, we refer to this as failed extrapolation. Therefore, the key to extrapolation failure is the output of invalid next-tokens, which is caused by anomalies in the model's output. We speculate that these anomalies arise because certain parts of the model receive input values that exceed their acceptable range.

Definition of the threshold H: For convenience in our analysis, we consider the outputs of the preceding layer In a multi-layer neural network, which serves as the input for the next. To ensure normal subsequent outputs, the values from the preceding layer need to remain within a reasonable range. This implies that there must be a bound on these values. We refer this observable boundary as the threshold H. This threshold can be either an upper bound or a lower bound. In our work, for the sake of theoretical analysis, we chose the lower bound as our threshold.

Definition of hidden state values o: the output of each layer in an LLM, which also serve as the input for the next layer.

Note that in our assumption, we assume that for specific dimensions in specific layers, there exists an acceptable range of input values (i.e., hidden state values o). The lower bound of this acceptable range serves as our threshold H. 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.

We apologize if it wasn't clear enough. We will add these refined definitions in the final version of the paper.

weakness 2:

Despite observing minor performance degradation in the 8k-11k token range for summarization tasks, our method consistently exhibits strong extrapolation capabilities across a diverse set of tasks, including passkey retrieval, language modeling, and summarization, as corroborated by Reviewers LfLf and bJoG.

We also noted that the performance in the 4k-6k range is good. We speculate that with a fixed extrapolation width param (E = 50), the 7k-11k range may spread the model's attention more thinly compared to the 4k-6k range. For the summary task, the model needs to pay closer attention to the dispersed information within the context, which requires a finer extrapolation width to minimize errors. We hypothesize that optimizing the extrapolation width could improve performance in the 7k-11k range.

Q1

The inequalities in our Assumption are not reversed since we consider H as the lower bound.

We will add clear definitions and explanations in the final version.

Q2

Dimensions 1 and 6, as well as dimensions 7 and 9 in the appendix (line 1246), were selected from the first 20 dimensions to showcase whether the extrapolation would fail using the observed threshold for NTK and Rerope. Note that different inputs activate different parts. Figure 2 used 'hello' as the probe word. Additionally, we demonstrated that dimensions 14 and 19 can be used to predict extrapolation failures using 'good' as the probe word (please refer to the uploaded PDF in Fig 7 and 8 for more detailed diagrams).

Dear Reviewer PZAX,

As Discussion Period progresses, we would like to know if our responses have addressed your concerns.

We believe our responses address all of your main concerns; if they do, we would very appreciate it if you raised your score accordingly.

If not, please let us know so that we can address any serious concerns that you might still have.

Once again, thank you very much.

Thanks for your work.

Weakness 1:

We add evaluations on Ruler, and the results are available in the uploaded PDF (see Fig 1, 2 and 4). These results indicate that our method also performs well.

Q1 & weakness 2:

We perform experiments on long context window "microsoft/Phi-3-mini-128k-instruct" model using Ruler dataset. The results show that our method can further extrapolate to 192k, indicating the effectiveness of our method. We show it in the uploaded PDF (see Fig 4).

Due to the inherent limitations of the phi3-mini-128k-instruct model itself when handling this task, even within a 128k window, the phi-3 model does not achieve 100% accuracy as shown in the right-side of Fig 4. Our method effectively extends its window based on the original model's capabilities, rather than improving its capability.

Q2

Although existing LLMs typically incorporate positional encoding (PE) components in their architecture, recent studies have shown that the NoPE (no positional encoding) method may better facilitate the extrapolation capabilities of LLMs. The choice of including or excluding the PE component in a LLM is critical, as it cannot be altered once the model is trained. We believe it is essential to theoretically model the extrapolation performance of both NoPE and PE. We discuss this part in our related work.

Due to a lack of suitable candidates, NoPE was excluded from our baseline. Aside from the 1B-size pre-trained model (Kazemnejad, Amirhossein, et al. "The impact of positional encoding on length generalization in transformers." Advances in Neural Information Processing Systems 36 (2024).), we did not find any other suitable candidates. However, this pre-trained 1B model’s capabilities are very limited, making it difficult for evaluating basic NLP tasks.

Additionally, by simply incorporating training-free weave PE into a PE-type LLM, we can achieve longer extrapolation. Through these efforts, we hope to demonstrate that choosing weave PE offers more advantages compared to NoPE.

Q3

Mesa-Extrapolation appends all previous tokens at the bottom of its attention matrix, allowing the last token to attend to all preceding tokens. When the input length significantly exceeds the model's maximum training window, it uses Stair PE to reuse positions and reduce errors further. The attention sliding window approach, on the other hand, discards tokens in the middle section. This method restricts the last token's attention to only the head and tail tokens, keeping the total attended token length within the maximum training window and directly using trained positions. This approach sacrifices the information from the middle tokens. In context-based tasks, the middle tokens often contain important information, and losing these tokens can prevent the model from providing accurate responses to user queries.

Dear Reviewer hjGM,

Thank you for your valuable suggestions and the time you’ve dedicated to our paper.

As the discussion period draws to a close, we would like to know if our responses have addressed your concerns.

Following your suggestion, we have conducted verification on phi3-mini-128k, which significantly enhance the impact of our approach.

If our responses have resolved your concerns, we would greatly appreciate it if you could raise your score to a clear "accept".

Once again, thank you very much.

Thanks for your work.

weakness

Thank you for your suggestions. Our method shows good extrapolation performance on accuracy-related tasks, but we observe slight variability in extrapolation performance within mid-length (8k-11k) in the summary task. We will adjust our claims accordingly in the final version.

Q1

Illustration of Mesa-Extrapolation. The left figure shows the Chunk-based triangular attention matrix (before SoftMax operation) of Mesa-Extrapolation while a sequence of length 13 is fed into an LLM. The right figure shows an example of PE and Stair PE. The Stair PE is used to weave the relative position equipped by Mesa-Extrapolation.

Q2

In sec 3.4, we designed a probe experiment by repeating a word (e.g. “hello”) N times as input to the model, where N is the input length (the details are provided in Appendix, line 1206-1209). The reason for designing the probe this way is to eliminate the influence of different tokens since both of the input token and its position affect the hidden state values.

In Figure 2, we use a vertical black dashed line to indicate the position of maximum training length of the model. In this case, it is 4k for llama2-7b-chat model. The hidden state value at this position is designated as the observed threshold and marked with a horizontal red dashed line. When the hidden state value exceeds the red dashed line as the position changes, it signifies that the hidden state value has surpassed the threshold, suggesting a failure in extrapolation after that position. We will add these explanations to the final version.

Q3

We used a different experimental setup primarily for the following reasons: Figure 2 examines extrapolation failures, which primarily occur when input length surpasses the maximum limit, irrespective of the underlying task. Using the "hello" probe can clearly help us analyze the relationship between position and the corresponding thresholds. We then verify these predictions with other tasks, such as those depicted in Figure 3 or language modeling tasks shown in Figure 4. The results indicate that the threshold observed aligns with the results from the passkey task and language modeling task, demonstrating the effectiveness of the theory we developed.

Nevertheless, we added additional results visualizing hidden state values using the same task and with the same setup as Figure 3 (see Fig 8 in attached PDF ).

Q4

Thanks. We fixed these typos.

Limitations

Due to limitations of resources, we have not yet validated our method at longer lengths. For instance, we have verified that phi3-mini-128k model can be extrapolated to at least 192k with our method. Beyond this length, memory crashes occur.

Dear Reviewer LfLf,

As the discussion time goes by, we would like to know if our responses have addressed your concerns.

If our responses have resolved your concerns, do you think our paper should now be a clear "accept"?

We are very grateful for your valuable suggestions. If you have additional concerns, please let us know so that we can address them and further improve the quality of our paper.

Once again, thank you very much.