Towards Infinite Length Extrapolation: A Unified Approach

TL;DR

This work proposes a unified Generalized Positional Encoding framework that reframes attention score modification as a multiplicative transformation plus additive bias, encompassing RoPE and ALiBi. It introduces Adaptive Positional Encoding (APE) with adaptive frequency and a multi-term decay bias, and provides theoretical conditions for infinite-context extrapolation, including convergent normalization and entropy boundedness, while preserving gradient positional sensitivity. The authors validate APE through case studies on TinyStories and a new LongTinyStories dataset, showing improved perplexity and attention behavior for long-range dependencies, along with analysis of memory and compute trade-offs. The approach offers a principled path toward longer context processing in transformers, with practical implications for long-document understanding and scalable language modeling.

Abstract

Large language models (LLMs) have revolutionized natural language processing, but their ability to process long sequences is fundamentally limited by the context window size during training. Existing length extrapolation methods often suffer from performance degradation or computational inefficiencies. We thereby use a unified framework that reinterprets positional encoding methods as a decomposition of the attention score into a multiplicative transformation and an additive bias. This perspective not only subsumes popular approaches such as relative position embeddings and attention-bias moderated approaches but also exposes their inherent limitations in handling long-range dependencies. To address these shortcomings, motivated by our framework, we introduce Adaptive Positional Encoding (APE), which leverages adaptive frequency modulation and an intricately designed decay bias that incorporates linear, logarithmic, and square-root terms. Our theoretical analysis establishes conditions for infinite-context extrapolation, ensuring that the softmax normalization remains well-defined over unbounded sequences while preserving long-distance correlations, entropy boundedness and gradient positional sensitivity. We substantiate our claims with an experimental case study on TinyStories dataset as well as a new synthetic dataset, \emph{Long Tiny Stories} featuring stories up to 32,000 words. Relevant code, dataset and model weights are available at https://anonymous.4open.science/r/Check-2DAD/.

Figures (7)

• Figure 1: Plots of Perplexity on Validation set of Tiny Stories Dataset vs Prompt lengths for RoPE, ALiBi and APE trained with context windows of 64, 128 and 256.
• Figure 2: Plots of Attention Entropy vs Prompt lengths for RoPE, ALiBi and APE trained with context windows of 64, 128 and 256
• Figure 3: A comparison of batched training, inference speed and memory usage of RoPE, ALiBi and APE.
• Figure 4: Perplexity vs Prompt lengths for RoPE and ALiBi trained with context windows of 64, 128 and 256 compared with APE trained only with context window of 64
• Figure 5: Plots of Perplexity on LTS (upto 5,000 words) vs Prompt lengths for RoPE, ALiBi and APE trained with context windows of 64, 128 and 256

Theorems & Definitions (30)

• Definition 3.1: Convergent Normalization
• Definition 3.2: Entropy Boundedness
• Definition 3.3: Long-Distance Correlation Preservation (LDCP)
• Definition 3.4: Gradient Positional Sensitivity (GPS)
• Theorem 3.5: Equivalence of Entropy Boundedness and Convergent Normalization
• Lemma 3.6: LDCP - Convergent Normalization Tradeoff
• proof
• Proposition 3.7
• Proposition 3.9
• Proposition 3.10