Source-linked AI summary
A Length-Extrapolatable Transformer
Yutao Sun, Li Dong, Barun Patra, Shuming Ma, Shaohan Huang, Alon Benhaim, Vishrav Chaudhary, Xia Song, Furu Wei
TL;DR
Transformers often fail when evaluated on sequences longer than those used in training, creating a need for length-extrapolatable models. The paper defines attention resolution and combines extrapolatable position embedding with blockwise causal attention; LEX performs strongly on both short and long texts. Its scope is causal language modeling, with added inference cost and assumptions in the position-encoding analysis.
Problem
Transformers commonly handle only in-distribution input lengths, although training across all possible lengths is infeasible.
Method
The paper defines attention resolution and improves it using extrapolatable relative position embedding and blockwise causal attention during inference.
Results
LEX Transformer achieves strong performance on both short and long texts, with perplexity continuing to decrease on long inputs while other methods fail to extrapolate or increase perplexity.
Takeaways & Limitations
Length extrapolation benefits from jointly designing position embedding and attention behavior rather than relying on position embedding alone.
Takeaways & Limitations
The work focuses on causal language modeling, XPOS adds about 6% inference cost versus absolute position embeddings, and its analysis assumes E(∠q) ≤ E(∠k).
Abstract
from arXiv · showhide
Position modeling plays a critical role in Transformers. In this paper, we focus on length extrapolation, i.e., training on short texts while evaluating longer sequences. We define attention resolution as an indicator of extrapolation. Then we propose two designs to improve the above metric of Transformers. Specifically, we introduce a relative position embedding to explicitly maximize attention resolution. Moreover, we use blockwise causal attention during inference for better resolution. We evaluate different Transformer variants with language modeling. Experimental results show that our model achieves strong performance in both interpolation and extrapolation settings. The code will be available at https://aka.ms/LeX-Transformer.
1 Introduction
Transformers struggle to process sequences longer than those seen during training, motivating length-extrapolatable designs. The paper introduces attention resolution, XPOS, and blockwise causal attention, with LEX Transformer showing strong performance on short and long texts.
- Transformers typically handle only in-distribution input lengths, making length extrapolation important for wider usage.
- Position modeling is difficult in parallel attention, and existing relative-position methods can trade off extrapolation against general performance.
- The paper defines attention resolution to measure a Transformer’s ability to recognize position and support length extrapolation.
- LEX combines extrapolatable position embedding with blockwise causal attention to improve length extrapolation.
- LEX Transformer achieves strong performance on both short and long texts in language-modeling experiments.
2 Design Principles of Transformers for Position Modeling
The paper frames position modeling around order sensitivity, translation robustness, and the ability to handle varying input lengths. It argues that extrapolation requires coordinated design of position embeddings and attention behavior.
- Position information is essential for sequence modeling because removing it reduces a Transformer to a bag-of-words model.
- A sequence representation should remain robust when padding translates the input positions.
- A suitable Transformer should handle arbitrary input lengths without retraining for longer downstream sequences.
- Learnable absolute embeddings cannot extrapolate, while Alibi addresses extrapolation with exponential attention decay but may lose long-term dependency information.
- Effective extrapolation requires more than position embedding: proper attention can prevent perplexity from exploding, but ideal models should continue benefiting from longer context.
3 A Length-Extrapolatable Transformer
LEX improves Transformer length extrapolation by defining attention resolution, modifying relative position encoding with decay, and using blockwise causal attention during inference. These designs target stable position modeling over long-term dependencies while retaining standard causal training.
- 3.2 Improve Resolution by Position Encoding: LEX introduces a relative position encoding that generalizes ROPE by adding exponential decay to the rotation matrix, explicitly optimizing attention resolution.The design selects ζ to maximize R(gζ) while avoiding additional position-bias calculations.
- 3 A Length-Extrapolatable Transformer: The resulting Length-Extrapolatable Transformer combines extrapolatable position embedding with blockwise causal attention to improve length extrapolation.The architecture is named the LEX Transformer.
- 3.1 Attention Resolution: Attention resolution measures how well attention scores recognize relative position, with monotonic score expectations preferred for increasing token distance.The metric uses expected attention scores s[n] at distance n and quantitatively evaluates position monotonicity.
- 3.2 Improve Resolution by Position Encoding: The position-encoding derivation assumes E(∠q) ≤ E(∠k) for generative models and computes attention expectations rather than the stronger upper bound used by prior work.This assumption is introduced by the language-model setting, while the prior inequality scaling is described as too strong to retain generality.
- 3.2 Improve Resolution by Position Encoding: Adding decay stabilizes attention expectations at large relative distances, reducing ROPE’s dramatic oscillation while preserving short-term fitting behavior.The decay primarily suppresses unstable dimensions associated with large rotation frequencies; short-term and long-term dependencies remain continuously divided.
- 3.3 Improve Resolution by Attention Mask: During inference, blockwise causal attention divides queries into l/2-length blocks that attend to their own and previous blocks, reusing key and value vectors.Training uses vanilla causal attention on shorter texts, while blockwise masking is applied to longer sequences during inference.
4 Experiments
Experiments evaluate LEX Transformer against alternative position embeddings on language modeling, interpolation and extrapolation, and attention resolution. LEX maintains strong perplexity as evaluation length increases, while BCA improves position distinguishability and supports extrapolation.
- Language Modeling: All Transformers converge to similar perplexity at lengths no greater than 1024, while XPOS achieves a stable 1–3 perplexity advantage.This result measures interpolation capability.
- Language Modeling: At lengths 2048 and 4096, XPOS perplexity decreases as sequence length increases, whereas other methods increase at length 4096.BCA is used with all position embeddings in this comparison.
- Language Modeling: LEX is the only method whose perplexity decreases as evaluation length increases across shorter and longer texts.The models are trained at length 1024 and evaluated at varied lengths.
- Measuring Resolution: XPOS improves position recognizability at training length, while BCA helps models distinguish positions in the length-2048 resolution ablation.Resolution is averaged across layers, and higher resolution indicates better context-token distinction.
- Ablation: Rotation is necessary for strong XPOS validation performance, and simple scaling alone performs worse than LEX.The ablation attributes the combination of rotation and decay to in-distribution and out-of-distribution ability.
- Blockwise Causal Attention: BCA prevents ROPE perplexity from exploding, while removing BCA raises XPOS perplexity by about 1 at 2048 and 8 at 4096.With BCA, XPOS can recognize positions under the blockwise constraint.
5 Related Work
Related work addresses long-sequence modeling through efficient attention and position representations, but this paper focuses on extrapolating from short training sequences to long evaluation sequences. It proposes XPOS to optimize attention resolution while retaining strong performance.
- Long-Sequence Transformers: Long-sequence Transformers address inefficient computation or memory use and the trade-off between performance and efficiency.These concerns motivate methods that alter attention computation or representation.
- Long-Sequence Transformers: Linear-attention methods use kernel-based or low-rank approximations and often improve efficiency while underperforming vanilla Transformers at regular lengths.Sparse-attention methods instead use structured sparsity to reduce computation.
- Length Extrapolation: This work targets extrapolation for language modeling: training uses short sequences with dense vanilla attention, while long-sequence capability is obtained during inference.The training efficiency is therefore kept unchanged relative to the stated baseline setup.
- Position Embeddings: Absolute sinusoidal, learnable, relative, rotary, and decay-based position embeddings provide different approaches to encoding position and length information.Relative-position methods include clipping, reparameterization, log-bucket scalars, rotation, and explicit exponential decay.
- Position Embeddings: Alibi supports length extrapolation through explicit decay but tends to perform worse than ROPE on contexts shorter than pre-training length.This exposes a performance trade-off between extrapolation and shorter-context performance.
- XPOS: XPOS is presented as a theoretically derived relative position embedding that optimizes attention resolution while supporting length extrapolation and strong performance.The paper combines this embedding with blockwise causal masking as its two extrapolation designs.
Limitations
The work focuses on causal language modeling, leaving integration with bidirectional attention as an additional effort. XPOS also adds about 6% inference cost compared with absolute position embeddings, while accelerating training convergence.
- The work focuses on causal language modeling rather than bidirectional attention settings such as masked language modeling.Integrating the proposed methods into bidirectional attention requires additional effort.
- Integrating the proposed methods into bidirectional attention, including masked language modeling, requires additional effort.
- About 6% inference cost is introduced by XPOS compared with absolute position embeddings, although training convergence accelerates.