Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Position as Probability: Self-Supervised Transformers that Think Past Their Training for Length Extrapolation

Philip Heejun Lee

TL;DR

The paper tackles the challenge of length extrapolation in neural sequence models by introducing PRISM, a probabilistic relative-position encoding implemented as a differentiable histogram update. PRISM maintains a distribution over relative positions and forms a position stream via sinusoidal embeddings, enabling accurate long-horizon reasoning and exact extrapolation on algorithmic tasks beyond training lengths, all within a self-supervised framework. By interpreting the histogram update as a non-stationary Hidden Markov Model/CRF and adding a copy branch for flexible anchoring, PRISM achieves state-of-the-art extrapolation across arithmetic, SCAN, and copy benchmarks with interpretable internal states. These results advance robust long-context reasoning in Transformers and open avenues for reliable algorithmic generalization in real-world applications.

Abstract

Deep sequence models typically degrade in accuracy when test sequences significantly exceed their training lengths, yet many critical tasks--such as algorithmic reasoning, multi-step arithmetic, and compositional generalization--require robust length extrapolation. We introduce PRISM, a Probabilistic Relative-position Implicit Superposition Model, a novel positional encoding mechanism that enables Transformers to extrapolate accurately up to 10x beyond their training length. PRISM learns continuous relative positions through a differentiable histogram-filter update, preserving position uncertainty via a probabilistic superposition rather than conventional deterministic embeddings. Empirically, PRISM achieves state-of-the-art length extrapolation, successfully generalizing to previously intractable sequence lengths across algorithmic benchmarks--including arithmetic (addition, multiplication), SCAN compositionality tasks, and complex copy variants derived from DeepMind's recent datasets. Our analysis demonstrates that PRISM's stochastic positional encoding maintains sharp and interpretable internal states, providing a theoretical basis for reliable length generalization. These results advance the goal of neural sequence models that remain algorithmically robust at lengths far exceeding their training horizon.

Position as Probability: Self-Supervised Transformers that Think Past Their Training for Length Extrapolation

TL;DR

The paper tackles the challenge of length extrapolation in neural sequence models by introducing PRISM, a probabilistic relative-position encoding implemented as a differentiable histogram update. PRISM maintains a distribution over relative positions and forms a position stream via sinusoidal embeddings, enabling accurate long-horizon reasoning and exact extrapolation on algorithmic tasks beyond training lengths, all within a self-supervised framework. By interpreting the histogram update as a non-stationary Hidden Markov Model/CRF and adding a copy branch for flexible anchoring, PRISM achieves state-of-the-art extrapolation across arithmetic, SCAN, and copy benchmarks with interpretable internal states. These results advance robust long-context reasoning in Transformers and open avenues for reliable algorithmic generalization in real-world applications.

Abstract

Deep sequence models typically degrade in accuracy when test sequences significantly exceed their training lengths, yet many critical tasks--such as algorithmic reasoning, multi-step arithmetic, and compositional generalization--require robust length extrapolation. We introduce PRISM, a Probabilistic Relative-position Implicit Superposition Model, a novel positional encoding mechanism that enables Transformers to extrapolate accurately up to 10x beyond their training length. PRISM learns continuous relative positions through a differentiable histogram-filter update, preserving position uncertainty via a probabilistic superposition rather than conventional deterministic embeddings. Empirically, PRISM achieves state-of-the-art length extrapolation, successfully generalizing to previously intractable sequence lengths across algorithmic benchmarks--including arithmetic (addition, multiplication), SCAN compositionality tasks, and complex copy variants derived from DeepMind's recent datasets. Our analysis demonstrates that PRISM's stochastic positional encoding maintains sharp and interpretable internal states, providing a theoretical basis for reliable length generalization. These results advance the goal of neural sequence models that remain algorithmically robust at lengths far exceeding their training horizon.

Paper Structure

This paper contains 71 sections, 47 equations, 23 figures, 4 tables, 4 algorithms.

Table of Contents

  1. Introduction
  2. Interpretability and analysis.
  3. Overview
  4. Related Work: Length Extrapolation in Transformers and other Sequence Models
  5. Recent Position–Extrapolation Landscape
  6. Positional ablation and randomization.
  7. Explicit index hints and similar hard-coded features.
  8. Hard-ALiBi (Jelassi et al., 2024) and other hashing-based retrieval approaches.
  9. Fixed or data–adaptive biases.
  10. Recurrent or looped computation.
  11. Language‐modeling extrapolation.
  12. Hand-labeled positional training data, self-training, and other curricula outside of standard self-supervised settings
  13. Broader surveys on length extrapolation.
  14. Summary.
  15. Definitions
  16. ...and 56 more sections

Figures (23)

  • Figure 1: Addition (CoT) extrapolation accuracy. PRISM vs. Baseline. Dashed line marks the training max length ($10$).
  • Figure 2: dyn_str_cpy extrapolation when trained up to length 5 (PRISM only).
  • Figure 3: dyn_str_cpy extrapolation when trained up to length 10. PRISM vs. Baseline. Dashed line marks the training max length ($10$).
  • Figure 4: SCAN-CoT extrapolation accuracy for test dataset. Refer to Appendix \ref{['sec:scan-len-dataset-distr']} for length distributions of train and test sets.
  • Figure 5: Reverse-String extrapolation for max‐train length $40$: PRISM vs. Baseline.
  • ...and 18 more figures

Theorems & Definitions (1)

  • proof