Title: Length Generalization via Auxiliary Tasks
Venue: NeurIPS 2025
Accepted: yes
Keywords: length generalization

theory
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\n**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.\n\nQUESTIONS:\nThis 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.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-21"}

comments: 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.

STRENGTHS AND WEAKNESSES:
Strengths:
* addresses a key problem facing transformers
* includes both theoretical analysis and experiments
* experiments show clear promise of the proposed method

Weaknesses:
* 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.
* 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.
* 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?

QUESTIONS:
See Weaknesses, especially the definition of the Auxiliary Tasks and the link between the experiments and the theorem.

LIMITATIONS:
yes
TITLE: Official Review
CLARITY: 2
IMPACT: 2
ORIGINALITY: 3
DATE: 2025-07-01

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n# 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.\n\nQUESTIONS:\n# 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?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 1, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\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?\n\nQUESTIONS:\n1) 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.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-05"}