Length Generalization via Auxiliary Tasks
Authors: Pranjal Awasthi
Anupam Gupta
Ravi Kumar
Venue: NeurIPS 2025 poster ()
Keywords: length generalization
theory
Primary area: theory
TL;DR: A new framework for analyzing and proving length generalization bounds.

Abstract: _Length generalization_, the ability of sequence models to
  generalize to sequences longer than those encountered during
  training, remains a key challenge for transformers,
  especially in tasks requiring algorithmic reasoning. Existing
  theoretical understanding of length generalization is limited, often providing
  only asymptotic results or focusing on specific problem classes or
  architectural variants, while empirical approaches frequently rely on
  ad hoc and often fragile techniques.

  In this work we introduce a novel framework for analyzing and
  proving length generalization bounds under specified, verifiable assumptions. A key
  outcome of the theory is the identification of a natural set of
  _auxiliary_ tasks, intricately related to the primary task structure,
  such that strong performance on these auxiliary tasks, alongside the
  primary task, provably guarantees length generalization within the framework.  This
  motivates a multi-task training procedure that explicitly optimizes
  performance on both the primary and the identified auxiliary tasks.

  Empirical evaluations on a variety of synthetic benchmarks
  known to be challenging for length generalization, including sequence
  sorting, and reversal, demonstrate that our proposed method yields
  significant improvements in generalization to substantially longer
  sequences.

Length Generalization via Auxiliary Tasks
Pranjal Awasthi
Anupam Gupta NYU anupam g@nuu..nuu......m edu
Ravi Kumar Google ravi Favi.k63@gmail.. com
Google
pranjaalaastth com
Abstract
Length generalization, the ability of sequence models to generalize sequences longer than those encountered during training, remains key challenge for trans- formers, especially in tasks requiring algorithmic reasoning. Existing theoretical understanding of length generalization is limited, often providing only asymptotic results focusing on specific problem classes or architectural variants, while em- pirical approaches frequently rely on ad hoc and often fragile techniques.
In this work we introduce novel framework for analyzing and proving length generalization bounds under specified, verifiable assumptions. A key outcome the theory the identification of natural set of auxiliary tasks, intricately related to the primary task structure, such that strong performance these auxil- iary tasks, alongside the primary task, provably guarantees length generalization within the framework. This motivates multi-task training procedure that explic- itly optimizes performance on both the primary and the identified auxiliary tasks. Empirical evaluations on variety of synthetic benchmarks known to be challeng- ing for length generalization, including sequence sorting, and reversal, demon- strate that our proposed method yields significant improvements in generalization to substantially longer sequences.
1 Introduction
Transformer architectures (Devlin et al. 2019 Vaswani et al.) 2017) have become the cornerstone of modern sequence modeling. achieving state-of-ther tf-the-rrt performance across a vast range of ap- plications. However, persistent and critical limitation is their often poor performance
length generalization- the ability to process sequences at test time that are significantly longer than those seen during training (Anil et al.) [202]] Kazemnejad al. [023] Press et al.) [2022] TTy al. (021). This deficiency particularly evident in tasks demanding algorithmic reasoning or extrapolation, such multi-digit arithmetic, symbolic manipulation, processing long documents (Saxton al.) 2019 Zhou et al.) 2024b). The failure to robustly generalize beyond the training length distribu- tion only limits practical applications, but also casts doubt whether these models truly learning underlying algorithms. or merely performing sophisticated pattern matching within a con- strained length range. Despite considerable research effort, achieving reliable length generalization remains a fundamental open problem (Nye et al. 2021 Saxton et al. 2019 Zhao et al. 2024).
Efforts address length generalization have yielded mixed results. Empirically, various techniques have been proposed, often involving modifications to positional encodings (PEs) (Anil et al.) 2022]
Saxton al. 2019) attention mechanisms. More recent works have found highly specialized PEs
like FIRE (Li et al. 2024) or position coupling (Cho et al. 20244] McLeish et al. 2024) essential for tasks like addition. However, it has remained difficult draw clear lessons from these works about when and why they work; moreover, the resulting length generalization can be fragile, highly sen- sitive factors like data formatting (c.g., reversed sequences, index hints for addition), scratchpad formats, random weight initialization, and training data order. This fragility raises the question of
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

whether these empirical fixes are merely learning task- shortcuts rather than inducing robust and generalizable algorithmic capabilities.
There have been several theoretical investigations into length generalization and how to boost per- formance (e.g., [Ahuja and Mansouri 2024 Huang al. 2025 (Sabbaghi al., [0224] Zhou et al.) 2024a)), but our understanding remains limited. The most relevant prior works suggest that spar- sity is a critical factor for length generalization (e.g., (Golowich et al. 2025 Huang et al. 2025
Sabbaghi et al. 2024)); ; we compare our work to these in Section
Our results. We propose a novel framework for studying length generalization transformers.
1. We present general theoretical framework to study length generalization. Rather than special- izing to particular tasks or training procedures, we aim for broad upplicability and focus on the class of sparse tasks: loosely speaking, these tasks where each output token depends on the presence of a small number of preceding tokens. (We formally define sparsity in Section
2. We show that robust performance at small training lengths implics generalization for such sparse tasks. Formally, our theoretical result considers input distributions { where D2 is distri- bution over length-k input sequences. We show that if the transformer model is "rrbustly" correct (with error most and large margin) on sequences of length n, then it has error N/n on quences of length N > n-i.e., the error increases only linearly in the sequence length. Here, being robustly correct at length n means that the error of the model is less than E, not only on sequences drawn from D.,, but also on suitable n-length random restrictions of sequences drawn from DN- (We present the notions of (N) pertuutuubaaso and robust correctness in Section 2.1
3. Our framework is not just descriptive, it also prescribes new training procedure. Indeed, the above results identify robust correctness as a desirable property: one way to ensure it at length n to actually train the model not only on the "native" distribution D,, but also on random n-length restrictions of -w-we call this robust training. It is well-known that training on just D, leads to poor length generalization, so our hope is that training on these random restrictions will lead robustly correct model that can
this is supported by our formal results.
Moreover, many common sparse tasks generate the next token by applying one of a small set of functions to (a small subset of) the preceding tokens. For instance, the sorting task requires repeat- edly finding the successor; likewise, the increment task repeatedly computes sums of two three digits. We treat these functions as auxiliary tasks and train the model jointly on the original and auxiliary tasks, drawing data from both D, and random n-length restrictions of Dx-
4. Finally, we empirically validate cluding commonly studied synthetic tasks such as sorting and reversal. We also introduce new tasks, including the SLiM task, which is aimed to capture toy language modeling scenario. Our results show that performing robust training at length n leads to good performance at length N > n.
approach by conducting experiments several tasks, in-
The structure of the paper is as follows: after reviewing related work in Section 1.1 we present the
theoretical framework in Section[] Our length generalization result is given in in Section 3] showing how robust correctness implies length generalization. Section Section contains our robust training process inspired by the theory, and presents the results of our experiments. Finally, Section contains concluding remarks and future directions.
1.1 Related Work
Empirical studies have consistently shown that transformers exhibit mixed suc- cess in length generalization (Duan et al., 2023] Jelassi ct al. 2023). While proficient at many sequence tasks, they often fail dramatically when tested on sequences longer than those seen during training, particularly on algorithmic tasks like multi-digit addition (Cho et al. 2024), sorting (Lec ct al. 2024), sequence reversal, parity checking (Zhou et al. 2024b), and ListOps (Nangia and Bow- man 2018). These tasks serve as crucial benchmarks because their underlying algorithms possess inherent recursive structures that models should ideally learn to apply, irrespective of length.
Empirical Work.
There is much work on the role of Positional Encoding (PE). Relative PEs were initially favored over absolute ones (Wang et al.) 2024a), but subsequent works gave conflicting evidence. Some found that removing explicit PFs entirely (NoPE) could yield better length generalization on certain downstream tasks (Wang et al. 2024a), whereas others showed that highly task-specific PE schemes,

---

such as FIRE (Li et al. 2024) or position coupling (assigning identical position IDs to structurally related tokens) (Cho al., 2024 McLeish et al. 2024), helped achieve length generalization on tasks like addition. Our work only uses NoPH, and it remains a fascinating question to see if it can
be combined with more sophisticated PEs.
Data formatting and representation also play critical role. Techniques like reversing the sequence order for addition (Lee al. 2024 Li et al. 2024), giving explicit index hints within input/output sequences, or using intermediate scratchpads or Chain-of-Thought reasoning steps (Merrill and Sabharwal 2024) can dramatically improve length generalization on specific tasks. Unfortunately, scratchpads are not a panacca-eega for the problems we consider, the baseline models (with scratch- pads) do not length-generalize to the extent that our proposed models (also using scratchpads) do. Index hints and position coupling offer considerable power, but require the input to be fed with hints during both training and inference: our goal here is to investigate approaches where the input does not have to be augmented with hints at inference time. (It would be interesting to study if these ideas can be combined.) Other empirical approaches include modifications to the attention mechanism itself (Fan al., 2025 or develop specialized training techniques (Lee et al. 2025).
Another potential approach improving length generalization is use curriculum learning (e.g.. (Pouransari et al. 2024)) and train on progressively more complex instances that are also longer. Our work once again orthogonal, and in principle, can be combined with curriculum learning; we leave this as an interesting future research direction.
Theoretical Work. Theoretical understanding of length generalization is an active area of re- search, with several perspectives emerging. One line of work aims to characterize the class of func- tions on which transformers length-generaizz The RASP-L Conjecture (Zhou et al.. 2024 posits that transformers length-generaliz precisely on tasks solvable by short programs in RASP (a lan- guage that models transformer computation (Weiss et al. (02T]); despite recent work (Huang al.] 2025 Yang and Chiang, 2024). this conjecture remains open. Some other works have studied sim- pler models like linearized attention (Sabbaghi et al. 2024) or average pairwise attention (Ahuja and Mansouri 2024), and show asymptotic length generalization for these classes. Huang et al. (2025) define the limit transformer and show length generalization asymptotically. [Li et al. (2025 con- sider the "vanishing variance" phenomenon where attention output variance decreases with length under certain assumptions. Many other works focus on showing expressivity of transformers, and do not consider issues of learnability. Our work focuses on giving general conditions under which the model exhibits length generalization, and gives quantitative bounds for finite lengths.
The work of Wang et al. (024b) proves that performing gradient descent on transformers success- fully learns sparse token selection (i.e., the task of choosing a sparse set of i out of tokens and taking their average); moreover it exhibits good length generalization, where the error scales as the square of the length increase. Although it relies on a specific training process, this result very much in the same spirit as ours. A concurrent work focusing sparsity is that of Golowich et al. (025),, who also propose a theoretical framework to study generalization via sparsity; they use a notion of predictive position coupling to achieve length generalization. This idea is orthogonal to ours, and would be interesting to see if they can be combined. Finally, formal language theory offers different perspective by investigating the classes of formal languages that transformers can recognize or generate (Strobl et al. 2024).
Previous Version. This paper substantially extends the preprint by Awasthi and Gupta 2023 which investigated task hinting as a mechanism for length generalization. Our theoretical analysis yields new robust training framework based auxiliary tasks component absent in the prior work. Furthermore, this new approach leads to improved empirical performance across a range of tasks. For instance, the increment task length-generalized (i.e., achieved > 90% test accuracy) by most a factor of 1.2 in the previous work, whereas it does so by at least a factor of 1.5 in ours.
2 A Framework for Length Generalization
In this section we introduce formal framework to analyze the conditions under which transformer models -engt--generalize on in Section 3
tasks; we then prove the generalization bounds
3

---

Let E and I' be finite input and output alphabets respectively. A sequence-to seeeenceet---eqeeee -e-euuunce task is defined by target function f Σ- - I where Σ and I denote the set of all possible finite sequences over the respective alphabets. The model, typically transformer, is trained on dataset Pmiin consisting of input-output pairs (x. y) where x E Σ" y f(r),, and the length of the input sequence [x] I n for some maximum training length n.
Length generalization refers the model's ability to perform well on input sequences æ with lengths N > n drawn from a test distribution Diss Formally, let M be a model trained on Duain- Let accuracy) of model M on distribution D. We say that M exhibits length generalization if Err (M, Disa) for N > remains low close to Err(M, Duain) for N = n and degrading gracefully as N increases rather than exhibiting sharp increase in error immediately beyond n. Length generalization contrasts with in-distribution generalization, which concerns performance on unseen examples (ar, y) where x ~ Dirain and hence |a| S n.
Err(M, D) denote the expected error (i.e.,
In this work, we focus on decoder-only transformers. Let .M = (E, D, K, Q, V) be a decoder-only single-layer transformer model with a single attention head and no MLP layer. Here E and D are the encoder and decoder matrices, and K, Q, V are the key, query, and value matrices respectively.
2.1 Modeling Single-Token Outputs
We first consider the problem of computing a function f: Σx + I mapping each input sequence = (I11 "f1,.....a0 .n)) to a single output symbol f(x) T. We assume the existence of special end- of-input symbol 1 & Σ U.. The transformer model M is given the augmented sequence and we use the representation 21 the 1 token after the attention layer to generate the output M(x)= = arg max dz), where dj is the decoder column corresponding to j E l.
I,
For n E Z4, consider distribution Dn on sequences of length from E. Let DD refer to this ensemble of distributions D,,, one for each n E Z+.
Definition 2.1 (Goodness). A model M is (e, 7)-oood at length with respect to distribution Dz, if
P Pa BJEr:fr((m) (71, dz) 2 (#1, (iid)) - -7 S
where dj is the vector mapping to the output symbol E T. Here we call ry the margin at length n.
Observe that if the model is (s, -goood for any value of y 2 0, then the probability of error is at most E; moreover, a large value of y indicates that the model is correct and confident whp.
Definition 2.2 (Sparse Functions). A function f Σ' + I is s-sparse if for each x E Σ', there exists some set A(x) C (I...., xx)} } of s positions such that for any subsequence x' of x containing the symbols at the positions in A(x), we have f(x') = f(x). We call A(x) the anchor set for x.
For example, the minimum and maximum functions are both 1-sparse functions, whereas the suc- cessor function (which, given an unsorted sequence æ. and number in this sequence, outputs the next higher number in x) is 3-sparse. We focus on s-sparse functions for the rest of this section.
Definition 2.3 (Perturbation). For N n, define the (N, pertuurbbaii of D' as the distribution over sequences obtained by first drawing a random sequence x from Dx, keeping the symbols in the anchor set A(x) and keeping each symbol in the N positions in [N] A(x) independently with
i probability N= 2 R. where s = [A(x)). The expected length of the resulting sequence is n + s.
We could have defined perturbations be random sequences of length exactly retailing a uniformly random set of size n from the N - positions; we focus on the i.i.d. case for simplicity of exposition. Note that in the typical case of s < n. the length of the subsequence obtained in the i.i.d. setting is (1 + (n + s) with probability 1 exp(-O(e2n)).
obtained by
Definition 2.4 (Robustness). A model M is (N, n)-rruustt w.r.. D with parameters (E, 7) model is (e, -goo with respect to Dn, and also with respect to the (N, n pertuurrrtioo of DN.
(E,7) if the
While our theoretical model considers the performance on sequences of length x (1+)) s) with high probability, our experiments are done on sequences of length n + s. This choice to consider i..d. sampling in our theoretical results is purely for analytical simplicity.

---

3 Length Generalization Bounds
We first consider the case of single-token outputs. Our main generalization bound here says that if the model is (N, n.) -robuus, then it is also good at length N (with suitably weaker parameters). Formally, let [[V]] denote the operator norm of the value matrix. Then we show:
Theorem 3.1 (Single-Token Generalization). For L, let M = (E, K, Q) be one-layer attention-only transformer: assume that the encoder matrix E and decoder matrix D have unit- length columns. Suppose . is (N)uus with parameters (E,7), where y > 2 (1 1/ac) for N S en. Then M is (O(s N/n),, -)good good at length N.
In particular, the probability of erring at length N degrades linearly in the sequence length. While the theorem does not depend on how the model is trained, or the length of sequences which this training should be done, our experiments are performed when training on length up to n + 8.
To prove Theorem 3.1 it will be convenient to have the following "anti- coneitconcentration" bound. (We use this only for the special case of independent random variables, but we prove for more general negative correlation property that captures sampling without replacement.)
Theorem 3.2 AntiiConncntration Bound). Suppose Zi...., ,,m are identical Bernoulli random
variables with [Z] q and satisfying the negative cylinder property that Pr[mer Zi = 0] s Teer Pr|Z; = 0) for all T S [m]. Let a1,...,,0 .0m be non-negative real numbers and suppose
>; a. 2 1. Define
Σ; 0; Zi (hence [[[] 2 q). Then, we have Pr[X 2 q/4] 2 q/2.
We can now return to the proof of our length generalization result for single-token outputs.
Proof of Theorem 3.7 Consider an input sequence x = (x1 of length N on which the model makes mistake. Recall that we feed the augmented sequence æ I to the model, which returns arg max d,) as the answer. (Here 21 is the embedding of the end-of-input token after passing through the model, and dj is the column of the decoder matrix corresponding to output symbol j E r.) Let j = fx) E be the correct answer on this sequence; since the model makes a mistake on æ, it outputs some other symbol j E J J¹*.
The attention structure of our model ensures that the representation of the 1 token must be
Ei exp((Kez,, Qe1)) Vez. EE exp(( Qe1))
EVENNY W:P;
211
EiE[N] uy
where ex; the encoding of the symbol xy and ey that of the end-of-input symbol 1, and we use v. Vez, and w, exp((Kez, Qel)) 0 to reduce notation. Define a (v,, d,, dzi).
Since each of the e,, d, vectors has unit length, we have the following upper bound on a;:
a (v, d,- d (Vez. dj- -d;t) S 2||
Now, since the model errs on x and outputs jt instead of jf f(x), we have
0 2 (71, d,. dgi) = EiEA(e) W,a; + EiENNAA(a) W,0;;
(2)
ΣiENN] wi
Suppose the model errs at length N with probability 5. We want to show the following claim: conditioned on making an error (say on sequence x.), the probability that the margin less than ? on the (N, pperturratto of'D '' is at least óp/ 2, where p = N= know this probability of small margin must be smaller than e, so 8 nn(2NN) < 6p/2 < E, so 8 < E (2N/n). proving the theorem.
a 2 N 2 } Indeed, by assumption,
Now we prove the claim. Define == terms in the numerator of (2) but outside the anchor set and let P := -NN] x p' thus contains either the indices of the non-negative terms in the numerator of (2) or indices that are in the anchor
: (i E
la;
aa < 0} to be the indices of the negative
set (hence retained in the perturbation with probability 1). We focus on this numerator and regroup
terms to get:
Σ w V iEP
Σ
w/ wss
(3)
'
5

---

Recall the definition of the (N, pertturbbbaiiio of the sequence: it retains all the s tokens that "force" the correct output j" (i.e., samples them with probability 1) and samples each other token independently with probability p. Let Y; be the indicator r.v. for the token with i E [N] being sampled, and let ă denote the (random) shorter subsequence obtained from æ: by restricting to the sampled indices (and y being the r.v. denoting the embedding of the I token on this input). Consider the outcomes of { for i E P", and let us condition on their values being Y; = yi this gives us
(21. 1o-ddy -dj)) = Σiep' Wiayys DiEq'' W;los/Y;
(4)
LiEp' Pr Wiyi + +iieg W;Y,
Now, we consider the sum S == Eigr 1400yi- If S S O, then we are done since it means the model made error length and hence the small margin claim holds vacuously. Therefore, focus on S > O. Now define a, := loll///SS for i E Q'; observe that ΣiEoy ag 2 because of (3). Hence the 04 values and the associated i..dd rv's Y, satisfy the conditions of Theorem [.2 which implies that Pr PrlE Wala;/Y, 2 Sp/4] 2 p/2. Moreover, Eieq w,Y, 2 0 always holds for any indicator r.v'' Y; in Q' Substituting these back into ( we get that conditioned on M making a mistake on æ, and for any conditioning Y; = for the r.V.'s in P", the following upper bound on the margin holds w.p. least p/2:
dy. - d (1 pp/4) EiEP W,Oyye
2|| (1- 4) S 2||1|
(1-1)-
1 O 4c
EiEp wayh
3.1 Extension Multi-Token Outputs
Our theoretical framework also extends naturally to capture the more realistic scenario where the output is a sequence of tokens as opposed to single token. Note that in this case one can view the problem of predicting each output token as an individual single-token output prediction problem with a corresponding auxiliary task. Hence, in the multi-token output setting we have a set of auxiliary tasks related to the main task, one per output position. Based on this we provide the following guarantee. We defer the full details to the Appendix.
be one-layer attention-only transformer; assume that the encoder matrix E and decoder matrix D have unit- length columns. Suppose M is (N, n) m-mmii-ttkke--rbbee with parameters (E,7), where ? 2|| (1 1/sc) N S cn. Then M is (O(e N/)), 0) -good at length N.
K, Q, Theorem 3.3 (Multi-Token Generalization). For 1, let (E,
3.2 A Training Recipe for Length Generalization
Our framework suggests a natural training recipe for inducing length generalization for sparse tasks.
1. Identifying auxiliary tasks: For a sparse function, we produce multi-token outputs, where each token at timestep I obtained by applying some "local" function Se small num- ber of previous tokens. Given the sparsity, there can only be a bounded number of such functions; moreover, for many common tasks, there are just a handful of functions being repeatedly applied, e.g., finding the minimum and successor in sorting, or doing single- digit addition in the increment task, or summing two tokens in the SLiM task, etc. These commonly-used functions for particular problem are natural auxiliary tasks for it.
2. Robust training for auxiliary tasks: Having identified these auxiliary tasks, we train the model to be (N, n)-rouust on auxiliary tasks (in addition to training on the main task). Specifically, we use the standard multi-task training where, during each gradient step, batch of data from either the main task or an auxiliary task is used (in round-robin manner) for backpropagation. Note that all training on n-llngth sequences.
3. Robustness implies length generulization: We can now use our theoretical guarantees, which say that if we are (N, -robuustt with good parameters, then length generalization happens (i.e., low error on N-length sequences).
4 Experiments
In this section we present empirical results evaluating our robust training recipe for length gener- alization. The goal is to assess its effectiveness in improving length generalization on challenging

---

Sorting. We consider sorting sequences of a partieular length m where each element is drawn uniformly at random from task corresponds outputting the sorted sequence. The natural auxiliary task is produce the successor of a given element in the output sequence. Figure [ shows an example of the sorting task and some well-known synthetic tasks and to validate the underlying theoretical assumptions in our framework. We begin by describing the synthetic tasks that we consider, and the corresponding auxiliary tasks. In each case we will provide example illustration of the main task and the cor- responding auxiliary task. When presenting the auxiliary task we will use a red colored token to represent the current token for which the next token output needs to be found. We will use blue colored tokens to highlight the corresponding relevant tokens for the auxiliary task. (More details about the experimental setup and auxiliary tasks appear in Appendix C
4.1 Synthetic Tasks
100}. Given an input sequence æ (#1,..... Zm L) the main
and corresponding auxiliary task. For the sorting task train on sequences of length up to 20 and measure length generalization for higher length sequences.
Figure 1: An example of the sorting main task and an auxiliary task.
Reversal. We consider the task of reversing a given sequence where each element is drawn uni- formly at random from f1.... 100).. Note that for the reversal task, a model has to crucially keep track of the position in the input of the current token that is being output. We consider the following task representation that makes this explicit and lets us design auxiliary tasks that depend on small number of relevant tokens: we explicitly append pos and vol tokens to keep track of the position information; see Figure for example. We train the reversal task on sequences of length up to 20.
pos
p00
pos
pos
pos
val
16
ooe
vell
21
pow
val
21
pise
vol
16
An example of the reversal main task and auxiliary task: to output value, we pay attention to the current position (which is 0), and to the corresponding pos O and val 16 in the input.
Figure 2:
Copy. This is similar the reversal task except that is required copy the input instead of reversing it; see Figure for an example. We use similar representation as with reversal and generate training data for the main task and the auxiliary tasks for sequences up length 20.
16
21
pos
0
val
16
pos
vall
21
pos
0
vel
16
pos
21
Figure 3: An example of the copy main task and an auxiliary task.
Increment. The increment task consists of adding to sequence of digits chosen from 0 to 9. As is standard for length generalization on arithmetic tasks, we consider producing the incremented number in reverse format (lowest significant digit first). Furthermore, along with the sum and carry tokens that us define sparse auxiliary tasks; see Figure for an example. We train the increment task for sequences of length up to 20.
again include pos token

---

pow po0
pde
sum
carry
sum
poe
val
pos
val
po6
sum
carry
200
sum
Figure 4: An example of the increment main task and an auxiliary task. Since we have carry of 0 at position O, and a value of at this position in the input, the output is 1.
Parity. In this task given sequence of n bits we are required to compute n parity bits where the ith output bit must equal the parity the first bits of the input. Similar the reversal and copy tasks, we also augment the training sequence with pos and val tokens to keep track of the current position information. Furthermore we also use the xor token to keep track the parity of the prefix computed so far; see Figure 5 for an example. We generate training data for the main task and the auxiliary tasks for sequences up to length
val
pos
vol
x0r
pos
0
Figure An example of the parity main task and an auxiliary task.
SLiM (Small Language Modeling). This synthetic task is inspired bylanguage modeling, and in particular, the notion of contextual lookup/recall. In a typical language modeling setting, each output token depends on a small number of prior tokens. However, these tokens may not always be the most recent ones, but may appear at arbitrary prior positions the input sequence. Furthermore, given the correct tokens to attend to, the rule for producing the next output token is more involved than a simple copy operation. We consider an idealized scenario to capture this behavior. the vocabulary {1, (1...., 100} that is randomly partitioned into two classes A and B. The sequence is seeded with set of four tokens, two chosen randomly from class and the other two randomly from class B. Given current output token from class A (class B, resp.), the next token is produced by finding the previous class A token (class B, resp.) in the sequence, and computing the sum of the example. Here, given input sequence, followed by se
consider
two values (modulo 100). See Figureoffor
a "query" token the auxiliary task is to look up the last input token r of the same type as and to output q + (modulo 100). We train the SLiM task on sequences of length up to 20.
78
37
Figure 6: An example of the SLiM main task and an auxiliary task. this example numbers from 1 50 belong class A and the numbers from 51 to 100 belong to class B. The token 33 pays attention to the previous token of the same class, i.e., to 4, and sums to output 37.
4.2 Training Details
A standard decoder-only transformer architecture used (unless otherwise specified), with details on the number of layers, heads, embedding dimensions, and activation functions provided in the Appendix. We trained each model without positional embeddings. The proposed multi-task training method used the same backbone architecture with an additional output head for the auxiliary tasks. For all the tasks, we train on sequences of length up 20. In order to generate the training data for the main task, we sample length n from [4, 20) at random and chose a random sequence of length n. (For the increment task we consider slightly non-uniform sampling where 20% the training data consists of sequences that end with 9.) Crucially, the training data for the auxiliary tasks is generated from (N, pertuurbations i.e., first drawing an input sequence larger length and then random subsampling the irrelevant tokens down to length (As mentioned above, we subsample

---

these tokens without replacement to get fixed length, instead of sampling each one independently.) The amount of training data used for the two upproaches is kept equal. so that the baseline model uses the same amount of data as the proposed model. Moreover, both the baseline and proposed models use scratchpads; the baseline models do not length-generalize even with scratchpads.
Models were trained using the AdamW optimizer with a cosine learning rate schedule and a batch size of 1024 (see Appendix). The baseline model was trained solely on the primary task; our model was trained with auxiliary tasks using multi-task training. The primary metric was the "exact match" accuracy, i.e., the fraction of test sequences for which the model's output matched the ground truth.
4.3 Results
Figure 7 summarizes the results across the synthetic tasks that we consider. The baseline models, while achieving high accuracy within the training length distribution (length at most nt), exhibit the typical sharp drop in performance for lengths beyond n. In contrast, the models trained with the auxiliary tasks maintain high accuracy far beyond n. For instance, on the sorting task, training lengths up to 20, the baseline model's 's accuracy drops below 1% for length 100, whereas our method achieves over 90% accuracy. Similarly strong gains observed for the other tasks as well.
100
100
60
60-
Jun Socal
60
80
50
100
18
22
26
Eequeaie kogtt
Saguiet
Copy
Incronent
300
100
Koot
60
03
AcCRust
20
20
1.
13
Segarsio imggh
Parity
SUM
300-
00
00
60-
00
20
Jeciaa Moecl
20
4,5
Figure 7: A comparison of the length generalization performance of the baseline training method vs. the proposed auxiliary task based method on various synthetic tasks. The z-axis represents the sequence length and the y-axis is the exact match accuracy. The shaded region represents the sequence lengths used during training and the non-shaded region represents the sequence lengths that the model does not see during training time.
We make few observations about the results. Note that the performance on sorting is good up to length 60 even for the baseline model. This consistent with similar observations made in prior works (Zhou et al. 2024a). The performance of the baseline model starts to drop sharply beyond

---

60 whereas the model trained with auxiliary tasks maintains high accuracy (more then 95%) even up to length 100. For tasks that involved carefully maintaining position information and indexing into a specific position in the input the performance of the baseline model drops sharply whereas the auxiliary task trained model has much greater robustness. Finally, the SL.M task, which involves learning simultaneously rule for attending to specific tokens and a rule for using the attended tokens is significantly harder for the baseline model and its performance drops sharply beyond the training length. We again observe that by incorporating the auxiliary tasks in the training procedure the model is robust for up to two times the training length.
These results strongly support the hypothesis that explicitly training on the auxiliary tasks suggested by our framework, and using perturbed distributions, effectively guides the model learn the inter- nal mechanisms necessary for length generalization.
4.4 Scaling Experiments
Figure 8: Scaling plots for the increment task. The left figure shows the actual accuracy at various lengths and the accuracy predicted by (5). for a model trained for length up to 20. The right figure shows the same plot for model trained for length up to 25.
In this section we further aim to validate the practical relevance of our framework. Note that our key conceptual contribution is to establish a connection between the performance of a model on auxiliary tasks with its ability to mantth-generalize Our main claim (Theorem 3.1) needs one to achieve a margin of 2|[ 1.) over the set of auxiliary tasks. However, this bound may be too pessimistic in real settings. In particular, analysis of a typical average case scenario reveals the following rule-of thumbb for length generalization:
(1
1/c) n
Pr [error at length N = c c-N]= ~ Pr aux. task margin S a
(5)
where a is problem dependent constant.
We then perform a series of scaling experiments where we train models for our proposed synthetic tasks at various training lengths and compare how well does proposed predicted length general- ization performance from (5) match the empirically observed accuracy. Figure Jshows one instance of this for the increment task trained at two different lengths (sce Appendix for other scaling exper- iments). We see that (5) is consistently good predictor of the length generalization performance across various values of N. These experiments help validate the connection between the theoretical assumptions made in our framework and achievable performance in practical settings.
shows
Conclusion
Our work aims to bridge the gap between thcoretical understanding and practical improvements in length generalization. By deriving auxiliary tasks directly from novel framework, we offer prin- cipled alternative to current mostly heuristic approaches. Limitations of our work include the need to further validate the assumptions and extend the framework to more complex tasks, particularly in natural language processing where defining clear algorithmic structures and relevant auxiliary tasks can be challenging. The specific theoretical models analyzed (e.g., single-layer transformers) are simplifications, and extending our results to deeper, more complex architectures is an impor- tant future step. Additional interesting directions include relaxing the theoretical assumptions and automatically discovering relevant auxiliary tasks.
10

---

References
Ahuja, K. and Mansouri, A. (2024). On provable length and compositional generalization. arXiv:2402.048//
Anil, C., Wu, Y., Andreassen, A., Lewkowyzz, A., Misra, V., Ramasesh, V., Slone, A., Gur-Ari, G., Dyer, E., and Neyshabur, B. (2022). Exploring length generalization in large language models. In
NeurIPS, pages 38546-38556.
Awasthi, P. and Gupta, A. (2023). Improving length-generalization in transformers via task hinting. arXiv 2310.00726.
Cho. H., Cha, J., Awasthi, P., Bhojanapalli, S., Gupta, A., and Yun, C. (2024). Position coupling: Leveraging task structure for improved length generalization of transformers. In First Workshop on Long-Context Foundation Models@ ICML.
Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. (2019). BERT: Pre-training of deep bidirec- tional transformers for language understanding. In NAACL-HLT (1), pages 4171-4186.
Duan, S., Shi, Y., and Xu, W. (2023). From interpolation to extrapolation: Complete length gener- alization for arithmetic transformers. arXiv:2310.11984,
Dubhashi, D. P. and Panconesi, A. (2009). Concentration of Measure for the Analysis of Randomized Algorithms. Cambridge University Press.
Fan, Y., Du, Y., Ramchandran, K., and Lee, (2025). Looped transformers for length generaliza- tion. In ICLR.
Golowich, N., Jelassi, S., Brandfonbrener, D., Kakade, S. M., and Malach, E. (2025). The role of sparsity for length generalization in transformers. In ICML.
Huang, X., Yang, A., Bhattamishra, S., Sarrof, Y., Krebs, A., Zhou, H., Nakkiran, P., and Hahn, M. (2025). formal framework for understanding length generalization in transformers. In ICLR. Jelassi, S., d' Ascoli, S., Domingo-Enrich, C., Wu, Y., Li, Y., and Charton, (2023). Length gener- alization in arithmetic transformers. arXiv:2306.1544
Kazemnejad, A., Padhi, I., Natesan Ramamurthy, K., Das, P., and Reddy, S. (2023). The impact of positional encoding on length generalization in transformers. NeurIPS.
Lee, N., Cai, Z. Schwarzschild, A., Lee, K., and Papailiopoulos, D. (2025), Self-improving trans- formers overcome hard and length generalization challenges. In ICML.
Lee, N., Sreenivasan, K., Lee, J. D., Lee, K., and Papailiopoulos, (2024). Teaching arithmetic to small transformers. In ICLR 2024.
Li R., Boduljak, G., and Zhou, J. (2025), On vanishing variance in transformer length generaliza- tion. 022807.
Li, S., You, C., Guruganesh, G., Ainslie, J., Ontanon, S., Zaheer, M., Sanghai, S., Yang, Y., Kumar, S., and Bhojanapalli, S. (2024). Functional interpolation for relative positions improves long context transformers. In ICLR.
McLeish, S., Bansal, A., Stein, A., Jain, N., Kirchenbauer, J., Bartoldson, B. R., Kailkhura, B., Bhatele, A., Geiping. J., Schwarzschild, A., and Goldstein, T. (2024). Transformers can do arith- metic with the right embeddings. In NeurIPS.
Merrill, W. and Sabharwal, A. (2024). The expressive power of transformers with chain of thought. In ICLR.
Nangia, N. and Bowman, S. R. (2018). ListOps: diagnostic dataset for latent tree learning. In NAACL 2018.
Nye, M., Andreassen, A. J., Gur-Ari, G., Michalewski, H., Austin, J., Bieber, D., Dohan, D.,
Lewkowyez, A., Bosma, M., Luan, D., Sutton, C., and Odena, A. (2021). Show your work: Scratchpads for intermediate computation with language models. arXiv:21 /2.00////
Panconesi, A. and Srinivasan, A. (1997). Randomized distributed edge coloring via an extension of the chernoff chermoff-hoeffding bounds. SIAM Journal on Computing, 26(2):350-368.
Pouransari, H., Li, C., Chang, J. R., Vasu, P. K. A., Koc, C., Shankar, V., and Tuzel, O. (2024). Dataset decomposition: Faster L.M training with variable sequence length curriculum. In NeurIPS 2024.
11

---

Press, O., Smith, N. A., and Lewis, M. (2022), Train short, test long: Attention with linear biases enables input length extrapolation. ICLR.
Sabbaghi, M., Pappas. G., Hassani, H., and Goel, S. (2024). Explicitly encoding structural symmetry is key to length generalization in arithmetic tasks. arXiv:2406.0 01895.
Saxton, D., Grefenstette, E., Hill, F., and Kohli, (2019). Analysing mathematical reasoning abili- ties of neural models. In ICLR.
Strobl, L., Merrill, W., Weiss, G., Chiang, D., and Angluin, D. (2024). What formal languages can transformers express? A survey. TACL, 12:543-561.
Tay, Y., Dehghani, M., Abnar, S., Shen, Y., Bahri, D., Pham, P., Rao, J., Yang, L., Ruder, S., and Metzler, D. (2021). Long range arena: A benchmark for efficient transformers. In ICLR.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polo- sukhin, I. (2017). Attention is all you need. In NIPS, pages
Wang, J., Ji, T., Wu, Y., Yan, H., Gui, T., Zhang, Q., Huang, X., and Wang, X. (2024a). Length generalization of causal transformers without position encoding. In ACL (Findings).
Wang, Z., Wei, S., Hsu, D., and Lee, J. D. (2024b). Transformers provably learn sparse token selection while fully-connected nets cannot. In ICML.
Weiss, G., Goldberg, Y., and Yahav, E. (2021). Thinking like transformers. In ICML.
Yang, and Chiang, (2024). Counting like transformers: Compiling temporal counting logic into softmax transformers. In COLM.
Zhao, L., Feng, X., Feng, X., Zhong, W., Xu, D., Yang, Q. Liu, H., Qin, B., and Liu, T. (2024).
Length extrapolation of transformers: A survey from the perspective of positional encoding. In EMNLP (Findings).
Zhou, H., Bradley, A., Littwin, E., Razin, N., Saremi, O., Susskind, J., Bengio, S., and Nakkiran, P. (2024a). What algorithms can transformers learn? A study in length generalization. In ICLR.
Zhou, Y., Alon, U.. Chen, X., Wang, X., Agarwal, R., and Zhou, D. (2024b). Transformers can achieve length generalization but not robustly. In Workshop on Understanding Foundation Mod- els@ ICLR.
12

---

NeurIPS Paper Checklist
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
See abstract and the introduction.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
See the conclusion section.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The au- thors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the ap- proach. For example, a facial recognition algorithm may perform poorly when image resolution is low images taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to ad- dress problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Ys]]
13

---

Justification: See section 3.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theo- rems.
The proof' can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be comple- mented by formal proof's provided in appendix or supplemental material. Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclu- sions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yss]
Justification: See Section 4 and the appendix.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, a No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary either make it possible for others replicate the model with the same dataset, or provide access to the model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all sub- missions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or a way to re- produce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case au- thors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
14

---

Answer: [No]
Justification: The code needs to through an internal approval process for release. We will submit post paper acceptance.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips..... cc/
Sowulidde//Coeeuuissssss/s/ss Policy) for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https:
//nipscc./ wtt.cc/ppubb. ppwplic/guides..... /ttee//guiddee/eeeubbbmmt (Cooeeobmiissseeeee ssiiisssiicycyy for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
See the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropri- ate information about the statistical significance of the experiments?
Answer:
[Yes]
Justification:
See the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (c.g., Normally distributed errors).
15

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should prefer- ably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
See the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g.. preliminary or failed experiments that didn't t make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification:
We have reviewed the code of ehtics and we follow it.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there is a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: There are no immediate societal impacts of the work.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no socictal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact spe- cific groups), privacy considerations, and security considerations.
16

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitiga- tion strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Question:
Answer: [NA]
Justification:
We don't 't deal with any sensitive or large scale data that has potential of
misuse.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by re- quiring that users adhere usage guidelines restrictions access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do require this, but encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: the appendix.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from particular source (e.g., website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms of use the pack- age should be provided. For popular datasets, batasets..paperswtt wwwpppppppwaia.ee wwwrswittthh.....c com /atatttt has
curated licenses for some datasets. Their licensing guide can help determine the li- cense of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
17

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documenta- tion provided alongside the assets?
Answer: [NA]
Justification:
No assets added.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmoel part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the pa- per include the full text instructions given to participants and sereenshots, applicable, well as details about compensation (if any)?
Answer: [NA]
Justification: No human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contri- bution of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, cura- tion, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: no human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equiva- lent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
18

---

Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: LLM was used to polish the writing and to generate draft ideas for the intro- duction.
Guidelines:
that the method development in this research does involve LLMs as any important, original, or non-standard components.
The answer NA
Please refer to our LLM policy http:////earrppssccc/ooeeeeeeeee for what should or should be described.
/LMM)
19

---

A Extension to Multi-Token Prediction
In this section we extend our previous analysis to the broader case of functions that map an input sequence to an output sequence, rather than a single token. In particular, we consider the problem of computing a function f M + I* mapping each input sequence æ = (x)...... En) to the output sequence ã ( Tm.. where m := m(n) is the length of the output depending before we append the 1. token to the input æ.
As
Note that as opposed to the single token case where we only have one output token conditioned on the embedding of the token, in the multi-token case we have m 1 output tokens. Hence, for any function f we can associate a sequence of single-token output functions (fi..... fm) where fi maps a to 31 and fi maps o 1. 31 o... - 8;
Analogously we will extend the definitions from Section 2 to the multi-token case. We will use 21 to denote the (last layer) embedding of the 1 token and 2; to denote the embedding of the token ø;-
Definition A.1 (Multi-Token Goodness). A model M is (s, -goood at length with respect to distribution D,, if
P PEND, AErrAAa (211 d,) 2 (21, dz,) - - s P Pa-Da 3i 2 2, 3jer.zzz (2--1. d,) 2 (2--1, d2,) -7] E,
where dj is the vector mapping to the output symbol E Here we call yy the margin at length n. Definition A.2 (Sparse Multi-Token Functions). function f : + r (fr...., ..m) s- sparse for each i {1....,m} the function fo is s-sparse.
E
Definition A.3 (Perturbation Set). For N > nI define the (N, n)-perturbaatin set of Dx as the product distribution 8 x... Sm(N) where each 8; an (N + -.)tttpreuuratioon of fi.
We are now ready to define what it means for a transformer model for multi-token prediction to be robust.
Definition A.4 (Multi-Token Robustness). A model M is mmuuii--iii-- w.r.t. D with parameters (8,7) if the model is (e, -mmmliii-token good with respect to D,, and also with respect to the (N, n) pperturbatio set of D
With the above definitions we have the following theorem.
Theorem A.5 (Multi-Token Generalization). For e >> 1, M = (E,,,K, be a one- layer attention-only transformer; assume that the encoder matrix E and decoder matrix D have unit-length columns. Suppose M is (N, n. min-mmlmii-ttie multi-token--rontt 2|[ - (1 - 1/4c) for N S cn. Then M is (O (N+m(NN)/)) -ood at length N.
with parameters (E,7), where >
Proof. The proof very closely mirrors the proof of the single token case. Consider an input sequence (I .N) of length N which the model makes mistake. Let p be the first output position where a mistake is made. In other words the model makes a mistake at fp. Recall that we feed the augmented sequence æo porio,,,...... ... Ap-1 to the model, which returns arg max,er (2p-1. 1, d,) as the answer. (Here 2p-1 is the embedding of up-1 after passing through the model, and d, is the column of the decoder matrix corresponding correct answer on this sequence; since the model makes a mistake on æ, it outputs some other symbol jt r, jt + jr Let N' = N + p < N + m(N) be the length of the sequence that is input to the model.
output symbol j Т.) Let j" S E I be the
The attention structure of our model ensures that the representation of ūp-1 token must be
EN exp((Kexs Qep-1)) Vex: Eie(Ny w.
2p-1 M exp((Kez, Qeg-1))
EiELNy W;
where en; is the encoding of the symbol x; and 6p-1 that of T.-1, and we use v; = Ven, and w; = eex((Ke,, Qep-1)) 1)) 2 to reduce notation. Define a; := (D,; dy. - dgt). ). Since each of the
e, dj vectors has unit length, we have the following upper bound a;;
a (vi dj. d = (Vez. dj. d < ||[
(6)
20

---

Now, since the model emrs on x and outputs jt instead of j" = f(z)), we have
2 (2p-1, -d1) EiEn(m) wiak; + Li A(a) waa;
EiE[Ny wi
where A(x) is the sparse set for the function fp. Now, suppose the model ems at length N' with probability 8. We want show the following claim: conditioned on making an error (say on se- quence at), the probability that the margin is less than on the (N', pnrrpuerurrbation of D is at least 6p/2, where p = N 2 No 2 1. Indeed, by assumption, we know this probability of small margin must be smaller than E, so ô n/(2N') I 6p/2 S E, so d S E. (2N' proving the theorem.
a < 03 to be the indices of the negative terms in the numerator of ( but outside the anchor set and let P == [N'] Q'. P' thus contains either the indices of the non-negative terms the numerator of (7) or indices that are in the anchor set (hence retained in the perturbation with probability 1). We focus on this numerator and regroup terms to get:
Now we prove the claim. Define Q' = fi
Σ WaO; S EE
Σ w. wjlass
8
'EQ'
Recall the definition of the (N', n) -)ppeuuurbaaiion of the sequence: it retains all the s tokens that "force" the correct output 5" (i.e., samples them with probability 1) and samples each other token independently with probability p. Let Y, be the indicator r.v. for the token o; with E [N'] being sampled, and let è denote the (random) shorter subsequence obtained from æ by restricting to the sampled indices (and ēp-1 being the r.v. denoting the embedding of up-1 token on this input). Consider the outcomes of {Y,) for i P', and let us condition on their values being Y; yı this gives us
dj. -d.) ) Eiep W,aKH - EiEgr wylas/Yi
Ever wilh + EiEg' W:Yy
Now, we consider the sum S := -;=Ee E W;004/s. I < 0, then we done since it means the model made an error at length n and hence the small margin claim holds vacuously. Therefore, we focus on S > O. Now define a; == wy lalll//SSS for i E Q'; observe that Eieg a; 2 because of (8). Hence the a, values and the associated i.i.d. random variables Y, satisfy the conditions of Theorem 3.2 which implies that Pr Pr|Exeq' W; W; 2 Sp/4] 2 p/2. Moreover, Exeq' w;Y, 2 0 always holds for any indicator r.v's Y, in Q'. Substituting these back into (.. we get that conditioned on M making a mistake on æ. and for any conditioning Y, y, for the r.v'' in P", the following upper bound on the margin holds w.p. at least p/2:
(1 - -//4) ΣiEr WoOKy 2||VI| (1- 7) s 2||1
(1- (1- 1 o 4e
(šp-1, dz. dri)
EiEp wilk
B Proof of Theorem 3.2
Proof of Theorem 3.2 By scaling the a;'s 's down if needed, let us assume that Σ, az 1, so that
E[X] = q. There are two cases:
1. Suppose there exists i' such that ar 2 9/4 Then
Pr[X 2 4/4 Z Pr|ay Zv 2 q/4] Pr(Zz 1 9/2.
2. Else, each a, < 4/4. Then define df = # X = aiZi, and X' = Ex Now we have a: [0, 1] and also X [0, but E X 4. We want to bound show that Pr[X'' 2 1] is large. Indeed, we can use concentration bound for negative cylinder random variables (Dubhashi and Panconesi 2009 Panconesi and Srinivasan 1997) to show that the probability of a "bad" event is:
Pr[X' 1 = Pr X' S E[X' //4
E
(1 - (1-1/4)
- (4-1))/(2-) s $1/e.
< exp
2
21

---

Hence,
Pr|X 2 R[x]/4 = Pr(x || 1/e 9/2
for any E [0, 1]. as claimed.
Ths completes the proof of Theorem 3.2
C Experiment Details
All our experiments are conducted using the Jax Flaxformer codebase and involve training decoder only transformer models from scratch. For training the baseline models, we partition the training dataset for the main task into batches specified size (see details below) and use the AdamW optimizer. We run our models for prespecified number of gradient steps. To set the learning rate, we start with a learning rate zero and linearly ramp up to prespecified base learning rate over 10 epochs. Following that we anneal the learning rate to zero using the standard cosine learning rate scheduler (with one cycle); see Table ! We train all our models for 200, 000 gradient steps (the in-distribution performance saturates much before that for all the tasks). Besides the four main tasks mentioned in main body, we also experiment with two additional tasks: the task of copying a sequence and the task of computing the parity of given binary sequence.
Parameter
Value
Embedding size (d)
1024
103 (for sorting, SLiM)
Vocabulary size (4)
200 (for reversal, increment, copy and parity)
Position embedding type
None
# Attention heads (h)
16
MLP inner dimensionality (d')
2048
Sequence length
512
Base learning rate
le-5
Optimizer
Adam
LR warmup
Linear for 10 epochs
LR decay schedule
Cosine, cycle with default parameters
Dropout
None
Activation
GELU
(Tor sorting)
Depth
(for reversal, copy, increment, parity, and SLiM)
Table I: Hyperparameters for the experiments.
D Dataset Details
Below we describe our example generation process for each of the synthetic tasks, for both main and auxiliary tasks. For the main task, generate 100M examples, except for SLiM, where we generate 1M examples. For all synthetic tasks, we generate 100M examples for the auxiliary task.
D.1 Sorting
For the main task first pick length value uniformly at random in the range [4, 20). Given l, with probability 0.8 we generate a length sequence by picking each of the numbers i.i.d from the vocabulary V {1...., 100}. With probability 0.2 generate the sequence from shortened vocabulary V C V. random subset of size f/2. This ensures that we train the models to do well inputs with repetitions.
To generate a training data point for the auxiliary task, we simply generate a sequence of length 100 at random following the same process as above and then pick output position i uniformly at random. We fix the relevant tokens needed to predict the token following position and subsample the rest with probability n/N.
22

---

D.2 Reversal
For the main task we first pick a length value uniformly at random in the range [4, 20). Given L, we generate a length sequence by picking the é numbers i.dd from the vocabulary {0, 10.1..... 100).
To generate training data point for the auxiliary task we simply generate sequence of length N = 60 at random following the same process as above and then pick an output position uniformly random. We fix the relevant tokens needed predict the token following position i and subsample with probability n n/N. We also augment the training sequence with pos and val tokens to keep track of the current position information. In order for the models to be able to handle lengths larger than the ones seen during training, we generate pos values by sampling a position uniformly at random in the range [0, 80] and incrementing by one for each new digit in the sequence.
D.3 Copy
For the main task we first pick a length value uniformly at random in the range [4, 20). Given I, we generate length i sequence by picking the l numbers i.idd from the vocabulary {0, ...., 100).
To generate a training data point for the auxiliary task, we follow the same process as in reversal.
D.4 Increment
For the main task we first pick a length value uniformly at random in the range [4, 20). Given f, with probability 0.8 we generate length sequence by picking the number at the most significant position uniformly at random from {L....,
.......) 9} and each of the remaining e numbers i.idd from the vocabulary V = (0, ... 9). With probability 0.2 we generate a random sequence that ends with a suffix of 9's where the suffix length is chosen to be uniform between [1, t/2). This ensures that the models learn the carry operation faithfully.
To generate a training data point for the auxiliary task, we simply generate a sequence of length N = 60 at random following the same process as above and then pick output position uniformly at random. We fix the relevant tokens needed to predict the token following position and subsample the rest with probability n /N. We also augment the training sequence with pos, val, sum, and carry tokens to keep track of the current position information. In order for the models to be able to handle lengths larger than the ones seen during training, generate pos values by sampling position at random in the range [0, 80) and incrementing it by one for each new digit in the sequence. We also follow the standard practice of producing the output in reverse order, i..,, least significant digit first.
D.5 Parity
For the main task we first pick length value l uniformly at random in the range [4, 20). Given é, we generate a length l sequence by picking the l numbers i.i.d from the vocabulary {0, 1). The output sequence corresponds to the sequential parity of every prefix of the input.
To generate a training data point for the auxiliary task, we simply generate a sequence of length N 60 at random following the same process as above and then pick an output position uniformly at random. We fix the relevant tokens needed to predict the token following position i and subsample with probability n n/N. In order for the models to be able to handle lengths larger than the ones seen during training, we generate pos values by sampling a position uniformly at random in the range [0, 80) and incrementing by one for each new digit in the sequence.
D.6 SLiM
For the main task we first pick a length l uniformly at random in [4, 20). We pick a seed sequence of length 4 by picking two random tokens from class A and two from class B, Given the seed sequence, we generate a length-f output sequence by applying the latent (deterministic) rule.
To generate a training data point for the auxiliary task we simply generate sequence of length N 60 at random following the same process as above and then pick an output position uniformly at random. We fix the relevant tokens needed to predict the token following position i and subsample with probability n/ n/N. So if the token at position i of type A, the auxiliary task to find the previous token of type A before position i, and output the sum of those two modulo 100.
23


Reviews:
[{"review_id": "swNqzWcoHe", "reviewer": "Reviewer_itBv", "summary": "The paper introduces a new methodology for training transformers for length generalization, based on identifying auxiliary tasks which should improve the length generalization on the main task when trained on jointly. The proposed method improves the length generalization  performance on s-sparse tasks such as sorting, reversal, increment, SLiM, copy and parity as shown by a significant margin. This method is inspired by the theoretical analysis provided by the authors, which shows that under certain assumptions, the generalization error decays linearly in N.", "questions": "1) Could the authors make the connection between the theoretical analysis and the proposed training procedure clearer? \n\nMy current understanding is that for s-sparse functions - which are functions that are basically determined by s positions in the input sequence - you can introduce a notion of $(N, n)$ perturbations. This is the distribution over sequences of length $n+s$, obtained by sampling sequences of length $N$, fixing the s positions of the s-sparse function and then sampling the rest. And now you have a notion of robustness connecting the probability of mistakes on $D_n$ and the $(N, n)$ perturbations of $D_N$. \n\nHow does this analysis imply directly the training methodology you suggest?\n\n2) Is there any connection between the proposed method and curriculum learning, for example the lenght curriculum proposed in [1]?\n\n3) In Figure 5, is the amount of data fixed for the 2 tasks (main and auxiliary)? Could the improvement come from the fact that the AuxTask model has been trained on more unique data?\n\n4) The choice of the auxiliary tasks seems too task-dependent: is there a general fix? How to adapt this strategy to general training on a mixture of tasks? How much is the gain sensitive to the choice of the auxiliary task?\n\n5) Could the authors explain what is the auxiliary task for SLiM? It is not clear from Figure 4.\n\n[1] Pouransari, Hadi, et al. \"Dataset decomposition: Faster llm training with variable sequence length curriculum.\" Advances in Neural Information Processing Systems 37 (2024): 36121-36147.", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n- The theoretical analysis provided by the authors is sound and is able to predict the empirically observed results in Figure 6\n- Across all synthetic tasks that the authors evaluate on, they obtain significant gains in performance as shown in Figure 5\n\nWeaknesses:\n- The main weakness of the paper is the limited empirical evaluation. The method proposed by the authors would be much stronger if it also worked in experiments on addition and multiplication, as well as experiments in more realistic architectures such as deeper networks (with MLPs) with an ablation over positional embeddings\n- It is not clear to me how this method could be extended to other tasks, such as ListOps, or addition. Could the authors give examples of what would be an auxiliary task in these cases?", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their encouragement and comments. We are glad that you liked the general framework, the fact that the experiments achieve the significant gains as suggested  by the theoretical model. Please find responses to your specific questions and comments below.  \n\n**Clarification of the Setup**: Your statement regarding the problem setup is correct: we focus on sparse functions, consider $(N,n)$-perturbations (which are sequences of length $n$ produced by sub-sampling $N$-length sequences), and train for both the auxiliary and main tasks on these sequences. To complete the picture, the steps are as follows:\n\n   a. **Identifying Auxiliary Tasks**: For a sparse function, we produce multi-token outputs, where each token at timestep $t$ is obtained by applying some “local” function $f_t$ to a small number of previous tokens. Given the sparsity, there can only be a _bounded_ number of such functions; moreover, for many common tasks, there are just a handful of functions being repeatedly applied, e.g., finding the minimum and successor in sorting, or doing single-digit addition in the increment task, or summing two tokens in the SLiM task, etc. These commonly-used functions for a particular problem are natural auxiliary tasks for that problem. \n\n  \n   b. **Robust Training for Auxiliary Tasks** : Having identified these auxiliary tasks, we train the model to be $(N,n)$-robust on auxiliary tasks (in addition to training on the main task). Recall that this training is on $n$-length sequences.  \n\n\n   c. **Robustness implies Length Generalization**: We can now use our theoretical guarantees, which say that if we are $(N,n)$-robust with good parameters, then length generalization happens (i.e., low error on $N$-length sequences).  \n\n\n **Relationship to other Approaches:** We feel the ideas in our approach and in curriculum training are orthogonal. The latter focuses on training on progressively more complex instances (e.g., the approach of [1] buckets the sequences by their lengths, and then trains on some carefully chosen mixture of these sequences). Our approach is to identify useful building blocks (i.e., the auxiliary tasks) that can improve the training performance, and then to perform robust training for them (in addition to the main task). We feel the two approaches could be fruitfully combined in the future.  We will add a remark to this effect in the revision.  \n\n\n**Data Usage Identical:** Yes, the amount of data is the same for the two approaches. So the baseline model is in total using the same amount of data as the proposed model (with data for main task + data for auxiliary task. On the point of unique data, it is an interesting comment: in fact, that is to some extent what our theory proposes—if we show more robust trajectories (which will be unique) to the model then it will length-generalize.  \n\n\n**Choice of Auxiliary Functions:** This is a great question. While the choice of auxiliary tasks currently *does* depend on the specific problem, we believe this is natural for problems with precise answers. Indeed, such problems require particular tasks to be done correctly and repeatedly, and hence it makes sense to robustly train for them. If we want to consider multiple problems simultaneously, we can definitely go beyond the naive approach of training on the auxiliary tasks for each of these problems—we can use a Boosting/Hedge-based approach, where we adaptively train the model on a small set of auxiliary tasks on which the model has the worst performance.  \n\nIt is a very interesting question about how to isolate these auxiliary tasks automatically: one way would be to have a “library” of “basis” functions for the “local” tasks, and then use these as subroutines in a general training setup where they can hopefully enable length generalization for a variety of downstream applications. Thank you for raising this point; we will add it to the discussion part of the revision.\n\n**Clarification of Auxiliary Task for SLiM:** For SLiM the auxiliary task is as follows: given a valid sequence $S$ of tokens (produced as in the original problem), and a specific “query” token $q$, the auxiliary task is to lookup the last token $r$ in the input sequence $S$ of the same type as $q$, and to output $(q+r)$ modulo $100$. Note that this **contextual lookup/recall** is precisely the task the SLiM has to repeatedly perform using the last output token as its query, which is why it was chosen as an auxiliary task. Moreover, since contextual lookup/recall is an essential part of language, we think of this as a first yet concrete step towards understanding length generalization in broader contexts."}, {"review_id": "dnj7tmjNyp", "reviewer": "Reviewer_hBbS", "summary": "This work addresses the problem of length generalization in small transformers from both a theoretical and empirical perspective. It argues that training on short sequences often leads models to rely on length-specific heuristics, failing to generalize to longer inputs. To mitigate this, the authors propose joint training on a main algorithmic task (e.g., sorting) and auxiliary tasks that isolate sparse, local subcomputations (e.g., predicting the successor of an element in sorted order). Their theoretical contribution establishes generalization bounds that grow linearly with input length, under a sparsity condition where each output token depends on a limited number of input positions. Experiments on structured synthetic tasks (sorting, reversal, Increment and small language modeling) show that this method enables generalization to inputs up to 10× longer than those seen during training, while standard training fails catastrophically. However, its application is limited to tasks with known structure, and real-world applicability remains untested.", "questions": "# Questions\n\n1. For tasks like language modeling or open-domain QA, where the latent task structure is not known, how might one design or discover auxiliary tasks automatically?\n2. Is there any experimental comparison to curriculum learning or random length scaling, which might also improve extrapolation by varying sequence lengths during training?", "limitations": "Yes", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 1, "strengths_and_weaknesses": "# Strengths\n1. Theoretical Grounding: The paper offers a rigorous generalization bound for sparse sequence-to-sequence functions showing test error grows linearly as a function of length.\n2. On a few synthetic tasks, jointly training a main task with an auxiliary task (that contains subcomputation) in a multitasking manner  enables length generalization where baselines collapse. \n\n# Weaknesses\n\n1. Significant overlap with Prior Work: The idea of improving length generalization through auxiliary task multitasking has clear precedent in \"Improving Length-Generalization in Transformers via Task Hinting\" (Chughtai et al., 2023), which uses similar tasks (sorting + successor) and shows that multitask training of a main algirthimc task with a relevant auxiliary task, say successor finding, significantly improves length generalization.                                                                                                                                                                                                                             This major overlap is not acknowledged (as far as reviewer can see) and limits the originality of the contribution.\n2. Synthetic-Only Evaluation: All experiments are on clean, synthetic tasks with known structure. The paper does not attempt even lightweight transfer to real-world tasks or semi-structured inputs, leaving open questions about it practical relevance.\n3. Manual Auxiliary Design: The auxiliary tasks require prior knowledge of the task structure (e.g., that sorting can be decomposed into comparisons or successors). It's unclear how to design such tasks in domains where the compositional structure is not obvious.\n4. Missing comparison to alternative methods such as curriculum learning, or architectural constraints, which have also been shown to improve length generalization.", "was_revised": true, "final_justification": "Significant overlap with prior work makes its core idea not novel. Connection to larger-scale models is not clear either. Authors rebuttal did not address those satisfactorily. I maintain my original score.", "rebuttal": "We thank the reviewer for their encouragement and comments. We are glad that you liked the theoretical justification of the training methodology we propose, and that the experiments showed significant gains when baseline training methods failed. \n\nPlease find responses and clarifications to your specific questions and comments below. \n\n**Relationship to Chugtai et al:** We were unable to locate a 2023 paper by \"Chugtai et al.\" on this topic and suspect the reviewer may be referring to Awasthi and Gupta [AG23], given the similar problem space. We sincerely apologize for this oversight in our initial submission and will, of course, add a detailed comparison to our Related Work section.\n\nOur work presents three strict improvements over [AG23]:\n\n1. **Theoretical Justification:** In [AG23], there was no formalization of why task hinting would lead to length generalization. We, on the other hand, provide a theoretical framework to show why auxiliary tasks help with length generalization. \n\n2. **Theory suggests new Training Recipe:** Our theoretical analysis leads to a robust training recipe, which was not present in [AG23].\n \n3. **Improved Performance via Robust Training:** The new robust training recipe leads to better empirical performance across a variety of tasks. For example, for the increment task [AG23] length-generalized (i.e., obtain more than 90% test accuracy) _only_ by a factor of at most 1.2.  In our work, however, we can length generalize to a factor of more than 1.5 for the same task.\n\n**Use of Synthetic Tasks:** Synthetic tasks are standard in the literature on length generalization as it is easy to clearly define the task for various input lengths ($n$). Moreover, we wanted to focus the investigation on tasks with precise answers, like sorting, increment, etc., where the effectiveness of length generalization is concretely and objectively measurable. In fact, the SLIM task is an attempt at a concrete problem capturing “contextual recall”, where each number looks for the previous number of the same “type” and sums them. \n\n**Choice of Auxiliary Functions:** This is a great question. While the choice of auxiliary tasks currently *does* depend on the specific problem, we believe this is natural for problems with precise answers. Indeed, such problems require particular tasks to be done correctly and repeatedly, and hence it makes sense to robustly train for them. If we want to consider multiple problems simultaneously, one can go beyond the naive approach of training on the auxiliary tasks for each of these problems by use a Boosting/Hedge-based approach, where we adaptively train the model on a small set of auxiliary tasks on which the model has the worst performance.\n\nIt is a very interesting question about how to isolate these auxiliary tasks automatically: one way would be to have a “library” of “basis” functions for the “local” tasks, and then use these as subroutines in a general training setup where they can hopefully enable length generalization for a variety of downstream applications.  Thank you for raising this point; we will add it to the discussion part of the revision.\n\n**Comparisons to Other Approaches:** In the revision, we will add the comparisons to the mentioned approaches.  Please note that we did not add those because the approaches are complementary to ours: curriculum training increases the complexity of the training data, whereas architectural constraints try to enforce additional structure to achieve their results. In our framework, our focus is quite different.\n\n**Extensions to “Fuzzy” Problems:** This is a great question. Our focus has been on “algorithmic” problems for which the auxiliary tasks can be naturally identified, and where length generalization can be precisely and objectively measured. Language modeling is a more amorphous task, and we feel that your question will require further research. E.g., one can imagine that for languages for which we do not have huge data corpuses, we should be able to use its grammar to define auxiliary tasks, and train the models to be better at handling intricate language constructs. (Our SLiM task tries to capture the kind of contextual recall typical in natural language.)\n\n**Further Experimental Comparisons:** We currently do not have experimental comparisons to curriculum training, since we consider it orthogonal to our focus. We will mention this in the “Limitations” section."}, {"review_id": "lEhrk2IS1S", "reviewer": "Reviewer_TSCa", "summary": "This paper addresses length generalization, a major challenge facing transformers. It shows that length generalization can be boosted by adding training data from an auxiliary tasks at longer lengths.", "questions": "See Weaknesses, especially the definition of the Auxiliary Tasks and the link between the experiments and the theorem.", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\n* addresses a key problem facing transformers\n* includes both theoretical analysis and experiments\n* experiments show clear promise of the proposed method\n\nWeaknesses:\n* The paper claims that the auxiliary tasks are derived from the novel framework of theoretical analysis (e.g. line 343), but I didn’t follow the link. Section 3.2 is quite unspecific and informal. The link to Theorems 3.1, 3.3 remains unspecified.\n* I found the definitions of the auxiliary tasks confusingly presented. Example: line 261: “produce the successor of a given element in a sequence” What exactly does this mean? I don’t understand the link to Figure 1. I presume the second row in the figure is the auxiliary task? But what exactly is the task? I had similar trouble understanding the other tasks.\n* While I did understand the SLiM task, I didn’t understand how it is a plausible model of language modeling. It is both very simplified and quite contrived. Is there a theoretical or conceptual justification for it?", "was_revised": true, "final_justification": "Based on the rebuttal, I trust the authors that they can address the Weaknesses that I raised. I have thus increased my score.", "rebuttal": "We thank the reviewer for their suggestions and comments. We are happy that you concur with the fact that this is a key problem on which progress is essential, and that you also see the promise of our approaches to getting better length generalization. \n\nPlease find responses and clarifications to your specific questions and comments below. \n\n**Link between Theoretical Results and Training Methodology**: Recall that we  focus on sparse functions. For this we consider $(N,n)$-perturbations (which are sequences of length $n$ produced by suitably sub-sampling $N$-length sequences), and train for both the auxiliary and main tasks on these sequences. Given a problem, we do the following\n\n   a. **Identifying Auxiliary Tasks**: For sparse problems with multi-token outputs, each token at timestep $t$ is obtained by applying some “local” function $f_t$ to a small number of previous tokens. Given this sparsity, there can only be a _bounded_ number of such functions; moreover, for many common tasks, there are just a handful of functions being repeatedly applied, e.g., finding the minimum and successor in sorting, or doing single-digit addition in the increment task, or summing two tokens in the SLiM task, etc. These commonly-used functions for a particular problem are natural auxiliary tasks for that problem. \n\n   b. **Robust Training for Auxiliary Tasks** : Having identified these auxiliary tasks, we train the model to be $(N,n)$-robust on auxiliary tasks (in addition to training on the main task). Recall that this training is on $n$-length sequences.  \n\n\n   c. **Robustness implies Length Generalization**: We can now use our theoretical guarantees, which say that if we are $(N,n)$-robust with good parameters, then length generalization happens (i.e., low error on $N$-length sequences).  \n\nWe hope this clarifies the setup. We will make this more explicit in the next version.\n\n**Clarification of Sorting Task:** For sorting, the auxiliary task is as follows: given an input unordered sequence $S$ of tokens (which in our case are numbers modulo 100), and a specific “query” token $q$, the auxiliary task is to lookup the smallest number in $S$ that is greater than $q$ (“successor”). Note that this “lookup” task is precisely what sorting repeatedly performs, since it outputs the minimum element, and then repeatedly outputs the successor of the last-output token. We will clarify this in the revision.\n\n\n**Motivation for the SLiM Task:** The motivation for the SLiM task was to serve as a simple example of _contextual recall_, where each number pays attention to the closest previous number of the same “type” and sums the two. This gives us a precisely defined task on which we can objectively measure length generalization, while capturing some of the essential properties of language (i.e., contextual lookup and then some transformation of the tokens looked up this way).  One could consider more sophisticated probabilistic versions, e.g., each token $a_t$ pays attention to multiple previous tokens having some types (that depends on $a_t$’s type), and then outputs the next token chosen from a probability distribution.  This is another step closer to natural language and we leave this for future work."}, {"review_id": "X8xBFsq1y8", "reviewer": "Reviewer_1jdG", "summary": "The paper focusses on a specialised set of sparse task setups, where the output token only depends on a small set of preceding tokens and studies these kind of tasks with a novel framework that enables length generalisation. The paper shows how good performance on smaller training lengths can lead to generalisation for such sparse tasks. For transformers, the error grows linearly in sequence length of the input. Thus that naturally leads to a recipe for training to help in length generalisation, which is to train on the native distribution of $\\mathcal{D}_n$ and also on random n-length restrictions of $\\mathcal{D}_N$ (the random restrictions here implies -- making the longer input with $N$ tokens smaller to length $n$, where all the $n$ tokens will be utilised to make the next prediction). This is the special key step being proposed here, as just training on $\\mathcal{D}_n$ is already known to not lead to strong length generalisation. So, if I understand correctly, then the idea is to train both on sparse as well as in a certain sense dense setups [referred to auxiliary tasks here] at smaller training lengths, and that should lead to length generalisation for the sparse setting at longer lengths. Finally several experiments are conducted to verify the theoretical claims and show the training recipe being successfully applied to a variety of tasks.", "questions": "This is not a question, more of a suggestion / comment. Currently the paper writing seems to be written in 2 styles. The empirical experiments are quite easy to understand, insights are also intuitive and is quite well written. The theory on the other hand while is quite trivial to understand once one gets past the notation, I feel there is room to make the things a bit more intuitive for the reader as well. There can be for instance a paragraph defining the proof sketch, before the actual proof starts, to help prepare the reader a bit better for how the proof is going to go. However, this is not quite a criticism, just an observation, as the experimental section was quite intuitive, but the theory took a bit of time to go through because of the definitions and heavy notation, which could be circumvented with some kind of short proof sketches. \n\n\nOverall I feel that this paper has quite a lot of potential and I really like the overall idea being proposed here. However, there needs to be some more interaction with prior literature, not just in terms of mentioning what has been proposed before, but in terms of how good the current method is compared to other things. Some more discussions about positional encodings, it's role and bit more exposition on how to actually identify the underlying auxiliary tasks in a general sense would be helpful. In case, auxiliary tasks cannot be found across setups, then that has to acknowledged in the paper somewhere.", "limitations": "yes", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n    - There is a clear practical insight that can be taken from theoretical findings, and a training recipe that can be applied readily.\n    - Empirical results shown in the paper support the theoretical claims being made in the paper.\n    - The core idea of having an auxiliary task (reduce the actual task by reducing the number of tokens to the ones that matter and that don't change the answer) and then train based on it seems quite useful and is grounds for further exploration on its own.\n\n**Weaknesses**\n  - Only No Positional Encodings have been considered. It's unclear what the inclusion of positional encodings will do, especially since, something like Absolute positional encodings requires longer positions to be also trained, so clipping the length to smaller lengths for the auxiliary tasks might make training slower. Similarly with Rotary positional encodings, the rotation angle etc. might require modifications to the training recipe proposed here. \n- Not all tasks that we can think of have an obvious auxilliary task that we can train them on. For example, math problems don’t have obvious auxilliary tasks that I can think of, thus although the insights from the paper are actionable, they are not easily applicable beyond across settings.\n- Empirical Comparison has only been made with the base model of their choice, discussion about other hints / tricks suggested by prior literature is not discussed. Thus, although training this way can be considered better than pure vanilla training, but it’s not clear, whether this is objective better than something like index hinting, providing scratch pads to the model etc.  For instance it was pointed out in the related section of the text -- that scratchpads are not universally effective, but neither is this approach as far as I see it. Thus for a fair comparison, there needs to be some kind of comparisons with scratchpads. Similar arguments can be made for other kinds of training recipes proposed in the literature before.", "was_revised": true, "final_justification": "The authors were able to clarify my questions. I like the proposed method and see considerable potential for the paper to improve through deeper engagement with prior work, experiments on other positional encodings, and a clearer delineation of where the method is applicable and where it is not. In light of the authors’ responses and the other reviews—particularly regarding the mention of Awasthi and Gupta on task hinting (as the proposed method is quite similar, and so therefore the current paper by the authors requires somewhat more contribution to be valuable contribution and not just a marginal one)—Therefore I feel the current version is not ready for publication, and needs some revisions to be a solid paper.", "rebuttal": "We thank the reviewer for their suggestion, encouragement, and comments. We are happy that you appreciated the clean idea behind our work: that sparsity naturally suggests a set of “core” tasks, which can then be trained on robustly, and this can ensure length generalization. We hope that this conceptual simplicity/message will be useful more broadly, beyond the specific problems in this paper.\n\nThank you also for the detailed and valuable suggestion: we will rephrase to present the proofs more transparently, and will also unify the presentation styles. We completely agree that it will make the ideas of the paper more accessible, and will help the reader. We will also add more connections and comparisons to previous work.\n\nPlease find responses and clarifications to the other questions and comments below. \n\n**The Role of Encodings:** Thanks for the question. We view the role of position encodings as complementary (and indeed, additive) to our proposed recipe. Typically in order to train a model with positional encoding for length generalization standard techniques such as a random position offset is used during training. If positional encodings are present, the same techniques will apply to our framework too.   \n     \n**Choice of Auxiliary Tasks:** It is a great question to design auxiliary tasks automatically for more real world settings such as natural language tasks and math tasks. One way would be to have a “library” of “basis” functions for the “local” tasks, and then use these as subroutines in a general training setup where they can hopefully enable length generalization for a variety of downstream applications.  \n    \n**Comparison to other Models:** Regarding the use of *scratchpads*, we would like to note that in all our experiments the baseline models use scratchpads and still do not length-generalize; this is consistent with what other literature on this topic has observed. *Index hinting* techniques such as position coupling are quite powerful for length generalization but rely on the right hints to be fed into the model at inference time as well. These are very strong hints; our focus was to make the model naturally learn to attend to the right tokens via robust training without such direct hints, since such hints will typically not be available at inference time. We are happy to include these methods for comparison in the next version of the paper."}]

Author rebuttal:

Meta-review: The paper provides novel theoretical framework proving length generalization results on a single layer Transformer. Specifically, the authors show that under an assumption of robustness to perturbation, any model with low error on "short" sequences will suffer moderately increasing error for longer sequences. Following this result, the authors show that this robustness assumption can be enforced during training by constructing an auxiliary task, such that when adding this auxiliary task to the training it improves length generalization capabilities on the original target task.

Reviewers appreciated the results and methods described in the paper, and overall found the theoretical framework novel and the underlying motivation to be important. One major criticism is the lack of proper reference to prior work, and specifically the work of Awasthi and Gupta [AG23], which empirically demonstrates the benefit of a very similar method of "task hinting" for length generalization. Additionally, the reviewers mentioned that there was no discussion on the effect of positional encoding and that the derivation of the auxiliary tasks is not clear. They also mention that there is some disconnect between the theoretical and empirical sections, and that the connection between them should be further clarified.

I believe that the theoretical framework and results introduced in this work are indeed valuable and do amount to significant contributions over prior results. However, the original manuscript has flaws as mentioned above, and in particular should have framed these nice theoretical contributions in the context of existing empirical works. I therefore recommend to accept the paper, assuming that the authors make the necessary changes to properly discuss prior works, as well as incorporating other feedback from the reviewers.
Decision: Accept (poster)