The Role of Sparsity for Length Generalization in LLMs

Thank you for your comments. Below we contrast our work to each of the papers you mentioned. We want to emphasize that with the exception of the 4th paper below (whose theory is different from ours), none of them provide any theoretical analysis, unlike our paper. Moreover, we wish to emphasize that our core method is neither presented in Figure 6 nor is sparse attention. Rather, our main contribution is to explain, using both theory and experiments, why having a sparse dependency structure in data can allow length generalization. We also note that it does not appear that the experiment presented in Figures 2 & 6 is present in any of the papers you pointed out.

1. *What is Wrong with Perplexity for Long-context Language Modeling?''* While the quantity $LSD_\theta(x_i)$ from this paper (which measures the difference in log-likelihoods of the current token using long and short context windows, and is essentially equivalent to what we call $L_{{long}} - {L}_{{short}}$) also plays a role in our paper, a key aspect of our experimental results is that of the tokens in the longer context window, a small number of them can be used to predict the current token $x_i$ nearly as well as all of the long-distance tokens. This sparse dependency pattern'' is not highlighted in this prior work.

2. ``DAPE V2: Process Attention Score as Feature Map for Length Extrapolation.'' While the main method proposed in this paper, DAPE, may be motivated by similar considerations regarding the sparse dependency structure of many tasks (such as associative recall), we emphasize that the content of this paper, which proposes a new empirical method based on applying an MLP to attention patterns so as to improve length extrapolation, is entirely different from ours.

3. ``Train short, test long: Attention with linear biases enables input length extrapolation.'' We already cite this paper (which introduces ALiBi). We would like to highlight one key difference between our message and that of the ALiBi paper: whereas ALiBi biases the model to put less attention on tokens far in the past, one of our main messages is that length generalization can occur in the (orthogonal) situation when the current token depends on a small number of tokens far in the past.

4. ``Mesa-Extrapolation: A Weave Position Encoding Method for Enhanced Extrapolation in LLMs.'' This paper has some theoretical analysis showing that certain transformers which fail to length extrapolate for NoPE or ALiBi can length extrapolate for some modified positional encoding schemes introduced therein (namely, Weave PE Extrapolation and Mesa Extrapolation), as well as supporting experiments. These contributions are entirely different from our own.

5. ``Sparser is faster and less is more: Efficient sparse attention for long-range transformers.'' Using sparse attention, as proposed in this paper, may be motivated by observing that the attention patterns of many actual transformers is sparse, which is closely related to the sparse dependency structure in the data which is the main focus of our paper. Nevertheless, our contribution is not to (re-)introduce sparse attention, but rather explain why such a sparse dependency structure suffices for length extrapolation using both theory and experiments.

Thank you for your positive comments!

Regarding the how the $k$ influential tokens are chosen: please see the discussion around Eq.~(20) in the appendix.

Thank you for your helpful feedback!

Regarding the various $L$ factors. Thank you for these great insights that allow us to improve our bounds! Though we did not attempt to optimize these factors in our submission, as you point out, one can indeed improve the dependence on $L$ and $\bar L$, which we discuss in further detail below in response to your questions. (For the purpose of interpreting our theoretical results as stated currently, one can think of $\bar L = O(L), \eta_L = O(L), \eta_{\bar L} = O(\bar L)$ (see Remark C.1), and $\delta \leq o(L^{-5})$, which implies a length generalization error (per Theorem 4.3) of $o(1)$.)

• Yes, the factor of $L$ in the error bound can be removed by adopting a minimax objective, as you suggest. (We remark that the objective in Eq.~(3) aligns more closely with what is done in practice, which is to mix over a range of context lengths during training.)

• Yes, the factor of $\bar L^2$ in the error bound can be replaced with $(\bar L / L_{{local}})^2$. The intuition here is that the proof is using $L_{local}$-locality to incorporate the longer context in ``chunks'' of size $L_{local}$. Since there is added error each time we add a chunk, the overall error will scale with $\bar L/L_{local}$. Thus, to minimize the error we want to maximize the size of the chunks. The largest possible value of $L_{local}$ which still satisfies the assumption of the theorem statement is $L_{local} = L/4$; by choosing this value of $L_{local}$, we can further improve the error dependence to $O((\bar L / L)^2)$. (Essentially, the current counterintuitive bound occurs because the technique used in the proof with values $L_{local} < L/4$ is a suboptimal approach to bound the error, which becomes less suboptimal as $L_{local}$ increases).

• Regarding the size of the parameters $\eta_L, \eta_{\bar L}$, please see Remark C.1, which gives reasonable conditions under which we have $\eta_\ell = O(\ell)$ for all $\ell$. Note that, for a fixed choice of distributions $Q_\ell^{pos}$, increasing the parameter $L_{local}$ can only make it easier to satisfy Assumption 4.2 for a given choice of $\eta_{\ell}$. (That said, if we maximize over distributions $Q_\ell^{pos}$ supported on sets $S^\star$ of a given locality $L_{local}$, then the worst-case value of $\eta_\ell$ may increase with increasing values of $L_{local}$.)

• We remark that the dependence on $\eta_L, \eta_{\bar L}$ in the overall error bound of Theorem 4.3 can also be removed under appropriate assumptions: for instance, if we make the distributional assumption of Remark C.1, then the instance of $\eta_L$ in Line 885 can be replaced with $O(1)$ since it is comparing between lengths $L$ and $L - 2L_{local} \geq L/2$. Under the same assumption, the instance of $\eta_{\bar L}$ in Eq.~(10) can also be replaced with $O(1)$ since it is comparing between lengths $\bar L$ and $L - 2L_{local}$; the former length $\bar L$ is larger and on the numerator, meaning that the distribution $Q_{\bar L}^{pos}$ is more ``spread out'' and will put less mass on any $S^\star$ under the assumption in Remark C.1.

• Summarizing, if we incorporate all of the optimizations discussed above, we can decrease our overall error bound from $O(\eta_L \eta_{\bar L} L \bar L^2 \cdot \delta)$ to $O((\bar L / L)^2 \cdot \delta)$, which does not decay with $L$. Hopefully this addresses your concern. We will update the paper to incorporate this discussion.

• We will add more details surrounding the equation around line 885.

Regarding details of PPC. The description of PPC can be found on page 7, though we will add a more detailed description. To answer your questions: when discussing PPC, we distinguish between the prompt and the response. During training time, the coupled position IDs for both prompt tokens and response tokens are fed to the model, and the model is trained to predict the coupled position IDs for the response tokens (as well as the actual response tokens). During testing time, the coupled position IDs are only needed for the prompt tokens, but not the response tokens (the model will predict the coupled position IDs for the response tokens).

Yes, training the model to predict position IDs adds an additional axis of length generalization. While the model sometimes makes errors at predicting response position IDs (which we count as an error in our plots), for many of the synthetic tasks that are studied, typically each coupled position ID actually only depends on a small number of previous tokens (e.g., many of the position IDs just increment by 1 at each step). Please see Sections E.2.1 and E.3.1 for the description of our choices of coupled position IDs for our synthetic experiments.

Regarding Assumption 3.2. Yes, when $k = 1$, relativity leads to a sort of ``bag of words'' model which is invariant to permutations of the sequence of tokens. We believe the interesting cases occur when $k > 1$.