Dear Reviewer PZAX,
Thanks for providing your further concerns.
## About Assumption 1
We suspect there is a misunderstanding on the definition of o. In our assumption, "o" represents the hidden state value, not the sequence length of the input. Please refer to our first response for the definitions of "o" and "H".
For example, consider an input sequence of length I with tokens denoted as below:
input_tokens = $[x_0, x_1, x_2, …, x_I]$,
where $x_i \in \mathbb{R}$
After transformation by the Transformer matrix, we obtain the query vector, key vector, and value vector as below:
Q = $[q_0, q_1, q_2…., q_I]$
K = $[k0, k1, k2…., k_I]$
V = $[v0, v1, v2, …v_I]$
where $q_i, k_i, v_i \in \mathbb{R}^{d}$, e.g d =4096 for Llama2.
Passing them to the Self-Attention part, we obtain that:
P = $Q^T K \in \mathbb{R}^{I \times I}$
O = $V \cdot Softmax(P)$ = $[O_0, O_1, O_2, ….., O_I]$
where $O_i \in \mathbb{R}^{d}$, and let $O_i = [o_1, o_2, ..., o_{d}]^T$.
Please note that $O_i$ corresponds to the hidden state values at position i. **We used $o$ to represent the hidden state value at a specific dimension**, that is, **$o_j$ is the j-th element of $O_i$**.
We assume that at a specific dimension (e.g. $o_j$ or dim j), there is a reasonable and acceptable range for this value. For simplicity, we have assumed a lower bound. If the value in this dimension falls below this lower bound, it could cause abnormal outputs in subsequent layers, ultimately leading to abnormal model behavior. This construction is intended to demonstrate that threshold-based criteria can be used to assess whether extrapolation is successful.
## About dimensions
We suspect that there might be some confusion about dimensions in our context and canonical basis vectors.
In our context, we use “dimension j” to describe the j-th element of the output vector $O_i$, as stated above . Therefore the concept of canonical basis vectors e_1 and e_{20} is not relevant in our context.
Here we only chose specific examples of dimensions to illustrate the phenomenon of extrapolation failures. We observed that **this behavior is consistent** across different inputs and tasks (passkey, language modeling etc.). By observing the corresponding threshold, we can predict the extrapolation capability of the LLM, which is limited by its maximum pre-training window, regardless of the input. Given the consistent results across various experimental conditions, we believe our findings have broad implications and are not data-specific.
## About the design of the chunk-based triangular attention
First, our assumption that the beginning tokens and surrounding tokens are important is well-established and has been used in many of the previous works, such as Streaming-LLM and LM-Infinite. And the reasoning is below:
**Beginning of the sequence:** 1) people often place instructions, task descriptions, etc., at the start of a sentence, which is related to the way prompts are structured. 2) a start token is typically added by default, which is an important position holder. For instance, Llama2 automatically adds a <bos> token at the beginning. We’ve also explained this in detail in the appendix (line 1206-1209).
**Surrounding tokens**: According to the principles of the attention mechanism, attention decreases as the distance increases, meaning the nearby tokens generally receive higher attention scores.
Focusing on the initial tokens and their neighbors has been validated by other previous works (e.g,. LLM-Streaming and LM-Infinite). When designing the chunk-based triangular attention, we also experimented it and found this design to be the most effective.
Note that Mesa-Extrapolation appends all previous tokens at the bottom of its attention matrix, allowing the last token to attend to all preceding tokens. In contrast, Streaming-LLM and LM-Infinite also discard middle tokens at the bottom of their attention matrix, inevitably resulting in loss of information from these tokens, leading to inferior performance.
If you still have any questions regarding the explanation above, please let us know so that we can provide further clarification.
Once again, thank you very much.