# Position Coupling: Improving Length Generalization of Arithmetic Transformers Using Task Structure

arXiv: https://arxiv.org/abs/2405.20671 · HTML: https://mlsanity.com/paper/2405.20671/

## Key facts

- Cited by 23 papers in the map (data through 14 Oct 2026).
- Builds on 42 of its 50 references that are in the map.
- First posted on arXiv 31 May 2024.

## Builds on (most cited first, 30 of 42)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [Gaussian Error Linear Units (GELUs)](https://mlsanity.com/paper/1606.08415/) (2016)
- [HuggingFace's Transformers: State-of-the-art Natural Language Processing](https://mlsanity.com/paper/1910.03771/) (2019)
- [GLU Variants Improve Transformer](https://mlsanity.com/paper/2002.05202/) (2020)
- [Root Mean Square Layer Normalization](https://mlsanity.com/paper/1910.07467/) (2019)
- [Solving Quantitative Reasoning Problems with Language Models](https://mlsanity.com/paper/2206.14858/) (2022)
- [LaMDA: Language Models for Dialog Applications](https://mlsanity.com/paper/2201.08239/) (2022)
- [Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation](https://mlsanity.com/paper/2108.12409/) (2021)
- [Gemini: A Family of Highly Capable Multimodal Models](https://mlsanity.com/paper/2312.11805/) (2023)
- [Language Modeling with Gated Convolutional Networks](https://mlsanity.com/paper/1612.08083/) (2016)
- [Convolutional Sequence to Sequence Learning](https://mlsanity.com/paper/1705.03122/) (2017)
- [Are Transformers universal approximators of sequence-to-sequence functions?](https://mlsanity.com/paper/1912.10077/) (2019)
- [The Impact of Positional Encoding on Length Generalization in Transformers](https://mlsanity.com/paper/2305.19466/) (2023)
- [Autoformalization with Large Language Models](https://mlsanity.com/paper/2205.12615/) (2022)
- [Neural Networks and the Chomsky Hierarchy](https://mlsanity.com/paper/2207.02098/) (2022)
- [Exploring Length Generalization in Large Language Models](https://mlsanity.com/paper/2207.04901/) (2022)
- [What Algorithms can Transformers Learn? A Study in Length Generalization](https://mlsanity.com/paper/2310.16028/) (2023)
- [Thinking Like Transformers](https://mlsanity.com/paper/2106.06981/) (2021)
- [Investigating the Limitations of Transformers with Simple Arithmetic Tasks](https://mlsanity.com/paper/2102.13019/) (2021)
- [Teaching Arithmetic to Small Transformers](https://mlsanity.com/paper/2307.03381/) (2023)
- [Randomized Positional Encodings Boost Length Generalization of Transformers](https://mlsanity.com/paper/2305.16843/) (2023)
- [Unveiling Transformers with LEGO: a synthetic reasoning task](https://mlsanity.com/paper/2206.04301/) (2022)
- [Tracr: Compiled Transformers as a Laboratory for Interpretability](https://mlsanity.com/paper/2301.05062/) (2023)
- [Low-Rank Bottleneck in Multi-head Attention Models](https://mlsanity.com/paper/2002.07028/) (2020)
- [Transformers Can Achieve Length Generalization But Not Robustly](https://mlsanity.com/paper/2402.09371/) (2024)

## Built on by (most cited first, 23 of 23)

- [Is Chain-of-Thought Reasoning of LLMs a Mirage? A Data Distribution Lens](https://mlsanity.com/paper/2508.01191/) (2025)
- [Self-Improving Transformers Overcome Easy-to-Hard and Length Generalization Challenges](https://mlsanity.com/paper/2502.01612/) (2025)
- [Extrapolation by Association: Length Generalization Transfer in Transformers](https://mlsanity.com/paper/2506.09251/) (2025)
- [Beyond Single-Task: Robust Multi-Task Length Generalization for LLMs](https://mlsanity.com/paper/2502.11525/) (2025)
- [Born a Transformer -- Always a Transformer? On the Effect of Pretraining on Architectural Abilities](https://mlsanity.com/paper/2505.21785/) (2025)
- [A Theoretical Framework for OOD Robustness in Transformers using Gevrey Classes](https://mlsanity.com/paper/2504.12991/) (2025)
- [To Infinity and Beyond: Tool-Use Unlocks Length Generalization in State Space Models](https://mlsanity.com/paper/2510.14826/) (2025)
- [Stabilizing Extrapolation in Looped Transformers via Learned Stochastic Stopping](https://mlsanity.com/paper/2606.29983/) (2026)
- [Language Models, Graph Searching, and Supervision Adulteration: When More Supervision is Less and How to Make More More](https://mlsanity.com/paper/2503.10542/) (2025)
- [On the Role of Computation in Reinforcement Learning](https://mlsanity.com/paper/2602.05999/) (2026)
- [On the Ability of Transformers to Verify Plans](https://mlsanity.com/paper/2603.19954/) (2026)
- [Exact Learning of Arithmetic with Differentiable Agents](https://mlsanity.com/paper/2511.22751/) (2025)
- [Pre-Generating Multi-Difficulty PDE Data for Few-Shot Neural PDE Solvers](https://mlsanity.com/paper/2512.00564/) (2025)
- [From Small to Large: Generalization Bounds for Transformers on Variable-Size Inputs](https://mlsanity.com/paper/2512.12805/) (2025)
- [Arithmetic OOD Failure Unfolds in Stages in Minimal GPTs](https://mlsanity.com/paper/2603.26828/) (2026)
- [Training Transformers as a Universal Computer](https://mlsanity.com/paper/2604.25166/) (2026)
- [Length Generalization with Log-Depth Recurrent Units](https://mlsanity.com/paper/2605.26035/) (2026)
- [Position as Probability: Self-Supervised Transformers that Think Past Their Training for Length Extrapolation](https://mlsanity.com/paper/2506.00920/) (2025)
- [On the Limitations and Capabilities of Position Embeddings for Length Generalization](https://mlsanity.com/paper/2510.04130/) (2025)
- [Convex Optimization of Spiking Neural Networks with Global and Approximation Guarantees](https://mlsanity.com/paper/2605.08022/) (2026)
- [Relative Positions Generalize, Absolute Positions Memorize: An Implicit-Bias Account of Length Generalization in Attention](https://mlsanity.com/paper/2607.18759/) (2026)
- [Beyond the Training Horizon: Mechanisms and Limits of Length Generalization in Looped Transformers](https://mlsanity.com/paper/2609.33144/) (2026)
- [Counting with Attention and LayerNorm: Norm Scaling and Representation Geometry](https://mlsanity.com/paper/2610.16653/) (2026)

## Abstract (arXiv)

Even for simple arithmetic tasks like integer addition, it is challenging for Transformers to generalize to longer sequences than those encountered during training. To tackle this problem, we propose position coupling, a simple yet effective method that directly embeds the structure of the tasks into the positional encoding of a (decoder-only) Transformer. Taking a departure from the vanilla absolute position mechanism assigning unique position IDs to each of the tokens, we assign the same position IDs to two or more "relevant" tokens; for integer addition tasks, we regard digits of the same significance as in the same position. On the empirical side, we show that with the proposed position coupling, our models trained on 1 to 30-digit additions can generalize up to 200-digit additions (6.67x of the trained length). On the theoretical side, we prove that a 1-layer Transformer with coupled positions can solve the addition task involving exponentially many digits, whereas any 1-layer Transformer without positional information cannot entirely solve it. We also demonstrate that position coupling can be applied to other algorithmic tasks such as Nx2 multiplication and a two-dimensional task.