The Role of Sparsity for Length Generalization in LLMs
Authors: Noah Golowich
Samy Jelassi
David Brandfonbrener
Sham M. Kakade
Eran Malach
Venue: ICML 2025 poster (2025)
Keywords: Length generalization
Transformers
Positional Encoding
Primary area: theory->domain_adaptation_and_transfer_learning
TL;DR: We introduce a formal framework for understanding the role of sparsity in length generalization in LLMs and provide empirical evidence for it.

Abstract: Training large language models to predict beyond their training context lengths has drawn much attention in recent years, yet the principles driving such behavior of length generalization remain underexplored. We propose a new theoretical framework to study length generalization for the next-token prediction task, as performed by decoder-only transformers. Conceptually, we show that length generalization occurs as long as each predicted token depends on a small (fixed) number of previous tokens. We formalize such tasks via a notion we call k-sparse planted correlation distributions, and show that an idealized model of transformers which generalize attention heads successfully length-generalize on such tasks. As a bonus, our theoretical model allows us to provide justifications for techniques to modify positional embeddings which have been introduced to improve length generalization, such as position coupling.

We support our theoretical results with experiments on synthetic tasks and natural language, which confirm that a key factor driving length generalization is indeed a ``sparse'' dependency structure of each token on the previous ones. Further, inspired by our theory, we introduce Predictive Position Coupling, a generalization of position coupling which trains the transformer to predict the position IDs used in a positional coupling approach.  Predictive Position Coupling thereby allows us to broaden the array of tasks to which Position Coupling can successfully be applied to achieve length generalization.

The Role of Sparsity for Length Generalization in LLMs
Noah Golowich Samy Jelassi? 2 David Brandfonbrener' Sham M. Kakade Eran Malach Abstract
popular technique recently has been to use scratchpads or Mhain--ofThhoueee (CoT) to perform various logic and reasoning tasks, and the length of the scratchpad can become very large, especially when combined with search or reinforcement leam- ing techniques (DeepSeek-Al et al., 2025; Team et al., 2025).
Training large language models to predict beyond their training context lengths has drawn much atten- tion in recent years, yet the principles driving such behavior of length generalization remain underex- plored. We propose new theoretical framework to study length generalization for the next-token predic- tion task, as performed by decoder-only transformers. Conceptually, we show that length generalization occurs as long as each predicted token depends on small (fixed) number of previous tokens. Our the- oretical model justifies certain techniques to modify positional emheddings which have been introduced to improve length generalization, such position coupling. We support our theoretical results with experiments on synthetic tasks and natural language, which confirm that a key factor driving length gener- alization is a "sparse" dependency structure of each token on the previous ones. Inspired by our theory, we introduce Predictive Position Coupling, which trains the transformer to predict the position IDs used in a positional coupling approach. Predictive Posi- tion Coupling thereby allows us to broaden the arryy of tasks to which position coupling can successfully be applied to achieve length generalization.
Unfortunately, transformers struggle to length generalize on even very simple arithmetic and logic tasks, such as computing parities, integer addition, and variable assignment (Anil et al., 2022; Kazemnejad et al., 2023). Indeed, only recently have transformer models successfuly been trained to -rngth-gennralize to contexts many times their training length on integer addition tasks, using a technique known as position coupling (closely related to Abacus embeddings) (Cho et al., 2024a;b; McLeish et al., 2024). For many other simple prob- lems, comparable length generalization remains challenge. In light of this uneven progress, natural question is whether there is more principled way of understanding length generalization. In particular, can we formally reason about what structural properties of data aid or inhibit length generalization? More- over, can this suggest architectural modifications, such as to positional embeddinggs, which improve length generalization?
Contributions. In this paper we give a positive answer to the above questions: First, we introduce a class of data distributions, namely those with sparse planted correlations (Definition 3.2), which, roughly speaking, captures the property observed in many synthetic and natural language tasks that each token to be predicted depends on a small number k of previous tokens, which we call the sparsity. We then establish the following:
1. Introduction
Enabling large language models (LLMs) to generalize to con- texts longer than their training context length has emerged as key problem in recent years. Indeed, many factors limit the context length of sequences which can be used during train- ing, including the increased computational cost of training on long sequences (Tay et al., 2022) as well as the fact that longer sequences may be less numerous in the training dataset. Nev- ertheless, many applications require LLMs to be accurate on extremely long context lengths at inference time: for instance,
As long as the sparsity does not grow with the sequence length, then a simple class of models generalizing attention heads (namely, sparse functional attention; Definition 3.3) has provable length generalization (Theorem 4.3). This result also relies on an additional assumption of locality on the hypothesis class.
While the locality assumption is typically violated in practice, we show that (a theoretical abstraction of) position coupling can remove the locality requirement (Proposition 4.4), thus offering theoretical justification for this recent technique.
'MIT EECS 'Harvard University 'Kempner Institute at Harvard University. Correspondence to: Noah Golowich <nzg@miteduus Sany Jelassi sielsssseelesssiee @ fas hww.aaara.dddee edu? David Brandfonbrener<davvii www.ddidddaaadddbbeeeeeeeeeee gmail.com>> Sham M. Kakade-<sham @ scas. hddv, Eran Malach
We perform experiments (Section 5) on synthetic and natural language data to support theorical conclusions: for instance, our experiments adjust the sparsity of several synthetic tasks, and we observe that length generalization improves monotonically with decreasing sparsity. For natural language data, we provide evidence that -engtth-nnerrliiing transformers indeed make accurate predictions using small
<cran
www.mmallcceggaa.cmm....... gmail.com...
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Length Generalization in Transformers
number of past tokens.
We discuss it next; additional techniques modify position IDs to achieve length generalization are discussed in Appendix A.2.
Inspired by our theory, we introduce a modification of positional coupling, Predictive Position Coupling, which, unlike positional coupling, works tasks for which the coupled position IDs are input-dependent. We show (Section 5.2) that Predictive Position Coupling enables significant length generalization on such tasks.
2.1. Position Coupling
The technique of position coupling (Cho et al., 2024a;b; McLeish et al., 2024) (similar to Abacus in McLeish et al. (2024)) works specifically for problems with structured imput where there is clear one-to-one relationship between certain tokens. In particular, it assigns each token in a sequence a particular position ID in a way so that tokens assigned the same position ID should have a (lask-dependent) "one-to-one correspondence". For example, to solve string reversal, i.e., predict the last I. tokens of the sequence X1, A..... X,, <SEP>, XL.... X1, then since the ith-to-last reversed token is equal to the ith input token, we feed the following position IDs: L.2.....
Related work. A few prior works have attempted to give theoretical justifications for the length generalization abilities of transformers (Zhou et al., 2023; Huang al., 2024; Sabbaghi al., 2024; Ahuja & Mansouri, 2024); we compare these works, as well as many others, to ours in Appendix A. The idea that sparsity and locality are important for length generalization has appeared in several of these works (Huang et al., 2024; Sabbaghi et al., 2024). However, none has provided the precise theoretical and empirical evidence which we believe are essential for isolating the importance of these notions.
3. Theoretical model
2. Preliminaries
In this section, we formally define our theoretical model. Our primary inspiration is the class of decoder-only transformers, which are trained to predict each next token as a function of the preceding ones in a given sequence. Accordingly, our theoretical framework will focus on the next-token prediction task, where we fix lengths L < L, and attempt to show that models trained on sequences of length S I. to predict a single token (interpreted as the (L + 1)th token in the sequence) also succeed at predicting single token (interpreted as the [L 1)th token) when tested on sequences of length L.
Our focus in this paper is on decoder-only transformers trained to predict the next token. We review the architecture of such transformers in more detail in Appendix B.2. In this section, we discuss a few aspects of the position encodings of transformers which have previously been proposed to improve length generalization.
Initially, Vaswani et al. (2017) proposed to introduce positional information in a transformer model using absolute positional embeddings (APE) (see also Gehring et al. (2017)). The APE assigns to each token a position ID (typically the position of that token in the sequence) and adds an embedding vector vAPE depending on to the embedding vector for the corresponding token. Recently, it has become more common to use position encodings which encode relative positions. In particular, a popular positional encoding technique in many large-scale open-source transformers is the rotary poxitional encoding (RoPE) (Su et al., 2021). which adjusts the computation of attention scores as follows: it multiplies each key embedding vector with position ID by a rotation matrix depending on i, and each query embedding vector with position ID j by rotation matrix depending on j. The effect of these rotations is that the attention score for a (key.que .qurr) pair with position IDs (i.j) depends only on i-j (and not i or jrj individually).
Fix a set V which denotes the set of tokens. V -~=UUaa -Uzo De denotes the set of all arbitrary-length sequences of tokens. For technical reasons discussed below, we model the problem of predicting the embedding (i.e., representation) of the next token, as opposed to the token itself, given a sequence of previous tokens. Accordingly, we fix label set y which is a convex subset of Rd for some d E N denoting the space of possible token embeddings, and a loss function 2:)'x, x + [0,1] with L(YYY) = 0 for all Y Ey. For simplicity, we fix an arbitrary norm || on IRd, and assume that diam()) < 1 with respect to this norm and that L(Y,Y') = ||| - Y'IJ. For a set Ω and k e N, we let Sets be the set of size-k subsets of Ω.
A distribution ensemble P is a sequence of distributions Pi,P2..... where for each E E N, Pa E A(V x x))) represents a distribution over labeled input sequences of length f. Let H. yCV denote a class of functions h: V y. The length generalization problem is as follows: we aim to choose a hypothesis h E H which enjoys small square loss Z(-) for inputs drawn from Py for some Le if we are only allowed to choose h based off of inputs of lengths less than L, i..., those drawn from P. for L.L.
An obstacle to length generalization in transformers is that the joint distribution of particular tokens and their position IDs seen in training sequences may not match that seen in the longer sequences test-time. To account for this discrepancy, a common empirical technique is to modify the position IDs at training or test time. As discussed in Section 1, the positional coupling technique, which has recently paved the way for major improvements in length generalization in arithmetic and logic tasks, plays a key role our theoretical and empirical results.
How can we use inputs drawn from such Pe to choose h?? A classic paradigm in leamning theory is empirical risk
2

---

Length Generalization in Transformers
minimization (see e.g., Shalev-Shwartz & Ben-David (2014)), which stipulates to choose hypothesis in H to minimize the average loss over samples drawn from P: for some values of L. However, as our goal is specifically to understand the out- of distrributtion generalization behavior from inputs of length <<I L to inputs of length L. we instead assume that h is chosen so as to minimize the population risk for inputs of length << L. This choice allows us to avoid having to consider the (in-distribution) generalization error for samples drawn from Pe (f < L). Formally, we define length generalization as follows:
sequences drawn from such distributions have most of their tokens drawn independently from some distribution μ AV)) (which we think of as a "background distribution" representing tokens not relevant for the task at hand), and a small number k of tokens drawn from some "planted" correlated distribution over k-tuples of tokens, denoted by Q in Definition 3.2 below. One should interpret these planted k tokens as having "relevant information" for the task of predicting the label Y. The particular location of these tokens is drawn independently (denoted O below). Formally, we have:
Definition 3.1 (Length generalization), For L.LEN, CE (0.1), we say that the class H. has (L,L,e)-lengel generalization with respect to the ensemble P if the hypothesis
Definition 3.2 (k-sparse planted correlations). Fix positive integer ENN. We say that a distribution ensemble P = (Pe)een has k-sparse planted correlations if there are distributions μ A()). O. E ((Setsx((4) for N. - E: A(V%), and function g" :0 =>> so that the following holds. For each E N, a sample (X,Y) ~ Pp may be drawn as follows: first, we draw S 2 Z~ Quoe, and we set:
d:=argmin NEN B-UBIILLL/.DD YNMP
E
E
(() (1)
satisfies EX (
For a hypothesis class H. YVV and 8 (0,1), we say that the ensemble P is 6-approximately H-realizable if there some < 5 for each E E. We say that P is H-realizable in the case d =0
Xs. =2, Z,
XINNViES" Y=g
h H so that (xxy),
While the assumptions in Definition 3.2 that (a) the remaining tokens X, for iorig S are i.d. from μa and (b) the tuple S which indexes the correlated tokens that "matter" is drawn from a fixed distribution 2.06 are not realistic, we emphasize that they are made to simplify the proofs and ensure a simple model which captures the salient features that enable length generalization. We leave generalizations, such as that the X, are drawn from Markov chain or Hidden Markov Model, for future work.
3.1. Distributional assumptions: sparse structure
Without distributional assumptions on Pe. achieving length generalization per Definition 3.1 becomes degenerate in the following manner, even if we assume H-realizability. For any fixed (realizable) choice of the distributions Pe for t < L. unless is identically equal to h², we can choose Py to have all its mass on some sequence X E y² for which (()((X) (X)))(X) >> O. This choice prevents the loss of h defined in (1) from being small, thus ruling out length generalization, as formalized in Proposition B.I.
An example of a distribution ensemble satisfying Definition 3.2 is the distribution of tokens in the k-gram retrieval task, of outputting the token following a certain k-gram in a given sequence X. In particular, the k-gram in question together with the following token form the "planted sequence"; see Appendix B. for a formal description.
While some works (e.g., Ahuja & Mansouri (2024); Huang et al. (2024)) have offered explanations for length generalization by showing that in fact, ht is leamed exactly under appropriate assumptions, there is ample empirical evidence (Zou et al., 2023; Wei et al., 2023; Andriushchenko et al., 2024) that transformers can eIT on worst-case inputs. Accordingly we ask: Are there distributional assumptions which enable us to establish length generalization in settings where the ground truth hypothesis h² may not be learned (nearly) exactly in the sense that supxev. L(hh" (()((X(x,"))) (X)) is small?
3.2. Defining the hypothesis classes
In this section, we define a hypothesis class which allows us to effectively model the salient aspects of transformers while permitting
obtain provable length generalization. In particular, we aim to satisfy the following two desiderata:
1. First, we would like there to be many ensembles P with sparse planted correlations which are realizable with respect to H.
The sparse attention patterns in many transformers trained on natural language (e.g.. Child et al. (2019); Tay et al. (2022)) suggest that for modeling a wide spectrum of natural language, one needs to attend to only a small number of previous tokens to predict each successive token. Inspired by this observation, we define the following class of distribution ensembles, namely those with sparse planted correlations. Roughly speaking,
Second, we would like the class H to capture simple transformers (i.c., a single attention head).
To motivate how to arrive at our function class starting from the above criteria, we recall the structure of a single attention head (see Appendix B.2 for a more complete presentation). It takes as input embedding vectors hi.....y, E Rdd which could be, for in- stance, the results of an embedding matrix multiplied by one-hot encodings of tokens X XI,..,X,, E V. For some matrices K Ve E
1 'The lower hound of T/2 in the interval [T./2, I.] from which is sampled is unimportant; any constant factor of will suffice.
3

---

Length Generalization in Transformers
Rdxd and a query vector qe RR, the attention head computes
attention head with fixed query vector as discussed above: in particular, for 5= (=(+), we choose go (0(5,,5 (5,x)) = (q.Kh.) where h, is the embedding of x,, and gi (x:) = V -h, The fact that (Sxxs) can also depend on the position indices S should be interpreted as allowing us to model positional embeddings (see Section 4.1 for more discussion). The following proposition formalizes these observations, verifying that the sparse group attention class indeed generalizes attention heads.
AttHeadk v.a (h) (hi,,.,hu))
04 Vh;
where a- softmaxx(qq........oo (q,, (q. Khn)). The softmax over linear functions of h,,,, allows AttHead to attend to individual positions E [L] whose tokens satisfy some property (e.g.. the embedding vector points in particular direction of Ied which could represent a certain meaning of the token). Suppose we were take H. to be the class of attention heads and try to satisfy Item I: we would need attention heads to be able to apply some function (g" in the context of Definition 3.2)
Proposition 3.1 (Informal version of Proposition B.2). The class of attention heads corresponding to a vocabulary V and embedding dimension d equal to the 1-sparse functional atten- tion class Hartn (geen gyal) for appropriate choices of Ghey graa.
to certain sub-sequence of k tokens (Xs- in Definition 3.2).²
Additional properties. We believe that the abstraction of sparse functional attention defined in Definition 3.3 may be of broader interest in obtaining a theoretical understanding of various properties of transformers. Nevertheless, in order to analyze length generalization, we need to further restrict the class Hattn ggaa)) in the following ways: first, we assume that the class gleyy is local, meaning that it only outputs finite (i.e., not negative infinity) scores on subsets S C for which max(S)- - min(()) is bounded, Second, we assume that gley only uses relative positional information, meaning that shifting S does not change (u(S.x)) for all go E gley xx V
Unfortunately, a single attention head as in (3) is not able to do so for large classes of g". when & > 1 and even if Xs. is sufficiently "distinguishable": for instance, AttHeadk, v. (q(h....,hy) is invariant to the order of hi...,h, (and thus of the tokens, in the typical case that h, is an em- bedding of X,), whereas g" (Xs-) ) may not be. Decoder-only transformers correct for this deficiency by (amongst other factors, such as positional embeddings) stacking multiple layers of attention heads. When doing so, embedding vectors h, at individual positions of higher layers can contain information about multiple tokens of the input sequence. Inspired by this property, we introduce the following model of an "idealized transformer" - instead of stacking multiple attention heads, it directly applies attention on subsets of k tokens:
Assumption 3.2 (ghey is local and relative). We introduce the following properties of the class ghey;
Fix Lioca E N. We say that gtey is Locarllocal if for all S E Setsk, (N) for which maxiesi - mingesi > Taocal, we have g0(Sx)==o for all go E
Definition 3.3 (Sparse functional attention class). Fix positive integer i (the sparsity) together with hypothesis classes gley consisting of functioos mapping (N x V)) - (-o)- and gual mapping y* y. We define the k-sparse functional attention class Hattn = Hattn (ghey, to be the class of all hgv.91 V +y, indexed by go E ghey g1 E gval, where*
2.
We say that gkev is relative if for any E and S SeStss(()) (N), x y it holds that go
(ix) for all go E e
The assumption that gley is relative is inspired by the fact that many common positional embeddings only use relative positional information. While the assumption of locality is strong, we remark that it is provably necessary (Appendix C.3) and can be relaxed using techniques such as positional coupling (Section 4.1). Throughout the paper, we will assume that the parameters Looaal are chosen so that Liocal k.
exp(ao (xp(g((5xx)) (Sxx)) 9 (xs)
hg (x))).,xI.)= =
Σ
(()))ss((L)) S ESetss ([L)) ex (S" x)
Intuitively, one should think of the "attention function"' (Sxx)) as computing attention scores based off of sets of i tokens. Sparse functional attention with k captures single
4. Theoretical results
In this section, we establish formal length generalization guarantees for sparse functional attention classes Hattn (Definition 3.3) for distribution ensembles P with sparse planted correlation (Definition 3.2). To do so, we need to make few assumptions the ensemble P. The first assumption ensures that the ensemble P is approximately Hattn geattireali..
"In the case of decoder-only transfommers, q is commonly taken to be equal to be a linear function applied to hy, though in this discussion we take it to be fixed along with the other parameters of the attention head.
*Notice that since S* is not given explicitly, there are certainly cases where it is information theoretically impossible to determine S" from X. However, we give several examples in Appendices B.1 and C.3 where the planted sequence Xs+ is sufficiently "distinguish- able" from the remaining tokens X,, & and in fact can indeed he computed with high probability.
Assumption 4.1 (Realizability). We assume that P is 8- approximately Hatm
(ghee ( . i.e., there is ht* E Hatt g so that LE E (C (X), Y) for all E.
for all S (e.)) then we ha (X1,.., to he [Baa g (xs).
*If go(S, x3)
'For S= (k) S Sti== (ir +1).
4

---

Length Generalization in Transformers
Our second assumption states that the distributions 2,eh in the context of Definition 3.2 have bounded coverage at all locations in the sense that changing& & or shifting S by A units does not significantly change 2P(((5))
tribution Q over tuples (S*, .bb)) with S E Selss ((f)) and the [f] - [e] so that the marginal of S under Quos-c is 2i and with probability 1 under the draw of (5* ,ébe) -pos-
1.
Liocal-
Assumption 4.2 (Coverage of ),²²) Fix a h-sparse distri- bution ensemble P specified by 1, (2P Qroc, as well as Liocal E N. We assume that for each é E N with 82 Liocal, there is some positive value 77c ERR so that for all L' E [L
2. Wiz, S* satisfy some additional conditions which, e.g- prevent coupling of non-planted tokens (i.e., outside of s*; Definition D.2).
A 10 O and S ESetsx 2er [S] S T1e and (S")
(5")
(S*) Smc.
With Definition D.2 in hand, we can now describe how hypotheses in a sparse functional attention class Hatti which may not satisfy locality, can be transformed into new hypotheses which will use information from the output of the "position coupling" Uo and will also satisfy locality. Roughly speaking, given a sample (X,S*, Y) ~ Pee we will join (i.e., "couple") all tokens X; for which (/n(t) is identical. We will denote the resulting distribution over sequences of "coupled tokens" by PC(P)). Moreover, we define a "posstiinn-couledd" hypothesis class PC[[(att]] whose hypotheses, when given sequences of coupled tokens as above, "unpacks" them and applies the corresponding hypothesis in Hattn In Appendix D, we formally define these notions and prove the following propo- sition, which shows that this procedure of position coupling can remove the locality requirement from Theorem 4.3:
where the first inequality is only required if S A E Setsy ((eT)
and the second inequality is only required if S E Setsk ([")).
Our main result states that for any classes gkey, gual and distribution ensemble P satisfying Assumptions 3.2, 4.1 and 4.2, the class H (ghey grall) enjoys length generalization with respect to the ensemble P.
Theorem 4.3 (Provable length generalization). Fix any k E 14, and consider any k-sparse functional attention class Hatt = Hatt (g gral) which is Liocat-local and relative (Assumption 3.2) for some Tiacal EN. Then for any k-sparse planted correlations distribution ensemble P (Definition 3.2) satisfying Assumptions 4. and 4.2, and any integers L,, L. for which Liocal L L and L 2 4Liocal, Hatte achieves (I. I, nL L. (I./ (L 57-length generalization with respect to the ensemble P.
Proposition 4.4 (Informal version of Proposition D.2). Hf is sparse functional attention class and P is an ensemble with k-sparse planted correlations which satisfies 8-approximate strong realizability (Assumptiin 1) and Assumption 4.2, then for L,L satisfying the conditions of Theorem 4.3, PC[Haatn] achieves (L, L, MLME, LL? -y-length generalization with respect the ensemble PC[P].
In Appendix C.3, show that if we remove any one of the main assumptions of Theorem 4.3, then length general- ization can fail to hold. Moreover, we remark that the error Mr. mE. LT22. 8 on distributions of length L can be improved: see Appendix C.2 for further discussion.
5. Experiments
4.1. Improving length generalization: positional coupling
Conceptually, we view the main takeaways of Theorem 4.3 and Proposition 4.4 to be the following:
As discussed above, one limitation of Theorem 4.3 is its reliance on Item 1 of Assumption 3.2, which not satisfied in actual transformers: indeed, transformers may attend to tokens which are very far apart, as such dependencies can occur in natural language. Is there way to mitigate this limitation?
(TI) First, an important factor enabling length generalization is that the label Y depends on only l tokens of the input (in the sense of Definition 3.2 and in particular the relation Y = g" (Xss) in (2)). In particular, the parameter k: (which we refer to informally as the sparsity) must be the same for both the lengths on which we train (namely, lengths & L) and the length L to which we attempt to extrapolate, and sufficiently small compared to the maximum training length L.
We now show how a theoretical abstraction of position coupling as discussed in Section 2.1 can allow us to remove this locality requirement. Roughly speaking, this abstraction of position cou- pling states that there is a joint distribution over (S* ,úx). where < gives a way of "rewriting" position indices, so that the "rewritten" set (S*t of indices in the planted set S* satis- fies the locality condition of Assumption 3.2. In particular, for each position ID E [e]. the value of (i) should be interpreted as its "coupled position ID" as discussed in Section 1. More precisely, we have (see Appendix D for the full definition):
(T2) Second, locality of the hypothesis class (per Item of Assumption 3.2) plays an important role as well: the maximum distance Liocal between tokens which "matter" in predicting the label Y must be the same for lengths & & L on which we train and the length L to which we extrapolate. Moreover, as this requirement is quite strong (and unrealistic for many problems of interest), one way to mitigate it is the technique of position coupling (per Proposition 4.4).
Definition 4.1 (Informal version of Definition D.2). Fix a distribution ensemble P with k-sparse planted correlations per Definition 3.2 (defined by μe, Q QuOC). A Liocal local position coupling of P is defined by, for each l ENN, a joint dis-
In this section, we evaluate these conclusions for synthetic and
5

---

Length Generalization in Transformers
natural language modeling data.
Results & discussion. The accuracy of the trained trans- formers ixx... for predicting the last (parity) token samples of various lengths & and sparsity values A: is shown in Figure la for Kaain = 10 and in Figure 3 for the remaining values of trrai-- Two observations are notable: first, when the training sparsity K is small enough (i.e., K S 8; see Figure 3), then Ku.... experiences near-perfect length generalization up to lengths of 500, for all values of the test sparsity hrest S train However, for test sparsity values Rest > Ktrain the performance of fuk.... deteriorates rapidly, for both in-distribution and out-of diit--f-dittibttionn lengths. This behavior is consistent with Theorem 4.3, which implies that good length generalization occurs as long as the sparsity does not change between the train and test distributions and sufficiently small as a function of the training length.
5.1. Length generalization for sparse parity
First, we discuss a simple setting which measures the degree to which length generalization in transformers reflects the require- ment from (T1) that the sparsity & be sufficiently small and not grow function of the sequence's 's length. In particular, consider the following sparse parity task:* given k NN the in- put is drawn from a distribution D.F; over length-2f sequences, where tokens at even nee-nnumereed positions are bits, and tokens at odd-numbered positions belong to some set Q2 Exactly & to- kens at odd nudbbeeree positions belong to some "special subset" S'' C22 and the goal is to find the parity of the l tokens imme- diately following them. See Appendix E.1.1 for precise details.
A corollary of this reasoning is that for fixed maximum training length I., the length generalization behavior with respect to the distribution ensemble (D should degrade at some point if the sparsity & is allowed to increase enough. (In particular, our theoretical guarantee on length generaliza- tion only holds for sufficiently small training sparsity, i.e., Kraiin S L/4 in the case of Theorem 4.3, though the particular threshold is likely different in the case of actual transformers). This behavior is mirrored in Figure la as well as Figure 3 in Ap- pendix E.1.2, where, for each fixed value of K E {10,12). larger values of & E [Krain] have worse length generalization despite all having near near-perfect in-distribution performance." 8
Experimental setup: data. For each value of Kerain E (4,6,8,10,12). we train a transformer hyg,. to predict the last token of samples drawn from Di. where E is sampled uniformly subject to the length of the sample satisfying 26- E [20, 50] and k ~ Unii([XKrain]]] evaluate the performance of each of the trained transformers hk on samples drawn from Dr.hess of length 24 E [20,500] and with sparsities heest E
We then
Experimental setup: model. Our model is based off of the GPT-NeoX (decoder-only) transformer (Andonian et al., 2023), and uses rotary positional embeddings (RoPE). To ensure nontrivial length generalization performance, we combined RoPE with PoSE (Appendix B.3). Full training and evaluation details may be found in Appendix F.
5.2. Scratchpad with Predictive Position Coupling
Many tasks, unlike sparse parity, have the property that the sparsity of the next-token prediction task (as formalized by, c.g., Definition 3.2) can be quite large and in fact increase from shorter to longer lengths. To help achieve length gener- alization, one approach which has been suggested in numerous existing works (Anil et al., 2022; Hou et al., 2024) is the chain- thoughtsscratchpa technique (Nye et al., 2022; Wei et al., 2022; Lewkowycz et al., 2022). It proceeds by computing tokens representing intermediate steps, such as intermediate parities when the task is to compute the parity of a sequence of bits. It has been observed that the scratchpad technique alone is insufficient to ensure length generalization (Anil et al., 2022; Dziri et al., 2023; Hu et al., 2024; Kazemnejad et al., 2023; Lan- chantin et al., 2023). However, number of recent works (Cho et al., 2024a;b; McLeish et al., 2024) have shown that the tech- nique of position coupling (see Section 2.1), when combined with scratchpad (and even when used on its own, when appro- priate), can allow significant length generalization on arithmetic tasks such as parity, addition, and multiplication. This devel- opment parallels theoretical results in Section 4.1, where
Remark 5.1. Numerous other modifications to position IDs have been proposed for length generalization, such as position interpolation and various enhancements, which typically modify the way the transformer uses the position IDs at inference time, often after a small amount of fine-tuning (see Appendix A.2). We stick with PoSE in this paper (when position coupling is not applicable) because of its simplicity and since (a) does not require modifying the transformer's com- putations inference time; and (b) the fine-tuning for position interpolation requires sequences of length given by the resting context length, which fails to lie in our framework where we assume that any amount of training on such sequences is not al- lowed. Understanding the role of sparsity and locality for length generalization in transformers which make these inference-time modifications to position IDs is left for future work.
"This task is slightly different from standard formulations of sparse parity, where the sparsely chosen positions not identitied as part of the input. We use this version so as to allow different set of positions to he chosen for each example.
'Reaall that fusu.. was trained by placing equal weights on all sparsities & S Krann. so in this respect the distributions Di for I < Krain were treated "on equal footing."
"This task is not well-suited for position coupling, since the output token depends on 2 input tokens, all of which have different position DDs.
6

---

Length Generalization in Transformers
Tean Sparsity
Siedlyn Suadity 24
HDn Depth 11
(a) Sparse parity
(b) Parity with scratchpad & PPC
(c) Variable asst. with scratchpad & PPC
Figure Length generalization for experiments in Scctions 5.1 and 5.2. Shaded areas (b), (c) represent 95% CIs over multiple training runs.
we showed that position coupling removes the stringent require- ment of locality, which is not satisfied in practical settings.
sequence of coupled position IDs by a uniformly random offset subject to the constraint that all shifted position IDs be at most Lmax- In particular, the sequence of coupled position IDs ranges from to E, then we would shift by a uniform element of [O.I,...,Lmar -}.. We remark that using RoPE with PoSE (as oposed PPC) performs significantly worse; see Appendix E.
One downside of position coupling is that it requires significant structure in the format of the data sequence, including the scratchpad: in particular, once the length is fixed, the coupled position IDs (see Section 2.1) in all tokens the transformer is trained to predict must be determined, as they must be fed into the model to predict each successive token. (For instance, for the string reversal example in Section 2.1, when L denotes the length of the string to be reversed, the sequence of position IDs given as in Section 2.1.) As consequence, the applications of position coupling are currently limited small number of synthetic tasks where the computation of each token in the sequence (including the scratchpad) can be coupled with fixed position of the input.
5.2.1. WARMUP: PARITY WITH SCRATCHPAD.
As a warmup. we first consider the task of evaluating the parity of a sequence of l bits with use of a scratchpad (and PPC) to compute intermediate parities (as in, e.g., Anil et al. (2022); Cho et al. (2024b)). The scratchpad is structured as follows: it has length E, and the ith token of is the parity of the first tokens of the input; thus, the final (fth) token of the scratchpad is the desired answer, i.e., the parity of all l input tokens. To measure the impact of sparsity, we consider the following modification of the standard scratchpad format, which we refer to as having a scratchpad with "jumps". For values of k E {1,2,3,4), we only write every kth bit the scratchpad, meaning that the ith bit of the scratchpad the parity of the first ki tokens of the input. In particular, the case k == corresponds to the standard scratchpad format. We also modify the position coupling IDs so that each position ID is repeated & times in the input sequence. Full details of the data and position coupling formats are in Appendix E.
Towards relaxing this limitation to allow sequences where we propose
the coupled position IDs can be Predictive Position Coupling (PPC). PPC functions follows:
1. It modifies the transformer architecture instead predict the coupled position ID for each next token. In particular, we add an additional output embedding module so that, at each step, the transformer predicts two IDs: the next token ID, as well as the coupled position ID for that token.
2. The transformer is trained to predict both the actual tokens, as well as the coupled position IDs for each token. We use the standard entsopy loss for both predictions, and weight them equally.
Results & discussion. Figure 1b shows the length general- ization behavior for our transformer model trained as discussed above, for "jump lengths" k e (1,2,3,4}. Note that the jump length & parametrizes the sparsity of the next-token prediction task: each token in the scratchpad depends on k+ previous tokens (namely, the immediately preceding one together with the tokens the input which correspond to that step of com- putation). While all 4 jump lengths attain near-per nur--eefect accuracy for in-distribution lengths, smaller values of k have superior performance on longer out-of- distribbbut lengths, consistent with our takeaway (TI). This fact is particularly notable in light of the fact that for jump length k, the scratchpad is shorter by a factor of k. which priori could presumably have made it easier to correctly predict all tokens and position IDs correctly (the metric measured in Figure 1b). Thus, at least in this task, spar-
3. At generation time, the two predicted tokens at each step are fed in as the position ID and token ID at the following position.
In our results, we report the fraction of examples on which the model correctly predicts all tokens and coupled position IDs.
Experimental setup: model. We use the same NeoX model as described in Section 5.1, with the following modifications to implement position coupling. Following (Cho et al., 2024a;b; McLeish et al., 2024), we use absolute position embeddings and fix some integer Lmax denoting the maximum possible length of test-time sequences. During training, we shift each
7

---

Length Generalization in Transformers
sity (as measured here by k) is more important determinant of length generalization than the length of the output sequence.
To understand length generalization in models beyond this base- line, we ask: given I T., can transformer trained natural language sequences of length L achieve smaller perplexity on sequences of length L than on sequences of length L? Previous works have shown a positive answer to this question (as will our experiments) when one uses appropriate modifications to the positional embeddings, namely PoSE (Zhu et al., 2023) or posi- tion interpolation (Chen et al., 2023; emozilla, 2023; Peng et al., 2023; Chen et al., 2024). In light of this, natural way to mea- sure sparsity is as follows: is there way to select a small num- ber l of the tokens more thun I. units in the past from the present token, so that, using just those k tokens together with the most re- cent L tokens, we can nearty recover the perplexity of the model on full sequences of length L.? Informally, we are asking if the dependence on tokens more than L units in the past sparse, in terms of beating the trivial baseline for length generalization.
5.2.2. VARIABLE ASSIGNMENT WITH SCRATCHPAD.
Next, we consider the more challenging task of variable assignment with a scratchpad, of which several slight variants have been studied in numerous prior works (Zhang et al., 2023; 2022; Lanchantin al., 2023; Peng et al., 2024).
Experimental setup: data overview. In the variable assign- ment problem, we fix a depth d, denoting the number of "hops" it takes to compute the final value. Roughly speaking, the goal as follows: given a sequence of "variable assignments" the form o' (where v,e/ are token IDs), together with starting variable Unit follow sequence of d assignments starting at In more detail, problem instance consists of some number of "chains", each of the form U1 <- 02202 +- "a3"": <- 1a+1 for some depth parameter d. The variable assignments in each of these chains are interleaved randomly; exactly one of the chains has 1/1 = Unit- The goal of the task to find the final variable va on the chain for which U/ Uiit- To do so, we use a scratchpad which writes out in full the chain starting with Vnnit together with position coupling, which matches each element of this chain on the scratchpad with its corresponding position in the input. The data format and position coupling scheme are presented in detail in Appendix E.3. We mention here that, because the chains are interleaved randomly, the position IDs used in the position coupling scheme depend on the particular imput instance. Thus, the use of Predictive Position Coupling is crucial in order be able feed in the correct coupled position IDs in the scratchpad at inference time." We train on input se- quences of length É [20,60), and test on lengths Le [20,200].
Experimental setup. We trained our models using the OLMo codebase (Groeneveld et al., 2024) on the C4 dataset (Raffel et al., 2019). We trained a transformer model fishort us- ing context length 64. 10 Our aim is to length extrapolate to a context length of L 2L = 128. We used the PoSE technique with Lmax L. and c= 2 chunks. We found that training from scratch using PoSE did not yield good performance at length L, so instead first trained the model without PoSE for roughly 70% of the iterations and on the final 30% of the iterations used POSE Precise hyperparameters may be found in Appendix F.
CE loss oitf between long & short contexts
Results. Figure Ic shows our results for the variable assignment problem with position coupling: the technique of Predictive Position Coupling allows good length generalization for test lengths up to 3 times the training sequence length, for depths d E 4,,3,4,5). In Figs. 5a to 5c (Appendix E.3.2), we present baselines that remove either just the position coupling or the scratchpad altogether, as predicted by takeaways (T2) and (TI), respectively, these modifications significantly harm length generalization. In particular, using RoPE with PoSE (with or without scratchpad) performs significantly worse than PPC.
CE from unmasking tokens
Figure 2: Cross eooss-yntropy loss differences obtained by unmasking & influential tokens (æ-axis) or all of the first 64 tokens (g-axis).
5.3. Sparsity & length generalization in natural language
Finally, we study the interplay between sparsity and length generalization in natural language. A trivial way to obtain generalization from sequences of length < L. to those of length L> L is to simply ignore tokens more than L units in the past.
Results & discussion. Table 5 confirms that the perplexity of
hishort on a context window of 128 is less than that on context window of 64, verifying that hyhort is able to utilize information
Simply feeding in the correct value as the next position ID would be "cheating" would be performing part of the relevant computation of the scratchpad.
We chose the context length to be so short in order to decrease the performance of the trivial baseline for length generalization which ignores all hut the most recent tokens.
8

---

Length Generalization in Transformers
beyond the shorter length-64 training context window. More- over, in Figure 2, we plot, for each of total of 500 evaluation data points i.F. with context window L= 128, the amount of decrease in cross cross-entropy loss for predicting the last token X of the sequence one achieves by the following two strategies:
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learnin. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
(a) Unmasking k = "influential" tokens at positions i 64 compared to masking out all tokens at positions i4 (this difference is the z-coordinate);
References
Abbe, E., Bengio, S., Lotfi, A., and Rizk, K. Generalization on the Unseen, Logic Reasoning and Degree Curriculum. In International Conference on Machine Learning. arXiv, June 2023, arXiv:2301.13105 [cs, stat].
(b) Unmasking all tokens positions i I 64 compared to masking out all tokens at positions i 64 (this difference is the g-coordinate).
In particular, the y-coordinate of each point should be inter- preted as the amount the model uses the first 64 tokens to improve the prediction the last token, and the r-coordinate of each point should be interpreted as the amount the model uses select of the first 64 coordinates (see Figure 6 for an illustration). As can be seen, the data points roughly clus- tered along the line 2 my=y, y. which indicates that (up noise) the model's prediction of the last token in the sequence is sparse in its dependence on tokens "far in the past". consistent with Sec- tion 5 which predicts that such sparsity allows habort to length- generalize. We discuss additional details and experiments (in- cluding how the & influentail tokens are chosen) in Appendix F.
Ahuja, K. and Mansouri, A. On Provable Length and Compo- sitional Generalization, February 2024. arXiv:2402.04875 [cs, stat].
Andonian, A., Anthony, Q. Biderman, S., Black, S., Gali, P.,
Gao, L., Hallahan, E., Levy-Kramee J., Leahy, C., Nestler, L., Parker, K., Pieler, M., Phang, J., Purohit, S., Schoelkopf, H., Stander, D., Songz, T., Tigges, C., Thérien, B., Wang. P., and Weinbach, GPT-NeoX: Large Scale Autoregressive Language Modeling in Py PyTorch, 2023.
Andriushchenko, M., Croce, F,, and Flammarion, N. Jailbreak- ing leading safety-aligned Ilms with simple adaptive attacks, 2024.
Acknowledgements
NG was supported in part by a Fannie & John Hertz Foundation Fellowship and an NSF Graduate Fellowship. We thank anonymous reviewers for very helpful feedback.
Anil, C., Wu, Y., Andreassen, A., Lewkowycz, A., Misra, V., Ramasesh, V., Slone, A., Gur-Ari, G., Dyer, E., and Neyshabur, B. Exploring Length Generalization in Large Language Models. In Neural Information Processing Systems. arXiv, November 2022. arXiv:2207.044 .04901 [cs].
Chen, G., Li, X., Meng, Z., Liang, S., and Bing, L. CLEX: Continuous length extrapolation for large language models. In The Twelfth International Conference on Learning Representations, 2024.
Chen, S., Wong, S., Chen, L., and Tian, Y. Extending Context Window of Large Language Models via Positional Interpolation, June 2023. arXiv:2306. 15595 [es].
Child, R., Gray, S., Radford, A., and Sutskever, I. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509 2019.
Cho, H., Cha, J., Awasthi, P., Bhojanapalli, S., Gupta, A., and Yun, C. Position Coupling: Leveraging Task Structure for Improved Length Generalization of Transformers, May 2024a. arXiv:2405.2067 [cs].
Cho, H., Cha, J., Bhojanapalli, S., and Yun, C. Arithmetic Transformers Can Length-Generalhe Generalize in Both Operand Length and Count, October 2024b. arXiv:2410. 15787 [cs).
9

---

Length Generalization in Transformers
Cohen- Wang, B., Shah, H., Georgiev, K., and Madry, A. Contextcite: Attributing model generation context. In The eiihty- Annual Conference Neural Information Processing Systems, 2024.
Xu, Z., Zhang, Z., and Zhang, Z. Deepseek-rl: Incentivizing reasoning capability in Ilms via reinforcement learning, 2025.
Delétang, G., Ruoss, A., Grau-Moya, J., Genewein, T., Wen- liang, L. K., Catt, B., Cundy, C., Hutter, M., Legg. S., Veness, J.. and Ortega, P. A. Neural Networks and the Chomsky Hierarchy. In International Conference on Learning Rep- resentations. arXiv, February 2022. arXiv:2207 .02098 [cs].
Csordás, R., Irie, K., and Schmidhuber, J. The Devil is in the Detail: Simple Tricks Improve Systematic Generalization of Transformers. In Conference on Empirical Methods in Natural Language Processing. arXiv, February 202 arXiv:2108.1 12284 [cs].
Dzini N., Lu, X., Sclar, M., Li, X. L., Jiang, L., Lin, B. Y., West, P., Bhagavatula, C., Bras, R. L., Hwang, J. D., Sanyal, S., Welleck, S., Ren, X., Eutinger, A., Harchaoui, Z., and Choi, Y. Faith and Fate: Limits of Transformers Compositionality. In Neural Information Processing Systems. arXiv, October 2023. arXiv:2305. 18654 [es].
Csordás, R., Irie, K., and Schmidhuber, The Neural Data Router: Adaptive Control Flow in Transformers Improves Systematic Generalization. In International Conference Learning Representations. arXiv, May 2021b. arXiv:2110.07733 [cs].
emozilla. Dynamically scaled rope further increases performance of long context Ilama with zero fine- tuning. Dynamically-Seaan Scaled- RoPE,, 2023. Accessed: 2025-01-28.
Dai. Z., Yang, Z., Yang, Y., Carbonell, J., Le, Q.. and Salakhutdinov, R. Transfomer-XLL: Attentive language models beyond a fixed-length context. In Korhonen, A., Traum, D., and Mårquez, L.. (eds.), Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 2978- 2988, Florence, Italy, July 2019. Association for Computational Linguistics.
https:/ /tte://giitt gigt.b...... com coon/emozii 1
Fan, Y., Du, Y., Ramchandran, K., and Lee, K. Looped Transformers for Length Generalization, September 2024. arXiv:2409.16647 [es].
De,,Seek-A,, Guo, D., Yang, D., Zhang. H., Song. J., Zhang, R., Xu, R., Zhu, Q. Ma, S., Wang, P., Bi, X., Zhang, X., Yu, X., Wu, Y., Wu, Z. F., Gou, Z., Shao, Z., Li, Z., Gao, Z., Liu, A., Xue, B., Wang, B., Wu, B., Feng. B., Lu, C., Zhao, C.,, Deng, C., Zhang, C., Ruan, C., Dai, D., Chen, D., Ji, D., Li,
Gehring, J., Auli, M., Grangier, D., Yarats, D., and Dauphin, Y. N. Convolutional sequence to sequence learning. In Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 1243-1252. PMLR, 2017.
E., Lin, F., Dai, F., Luo, F, Hao, G., Chen, G., Li, G., Zhang, H., Bao, H., Xu, H., Wang, H.,, Ding, H., Xin, H., Gao, H., Qu, H., Li, H., Guo, J., Li, J., Wang. J., Chen, J., Yuan, J.,
Groeneveld, D., Beltagy, I,, Walsh, P., Bhagia, A., Kinney, R., Tafjord, O., Jha, A. H., Ivison, H., Magnusson, I., Wang. Y., Arora, S., Atkinson, D., Authur, R. Chandu, K. R., Cohan, A., Dumas, J., Elazar, Y. Gu, Y., Hessel, J., Khot, T., Merrill, W., Morrison, J., Muennighoff, N., Naik, A., Nam, C., Peters, M. E., Pyatkin, V., Ravichander, A., Schwenk, D., Shah, S., Smith, W., Strubell, E., Subramani, N., Wortsman, M., Dasigi, P., Lambert, N., Richardson, K., Zettlemoyer, L., Dodge, J., Lo, K., Soldaini, L., Smith, N. A., and Hajishirzi, H. Olmo: Accelerating the science of language models, 2024.
Qiu, J. Li, J., Cai, J. L., Ni, J., Liang, J., Chen, J., Dong, K., Hu, K., Gao, K., Guan, K., Huang, K., Yu, K., Wang. I.,, Zhang, L., Zhao, L., Wang, L., Zhang, L., Xu, L., Xia, L., Zhang, M., Zhang, M., Tang. M., Li, M., Wang. M., Li, M., Tian, N., Huang, P., Zhang, P., Wang, Q.. Chen, Q. Du, Q. Ge, R., Zhang, R., Pan, R. Wang, R., Chen, R. J., Jin, R. L., Chen, R., Lu. S., Zhou, S., Chen, S., Ye, S., Wang. S., Yu, S., Zhou,, S., Pan, S.. Li, S. S., Zhou, S., Wu, S., Ye, S., Yun, T. Pei, T., Sun, T., Wang, T. Zeng. W., Zhao, W., Liu, W., Liang, W., Gao, W. Yu, W., Zhang. W., Xiao, W. L., An, W., Liu, X., Wang. X., Chen, X., Nie, X., Cheng, X., Liu, X., Xie, X., Liu, X., Yang, X., Li, X., Su, X., Lin, X., Li, X. Q.. Jin, X., Shen, X., Chen, X., Sun, X., Wang, X., Song. X., Zhou, X., Wang. X., Shan, X., Li, Y. K., Wang, Y. Q. Wei, Y. X., Zhang, Y., Xu. Y., Li, Y., Zhao, Y. Sun. Y., Wang, Y. Yu. Y., Zhang, Y., Shi, Y. Xiong, Y. He, Y., Piao, Y., Wang. Y., Tan, Y. Ma, Y. Liu, Y. Guo, Y. Ou, Y., Wang, Y., Gong, Y., Zou, Y., He, Y., Xiong, Y., Luo, Y., You, Y. Liu, Y., Zhou, Y., Zhu, Y. X., Xu, Y., Huang, Y. Li, Y., Zheng, Y., Zhu, Y., Ma, Y., Tang. Y., Zha, Y., Yan, Y. Ren, Z. Z., Ren, Z., Sha, Z., Fu, Z., Xu, Z. Xie, Z. Zhang, Z., Hao, Z., Ma, Z. Yan, Z., Wu, Z., Gu, Z.,
Hahn, M. and Rofin, M. Why are Sensitive Functions Hard for Transformers? In Annual Meeting of the Association for Computational Linguistics. arXiv, May 2024. ar arXiv.2402.09963 [cs].
Han, C., Wang, Q. Peng, H., Xiong, W., Chen, Y., Ji, H., and Wang, S. Lm-infinite: Zero-shot extreme length generalization for large language models, 2024.
Hou, K., Brandfonbrener, D., Kakade, S., Jelassi, S., and Malach, E. Universal Length Generalization with Turing Programs, July 2024. arXiv:2407.083: 0000000..33100 [cs].
Zhu, Z_, Liu, Z., Li, Z.,, Xie, Z., Song, Z., Pan, Z., Huang, Z.,
10

---

Length Generalization in Transformers
Hsieh, C.-P., Sun, S., Kriman, S., Acharya, S., Rekesh, D., Jia, F., Zhang, Y., and Ginsburg, B. RULER: What's the Real Context Size of Your Long-Context Language Models?, August 2024. arXiv.2404.0664 [cs].
Lee, N., Sreenivasan, K., Lee, D., Lee, K., and Papailiopou- los, D. Teaching Arithmetic Small Transformers. In ICLR. arXiv, July 2024. (iXiv:2307.03381 [cs].
Lee, N., Cai, Z., Schwarzschild, A., Lee, K., and Papailiopou- los, D. Self-improving transformers overcome easy-to-hard and length generalization challenges, 2025.
Hu, Y., Tang, X., Yang, H., and Zhang, M. Case-based or rule-based: how do transformers do the math? In Proceedings of the 41st International Conference on Machine Learning. ICML'24. JMLR.org, 2024.
Lewkowyez, A., Andreassen, A., Dohan, D., Dyer, E., Michalewski, H., Ramasesh, V., Slone, A., Anil, C., Schlag, I,, Gutman-Solo, T., et al. Solving quantitative reasoning problems with language models. arXiv preprint 2.Ii:2200.118858, 2022.
Huang, X., Yang, A., Bhattamishra, S., Sarrof, Y., Krebs, A., Zhou, H., Nakkiran, P., and Hahn, M. A Formal Framework for Understanding Length Generalization in Transformers, October 2024. 10X0v2440.021400 [cs].
Li, S., You, C., Guruganesh, G., Ainslie, J., Ontanon, S., Zaheer, M., Sanghai, S., Yang, Y., Kumar, S., and Bho- janapalli, S. Functional Interpolation for Relative Positions Improves Long Context Transformers. In International Conference on Learning Representations. arXiv, March 2024. wrXi:2331004418 [cs].
Huang, Z., Liang D., Xu, P., and Xiang, B. Improve transformer models with better relative position embeddings. In Cohn, T. He, Y., and Liu, Y. (eds.), Findings of the Association for Computational Linguistics: EMNLP 2020, pp. 3327-3335, Online, November 2020. Association for Computational Linguistics.
Liu, B., Ash, J. T., Goel, S., Krishnamurthy, A., and Zhang, C. Transformers Learn Shortcuts to Automata. In International Conference on Learning Representations, arXiv, May 2022. arXiv:2210. 10749 [cs, stat).
Hupkes, D., Dankers, V., Mul, M., and Bruni, E. Composition- ality decomposed: how do neural networks generalise? In Journal of Artificial Intelligence Research. arXiv, February 2019. MxXiv1198.000 08351 [cs, stat].
Liu, B., Ash, J. T., Goel, S., Krishnamurthy, A., and Zhang, Exposing Attention Glitches with Flip-Flop Language Modeling. In Neural Information Processing Systems. arXiv, October 2023. arXiv:2306.00946 [cs].
Jelassi, S., d' Ascolii S., Domingo-Enrinn C., Wu, Y., Li, Y., and Charton, F. Length Generalization in Arithmetic Transformers, June 2023. arXiv:2306. 15400 [cs].
Lu, Y., Yan, N., Yang, S., Chiu, J. T., Ren, S., Yuan, F, Zhao, W., Wu, Z., and Rush, A. M. A controlled study on long context extension and generalization in llms, 2024.
Jelassi, S., Brandfonbrener, D., Kakade, S. M., and Malach, E. Repeat After Me: Transformers are Better than State Space Models Copying. In International Conference on Machine Learning. arv, February 2024. arXiv:2402.01032 [es].
Marsden, A., Dogariu, E., Agarwal, N., Chen, X., Suo, D., and Hazan, E. Provable Length Generalization in Sequence Prediction via Spectral Filtering, November 2024. 2xXiv:241.00108: [es].
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, 2024.
McLeish, S., Bansal, A., Stein, A., Jain, N., Kirchenbauer, J., Bartoldson, B. R., Kailkhura, B., Bhatele, A., Geiping, J., Schwarzschild, A., and Goldstein, T. Transformers Can Do Arithmetic with the Right Embeddings, May 2024. arXiv:2405.1733 [cs].
Kalavasis, A., Zadik, I., and Zampetakis, M. Transfer Learning Beyond Bounded Density Ratios, March 2024. arXiv:2403.1 11963 [cs, math, stat].
Kazemnejad, A., Padhi, L., Ramamurthy, K. N., Das, P., and Reddy, S. The Impact of Positional Encoding on Length Gen- cralization in Transformers. In Neural Information Process- ing Systems. arxiv, November 2023. arXiv:2305.1 19466 [cs).
Nagarajan, V., Andreassen, A., and Neyshabur, B. Understand- ing the Failure Modes of Out-of-Distribution Generalization. International Conference on Learning Representations. arXiv, April 2020. arXiv:2010.15775 [cs, stat].
Ke, G., He, D., and Liu, T.-Y. Rethinking positional encoding in language pre-training, 2021.
Nye, M., Andreassen, A. J., Gur-Ari, G., Michalewski, H., Austin, J., Bieber, D., Dohan, D., Lewkowyzz, A., Bosma, M., Luan, D., Sutton, C., and Odena, A. Show your work: Scratchpads for intermediate computation with language models. In Proceedings of the 10th International Conference on Learning Representations, 2022.
Lanchantin, J. Toshniwal, S., Weston, J. E., Szlam, A., and Sukhbaatar, S. Learning to reason and memorize with self-notes. In seirth- Conference on Neural Information Processing Systems, 2023.

---

Length Generalization in Transformers
Olsson, C., Elhage, N., Nanda, N., Joseph, N., DasSarma, N., Henighan, T., Mann, B., Askell, A., Bai, Y., Chen, A., Conerly, T., Drain, D., Ganguli, D., Hatieel-Doddss, Z., Hemandezz D., Johnston, S., Jones, A., Kernion, J., Lovitt, I.,, Ndousse, K., Amodei, D., Brown, T., Clark, J., Kaplan, J., McCandlish, S., and Olah, C. In-context learning and induction heads, 2022.
Shalev-Shwartz, S. and Ben-David, S. Understanding Machine Learning: From Theory to Algorithums. Cambridge University Press, 2014. ISBN 978-1-107-05713-5.
Shaw, P., Uszkoreit, J., and Vaswani, A. Self-attention with relative position representations. In Walker, M., Ji, H., and Stent, A. (eds.), Proceedings of the 2018 Conference of the North American Chapter of the Association for Compula- tional Linguistics: Human Language Technologies, Volume 2 (Short Papers), pp. 464-468, New Orleans, Louisiana, June 2018. Association for Computational Linguistics.
Ontanon, S., Ainslie, J., Fisher, Z., and Cvicek, V. Making Transformers Solve Compositional Tasks. In Proceedings of the 60th Annual Meeting of the Association for Computa- fional Linguistics (Volume I: Long Papers), pp. 3591-3607, Dublin, Ireland, 2022. Association for Computational Linguistics.
Shen, R., Bubeck, S., Eldan, R., Lee, Y. T.,, Li, Y., and Zhang. Y. Positional Description Matters for Transformers Arithmetic, November 2023. 1xXx:2311.11477 [es].
Peng, B., Quesnelle, J., Fan, H., and Shippole, E. YaRN: Effi- cient Context Window Extension of Large Language Models. In International Conference on Learning Representations. arXiv, November 2023. arXiv:2309.00) [cs].
Su, J., Lu, Y., Pan, S., Murtadha, A., Wen, B., and Liu, Y. RoFormer: Enhanced Transformer with Rotary Position Embedding. In Neurocomputiing arXiv, November 2021. arXiv:2104.098 [cs].
Peng, B., Narayanan, S., and Pupadimitrion, C. On limitations of the transformer architecture. In First Conference Language Modeling, 2024.
Tay, Y., Dehghani, M., Bahri, D., and Metzler, D. Efficient transformers: A survey. ACM Comput. Surv, 55(6), December 2022. ISSN 0360-0300.
Phuong, M. and Hutter, M. Formal Algorithms for Transformers, July 2022. arXiv:2207.0922 [cs].
Press, O., Smith, N. A., and Lewis, M. Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation. In International Conference on Learning Representations. arXiv, April 2021. arXiv:2108.1244 12409 [es].
Team, K., Du, A., Gao, B., Xing. B., Jiang. C., Chen, C., Li, C., Xiao, C., Du, C., Liao, C.,, Tang, C., Wang, C., Zhang, D., Yuan, E., Lu, E., Tang, F., Sung, F,, Wei, G., Lai, G., Guo, H., Zhu, H., Ding, H., Hu, H., Yang, H., Zhang, H., Yao, H., Zhao, H., Lu, H., Li, H., Yu, H., Gao, H., Zheng, H., Yuan, H., Chen, J., Guo, J.. Su, J., Wang. J., Zhao, J. Zhang, J., Liu, J., Yan, J., Wu, J., Shi, L.., Ye, L.,, Yu, L.. Dong, M., Zhang, N., Ma, N., Pan, Q. Gong, Q. Liu, S.. Ma, S., Wei, S., Cao, S., Huang, S., Jiang, T., Gao, W., Xiong, W., He, W., Huang, W., Wu, W. He, W., Wei, X., Jia, X., Wu, X., Xu, X., Zu, X., Zhou, X., Pan, X., Charles, Y., Li, Y., Hu, Y., Liu, Y., Chen, Y., Wang, Y., Liu, Y., Qin, Y., Liu, Y., Yang, Y., Bao, Y., Du, Y., Wu, Y., Wang. Y., Zhou, 2., Wang, Z. Li, Z., Zhu, Z., Zhang, Z., Wang, Z., Yang, Z., Huang, Z., Huang, Z., Xu, Z.,, and Yang. Z. Kimi k1.5: Scaling reinforcement learning with Ilms, 2025.
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 unified text-to-text transformer. arXiv e-prints, 2019.
Rozière, B., Gehring, J., Gloeckle, F., Sootla, S., Gat, I., Tan, X. E., Adi, Y., Liu, J., Sauvestre, R., Remez, T., Rapin, J., Kozhevnikov, A., Evtimov, I., Bitton, J., Bhatt, M., Ferrer, C. C., Grattafiori, A., Xiong, W., Défossez, A., Copet, J., Azhar, F., Touvron, H., Martin, L., Usunier, N., Scialom, T., and Synnaeve, G. Code Ilama: Open foundation models for code, 2024.
Ruoss, A., Delétang, G., Genewein, T., Grau-Moya, J., Csordás, R., Bennani, M., Legg. S., and Veness, J. Randomized Posi- tional Encodings Boost Length Generalization of Transform- ers. In Annual Meeting of the Association for Computational Linguistics. arXiv, May 2023. arXiv:2305. 16843 [es, stat].
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. Attention Is All You Need. In Neural Information Processing Systems. arXiv, August 2017. arXiv:1706.03760 [es).
Wang, 2., Wei, S., Hsu, D., and Lee, J. D. Transformers Prov- ably Learn Sparse Token Selection While cally-Connected Nets Cannot. In International Conference on Machine Learn ing. arXiv, June 2024. arXiv:2406.0689) [cs, math, stat].
Sabbaghi, M., Pappas, G., Hassani, H., and Goel, S. Explicitly Encoding Structural Symmetry is Key to Length Generaliza- tion in Arithmetic Tasks, June 2024. arXiv:2406.01 [cs].
Sanford, C., Hsu, D., and Telgarsky, M. Representational Strengths and Limitations of Transformers. In Neural Information Processing Systems. arXiv, November 2023. arXiv:2306.02880 [cs, stat].
Wei, A., Haghtalab, N., and Steinhardt, J. Jailbroken: How does LLM safety training fail? In seveth- Conference on Neural Information Processing Systems, 2023.
12

---

Length Generalization in Transformers
Wei, J., Wang, X., Schuurmans, D., Bosma, M., Ichter, B., Xia, F., Chi, E., Le, Q-, and Zhou, D. Chain-of-thought prompting elicits reasoning in large language models. In Proceedings of the 36th Conference on Neural Information Processing Systems, 2022.
Weiss, G., Goldberg, Y., and Yahav, E. Thinking Like Transformers. In International Conference Machine Learning. arXiv, July 2021. arXiv:2106.06981 [cs].
Wu, T., Zhao, Y., and Zheng, Z. Never Miss A Beat: An Efficient Recipe for Context Window Extension of Large Language Models with Consistent "Middle" Enhancement, June 2024. arXiv:2406.07138 [cs].
Zhang. C., Raghu, M., Kleinberg, J., and Bengio, Pointer value retrieval: A new benchmark for understanding the limits of neural network generalization, 2022.
Zhang, Y., Backurs, A., Bubeck, S., Eldan, R., Gunasekar, S., and Wagner, T. Unveiling Transforeer with LEGO: a syn- thetic reasoning task, February 2023. arXiv:2206.0430 [es].
Zhou, H., Bradley, A., Littwin, E., Razin, N., Saremi, O., Susskind, J., Bengio, S., and Nakkiran, What Algorithms can Transformers Learn? A Study in Length Generalization. In International Conference on Learning Representations. arXiv, October 2023. urXiv22310116022 [cs, stat].
Zhou, Y., Alon, U., Chen, X., Wang. X., Agarwal, R., and Zhou, D. Transformers Can Achieve Length Generalization But Not Robustly, February 2024. #rXiv:2402.09371 [es].
Zhu, D., Yang. N., Wang, L., Song, Y., Wu, W., Wei, F, and Li, S. PoSE: Efficient Context Window Extension of LLMs via Positional Skip-wise Training. International Conference on Learning Representations, arXiv, February 2023. arXiv:2309.1044 [cs].
Zou, A., Wang, Z., Carlini, N., Nasr, M., Kolter, J. Z., and Fredrikson, M. Universal and transferable adversarial attacks on aligned language models. arXiv preprint 15X1v:2307.1504,, 2023.
13

---

Length Generalization in Transformers
A. Related work
A.I. Theoretical guarantees for length generalization
Zhou et al. (2023); Huang et al. (2024) propose that transforees will singth-gnneraliz at tasks which can be solved using a short RASP program (Weiss et al., 2021). Theorem of Huang et al. (2024), however, assumes that the transformer is chosen so to minimize a certain regularizer tailored specifically for length generalization, and their result is asymptotic in nature (in contrast ours). Sabbaghi et al. (2024) show that using gradient flow 1-layer linear transformer model with relative position embeddings on a simple linear regression task will converge absolute position embeddings will fail to cength-geeeraiiz simple abstraction of attention heads that sums up pairwise dependencies of tokens experiences length generalization. Unlike our theoretical results, those of Sabbaghi et al. (2024); Ahuja & Mansouri (2024) are asymptotic in nature and cannot capture the (most common) case of seftmux attention head. Moreover, they proceed, roughly speaking, by showing that the learned model equal to the ground-trute model on the entire domain, which is generally not the case with actual language models (Zou et al., 2023; Wei et al., 2023; Andriushchenko et al., 2024). As our theoretical setup incorporates distributional assumptions, it establishes length generalization in cases where the learned model is only accurate on in-distribution inputs, which we believe to be more realistic. Moreover, all of the preceding works generally apply only to specific classes of transformers or linear classes; in contrast, our framework, while being able to capture an attention head, is significantly more general in that we allow arbitrary, potentially nonlinear, function classes (see Definition 3.3), thus meaning our results may have significance for other (non- nnon-rransformer classes of models as well.
model which length-generalizess whereas the use of Ahuja & Mansouri (2024) show that a model class defined as a
Hou et al. (2024) proposes to use Turing programs, scratchpad strategy inspired by Turing machines, to achieve length generalization an arra of tasks. Their theoretical results, however, are only representational in nature, showing that transformers can represent Turing programs without accounting for what models algorithms such as risk minimization will actually lear. Hahn & Rofin (2024) offer reasons why transformers struggle to length-generalize on parity, based on their sensitivity bias.
Further from our own work, Marsden et al. (2024) obtain provable guarantees for length generalization in the context of dynamical systems. Wang et al. (2024) prove that gradient descent on transformers can provably learn a task known as the sparse token selection task (Sanford al., 2023), which bears resemblance to Definition 3.2 (indeed, a slight modification of the task defined in Section 2.1 of Wang al. (2024) is in fact special case of distribution with sparse planted correlations (Definition 3.2)).
A.2. Modifying positional embeddings for length generalization
In addition to position coupling (Cho et al., 2024a;b) (and the closely related Abacus embeddings (McLeish et al., 2024)), which a focus in our paper, several other techniques have been developed to modify position embeddings during training and/or inference time to improve the length generalization performance of transformers. Zhu al. (2023) developed the positional skip-wise technique (PoSE; see Appendix B,3), which was later refined by modifying the distribution of position IDs used at training time in Wu et al. (2024). We remark that PoSE is conceptually similar to randomized position embeddings (Ruoss et al., 2023), which trains using a random set of position IDs from the test-length context window (with no guarantee on the contiguity of these IDs, unlike PoSE).
Another popular strategy to extend the context length of transformers which has seen traction for models at larger scales, such as as Code Llama (Rozière et al., 2024), is position interpolation. This technique scales down (" ("interpolates") the position IDs in the longer test-length context window match the length of the shorter training-length context. Several such interpolation strategies have been proposed, including the canonical choice of linear interpolation (Chen et al., 2023), as well NTK-RoPE (emozilla, 2023), YaRN (Peng et al., 2023), and CLEX (Chen et al., 2024); the latter strategies adjust the amount of interpolation done on a per-frequency basis, with different RoPE frequencies receiving different interpolation scales. One downside of these position interpolation strategies is that they generally require some fine-tuning on the longer test-length sequences in order to effectively use the longer context windows (c.g.. to achieve decreased perplexity on longer sequences than those seen during training). Such fine-tuning complicates the theoretical setting of length generalization, where it is typically assumed that any amount of training on the longer test-length sequences is not allowed. We remark, however, that these position interpolation techniques can be combined with PoSE (Zhu al., 2023; Wu al., 2024); exploring such combinations in the context of our experiments is interesting direction for future work.
Finally, we remark that some strategies, such as LM-Infinite (Han et al., 2024) and Self-Extend (Jin et al., 2024), have been proposed to adjust the attention mechanism at inference time so as to achieve length generalization without any fine-tuning, though their performance lags somewhat (Lu et al., 2024).
14

---

Length Generalization in Transformers
A.3. Empirical evaluations & explanations of length generalization
A number of papers have offered empirical evaluations and comparisons of various techniques for length generalization. Anil et al. (2022) studied few simple length generalization tasks, such as parity and variable assignment, and found that techniques such as finetuning and using scratchpad did not lead to much length generalization. Kazemnejad et al. (2023) compared the performance of various positional encoding techniques for length generalization, and found that NoPE performed best (though their analysis did not account for techniques such as position interpolation and PoSE which can significantly improve length generalization for, c.g., RoPE). Lee et al. (2024) observed that length generalization is difficult for transformers on arithmetic tasks. Finally, Lu et al. (2024) performed a systematic study comparing various approaches to extend the context length of LLMs, including various types of position interpolation.
Ontanon al. (2022); Dziri et al. (2023); Hupkes et al. (2019) study the out-of-distribution performance of transformers by focusing on compositional generalization, which refers to the ability of transformers to compose individual tasks found in the training data. There is also extensive work more broadly on out-of-distribution generalization (Nagarajan et al., 2020; Abbe et al., 2023; Kalavasis et al., 2024).
Finally, a number of works Hupkes et al. (2019); Liu et al. (2023); Delétang et al. (2022); Zhang et al. (2023); Hsieh et al. (2024)) introduce new benchmarks and datasets for studying length generalization and more broadly the performance of LLMs on long contexts.
Additional techniques for length generalization. Many papers have introduced new types of positional encoding schemes with the hope of improved length generalization, including relative positional embeddings (Shaw et al., 2018; Dai et al., 2019; Huang et al., 2020; Ke et al., 2021), ALiBi (Press et al., 2021), Hard-Alibi (Jelassi et al., 2024), and FIRE (Li et al., 2024).
Many other types of techniques have been proposed, such as priming the training dataset with a few long sequences (Jelassi et al., 2023), modifying the format of the input (Shen et al., 2023; Hu et al., 2024) such as by using a scratchpad (Lanchantin et al., 2023), architectural changes (Csordás et al., 2021b; Fan et al., 2024), and combining several such techniques (Csordás et al., 10 Liu et al., 2022; Zhou et al., 2024). Moreover, it was recently shown that length generalization can be achieved via a self-improvement framework (Lee et al., 2025), which trains the model on sequences of increasing length, which are labeled by previous versions of the model. We remark that such techniques lie outside the scope of this paper, as we require that the context length be bounded above by L throughout training. Moreover, such techniques still may suffer from the computational issues that plague longer context lengths.
B. Additional Preliminaries
B.I. Additional context on distributional assumptions
The following proposition rules out length generalization in the absence of distributional assumptions.
Propasition B.1. Fix any L.L with L< L, and a hypothesis class H VVV. Let PI....,PL be H-realizable distributions. realized by h² Suppose thut e > 0 and in defined in (I) satisfies supxevt. L(h² (()*Xx)(())) > e. Then there is an ensemble P extending PI. so that H does not have (lengh) generalization with respect to P.
Below, we give a simple example of a distribution ensemble having k-sparse planted correlations (per Definition 3.2).
Simple example: k-gram retrieval. Suppose that V -= [N] for some N N. and we want to model the "k-gram retrieval" task of outputting the token following a certain k-gram in the sequence X, which is closely related to the notion of induction heads (Jelassi et al., 2024; Olsson et al., 2022). Formally, for fixed length e, the tuple (X, (V is drawn as follows: first, draw Unif(y))=* independently, then choose a random e ~ Unif((e- - 21: 1]), and set X := X and
X,_x XI...,X,, Y: of ((x+4++) where v-V is some embedding function. In particular, the tokens X specify the particular k-gram we want to query. This distribution falls into the framework of Definition 3.2 by letting μ Unif(V), letting S* be uniform over length-(2k-o) +1) sequences (E'+ (L'+, L'k + - k+ L...,() where E' ~ Unif ([e- 2k -1), and Q be uniform over sequences
(Z.....,Z,, Z141 Z Zi.)..ZZ)) where "t...,Za++ 1 Un) Finally, g" (Zn (ZZ....,ZZz+4 = (Zz+1).
B.2. Transformers
In this section, we review the basic definitions and notations regarding the transformers architecture. We refer the reader to Phuong & Hutter (2022) for a more extensive overview.
15

---

Length Generalization in Transformers
Attention heads. We begin by defining atlention heads. Consider matrices K, Q, V ER and O where denotes the embedding dimension and h denotes the attention dimension. Let 0 == (K. Q V 0) denote the full set of parameters for an attention head. The corresponding attention head AttHead ) + R takes as input sequence h (hi,..,.r,) ,,,.) of L embedding vectors corresponding to individual tokens, as well as a query embedding vector Equery E R and outputs the embedding of the next token, (ttHeadg(h, by first computing:
q- Q.uqarry
kg K.,,,
v 0+ ( (q.k)))) WEELL,
where typically o =oftmaax Then, the attention head outputs AttHeady (eetaadg(hhhuuee (h. = 0 a
Masked self-attention. We consider transformers which perform masked self-attention, meaning that they aim to predict each token in sequence by attending to previous tokens in the sequence. Moreover, in doing so, multiple attention heads are combined into an attention layer, of which multiple are combined into a transformer. First, we describe an attention layer: for some m E N, we consider m attention heads parametrized by 01..-,0m. together with a multilayer perceptron MLP : :R. - TRd, with 2 layers. Letting 0 (LL...,....MLP)) we let the attention layer AttLayer, : (124) + (IR)) be the mapping defined as follows: for ie [4] and h (h1,..,h,,) E (
AttLayer (h), MLP +y) y Σ AutHeado, (h).
Note that AttLayer, (h), depends only on the embeddings up to index i, i... h..... h, Finally, given a collection 0= ((0())...,))))))) We) denoting parameters for each of A e N layers together with an embedding matrix W. Wd VI, we define Transformero: pL (R by, for h= (hi,....h.) E (Rd)L.
Transformer X)) Att AttLaayrrga W.-X):
where X is interpreted as a matrix in P whose columns are the one-hot encodings of each of the tokens of X. To produce a sequence of tokens, typically Transformer composed with an unembedding matrix W E R/ xd to each position, and a distribution for each position may be obtained by applying a softmax.
B.3. Positional Skip-Wiss (PoSE) Training
The Positional skip-wise (PoSE) technique (Zhu et al., 2023) (see also (Wu et al., 2024)) aims to ensure that: (a) the position IDs used during training cover all posible position IDs mmx. that could be observed test-time (where Lmax is the maximum length of test-time sequence), and (b) the differences between different position IDs seen in training sequences of similar magnitude to that seen in testing sequences. To do so, fix an integer denoting number of chunks, and given sequence of tokens X= (X.)..X.,, during training, we partition X into e contiguous chunks (i.e., subsequences) and assign to each chunk a random contiguous sequence of position IDs so that the first position ID of each chunk is greater than the last position ID of the previous chunk. At test time, one simply uses the true position IDs, namely (1, corresponding to a sequence X of length L.
The precise schemes to partition into chunks and assign position IDs that we use are as follows: in all of our experiments we take the number of chunks to be c=2: during training, we split a sequence X= (X,)..,Xx) into two parts by choosing a uniformly random position to split at. We then choose 2 integers Jo,Ji ~ (ni [[). and let the position ID of the first chunk begin min(Jo.JJ } and the position ID of the second chunk begin at max{ is essentially the same as the scheme used in Zhu et al. (2023) with all e = 0.
+ +1, where & is the length of the first chunk. This
B.4. Sparse functional attention captures a single attention head
In this section, we prove Proposition B.2, which shows that the k-sparse functional attention class in Definition 3.3 is sufficiently general so as to capture the class of single attention heads (as defined in Appendix B B.2). First we formally specify the relevant class of attention heads; to ensure that the class maps from sequences of tokens (as defined in Definition 3.3) as opposed to sequences of embedding vectors (as defined in Appendix B.2), we include an embedding matrix in the definition:
Definition B.1 (Single attention heads). Fix d EN and finite set V denoting the vocabulary. We define the class AttHeads which consists of mappings/ h: V + as follows. It is the set of all mappings from X E 4² an embedding vector as follows:
X AttHHadd (W. W.X..bqueryy.
16

---

Length Generalization in Transformers
where W ranges over all matrices W Rdxx Nquery ranges over all vectors hquery Rd. and 0 ranges over all tuples 0 = (K, Q V. V.O), where K Ve RRxd and O Raxh. Moreover, matrix X KER' V XL whose columns are the one-hot vectors of the individual tokens of the sequence X.
slightly abusing notation by interpreting X the
Proposition B.2 (Formal version of Proposition 3.1), Fix any d, E N and finite set V. Then there are function classes Gley consisting of functions mapping NxV- (BV) (Bo)- and gual consisting of functions mapping V so that AttHeadsd,h.v is equal to the I-sparse functional attention class Hatio ggeyygraa
Proof. We let Gkey be the set of all mappings 9n : NNv - R which only depend on the V-component of the input, and gval be the set of all mappings 91 :.:-V V
Now fix any 0= Q.V.). queery. We indexing some element of AttHeadsa,,.v For any L L given X E R whose columns represent the one-hot encoding vectors for some length-/. sequence, we have
AttHeady(W) W.W...K X, Qquer) - E.E. 04 VW.X;
where a, a;=softmax( CeciiiaaxxQQuuueee QQu kW.X,) ......... QQuqeer KW.Y iE iE[L].
Define go Ghey by, for X
90(1,)) (Q Qqerry. K KW
(where conflate xE V and its one-hot encoding vector in R and define 9i E by
91 (x) = O.W..
is now clear that the mapping X AttHeado (W. hquer) is identical to the 1-sparse functional attention mapping go defined Definition 3.3.
Moreover, since have assumed h 2 d, it is evident that any hoo.on for do ghey, 91 gual may be realized as X AttHeado (We. (....auuu) for some f. qquu We as above.
C. Proofs for length generalization
First, we make the following remark on Assumption 4.2.
Remark C.1 (On Assumption 4.2). To help interpret Assumption 4.2, note first that order for the distribution ensemble P to be realizable (Assumption 4.1) by a class ( satisfying lsscarllccality (Assumption 3.2), will typically the case that m-x minke Licca with probability at least 1-ô for S* E Q for any L E N. Thus, natural choice for Q is to fix some distribution 2 over sets in Scts,, ([Local)) and let 2,00 be the distribution of random shift of a sample from O Liocall) Formally, 2050 is the distribution of the shift S A, where S ~ Qpce and A Unif( (0..., Liocal )). It is straightforward
to see that such a construction ensures that Assumption 4.2 is satisfied with m)e = l. More broadly, any value rje S poly(f) leads to interesting conclusions in the context of our results, so we interpret Assumption 4.2 as being fairly mild.
C.1. Proof of Theorem 4.3
Proof of Theorem 4.3. Write Hattn : Hattn Hake Gual). By Assumption 4.1, we have E IX ((h, o for some hn* E Hattn, Let h Hattn be chosen according to (1), so that we can write h for some 90 E ghey . Evaal, Thus, for each ée [L/ we have
I5> E (X.Y)up
[C(h(X),Y))
(4)
Eseset (10) co
E (X.Y)>P
Y
E (xrEseteu Cxp( (S.
E EseSetsa exp exp (S 31 (Xs)
9" (Xxs))
(5)
(XPP
Σg St E ESetss (1)) exp Xsz))
17

---

Length Generalization in Transformers
where O* Di* 3/ is the function guaranteed by Definition 3.2. By Assumption 4.2, for any S Setsy, (L)), we have 20-2tou ) nrr and Lngg
then we use the convention that L (S*) 0.) At various points later in the proof, we will need to use (5) for t= I. Lecall or t- toocall: order for this to be valid, we need L 2-2 L/2, i.., L ALocaa, which is ensured by the hypotheses of Theorem 4.3.
(1()") S y/r. (Recall that if, e.g., S & Setsa((L-Laae
21(3")
Thus, by the definition of the k-sparse planted structure (Definition 3.2; in particular, are using the fact that X, p for all i*), we have that, for any X XEVL,S ,S' ESets ([L),
PL-z2 21ica2 3T>.ITT L (X,, Smm..
PL-21oca (X1:L- S² 1 Luzal 2h (1))
Smn. (6)
(x,5*))
PL Lonca (X,5*)
follows that
(Xs+) Esesecs (LL) (xp(0v(S.Xs)) (Xs)
TL (L5)
>
E (X (1 Wor Ja 2iipal
E (1x (X+)))
21 the map 821
([eSectu(L--Laa) 9' ('(x)) Esesetar (IL. Lana (mcll exp (5 (Xs)
Mh (L.)
E (X122 26 S" yo 24ypoal
(8)
Esresettu ( L-- Lmcal) exp (0o (5,
The triangle inequality then gives
Esesorar (XP Eseesaas (s))
Esesetex (ncal) exp (xx s) (Xs)
271-((00) (LO)>
E (X.:) 21 PwPp shuooll X1.-21 map oa
Es ((L- exp x
99)
We define the following random variables, where (X1:1-2laca SS) ~ PL 2naaar X, ji j> I. 2 locals and 12
Ao=
Σ
exp(Oo(S.Xs)) (Xs);
Bo:
Σ
exp(0o(S" xs))
SESets ((L. 2boal)
S'ESets, 21mca)
A;=
Σ
exp(in(S,Xs)) (Xs).
([ +ilmaal) Setay (LL 2Lacr (t ((=a))
B,
Se ([L- (|-- 2 Looat Lonal) SeStss (LL (o--2borrti1)/) exp(oo(S,Xs)).
Σ
Note that Ag.A ERd are vectors, and B,, B R are scalars. We have that
Ao
Ao A
E A Ao+ A1 Bi+ Bu+
E
2 (L5),
22n (I.).
Bu
Bu+BI
where the first inequality uses (5) with 8 T, 2 Aocal, (8) and the triangle inequality, and the second inequality uses (9). Let Zo denote the collection of random variables XI.I- X 2 and for i 2 1, let Z, denote the collection of random variables Then Zo. 21,.. are independent, and Z1, Z2... are identically distributed (since we have PI-2logg)) ). Since (S. = - Do for any S for which max(5)- (min(S))))))) >>>>>>>>> Liocal (Item of Assumption (X1-L- 2i:LL21pp!!'' )~~ ~~
X1, 2i--aat 2-
3,2), A; can each be written as functions of Z, for each i 120 O. Moreover, by Item 2 of Assumption 3.2, this function does not depend on for fori21 1. By Lemma C.1, it follows that, for any 12 -2,
E EseSetay (LL (L-2L (xp () (Xs) Esesery ( ('L ((++ exp( (S,Xs)) (. (Xs) Es'
B (X1.L 24 S" PuP atroes X,mN 482L 2lacal
Eseest @Seta (IL )لmcaa) (xp(0o
Es E ([L+1 exp(Co ( (s))),)))
<6nL. (t+))
A symmetric argument establishes that for any t2 0,
Ese Setuy (IL - 2lpre)) exp S (SXxs)) (Xx)
Σ Eseeets (I- (-ttreert -2 exp exp(gn S,Xx)) (Xs)
E (XLL 2brpol Pu 24baal Xymup vi50
Es E'eesets, eBets, (L-- (1 exp(6n( (S' ()
Es Estesets EBets, (I-t 1, 1 2 exp((n (S' .xss))
(T5).. t.
18

---

Length Generalization in Transformers
Combining the two displays above with Lemma C.2 (and again using Item 1 of Assumption 3.2), see that for any 4. 20,1 -2,
Esesets, (lL 21 (x ( 99
(I-Stts 10
E 1X 20--2tunge 2tooed S8t P- PL 2o Xy v350.1>
Extesma ESnta (It (l exp (s))
Est 1o exxpx (gr))
(4
(10)
Let fo.t denote the distribution of X E (VT+tnaaa(ttttt and S* Sctsk ([L+ Liocal Lioal (fo + +t))) defined as follows: first, draw (Xi:L-Lows S** S PL-2 as well as Xj ~ for all j /0, j> - iioca. and then set X, j fo for jE t (1)+ and 8 +0 Then (10) together with (5) (with 8= L L- 2 gives that, for all to 2 0,4 2 -2
Esesota (lsh
(Xs)
E Pr 20
(Xs))
(0mm-(68) (Tm)
(II)
By Definition 3.2 and Assumption 4.2, for any X E v and S E (eL))((L)) with maxxs"* < Lical, there are some tn 2 0,11 2 -2 satisfying to to 1 +2= (L- L)/ Loocal Il so that
Pr. (X
2T((5)
Sm Noting that we have L = L + Lto + t1), using (12) and summing (11) over all possible values of to, t for which to (L- Llocal gives that
I 2L-2 -S+ Llocal)
(12)
2a (X+)
(()) (x) (L)) exp
0- (Xss)
(s:))
<0 (L0)
Using that ("(Xx))=X under (X PL conclude that
[(((())))) sO(nmm LICoaa -25).
E (X,Y)-P).
Lemma C.1. Suppose that Zo Zi Zz, ... are independent random variables with sample space Q. and that Zi, Zz ... identically distributed. Furthermore suppose fo Ω Rd and 9. 9.900 Q R are measurable functions, and let 11, Zi-i for each i 1, and Ao Zo) (o
(7). Suppose >> O satisfies
Ao
Ao
Ao A
Ao+ A Bo B B2
E
Se,
E
SE.
(13)
Bo
Bo
B
By + B1
Then for any i21, 1, we have
Ao
Ao +
E
33E..
Bo
Bo
B +
B
Proof. We use a hybrid argument. Fix any i2 2, and define = =1= (Z),, We have
Ag
A1
E
Bo
Bo+ B, B
An
Ao A
Ao+ An+ A
TR
+E
$26,
Bu
Bo
+ B(
Bo+ B
Bo B + B
"In particular, any choice of to satisfying to Tobcal E Setsi (T. (T suffices.
19

---

Length Generalization in Transformers
where the final inequality follows from (13) together with the fact that Z2,22 are iid. By the triangle inequality, it follows that
An + A
Ao+ A1 + A.
E
3e.
Bo+BI
Bo+ B
Fix any j22. Summing over j and using Lemma C.2, we see that
Ao
E
<3c- Cj--(j-1). 1).
Bo+B1
BI+
final application of the triangle inequality yields the conclusion of the lemma.
Lemma C.2. Let LEN and a.q.
E Rd be vectors for 1 isis L and b,p; E eRyo be positive real numbers for 1 1sis L. Suppose
that for each iE [4].
gita Eç. Then
P.+
v Σ 1= 2i=1
to
E/-
Ei=
Pi +b
Proof. For each i, we have:
p.a bq:
a a
=
Sec
b(p; +b) b)
ps pb+bl
Multiplying both sides by bbp (which is positive), we see that
Jnna-ubl be
Summing both sides over i= 1 to
qi Σ paa ab
<b
(p)+
Let L S= and -
-i=1 qis so | Seb(S+b).
We need to bound
To do so, we compute
S+b
a V+a
WV
<MS+b)
b S+b
b(S+b)
B(S+b)
which gives the desired bound.
C.2. Improving the length generalization error
Below we discuss some ways that the error bound of O (T. (.5) in Theorem 4.3 can be improved:
First, the factor of can be removed instead of the training objective (I) we choose
h h argmin NEH RE(L/2..1 (C))(
E
(In particular, we take the worst length E [L) as opposed to an average length.) This holds because the factor of is picked up in the first step of the proof (4) where we bound the in in-distribution loss for each É E [L) We opt to use the formulation in (1) in the formal statement of Theorem 4.3 since more closely correponds to what done in practice.
When applying Theorem 4.3 it is without loss to take Liocal L/4, since making Liocal larger only make Assumption 4.2 easier to satisfy. Thus, the term (LoLa will be bounded as ((LL/L)²).
If we take 2P00 to be as in Remark C.1, then we can use the shift-invariance of the distributions 2,e0 to remove the term Ti from the error bound as follows. The term M7 in (6) may be replaced with (1)) since L-2L 2 L/2. Morcover, the instance 17F, in (12) can be replaced with O(1) since L > L LLccal
If we incorporate all of the above optimizations, then we obtain that (under the assumption on 2 introduced in Remark C.1), the error bound in Theorem 4.3 can be replaced with (T..T..2..).
20

---

Length Generalization in Transformers
C.3. On necessity of the assumptions
In the below proposition, we establish that each of the assumptions of Theorem 4.3 is necessary in that none of them can be individually removed.
Propasition C.3. There is constant >>> O so that the following holds. For any of the assumptions: (a) Item I of Assumption 3.2, (b) Item 2 of Assumption 3.2, (c) Assumption 4.1 with 800, (d) Assumption 4.2 with he ((f), and any integer LEN, there a 2-sparse functional functtioaa attention class Hattn = (gley gual) and a 2-sparse planted correlations distribution ensemble P which satisfy each of the other 3 assumptions with Liocal =2 2 but which does not have (L,L,c)-lengge generalization for any L 2L.
Proof. We show in turn that each of the assumptions cannot removed. For all of the examples, we set k= and V (1,2).
We let - [0,1] be the unit interval, and let -(-,
Removing locality. We first show that
(Item of Assumption 3.2) cannot be dropped. We consider the sparse function attention class Hattn = Hatto (Ghey, (vvly, where G are defined as follows:
= = (000.001 }. defined as follows:
:x=(2,2) or max(S)-min)
(S,x)
:otherwise
:x (2,2)
9o1 (S,x)
-00
:otherwise.
Gual consists of the single function 91 11:1 defined as 91 (2,2)= = and (x) = for all (x1xx2) (2,2).
Fix any integer L. We define a distribution ensemble with k-sparse planted correlations, as specified by distributions e ( 2 (A(Sess((())) for each & E,N, and e A(V)), well function g" : >V, as follows:
is the point mass on {1}.
Quoc is the point mass on {(2.2)]
Opes is uniform over the sets {1, 2}. (2, 3), ... [e 1, f if l < L; and is uniform over the sets ()..... (ff
We g" e as defined above.
is immediate from the above definitions that Hattn is relative (i.e., Item of Assumption 3.2). Next, realizability (Assumption
4.1) is clearly satisfied the function gon.gn achieves 0 loss. Finally, coverage (Assumption 4.2) is readily verified by taking Me 28 and Loocal 2. (In particular, for S* {a,a+ L) E ((e) and l' < L, we have that S* A ((")) for all choices of A and so we do not need to satisfy
< (such an inequality does not even make sense).
Finally, note that the risk minimization procedure which returns h h= haxa.a satisfies (1) with the chosen value of L and error 5=0, O. However, is readily checked that å experiences error 12(1) on the distribution Pe for any 222L.
Removing the relative assumption. Next we show that Item of Assumption 3.2 cannot be dropped. We consider the sparse attention class defined by the following gkey gual.
gley = (900.9011 }. defined as follows:
(22) and mai(S))min(())-
(Sx(S
:otherwise
:x= and max(s)-min(St =1
9o1 (Sxx)
0
-Sinn mmmin(s)22mmxxx) 2.max(5)>>>>>
xC
:otherwise.
21

---

Length Generalization in Transformers
gual consists of the single function 01 1/>+3/ defined 91 (2,2) and 91 (x1 -0 for all (X1,X2) + (2,2).
Fix an integer L. We define an ensemble P specified by 1422z Q
f .g" as follows:
is the point mass on {1}-
Quoc the point mass on {(222)
Opo is uniform over the sets (1.2),
g" - Eg
is immediate from the above definitions that H/attn = (gley gual) Locca-llocal (Item of Assumption 3.2) with Libcal =2 Realizability (Assumption 4.1) is satisfied with hi* = -0 Coverage (Assumption 4.2) satisfied with Liocal 11 and ryp = l.
The risk minimization procedure which returns om satisfies (I) with the chosen value of and error d = 0, but h experiences error 12(1) on the distribution Pg for any & L+1.
Removing realizability. Next, we show that Assumption cannot be dropped. We consider the sparse attention class defined by the following Gley gyal:
gley = [9100.901 J. defined as follows:
mmax($)-min$$}=== and x= (2,2)
90n
-00
otterwisse
-max(5)) =2 and (2,2)
901 (S,x)
:othewiise.
Fix an integer L. We define an ensemble P specified by the following 20 2,
Q "g²
j is the point mass on {1}-
Quoc the point mass on {(2.2))
P05 is uniform over the sets
if For L, we define:
(, = ... - -P--(()) 10(2-1)
9
((1.3))= (=
(14)
10(e-2)
"g"=91 gaal
It is immediate from the above definitions that H = (glen is Llocalllocal and relative (Assumption 3.2) with Liocal =2 Moreover, coverage (Assumption 4.2) satisfied with -
2.
However, note that the hypothesis h h9p0.. achieves expected loss on the distributions P; for l. S L, but 2(1) expected loss on the distributions P, for < L. 12 Moreover, this example shows that we cannot even hope for an agnostic guarantee whereby h is close to the best-in-class hypothesis, since for any & L, we in fact h3
E (xxY),,,((CCCe Y))
E
[Chh +)(+)))
(15)
(X Y)~P:
Note that for X Pe for / L, will have that (S, do for with nonzero probability; here, recall from Definition 3.3 legoo an (X) = [[]]]][[]]] Eseseta 91 91 (Xs).
13 Ennsuring that (15) holds the reason that we have O put more mass on sets with max(S)- (5)mmi((S) =2 in (14).
22

---

Length Generalization in Transformers
Removing coverage. Finally, we show that Assumption 4.2 cannot be dropped. We consider the sparse attention class defined by the following ghey, gual:
gley (9900-90) ), defined as follows:
(2,(2,) and
900 (S,x)
-00
:otherwise
(x,(2,2) and max(5)-min(5)
9o1 (S,x)
:otherwise.
Gual consists of the single function gr: 91:1= defined as gl (2,2)=1 and (X1,X2) =00))000 for all (x1,x)) /
Fix an integer L. We define an ensemble P specified by the following 2, 2
is the point mass on {1}.
Quec is the point mass {(2.2)).
cas uniform over the sets
if and uniform over the sets (...). L.
Lf
g" 91 Eg
is immediate from the above definitions that 7(2tt= H (Cley is Loccalllocal and relative (Assumption 3.2) with Liocal = 2. Moreover, realizability (Assumption 4.1) is satisfied by noting that 1900 has 0 loss.
However, note that the hypothesis h achieves expected loss on the distributions Py for es L, but 2(1) expected loss on the distribution Pe for L.
D. Theoretical analysis for position coupling
In this section we formally state and prove Proposition 4.4. We begin with some notations. Given set Q C: N and i Ωt let rinkg(1) denote the rank of i in 22 i.e., the position of i in Ω when the elements of Ω are sorted. Next for S E Sotss ([4), S= (k) and vhe [ - [L], we let vi((S) Setsy ([f) to denote the set (EE
First, we discuss some additional motivation for Proposition 4.4.
Remark D.1 (Limitations of Assumption 3.2). The combination of Assumptions 3.2 and 4.1 leads to the following restriction on the s-spas distribution ensemble P: in typical examples (modulo some degenerate ones where, e.g., functions g'al are constant), in order to satisfy realizability (Assumption 4.1), we will need the low-loss hypothesis h² Hgatte ggvely to be of the form h = for some I E Ghey which "selects out" the planted set S* and gi Egual which correctly evaluates the label Y given Xs+; formally, for each & N, with high probability under (X,S* (Y)
go (S,
g -S YS+S*
and of (Xs-)= -YY (We formally call this property strong realizability in Assumption D.1.) But by Item of Assumption 3.2, this means that - (S*((S)) S Liocal with high probability over the draw from Pr. Ideally, we would like establish results for planted k-sparse ensembles P for which the planted set S* is not local in this sense; this is accomplished by Proposition D.2.
We now formally define the notion of a local positional coupling, which expands out the technical details omitted in the informal version in Definition 4.
Definition D.2 (Local position coupling: Formal version of Definition 4.1), Fix distribution ensemble P with k-sparse planted correlations per Defimition 3.2 (defined by μ. QPes (oC A Liocarllccal position coupling of P defined by, for each { N, a joint distribution Qpas c over S* E Setsy ([f) and mapping ehe [[] + [e] so that the marginal of S* under Qpee-s is Q and with probability under the draw of (S* aht) co
1. max(v>((**) -ii min(()e(SS)))) Trocal-
23

---

Length Generalization in Transformers
2. For each i 1 (1) = (Iee., Indices not in S* are not coupled.)
3. For some fixed R (independent of E), [= and the distribution of the tuple ((rank, 8 (s-)(()) E ( (EV) does not depend e. 14
Position coupling operation on functional attention classes. We say that mapping 90 (N V) RU (-) (e.g., an element of Ghey is apsition-indepenneee if (S, depends only on X (and not S). For such go, we will write 9o: Vbc - We say that ghey is position-indeeeendee independent all gu E Ghey are position-indepennln independent.
We define a new vocabulary Vi= = (V |k)s and "position-coupled" versions, PC[go] : (N x D&)* RUU {-c) and PCIgi]:V V +y of mappings go: go: - Ruf -00 VK +y, as follows. First, given X V let
expand(X)EV>
be defined as follows: for each E [R] we write X, (Onn Za). : (Jire Zin) for some 2 0, Ji1.-...an E [e] and Zin V. Then we let expand(S,X. be the tuple consisting of all pairs (ju Z, (..)....(iinn (iinZinn Zin) across all ie [k:] where the indices Jm are sorted in increasing order across all i and all duplicates in the indices are removed. We next define
mmax(5) mmxx())mmm($S) Liccal
PC[go
:expand(X) (a) Vk otherwise,
(16)
(o(expand(X))
so that PC[1o] maps (Nx D4.** Roo and
91 (expand(XX)
:expand(S) e V
PC]gi(()::
[
(17)
otherwise,
where yo E is an arbitrary element; thus PC[g1] maps D y. For a sparse functional attention class 7/ = a g gual). we define
PC[ [NPCCm))PCCm)) : E Sn E
Note that hypotheses in PC[Hattt)] map V +..
Position coupling operation on distribution ensembles. Next, we define a "position coupling" operation which modifies distribution ensemble equipped with local position coupling (per Definition D.2) produce sequences in Vi2 with tokens coupled between positions.
Fix a distribution ensemble (Pe)ee with k-sparse planted correlations, which is specified by the tuple (n, 2P Q Given a sequence X E V a set S E Sotss,((()), and function the : [e]] - [C], we let PC(X,S] be the distribution over sequences X V defined as follows: for i [f]. the distribution X; is given by
=(rank ( (s) ())x,, jew
: (S))
X,
= (1,X m, (n)) Uni
(18)
Uni
:otherwise,
In words, we are "rearranging" the tokens of X by grouping together (i.e.. "coupling") tokens indexed by S which are mapped by ah to the same index. Moreover, for positions i not in (S)) we take the token at the smallest position of - and for positions i not in the image of v/z, we sample X; from Unif x μ.. We let PC[P:] denote the distribution of PC[X,S*] where (
Strong realizability. Finally, we introduce the following slight strengthening of realizability:
For Z N and E ri) is the position of in f when the elements of Q sorted in increasing order, Appendix D. Item technical condition which is satisfied by most practical position coupling schemes which have heen proposed.
24

---

Length Generalization in Transformers
Assumption D.1 (Strong realizability). Write Hat = (gheey gral). We say that an ensemble P is d-strongly approximately Hatto Haatt--realizablee if there are 96 g ot Egaal, for each & LEN, with probability 1 1-d ō under (X,5,,))~7
g Xs-)> (6(S.Xs)= VS
and o (Xs-)=
Using the fact that the diameter of Y is at most 1, is straightforward that d-strong approximate realizability implies d-approximate realizabillity.
Position coupling removes the need for locality. The below proposition confirms that applying the position-coupling operation distribution ensemble Pe with k-sparse planted correlations (Definition 3.2), to yield PC[Pt], and applying the position coupling operation PC to sparse functional attention class Hattt to yield PC[Hattm]] can effectively remove the requirement for locality to hold (per Assumption 3.2) in order to obtain provable length generalization:
Proposition D.2 (1 (Length generalization without locality from position coupling), Suppose that P= (Pe)aeN v is an ensemble with k-sparse planted correlations (Definition 3.2) defined by distributions (p, Dp ) for some Liocal E N is amenable to ppsition coupling (Definition D.2). Moreover suppose that the distribution of (Se(S))) for (S* zbe) Q satisfies the coverage assumption in Assumption 4.2 for some values of The > 0, EEN.
Then for any sparse functional attention class Hattn = Hattn (ghes, gral) for which Grey is position-iideeeendeeee and P 5-ap 5-approximately strongly Hattt--rellizall realizable, and any integers L,L [Lmadd for which Liocal L- L and L 2 ALiocal, the class PC[Hatii] achieves (L.,L.MLM. -LT.2..) -..22-))-engthh generalization with respect to the ensemble PC(P).
Proof of Proposition D.2. We will verify that the sparse functional attention class PC[ together with the ensemble PG [P] satisfy the requirements of Theorem 4.3 with respect to the vocabulary V defined above.
First, it is straightforward to see that PC [P] is itself k-sparse planted correlations distribution ensemble: for length e, the planted set is given by (S*) for S* Pi, and for all i S", the distribution of X, (for (<))( is an independent sample from Uni x μ.. (Here we have used the second and third cases in (18) together with Item 2 of Definition D.2.) We have additionally used Item 3 of Definition D.2, which ensures that the distribution of (xx) does not depend on é (recall from Definition 3.2 that Q cannot depend on E).
Next we verify each of the three assumptions made in Theorem 4.3:
Assumption 3.2 is satisfied by the first line of (16) (which ensures Liocar and since PS[g)](S.X) does not depend on S for sets S satisfying max(S) -min mmn(()) S Liccal (again by (16).
To show that Assumption 4. 1 is satisfied in the sense that PC[P) is 6-approximately Preiiattt--eeallzallle, we use that P is d-strongly approximately Hattn rcatt--calizabbee In particular, choose G0 E eC gf E egual per Assumption D.I. Consider a sample (X*55*) PC(Ppj) We can write 5* (S* and X as in (18) for sample (X,S* abe)~~Pet Pee. Thus, with probability 1-8 under (X,5+)~ (CCP)), we have from Eqs. (16) and (18) that
PC X 15 Xs.)>-o..
Moreover, in order to have PC (S,Xs)> ooo for some S E Setss ((I), we need (under the 1 - 5 probability event above) that expand(Xs) s) x== here we have used that dū E Egeyy is position-independent. For any such set S, it holds also that PC]gi](Xs) = of (expand(Xs)) - - (Xs-), which is equal to the label Y on this event.
Assumption 4.2 is satisfied on account of the assumption from the proposition statement on the distribution of (S* for (S* abt 2
Remark D.3 (Theoretical justification for PoSE), It is natural to wonder if our theoretical framework allows us to justify other tech- niques used to induce length generalization, such as PoSE (Appendix B.3), in a sense akin to Proposition 4.4. At a high level, PoSE is adjusting the distribution of position IDs during training, so that IDs always seen farther apart than the training context window at test time may nevertheless be observed in the same training instance. In other words, a model trained as such should "interpret" a greater range of sets S as satisfying the locality requirement of Assumption 3.2. Thus, we conjecture that this adjustment allows us to remove the locality requirement Assumption 3.2 as well; we leave a formal proof of this fact as an intriguing direction for future work.
25

---

Length Generalization in Transformers
E. Additional experimental details: synthetic data experiments
For our experiments with synthetic data (Sections 5.1 and 5.2), we used the GPT-NeoX decoder-only transformer model with a causal attention mask. The transformer had 12 layers, 16 heads, and an embedding dimension of 1024. We used the Adam optimizer with learning rate 5- 10- 5 and weight decay parameter equal to 0.I; moreover, the experiments use a linear learning rate scheduler with 300 warmup steps. All shaded areas in figures represent 95% confidence intervals, computed with respect to multiple training runs (the precise number of training runs for each experiment is noted in the following subsections).
Remark E.I. One might wonder why we emphasize Item I but not Item 2 of Assumption 3.2 in takeaway (T2). In fact, the conceptual message of Item 2, namely that position IDs only influence attention scores by their relative information (i.e., the difference between different positions) already captured by the fact that is common to use relative positional embeddings and variants (e.g.. RoPE (Su et al., 2021), FIRE (Li al., 2024)) in many open-soures transformer architectures. Due to the success of such embeddings, in this sense the constraint imposed by Item 2 can "come for free".
E.1. Sparse parity (Section 5.1)
E.1.1. HYPERPARAMETERS AND DATA FORMAT.
Hyperparameters. The hyperparameters specific our sparse parity experiments are shown in Table 1.
Parameter
Value
Minimum training length
20
Maximum training length
50
Minimum testing length
20
Maximum testing length
500
Training batch size
64
Number of training steps
50,000
Number of testing points Vocabulary size N Maximum position ID (Lmax)
512
20
510
Positional embedding
RoPE
Number of training runs
Table Hyperparameters for sparse parity experiments (Section 5.1).
Data format. For some positive integer N, we consider an alphabet V= satisfying & < we consider the task of evaluating the parity of a particular set of out of & tokens wven-numereed positions. The particular set of k tokens is defined as those tokens for which the directly preceding token (which must belong to [N) is <N/2. We remark that this task makes sense even for N 2.
Given integers k,
Formally, the task is defined as follows. For each f,k L, we let D denote the distribution over sequences of the form
<BOS> (1:6, <SEP
(19)
where (J1,...,Jt) E [V]] chosen uniformly subject to the constraint that exactly is of the values Ji.....Je are S N/2 and the remaining (- l values are > N/2. Moreover, by....bb ~ Unif and b* = Dus.gnph We consider the task of predicting the final token t
This task fits into the framework discussed in Section 3 as follows: Py is the distribution over tuples (X, Y) distributed as in (19) where X consists all tokens except the final one 4rr and Y = assumption that the l values of for which J/ N/2 are all within Lical of each other (for some choice of Liocal), then Assumption 4.2 is satisfied for π/ Off) as long as f2 22 [ooca)), and Assumptions 3.2 and 4.1 are both satisfied for natural choices of the sparse functional attention class Hattm
4,.." 15 It is straightforward to verify that, under the additional
15 T'eehhiially, the length of is 28 instead f, but this discrepancy is immaterial.
26

---

Length Generalization in Transformers
Train Sparsity
Train Sparsity
1.
Test Spansily
1.0
Tist Sparsity -6
Best Sponsty -a
Telt Spansiy -0
0.9
Spansty
Spansly
0.8
Accuracy
Accuracy
0.7
0.7
Spansty
Telt Spanity
led Spansiy
Test Sparsity 30
0.5
Tel Spanity 12
Sparsity
10c
200
300
400
500
10c
200
300
400
500
Length
Length
Train Sparsity
8
Train Sparsity 12
1.0
1.0
0.8
0.8
Accuraty
Accurary
Test Sparsity
-
Test Sparsity
0.6
Test Spersily
0.6
Test Spersily
Test Sparsity
Test Sparsity
Sparsity
Spariity
Sparsity
Sparsity
Test Sparsis sity- 16
Text Sparsity 16
0.4
0.4
Length
Length
Figure 3: Length generalization for the sparse parity task for values Kuain E (4,6,8,12) (see Figure la for Kuain 10).
1.2 ADDITIONAL. RESULTS
The analogous plots to Figure Ia for Kirain E [4,6,8,12) are shown in Figure 3. The same patterns for Kirain = 10 are evident for each of these other values of Ktrain (except that the values of Kerain E {4,6) are sufficiently small so that there is essentially perfect length generalization for each value of l & rrain).
E.2. Parity with scratehpad (Section 5.2.1)
E.2.1. HYPERARAMETERS AND DATA FORMAT
Hyperarameters.
The hyperparameters specific to our
experiments are shown in Table 2.
Data format. For each length E N, we train using data of the following format, which computes the parity of bits by successively computing the parity of the first bits for each i= l.
<BOS> by., <SEP> By....bbe
(20)
where b1....,bu ~ 10 are independent and 6/ +) by for E [.]. The model is trained to predict all tokens after the <SEP> (namely, the scratchpad). For position coupling, we use position IDs of 0,1,..,,,0.11...,, (which are shifted by random offset during training discussed above). Note that in this task, the position IDs used in position coupling in fact are determined solely by the length of the sequence.
We train on & for which the combined length of input and scratchpad satisfies 2 [10,40] and test for l satisfying 26 [10,200).
E.2.2. ADDITIONAL RESULTS
In Figure 4, we show the length generalization behavior for our experimental setup with modifications that (a) remove Predictive Position Coupling and simply use absolute position embeddings with random shift during training time, and (b) use RoPE with
27

---

Length Generalization in Transformers
PoSE. Both of these modifications significantly harm length generalization behavior.
Parameter
Value
Minimum training length
10
Maximum training length
40
Minimum testing length
10
Maximum testing length
200
Training batch size
64
Number of training steps
50,000
Number of testing points
192
Maximum position ID (Lmax)
210
Positional embedding
Absolute Position Embeddings (leared)
Number of training runs
10
Table 2: Hyperparameters for parity-with-scracchpad
expp:peeeeeecccccceeceeeeeeteete experiments (Section 5.2.
Accurney
0.-
0.4
75
100
125
150
175
25
50
75
300
125
150
175
Length
Lengthh
(a) Absolute positional embeddings with random shift
(b) RoPE with PoSE
Figure Length generalization figure for parity-with--ae scratchpad experiments (Section 5.2.1), with modifications that replace Predictive Position Coupling. Length generalization behavior is significantly worse than that in Figure 1b.
E.3. Variable assignment with scratchpad (Section 5.2.2)
Parameter
Value
Minimum training length
20
Maximum training length
60
Minimum testing length
20
Maximum testing length
200
Training batch size
32
Number of training steps
50,000
Number of testing points
192
Vocabulary size (number of variables)
200
Maximum position ID (Lmax)
230
Positional embedding
Absolute positional embedding (learned)
Number of training runs
3!7
Table Hyperparameters variable assignment with scratchpad experiments (Section 5.2.2).
Foror the ablation studies in Figure 4, we only used 3 training runs.
1he ablation study in Figure 5h uses training runs.
28

---

Length Generalization in Transformers
E. HYPERPARAMETERS AND DATA FORMAT
Hyperparameters. The hyperparameters specific to our variable assigmment-wiih--ccrcchhhp experiments are shown in Table 3.
Data format. We now describe in detail the distribution of data for the "variable assignment with scratchpad" problem discussed in Section 5.2.2. To sample problem instance, we first draw some number w of chains, where a chain a sequence of d variable assignments of the form: e1 02;12 +- "3j"" ;Wa <- vay1- We sample uniformly random sequence of such chains with distinct variable IDs, where chain j is denoted vJ,1 f- 1j,2 f- " f- Vj,++1 We interleave these chains randomly (but keep the order within each chain preserved). We let Vinit = U1 and begin the input sequence with V1, (so that it is possible for the model to determine which chain to follow). Finally, the sequence concludes with a scratchpad of the form V1,1 < 01,2 2,,2 f U1,3; " < 01,4+1, evaluating the variable assignments along the correct chain. Overall, an example sequence consists of the following token IDs (here, +,, <SEF > <SEP>> are separator tokens):
<SEP2>,
1 U12...u2..... 01,41
(21)
Here we have used T (2), E [w]] to denote the chain index corresponding to each step in the input sequence, and (2).. E [d] to denote the position within its chain corresponding to each step in the input sequence. As an example with 2=2 and w= 2, we might have the sequence:
<BOS> 21,1, <SEP, 2,,,,,,,,,,,,,0,,, 22,,1111111 + 21,2,11,,, 11,3, 01 U2
The model is trained predict all tokens after the <SEP>>, For a fixed depth d and desired length of the input, we choose as large as possible so that the length of the sequence between <SEP,> and <SEP₂> is at most f.
For position coupling, the first part of the example (before <SEP2>) uses the tne position IDs 1, 1,2,3....,f, where & denotes the num- ber tokens before <SEP>>. The final part of the example (after <SEP>>) uses position IDs 11,44 111.11++220 2, + where for j E [d], 1 denotes the position ID assigned to V1,j in the first part of the example. Summarizing, the position IDs corresponding to the example (21) are:
Note that, because II, depend on the (random) functions x,a, it is crucial that the model learns to predict these coupled position IDs, so that they can be passed as the next position ID.
E.3.2. ADDITIONAL RESULTS
Figure 5 shows the length generalization behavior for our experimental setup with modifications that (a) remove Predictive Position Coupling and use absolute positional embeddings with random shift during training time, (b) use RoPE with PoSE (keeping the scratchpad), and (c) use RoPE with PoSE and remove the scratchpad. Notice that Predictive Position Coupling (Figure Ic) greatly outperforms all 3 of these modifications.
F. Additional experimental details: Natural Language Experiments
Additional experimental details. Hyperparameters used for our training procedure for the model hasoors discussed in Section 5.3. The model /atort with context length T= 64 discussed in Section 5.3 was trained for 34000 steps without POSE, and 15000 steps with POSE. We additionally trained model hlong with context length L 128 for 48000 steps, whose hyperparameters are identical to those of hahort except the batch and microbatch sizes are halved. Altogether, both models P were trained on roughly 25B tokens, and each has 1.3B parameters.
Formal setup for results. Let hishore (X; X1s-1) denote the probability Asaor assigns to token X; given the context X X For each test example X - (X, (.)), we first computed a set of = "influential" tokens for predicting X,, chosen amongst the first tokens, by finding the values of j L 1] which minimize
hanorr K1gi
(22)
29

---

Length Generalization in Transformers
10
Death
1.0
Douth
Arrumey
0.6
0zpth =2
Dephh =3
Depth
Dupth =5
0.0
0.0
Length
(a) Absolute positional embeddings with random shift
(b) RoPE with PoSE
00
Accursty
Depth
Desth
25
30
75
100
125
150
375
Length
(c) RoPE with PoSE & No Scratchpad
Figure Length generalization figure for variable assignment experiments (Section 5.2.1). Figs. Sa 5c show length generalization behavior with modifications that replace Predictive Position Coupling: length generalization significantly worse than that in Figure Figs. 5a and 5c report full-string accuracy on scratchpad; Figure 5h (which does not use scratchpad) reports accuracy of model predicting the answer token. Moreover, consistent with takeaway (T1) appears that, even when restricting to RoPE with POSE, having a scratchpad is superior to not having one (even when accuracy in the scratchpad case is measured by full-string correctness, as in Figure 5c), though there significant variance between different training runs.
i.e., where token Xj is masked out all attention computations. Let Ji,....Jg denote these is tokens. We then compute the following three negative boglliielhooods:
Labort = loghetort (X (i-L:i--)
Liong loghshor (X,, X 1,i--1)
Lsport,sparee loghaton x XJ
where compute Sparet mask out all tokens in the first L- tokens except those at positions JR,...JR. In words, Glong-Cshort, C denote the cross-entropy losses for predicting X;, when the (a) full context X is used, (b) the L most recent tokens X i-LL--11 are used, and (c) the I. most recent tokens as well as X are used, respectively. We illustrate the quantities Lehort Clong Llong Spaart for particular choice of L,L.Jı.t in Figure 6.
Results & discussion. In Figure 2, we plotted the pairs (Leshort Liong Lasvort Lashort for each example X. As can be seen, the quantity Lahort - Liong, which may be interpreted as the amount the model hthhart uses lokens X1-L-L-1 to improve the prediction of X i,, is roughly equal (up to noise) to Lahort - fashort which may be interpreted as the amount the model uses tokens X., to improve the prediction of X,, This provides evidence for the hypothesis that the model's prediction of the Lth token is sparse in its dependence on tokens "far in the past", and is thus consistent with (TI) which predicts that such sparsity allows the model fishort to annth--generaivee
is natural wonder whether the behavior seen in Figure 2 occurs even absent the of methods such as PoSE extend the context
18 This is essentially equivalent to "leave-one-out" baseline Cohen- Wang et al. (2024).
30

---

Length Generalization in Transformers
Hyperparameter
hihort
hoog
Model Dimension
2048
2048
MLP Hidden Size
8192
8192
Number of Heads
32
32
Number of Layers
24
24
Position Embedding
RoPE (w/ PoSH)
RoPE
Activation Type
GeLU
GeLU
Vocabulary Size
32000
32000
Tokenizer
Llama-2-76
Llama-2-7b
Batch Size
8192
4096
Microbatch (per-device) size
256
128
Training Context Length
64
128
Number of training steps
49000
48000
Optimizer Learning Rate
AdamW (w/o weight decay)
AdamW (w/o weight decay)
10-3
10
Table 4: Hyperparameters for our natural language modeling experiments (based off Mo codebase with dataset).
(2
Lioegi
Labrrti
X
X₂
X
X
X;
X
X;
X;
Caportaparee
X
Xa
X
X
X
X
X
X
Figure Illustration of the computation of E Lanort : spane sparce with L ===88=222 (i.e., not attended to) while those shaded blue are not masked (i.e., are attended (o).
2,J2 4. Tokens shaded gray masked
Train Context Length (Model)
Eval Length 64
Eval Length 128
64 (hstart
17.59
14.74
128 (hrong)
17.62
14.49
Table 5: Perplexities for (a) model trained on context length of 64 (top): and (b) model trained on context length of 128 (bottom).
length. Accordingly, also trained a model frong in the same manner as fishort but with context length L- 128 (and without PoSE). In Figure 7, we observe that similar positive correlation seen (a) for the model hlong trained on full contexts length L (Figure 7a), as well as for (b) the model hvong trained on contexts of length only L but where the indices J1.A are chosen as in (22) with respect to hypore (as opposed to hyghart Figure 7b), and conversely, the model fuhhort when the indices J1:k are chosen in (22) with respect hrong (Figure 7c). These observations suggest that the property that small number & of tokens at positions in [L-L- 1-L1)) 1] can nearly recover the cross-entrope loss at position obtained by training on all of the tokens at positions in [T.- [- may be more of prop- erty of the data distribution, opposed to particularity of any particular language model variant such h In particular, it indicates that some property like Definition 3.2 indeed governs the structure of the task of predicting the next token on long contexts.
31

---

Length Generalization in Transformers
bettoeen sbott
dif unmanting tokens
loes from unmaicing tosess
lnos from unmesting tobens
(a) hoong CE loss; Jus chosen from htong. (b) huong CE loss; Jusk chosen from hahort (c) Pusto CE loss; Jis chosen from hiong.
Figure 7: Modification of the plot of Figure where cross-entropy of hiong used instcad (Figs. 7a and 7b), and where the opposite model from the one being evaluated is used to generate the indices JI&s (Figs. 7b and 7c).
32


Reviews:
[{"review_id": "wXzKHYb0L9", "reviewer": "Reviewer_2aUD", "summary": "This work mathematically studies the context length generalization for a single next-token prediction task. It obtains a length generalization error upper bound when the task satisfies certain sparsity (i.e., only a subset of tokens in a context is important to solve the task), locality (i.e., the important tokens must be close to each other), and regularity (i.e., distributions of important token positions do not change a lot with respect to the context length). While locality might be difficult to be satisfied for general next-token prediction tasks, this work also shows that a task-specific position ID re-assignment, Position Coupling, can resolve the issue.\n\nOn the empirical side, first, the general theoretical findings (without position coupling technique) are verified on the sparse parity task. To develop the idea of position coupling and mitigate its limitation, a variant called Predictive Position Coupling (PPC) is proposed; the efficacy of PPC is tested for parity (with scratchpad) task and variable assignment task. Lastly, the interplay between sparsity and length generalization capability is studied for the natural language dataset.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your helpful feedback! \n\n**Regarding the various $L$ factors.** Thank you for these great insights that allow us to improve our bounds! \nThough 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. \n(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)$.)\n\n- 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.)\n\n- 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).\n\n- 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}$.)\n\n- 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.\n\n- 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. \n\n- We will add more details surrounding the equation around line 885.\n\n**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).\n\nYes, 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.\n\n**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$.", "claims_and_evidence": "The theoretical claims (mainly, Theorem 4.3 and Proposition 4.4) are clearly supported by their proofs provided in Appendices C and D. The authors’ main insight is also tested by several experiments. However a particular claim that “PPC relaxes the downside of position coupling” doesn’t seem to be well supported by any evidence because of lack of PPC’s implementation details: I hope to see a (sub-)section of appendix dedicated to more detailed descriptions about the implementation, training, and evaluation methods using PPC.", "methods_and_evaluation_criteria": "The paper has chosen several appropriate tasks to examine and verify their claim about crucial factors in achieving length generalization.", "theoretical_claims": "I checked almost every math in all proofs. Most of them seem “correct”, but I mainly am worried about the significance (or tightness) of the given error bound. Below, I list all my concerns about the theoretical parts of the paper.\n\n- The $L$ factor in the bound looks like appearing due to the very first step of the proof: when bounding the risk of $\\hat{h}$ for trained lengths $\\ell \\le L$ by $L\\delta$. This seems like an artifact of the analysis, especially caused by $\\hat{h}$ defined with an argmin of an expectation. Why don’t we re-define it as a minimax-type estimator, to consider the case where the in-distribution (ID) generalization is already done quite well? For example: $\\hat{h} = \\arg\\min\\_{h\\in \\mathcal{H}} \\max\\_{\\ell\\in [L/2, L]} \\mathbb{E}\\_{({\\bf X},{\\bf Y})\\sim P\\_\\ell} [\\mathcal{L}(h({\\bf X}),{\\bf Y})].$ In this case, we have a $\\delta$-bound instead of $L\\delta$-bound:\n    \n    $$\n    \\begin{align*} \\mathbb{E}\\_{({\\bf X},{\\bf Y})\\sim P\\_\\ell} [\\mathcal{L}(\\hat{h}({\\bf X}),{\\bf Y})] &\\le \\max\\_{\\ell'\\in [L/2, L]} \\mathbb{E}\\_{({\\bf X},{\\bf Y})\\sim P\\_{\\ell'}} [\\mathcal{L}(\\hat{h}({\\bf X}),{\\bf Y})] \\\\\\\\ &\\le \\max_{\\ell'\\in [L/2, L]} \\mathbb{E}\\_{({\\bf X},{\\bf Y})\\sim P\\_{\\ell'}}  [\\mathcal{L}(h^\\star({\\bf X}),{\\bf Y})] \\\\\\\\ &\\le \\delta.\\end{align*}\n    $$\n    \n    I am a bit concerned about the $L$ factor because the current theorem tries to claim that the error bound may increase as we elongate the trained sequences. It is quite opposite to a common empirical observation that the length generalization capability grows as we train a model on longer sequences. But good news, this issue might be able to be relaxed with a different definition of $\\hat{h}$! Still, the parameter $\\eta_L$ may scale with $L$ in worst case, so the problem is not completely gone.\n    \n- I don’t think the regularity assumption (Ass. 4.2) immediately implies, together with Def 3.2, that the density ratio bounds in math equations roughly in line 885, which is just above Equation (5). I’d be happy if a slightly more detailed derivation is shown in the proof (like in Equation (10)).\n- Typo: In a math equation (see line 885; above Equation (5)), $\\mathcal{P}_{L-L_{\\sf local}}(\\mathbf{X}_{1:L-L_{\\sf local}}, S^{\\star})$  seems correct.\n- There seem to be several errors in indices, especially after Equation (8): $t_0 + t_1$ or $t_0+t_1+2$?\n- I don’t think the risk bound provides any useful insight about the $L_{\\sf local}$-dependency. By merely following the proof, it seems that the error bound is proportional to $(\\bar{L} / L_{\\sf local})^2$. But is that all? I don’t think so, because it is difficult to understand that the length generalization capability may degrade as the locality gets better (i.e., $L_{\\sf local}$ gets smaller). In fact, $\\eta_{\\bar{L}}$ also depends on $L_{\\sf local}$. In my opinion, the theoretical analysis can be stronger by providing an appropriate discussion regarding these.\n- One of my biggest concern is the scale of $\\eta_{\\bar{L}}$. It can be extremely large as it bounds the ratio between two distributions relevant to a very short context length (namely, $L-2L_{\\sf local}$) and a large one (namely, $\\bar{L}$). If this is the case, can we say that the provided error bound is not vacuous?\n- Moving on to the necessity of assumptions (Appendix C.2): I guess the “provable necessity” of the introduced assumptions is an over-claiming since Appendix C.2 does not guarantee that we will fail to discuss length generalizability **every time** we violate one of them. A tone-downed claim might work.\n- Minor comment in Def 3.3: I guess $\\mathbb{N}^k\\times \\mathcal{V}^k$ is correct, instead of $(\\mathbb{N}\\times \\mathcal{V})^k$.", "experimental_designs_or_analyses": "I checked the experimental designs and their results. The experiments mostly support the main insight of the paper well, while missing some minor details or discussions.\n\n- I can’t find whether PPC doesn’t require the task structure at all. Does the proposed method never require the coupled position IDs in test time? How about in training time? Moreover, doesn’t the training on the prediction of coupled position IDs cause an additional axis of length generalization? If it does, why does the position ID prediction work well even at evaluation time, while the prediction of new next position ID might depend on every previously assigned position IDs (which does not seem to be a ‘sparse’ task)?\n- Although not being a *necessary* analysis, “Needle in a Haystack” benchmark might be an extreme case where the authors’ claim should hold. However, they did not discuss the relationship between this task and their claim in the paper. Note that, in fact, several LLMs fail in this benchmark.", "relation_to_broader_scientific_literature": "The key contribution of this work is highly related to scaling up the sequence-to-sequence models (like language models based on decoder-only Transformers).", "essential_references_not_discussed": "I can’t find any missing essential references.", "other_strengths_and_weaknesses": "**Strengths**\n\n- I love the paper explaining why position coupling can improve length generalization; it removes the necessity of locality assumption, allowing more tasks to be length-generalizable. It aligns well with the intuition on the benefit of position coupling explained in Cho et al (2024b), adding mathematical rigor significantly.\n\n**Weaknesses**\n\n- The relativity of the class $\\mathcal{G}^{\\sf key}$ is a bit confusing. In case of $k=1$, the definition of relativity will assign the same attention score to every token; is this realistic? Also, I hope the paper will explain more clearly about why the Item 2 of Assumption 3.2 models the usage of “relative position information”.\n- Be consistent with the jargon: there are both “position coupling” and “positional coupling” in the main text. It’d be better to choose only one of them.", "other_comments_or_suggestions": "See above.", "questions_for_authors": "Authors can find my questions above (sorry for not numbering them…). I am willing to update my evaluation if all the concerns get resolved.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xIgFjNL6HL", "reviewer": "Reviewer_J64p", "summary": "This paper studies why LLMs sometimes can or cannot generalize to input sequences longer than those they were trained on. The authors propose a theoretical framework centered on the concept of \"sparsity\": the idea that good length generalization occurs when each predicted token depends on only a small, fixed number of previous tokens rather than the entire context.\n\nThe paper shows that \"locality\" (tokens being close together) is also important for length generalization, but this requirement can be mitigated using position coupling techniques, and it introduces \"Predictive Position Coupling\", an extension of position coupling that allows the model to predict the position IDs for tokens dynamically rather than having them fixed.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive comments!\n\nRegarding the how the $k$ influential tokens are chosen: please see the discussion around Eq.~(20) in the appendix.", "claims_and_evidence": "The claim that natural language exhibits the sparse planted correlation property is only partially demonstrated. Their experiments show promising evidence, but a more comprehensive analysis of different types of linguistic dependencies and their relationship to length generalization would be needed to fully support this claim.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria seem sound.", "theoretical_claims": "- Definition 3.2 (k-sparse planted correlations) - the assumption that non-relevant tokens are i.i.d. from a background distribution is acknowledged as unrealistic. This simplification might affect how well the theory translates to real-world language data.", "experimental_designs_or_analyses": "- The paper doesn't fully explain how the \"k influential tokens\" in the natural language experiments are selected. This is crucial for interpreting the results in Figure 2.\n- The experiments don't thoroughly explore how the relationship between sparsity and length generalization might change with model scale, which is important for practical applications.\n- While the paper includes confidence intervals for synthetic tasks, a more thorough statistical analysis would strengthen the conclusions, particularly for the natural language experiments.", "relation_to_broader_scientific_literature": "In the broader landscape of LLM research, this paper helps bridge the gap between the theoretical understanding of transformers and practical techniques for improving their capabilities, particularly regarding handling long contexts efficiently.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "other strenghts:\n\n- The introduction of Predictive Position Coupling is an important contribution that extends position coupling to more general settings where the mapping between positions is input-dependent. This broadens the applicability of position coupling techniques.\n- The findings have clear practical implications for designing models that generalize better to longer sequences, suggesting that architectures encouraging sparse dependencies might perform better in length generalization tasks\n\nOther weaknesses:\n- While the paper includes experiments on natural language data, this section is less developed than the synthetic task experiments. The measurement of sparsity in natural language contexts could be more thorough, with more detailed analysis of how different types of linguistic dependencies affect length generalization.", "other_comments_or_suggestions": "this is just a suggestion, but i personally find a bit uncomfortable to read the paper given that all the references are in a very bright blue. i'd suggest changing the color in a less bright one.", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "N/A"}}, {"review_id": "ziwc0GHKhX", "reviewer": "Reviewer_KgQD", "summary": "This work suggests that length generalization occurs as long as each predicted token depends on a small (fixed) number of previous tokens. This work also conducts experiments on synthetic tasks and natural language.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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. \n\n1. *``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.\n\n2. *``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. \n\n3. *``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. \n\n4. *``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.\n\n  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.", "claims_and_evidence": "Yes, the experiment results support the claims and evidence", "methods_and_evaluation_criteria": "Yes, the syntheic task and nature language part results make sense.", "theoretical_claims": "I have checked the proofs.", "experimental_designs_or_analyses": "I have checked the soundness/validity of experimental designs and analyses", "relation_to_broader_scientific_literature": "This work finds that length generalization occurs as long as each predicted token depends on a small (fixed) number of previous tokens.", "essential_references_not_discussed": "It seems that the key contribution is not new, which is discussed in the following: \n\n[1] Fang, L., Wang, Y., Liu, Z., Zhang, C., Jegelka, S., Gao, J., ... & Wang, Y. (2024). What is Wrong with Perplexity for Long-context Language Modeling?. arXiv preprint arXiv:2410.23771.\n\n[2] Zheng, C., Gao, Y., Shi, H., Xiong, J., Sun, J., Li, J., ... & Li, Y. (2024). DAPE V2: Process Attention Score as Feature Map for Length Extrapolation. arXiv preprint arXiv:2410.04798.", "other_strengths_and_weaknesses": "Weakness: \n* The key contribution of the work is that length generalization occurs as long as each predicted token depends on a small (fixed) number of previous tokens. **However, it seems that the contribution is not new**, and it has already been discussed in the previous works [1,2]. Therefore, the author should clearly provide the difference between this work and previous work [1,2]\n* Also, the core method, presented in Figure 6, is just the sparse attention, which is well-discussed in previous works [3,4].\n* Therefore, the author should clearly distinguish their contribution with the previous works.\n\n\n\n[1] Press, O., Smith, N. A., & Lewis, M. (2021). Train short, test long: Attention with linear biases enables input length extrapolation. arXiv preprint arXiv:2108.12409.\n\n[2] Fang, L., Wang, Y., Liu, Z., Zhang, C., Jegelka, S., Gao, J., ... & Wang, Y. (2024). What is Wrong with Perplexity for Long-context Language Modeling?. arXiv preprint arXiv:2410.23771.\n\n[3] Ma, X., Liu, Y., Liu, J., & Ma, X. (2025). Mesa-Extrapolation: A Weave Position Encoding Method for Enhanced Extrapolation in LLMs. Advances in Neural Information Processing Systems, 37, 81389-81436.\n\n[4] Lou, C., Jia, Z., Zheng, Z., & Tu, K. (2024). Sparser is faster and less is more: Efficient sparse attention for long-range transformers. arXiv preprint arXiv:2406.16747.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "The key contribution of the work is not new and has already been discussed in previous works. The authors have to give a response to such limitation.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The reviewers unanimously endorse this paper, with two reviewers raising their scores after the rebuttal phase. The current scores stand at 4,4,3.

Initial concerns regarding the contextualization of the paper's contributions in relation to existing literature were successfully addressed during the rebuttal. The reviewers provided valuable suggestions for improving the tightness of the theoretical bounds, which the authors have acknowledged and explained in detail how they will incorporate these improvements in the camera-ready version.

I recommend acceptance.
Decision: Accept (poster)