Title: Extrapolation by Association: Length Generalization Transfer In Transformers
Venue: NeurIPS 2025
Accepted: yes
Keywords: Language models

LLM

Transformers

Length Generalization

Transfer
Abstract: Transformer language models have demonstrated impressive generalization capabilities in natural language domains, yet we lack a fine-grained understanding of how such generalization arises. In this paper, we investigate length generalization—the ability to extrapolate from shorter to longer inputs—through the lens of \textit{task transfer}. We find that length generalization can be \textit{transferred} across related tasks. That is, training a model with a longer and related auxiliary task can lead the model to generalize to unseen and longer inputs from some other target task. We demonstrate this length generalization transfer across a diverse suite of algorithmic tasks, including arithmetic operations, string transformations, and maze navigation. Our results show that transformer models can inherit generalization capabilities from similar tasks when trained jointly. Moreover, we observe similar transfer effects in pretrained language models, suggesting that pretraining equips models with reusable computational scaffolding that facilitates extrapolation in downstream settings. 
Finally, we provide initial mechanistic evidence that length generalization transfer correlates with the re-use of the same attention heads between the tasks. Together, our findings deepen our understanding of how transformers generalize to out-of-distribution inputs and highlight the compositional reuse of inductive structure across tasks.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper addresses the length generalization problem, a challenge facing transformers. The paper finds that length generalization can transfer between tasks: when training two tasks at different input lengths, the model may end up generalizing to the longer length on both tasks.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\n* 1 Addresses an important challenge facing transformers\n* 2 Provides clear experimental promise of the proposed approach on algorithmic tasks\n* 3 Provides diverse experimental results: algorithmic tasks, benefit from pretraining, mechanistic analysis\n* 4 The paper is easy to read\n\nWeaknesses:\n* 1 The paper argues that length generalization transfer happens for “related” tasks (e.g., line 6), but it remains unclear what kind of relatedness here is needed. There is no theoretical (or even principled conceptual) understanding of how one would select a useful auxiliary task. The selection of auxiliary tasks is heuristic and exploratory. Hence, the paper provides no deeper insights for guiding the future design of auxiliary tasks. [UPDATE: the authors have some initial ideas based on MI, though no deeper principled insights are promised]\n* 2 I found Figure 9 unconvincing. The caption claims that “All three metrics closely align” but other than maxima at 2000 and 16000 steps, I found little obvious alignment between the metrics. Also, (even if the alignment between the metrics were more obvious), the conclusion that “successful length generalization transfer accompanies shared attention mechanisms between tasks” is just based on correlational evidence [there is no causal evidence that attention map/ablation map similarity is causal to the generalization], and should be qualified accordingly. [UPDATE: the authors promise to address this]\n* 3 In Figures 3 and 4, the training lengths vary (16+32 in Figures 3a,3b,4a,4c; 8+16 in Figure 3c, and similarly variable in Figures 5, 6). What is the motivation for this? Why not report unified results with the same lengths?  [UPDATE: the authors have explained the motivation]\n\nQUESTIONS:\nQuestions:\n* see Weakness #3\n* In Figure 3c, it’s presupposed that when training Reverse Add up to length 8 on its own, it won’t length-generalize (which is not shown in Figure 3d). Is this the case?\n* The paper shows transfer for related tasks, but not lack of transfer in the case of unrelated tasks. Have the authors tried to pair up unrelated or less-related tasks to show that indeed “related” tasks (line 6) are needed for transfer to happen?\n\nMinor:\n* line 245: Figure 10 → Figure 9\n* line 180, “or” → “For”\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper investigates whether being trained on longer instances of an auxiliary task helps with length generalization on a main task. The authors find that it generally does, but only when the tasks are related and when the length difference is not too stark. They also find evidence that the improvement in generalization performance is associated with an increased sharing of mechanisms (as indicated by attention maps and ablation maps) between tasks.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\nI really liked this paper. The idea is clear and elegant, the results are convincing, and I appreciated that the authors studied (relatively small) pretrained language models as well. I also thought the paper was well written and I thought the schematic figures indicating the authors' hypotheses were very nice. The limitations section helpfully acknowledges questions that this paper does not address and highlights potential for future work. Finally, I thought it was good (and important) that the authors studied tasks from several domains.\n\n**Weaknesses**\n\nPrimarily, I think Section 6.2 is currently lacking. It is both a little difficult to assess the extent to which ablation and attention difference are similar over time (other than them both going down) and there are many confounding factors, as many quantities vary which the training epoch (e.g. training set performance, in-distribution generalization, convergence to more consistent solutions from potentially more random initializations). I therefore think the evidence provided for more length-generalizing networks having more mechanisms in common is pretty weak. I think it would be interesting to compare the ablation and attention difference between different random seeds that exhibit different length generalization or between the networks exhibiting different length generalization in Figure 8. Finding an association between ablation/attention difference and length generalization across different networks would be much stronger evidence in my opinion.\n\nAs a more minor comment, I think it would be nice to connect the limitations section (which, as noted above, was generally good) to prior work that could be useful in addressing these questions.\n\nQUESTIONS:\nAs noted above, I really liked this paper. My primary concern is the evidence in Section 6.2. I think a more extensive evaluation of the hypothesized association between mechanism sharing and length generalization would make the paper considerably stronger (whether or not this hypothesized association turns out to be true). In my comments above, I have provided one example for what such an analysis could look like. If this is not possible in the rebuttal period, I think it would be good to more strongly acknowledge the limitations of the current approach.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper uncovers a novel finding adding to our understanding of transformer models: length generalisation can transfer across similar tasks in a variety of data, task, and model settings. They validate this on several task families, show instances of lack of generalisation due to task dissimilarity, show the finding holds on pretrained models, and provide interpretability analysis showing cross-task transfer is related to shared underlying mechanisms.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n- This is a really interesting and novel finding that adds to our understanding of language models. The overall idea is well-motivated and the paper is very nicely written, and the analyses (generalisation gap, different checkpoints, control task transfer) is comprehensive. Overall, I am very convinced by the narrative presented by the paper.\n- I appreciate the inclusion of details of mechanistic results in sec. 6.2 and the relevant appendices, which elucidate why this happens.\n- Further experiments comparing NoPE and RoPE are great to see; while these are in the appendix, this is important (along with the control task failures) to show that the finding is not merely about better learning of how to use RoPE but something more fundamental about cross-task transfer.\n\nWeaknesses\n- None apparent.\n\nQUESTIONS:\n- How were these tasks chosen? For the synthetic tasks, is it possible to formalise why task transfer is happening? (E.g. is the [C-RASP](https://arxiv.org/abs/2506.16055) program for these tasks sharing some key operations?) This isn't necessary for the paper as it stands, but I'm curious if there is some intuition on the authors part as to when generalisation does or does not occur across tasks.\n- Can you explain figure 9 more? Particularly the leftmost subplot doesn't seem to be showing a strong trend. Maybe some error bars could help guage the significance of the correlation in the three metrics.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 5, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper investigates a novel mechanism for improving length generalization in Transformers, termed \"extrapolation by association.\" The central hypothesis is that a model can learn to generalize to longer, unseen input sequences for a \"main task\" if it is jointly trained on a related \"auxiliary task\" that includes examples of those longer sequences. The authors provide empirical evidence for this phenomenon, which they call \"length generalization transfer,\" across a diverse set of algorithmic domains: arithmetic (e.g., reverse add), string manipulation (e.g., capitalize-reverse), and maze navigation (e.g., DFS trace).\n\nThe key findings are:\n1. Co-training with a longer, related auxiliary task consistently enables the main task to extrapolate beyond its training length, a capability it lacks when trained in isolation.\n2. This transfer is contingent on the structural similarity between tasks, as control experiments with unrelated tasks show no such benefit.\n3. The effect is also observed when fine-tuning pretrained language models (SmolLM), and the capacity for generalization improves with the extent of pretraining, suggesting that pretraining builds a reusable \"computational scaffolding.\"\n4. Initial mechanistic analysis suggests that successful transfer correlates with the reuse of the same attention heads for both tasks, indicating a sharing of internal computational circuits.\n\nSTRENGTHS AND WEAKNESSES:\nStrength: \n- The paper introduces a novel and significant research direction for length generalization, shifting focus from purely architectural solutions to the training paradigm itself.\n- The experimental validation is thorough and well-designed, using three distinct task domains which strongly supports the generality of the core claim.\n- The inclusion of control experiments with unrelated tasks provides compelling evidence that the observed transfer is due to task relatedness and not merely exposure to longer sequences.\n- The investigation extends to pretrained models, bridging the gap between findings on synthetic tasks and their potential relevance to large language models.\n- The mechanistic analysis using attention head ablation maps provides valuable, albeit initial, insight into how the transfer might be occurring at the circuit level.\n\nWeakness: \n- The primary weakness is the reported instability of the training dynamics; the success of the transfer is highly sensitive to random model initialization, which may limit the reliability and practical applicability of the phenomenon.\n- The crucial concept of \"task relatedness\" is defined intuitively and only analyzed post-hoc; the work lacks a predictive, a priori metric to determine which tasks will facilitate transfer.\n- The experiments are confined to relatively simple, deterministic algorithmic tasks, and it remains unclear how these findings would translate to more complex, ambiguous, or knowledge-intensive domains like natural language.\n\nQUESTIONS:\n1. The instability of transfer due to random initialization (Sec A.1) is a significant concern, as it suggests the phenomenon might occur by chance rather than being a reliable property. Could you analyze the properties of the model initializations that lead to successful transfer versus those that fail? For instance, are there discernible statistical differences in initial weight norms or other properties? A positive result here would greatly increase my confidence in the robustness of your findings and could change my score.\n2. The concept of \"task relatedness\" is central to your claims but is defined intuitively. Your post-hoc analysis showing a correlation with circuit sharing is insightful, but it doesn't help in predicting which tasks will work together beforehand. Have you considered any a priori metrics for quantifying task relatedness (e.g., based on algorithmic complexity, information theory, or embedding similarity of task descriptions) that might predict the likelihood or effectiveness of length generalization transfer?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-07-04"}