Research map / Architectures and efficiency / Attention and transformers
Position Coupling: Improving Length Generalization of Arithmetic Transformers Using Task Structure
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes position coupling, which assigns identical position IDs to semantically relevant tokens (e.g., digits of equal significance), enabling decoder-only Transformers trained on 1-30 digit addition to generalize to 200-digit addition, with supporting theory on 1-layer constructions.
Problem
Transformers often fail to length-generalize even on simple arithmetic such as integer addition, learning only up to the training sequence length.
Approach
Embed task structure into absolute positional encodings by giving the same position IDs to tokens that are semantically relevant to each other (e.g., digits of the same significance across operands and sum), with random starting position IDs during training.
Key results
- Position coupling with models trained on 1-30 digit additions achieves over 95% median exact-match accuracy up to 200-digit additions (generalizable length 200 for 1-30 and 1-40 training).
- A 1-layer 2-head Transformer with coupled positions is theoretically constructed to add integers with operand length up to 2^floor((d-17)/2)-2, while a 1-layer NoPE Transformer's max EM accuracy is 6.2% for 3-digit addition.
- Scaling to training length 160 yields strong exact-match accuracy for 500-digit additions.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Train on longer operand lengths (ablation) | Position coupling for addition · Addition | generalizable length (digits): 70 (1--10) -> 135 (1--20) -> 200 (1--30 and 1--40) | Helped |
| Increase depth to 6 layers (ablation) | Position coupling for addition trained on 1--30 digits · Addition | length generalization performance | Hurt |
| Use zero-padding (ablation) | Position coupling for addition with a 1-layer 4-head model · Addition | length generalization performance | Helped |
| Reverse the response or all numbers (ablation) | Position coupling for addition with a 1-layer 4-head model · Addition | length generalization performance | Helped |
| Remove reversing (ablation) | Position coupling for addition with a 1-layer 4-head model · Addition | in-distribution performance | Hurt |
| Add zero-padding and reversing (ablation) | Position coupling for addition with a 6-layer 8-head model · Addition | performance | Hurt |
| Use position coupling instead of NoPE or random-start APE | Decoder-only Transformers for addition · Addition | length generalization | Helped |
| Use index hinting instead of NoPE or random-start APE | 6-layer 8-head Transformer for addition · Addition | generalizable length (digits): does not extend further than 50; complete failure starting from length 70 | Helped |
| Use shared rather than distinct positional embedding layers (ablation) | Position coupling for minesweeper generator · Minesweeper generator | generalizable board width: 10×10 -> 12×12; over 98\% accuracy on 12×12 and near 90\% on 14×14 | Helped |
| Combine position coupling with RoPE (ablation) | 12-layer Transformer for addition · Addition | length generalization: up to operand lengths 100 | Helped |
Surprising findings
- A 1-layer NoPE Transformer cannot fully solve addition: maximum achievable EM accuracy falls to 6.2% at 3 digits, 1% at 4 digits, and 0.13% at 5 digits.
- Deeper 6-layer models show slightly degraded performance under position coupling, attributed to implicit bias toward shortcuts.
Limitations
- Requires manually designed coupling rules based on known task structure; not applicable to implicit-structure NLP tasks.
- Generalization to multi-summand addition and variable-length multiplication left to future work.
- Results for 6-layer models show slight degradation; median-based generalizable length threshold of 95% is stricter than some prior reports.
Experimental setup
Datasets: synthetic arithmetic data (balanced-digit sampled addition, multiplication, copy/reverse, minesweeper generator)
Benchmarks: 5-200 digit addition exact match (in-house) N x 2 multiplication (in-house) minesweeper generator (in-house) copy/reverse (in-house)
Metrics: exact-match accuracy generalizable length
Baselines: NoPE random-start APE index hinting RoPE Abacus
Builds on: learned absolute positional embedding index hinting NoPE reversed format
Closest papers by meaning
- Arithmetic Transformers Can Length-Generalize in Both Operand Length and Count (ICLR 2025 · similarity 0.92)
- Transformers Can Do Arithmetic with the Right Embeddings (NeurIPS 2024 · similarity 0.91)
- A Formal Framework for Understanding Length Generalization in Transformers (ICLR 2025 · similarity 0.89)
- Understanding Addition in Transformers (ICLR 2024 · similarity 0.88)
- Looped Transformers for Length Generalization (ICLR 2025 · similarity 0.87)
- Assign and Add: A Mechanistic Study of Compositional Arithmetic (NeurIPS 2026 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.