# Improving Length-Generalization in Transformers via Task Hinting arXiv: https://arxiv.org/abs/2310.00726 ยท HTML: https://mlsanity.com/paper/2310.00726/ ## Key facts - Cited by 13 papers in the map (data through 14 Oct 2026). - Builds on 25 of its 32 references that are in the map. - First posted on arXiv 1 Oct 2023. ## Builds on (most cited first, 25 of 25) - [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018) - [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014) - [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020) - [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023) - [Chain-of-Thought Prompting Elicits Reasoning in Large Language Models](https://mlsanity.com/paper/2201.11903/) (2022) - [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022) - [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016) - [ALBERT: A Lite BERT for Self-supervised Learning of Language Representations](https://mlsanity.com/paper/1909.11942/) (2019) - [Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation](https://mlsanity.com/paper/2108.12409/) (2021) - [Show Your Work: Scratchpads for Intermediate Computation with Language Models](https://mlsanity.com/paper/2112.00114/) (2021) - [Progress measures for grokking via mechanistic interpretability](https://mlsanity.com/paper/2301.05217/) (2023) - [PaLI: A Jointly-Scaled Multilingual Language-Image Model](https://mlsanity.com/paper/2209.06794/) (2022) - [Emergent World Representations: Exploring a Sequence Model Trained on a Synthetic Task](https://mlsanity.com/paper/2210.13382/) (2022) - [Multi-Task Learning with Deep Neural Networks: A Survey](https://mlsanity.com/paper/2009.09796/) (2020) - [Emergent Linear Representations in World Models of Self-Supervised Sequence Models](https://mlsanity.com/paper/2309.00941/) (2023) - [Towards Generalist Biomedical AI](https://mlsanity.com/paper/2307.14334/) (2023) - [Exploring Length Generalization in Large Language Models](https://mlsanity.com/paper/2207.04901/) (2022) - [Saturated Transformers are Constant-Depth Threshold Circuits](https://mlsanity.com/paper/2106.16213/) (2021) - [Unveiling Transformers with LEGO: a synthetic reasoning task](https://mlsanity.com/paper/2206.04301/) (2022) - [Goat: Fine-tuned LLaMA Outperforms GPT-4 on Arithmetic Tasks](https://mlsanity.com/paper/2305.14201/) (2023) - [Generalization on the Unseen, Logic Reasoning and Degree Curriculum](https://mlsanity.com/paper/2301.13105/) (2023) - [Auto-Regressive Next-Token Predictors are Universal Learners](https://mlsanity.com/paper/2309.06979/) (2023) - [Length Generalization in Arithmetic Transformers](https://mlsanity.com/paper/2306.15400/) (2023) - [The EOS Decision and Length Extrapolation](https://mlsanity.com/paper/2010.07174/) (2020) - [Location Attention for Extrapolation to Longer Sequences](https://mlsanity.com/paper/1911.03872/) (2019) ## Built on by (most cited first, 13 of 13) - [What Algorithms can Transformers Learn? A Study in Length Generalization](https://mlsanity.com/paper/2310.16028/) (2023) - [Transformers Can Achieve Length Generalization But Not Robustly](https://mlsanity.com/paper/2402.09371/) (2024) - [Causal Language Modeling Can Elicit Search and Reasoning Capabilities on Logic Puzzles](https://mlsanity.com/paper/2409.10502/) (2024) - [Position Coupling: Improving Length Generalization of Arithmetic Transformers Using Task Structure](https://mlsanity.com/paper/2405.20671/) (2024) - [Why Softmax Attention Outperforms Linear Attention](https://mlsanity.com/paper/2310.11685/) (2023) - [Extrapolation by Association: Length Generalization Transfer in Transformers](https://mlsanity.com/paper/2506.09251/) (2025) - [The Expressibility of Polynomial based Attention Scheme](https://mlsanity.com/paper/2310.20051/) (2023) - [Beyond Single-Task: Robust Multi-Task Length Generalization for LLMs](https://mlsanity.com/paper/2502.11525/) (2025) - [To Infinity and Beyond: Tool-Use Unlocks Length Generalization in State Space Models](https://mlsanity.com/paper/2510.14826/) (2025) - [Towards Sampling Data Structures for Tensor Products in Turnstile Streams](https://mlsanity.com/paper/2510.03678/) (2025) - [Dissecting Linear Recurrent Models: How Different Gating Strategies Drive Selectivity and Generalization](https://mlsanity.com/paper/2601.12598/) (2026) - [Block-level Text Spotting with LLMs](https://mlsanity.com/paper/2406.13208/) (2024) - [Frontier Language Models Struggle to Copy: Text Can Be Better Viewed in 2D](https://mlsanity.com/paper/2607.16072/) (2026) ## Abstract (arXiv) It has been observed in recent years that transformers have problems with length generalization for certain types of reasoning and arithmetic tasks. In particular, the performance of a transformer model trained on tasks (say addition) up to a certain length (e.g., 5 digit numbers) drops sharply when applied to longer instances of the same problem. This work proposes an approach based on task hinting towards addressing length generalization. Our key idea is that while training the model on task-specific data, it is helpful to simultaneously train the model to solve a simpler but related auxiliary task as well. We study the classical sorting problem as a canonical example to evaluate our approach. We design a multitask training framework and show that task hinting significantly improve length generalization. For sorting we show that it is possible to train models on data consisting of sequences having length at most 20, and improve the test accuracy on sequences of length 100 from less than 1% (for standard training) to more than 92% (via task hinting). Our study uncovers several interesting aspects of length generalization. We observe that while several auxiliary tasks may seem natural a priori, their effectiveness in improving length generalization differs dramatically. We further use probing and visualization-based techniques to understand the internal mechanisms via which the model performs the task, and propose a theoretical construction consistent with the observed learning behaviors of the model. Based on our construction, we show that introducing a small number of length dependent parameters into the training procedure can further boost the performance on unseen lengths. Finally, we also show the efficacy of our task hinting based approach beyond sorting, giving hope that these techniques will be applicable in broader contexts.