AI
6 min read

Attention with Linear Biases (ALiBi): Mathematical Foundations, Static Positional Biases, and Context Extrapolation in Transformers

Standard Transformer architectures rely on positional encodings to inject sequence order into permutation-invariant self-attention operations. Early models utilized absolute sinusoidal encodings or learned position embeddings, while subsequent iterations introduced relative positional embeddings and Rotary Position Embedding (RoPE). However, models trained with standard absolute or rotary encodings historically struggled with length extrapolation: evaluating sequences longer than the pre-trainin

Attention with Linear Biases (ALiBi): Mathematical Foundations, Static Positional Biases, and Context Extrapolation in Transformers

Standard Transformer architectures rely on positional encodings to inject sequence order into permutation-invariant self-attention operations. Early models utilized absolute sinusoidal encodings or learned position embeddings, while subsequent iterations introduced relative positional embeddings and Rotary Position Embedding (RoPE). However, models trained with standard absolute or rotary encodings historically struggled with length extrapolation: evaluating sequences longer than the pre-training window often led to immediate perplexity spikes and model failure.

Attention with Linear Biases (ALiBi), introduced by Ofir Press, Noah A. Smith, and Mike Lewis at ICLR 2022 in their paper "Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation", removed positional embeddings from token representations entirely. Instead, ALiBi injects static, non-learned linear distance penalties directly into the pre-softmax attention score matrix. This simple modification allows language models to train on short sequence lengths (such as 1,024 or 2,048 tokens) and extrapolate zero-shot to substantially longer contexts at inference time without architectural modifications.


The Extrapolation Failure Mode in Positional Encodings

In standard multi-head self-attention, query and key vectors for tokens at positions ii and jj compute a scalar compatibility score:

Score(i,j)=qikj⊤dk\text{Score}(i, j) = \frac{q_i k_j^\top}{\sqrt{d_k}}

When position is injected via absolute embeddings (xi+pix_i + p_i), the inner product expands into four terms: content-content (xiWQWK⊤xj⊤x_i W_Q W_K^\top x_j^\top), content-position (xiWQWK⊤pj⊤x_i W_Q W_K^\top p_j^\top), position-content (piWQWK⊤xj⊤p_i W_Q W_K^\top x_j^\top), and position-position (piWQWK⊤pj⊤p_i W_Q W_K^\top p_j^\top). When the inference context length exceeds the maximum pre-training index LL, the model encounters unseen positional vectors or out-of-distribution frequency combinations, causing attention distributions to collapse.

Relative position approaches, such as T5 bias (Raffel et al., 2020), map relative distances i−ji - j to learned scalar buckets added to attention logits. While T5 biases support relative distances, learned buckets saturate at a fixed distance boundary and require additional parameters and memory overhead during training and inference.

Attention Matrix Distance Penalty Slopes

Mathematical Formulation of ALiBi

ALiBi removes positional embeddings from input tokens entirely (hi=xiWEh_i = x_i W_E). Instead, for query position ii and key position jj (where j≤ij \le i in causal language models), ALiBi modifies the attention logits as follows:

AttnLogith(i,j)=qikj⊤−mh⋅(i−j)\text{AttnLogit}_h(i, j) = q_i k_j^\top - m_h \cdot (i - j)

Where:

  • qi,kj∈Rdkq_i, k_j \in \mathbb{R}^{d_k} are the standard query and key projections.
  • (i−j)≥0(i - j) \ge 0 represents the non-negative relative distance between the query token and the preceding key token.
  • mh∈R+m_h \in \mathbb{R}^+ is a head-specific, non-learnable scalar slope fixed prior to training.

The full self-attention weight distribution for head hh is then calculated via the standard softmax function:

αi,j(h)=exp⁡(qikj⊤dk−mh(i−j))∑ℓ=1iexp⁡(qikℓ⊤dk−mh(i−ℓ))\alpha_{i,j}^{(h)} = \frac{\exp\left( \frac{q_i k_j^\top}{\sqrt{d_k}} - m_h (i - j) \right)}{\sum_{\ell=1}^i \exp\left( \frac{q_i k_\ell^\top}{\sqrt{d_k}} - m_h (i - \ell) \right)}


The Geometric Sequence of Head Slopes

The slope values mhm_h are not learned during training. Instead, they are fixed according to a geometric sequence across the HH attention heads.

For a model with HH attention heads where HH is a power of two, the slope for head h∈{1,2,…,H}h \in \{1, 2, \dots, H\} is defined as:

mh=2−8H⋅h=(2−8H)hm_h = 2^{-\frac{8}{H} \cdot h} = \left(2^{-\frac{8}{H}}\right)^h

For example, in an 8-head attention layer (H=8H = 8), the slopes are:

m∈{121,122,123,124,125,126,127,128}={12,14,18,116,132,164,1128,1256}m \in \left\{ \frac{1}{2^1}, \frac{1}{2^2}, \frac{1}{2^3}, \frac{1}{2^4}, \frac{1}{2^5}, \frac{1}{2^6}, \frac{1}{2^7}, \frac{1}{2^8} \right\} = \left\{ \frac{1}{2}, \frac{1}{4}, \frac{1}{8}, \frac{1}{16}, \frac{1}{32}, \frac{1}{64}, \frac{1}{128}, \frac{1}{256} \right\}

For a 16-head attention layer (H=16H = 16), the ratio between consecutive slopes is 2−8/16=2−0.5≈0.70712^{-8/16} = 2^{-0.5} \approx 0.7071, yielding slopes from 2−0.5≈0.52^{-0.5} \approx 0.5 down to 2−8=12562^{-8} = \frac{1}{256}.

When HH is not a power of two, ALiBi constructs slopes by taking the geometric sequence for the largest power of two less than HH (2⌊log⁡2H⌋2^{\lfloor \log_2 H \rfloor}), and interpolating intermediate slopes from the sequence of 2⋅2⌊log⁡2H⌋2 \cdot 2^{\lfloor \log_2 H \rfloor} with a stride of 2.

This geometric distribution assigns distinct functional roles to different attention heads:

  1. Steep Slopes (e.g., m=1/2m = 1/2): Penalize distant tokens heavily, enforcing strong local recency bias and local syntactic parsing.
  2. Gentle Slopes (e.g., m=1/256m = 1/256): Impose minimal distance penalties, allowing the query-key semantic similarity qikj⊤q_i k_j^\top to dominate across extended sequence lengths.

Multiplicative Weight Decay via Softmax Mechanics

Subtracting a linear penalty mh(i−j)m_h (i - j) in logit space translates to an exponential decay factor when passed through the softmax operator:

exp⁡(qikj⊤dk−mh(i−j))=exp⁡(qikj⊤dk)⋅exp⁡(−mh(i−j))=exp⁡(qikj⊤dk)⋅(e−mh)i−j\exp\left( \frac{q_i k_j^\top}{\sqrt{d_k}} - m_h(i - j) \right) = \exp\left( \frac{q_i k_j^\top}{\sqrt{d_k}} \right) \cdot \exp\left(-m_h(i - j)\right) = \exp\left( \frac{q_i k_j^\top}{\sqrt{d_k}} \right) \cdot \left(e^{-m_h}\right)^{i - j}

For two key candidates jj and kk with identical dot-product affinities (qikj⊤=qikk⊤q_i k_j^\top = q_i k_k^\top), the ratio of their attention weights simplifies to:

αi,j(h)αi,k(h)=exp⁡(−mh[(i−j)−(i−k)])=exp⁡(−mh(k−j))\frac{\alpha_{i,j}^{(h)}}{\alpha_{i,k}^{(h)}} = \exp\left( -m_h \left[ (i - j) - (i - k) \right] \right) = \exp\left( -m_h (k - j) \right)

Each additional token of separation reduces the unnormalized attention score by a constant factor of e−mhe^{-m_h}. Because this decay is smooth, monotonically decreasing, and defined for all integer values of i−j≥0i - j \ge 0, the mathematical form remains valid regardless of how large i−ji - j becomes at test time.

Sequence Length Extrapolation Perplexity Curves

Empirical Validation and Length Extrapolation

In the original experiments conducted by Press et al. (2022) on the WikiText-103 benchmark:

  1. Extrapolation Capability: A 1.3B parameter model trained with ALiBi on sequence lengths of L=1024L = 1024 maintained or improved validation perplexity when evaluated on context lengths of Ltest=2048L_{\text{test}} = 2048 (perplexity dropped from 18.52 to 16.44 due to increased context) and Ltest=3072L_{\text{test}} = 3072. In contrast, sinusoidal position encodings saw perplexity explode from 18.70 at L=1024L=1024 to over 100 at L=2048L=2048, and Rotary Position Embeddings (vanilla RoPE) saw perplexity degrade when evaluated beyond training length.
  2. Compute Efficiency: Because ALiBi models can be trained on shorter sequences without sacrificing long-context inference performance, pre-training compute can be reduced. Training a model on L=1024L=1024 and deploying at L=2048L=2048 saved approximately 90% of the training time required to train directly on L=2048L=2048.
  3. Parameter and Memory Overhead: ALiBi introduces zero learned parameters and avoids allocating positional embedding tables. The distance matrix is symmetric and static, allowing it to be compiled directly into fused kernel operations such as FlashAttention without separate memory bandwidth round-trips.

Production Adoption: BLOOM and MPT

ALiBi was selected as the core positional mechanism for several prominent open-weight foundation models:

  • BLOOM (BigScience, 2022): The 176-billion parameter multilingual model adopted ALiBi across its 70 layers and 112 attention heads, enabling multilingual zero-shot length extrapolation without training on prohibitive sequence lengths.
  • MPT-7B and MPT-30B (MosaicML, 2023): MosaicML deployed ALiBi across the MPT architecture, training on 2,048 tokens and demonstrating zero-shot context evaluation up to 65,536 tokens on synthetic retrieval benchmarks.

Architectural Trade-offs and the Shift to RoPE Extensions

While ALiBi delivers zero-shot extrapolation and computational simplicity, subsequent research identified specific trade-offs:

  1. Locality Bias vs. Global State Tracking: The fixed exponential decay penalty intrinsically prioritizes recent tokens. For needle-in-a-haystack tasks or complex associative recall requiring uniform attention over distant tokens regardless of position, heavy distance penalties can suppress retrieval accuracy relative to rotary embeddings.
  2. Evolution of RoPE Interpolation: Early RoPE models struggled with extrapolation because high-frequency rotation dimensions became out-of-distribution at Ltest>LtrainL_{\text{test}} > L_{\text{train}}. However, the development of Positional Interpolation (PI), Neural Tangent Kernel (NTK)-aware scaling, YaRN (Yet another RoPE extensioN), and LongRoPE enabled RoPE-based models to extend contexts from 4,096 to over 128,000 tokens while maintaining precise positional phase relationships.

Despite the widespread adoption of RoPE in modern LLaMA, Mistral, and Qwen architectures, ALiBi remains one of the cleanest demonstrations of architectural inductive bias in transformer design: replacing complex learned position representations with a simple, static logit penalty that scales across sequence dimensions.


Sources

  • Press, O., Smith, N. A., & Lewis, M. (2022). Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation. International Conference on Learning Representations (ICLR 2022). arXiv:2108.12409
  • BigScience Workshop. (2022). BLOOM: A 176B-Parameter Open-Access Multilingual Language Model. arXiv:2211.05100
  • MosaicML. (2023). MPT-7B: A Commercial-Use Open-Source Language Model with 65k Context. MosaicML Technical Blog
  • Raffel, C., et al. (2020). Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer. Journal of Machine Learning Research (JMLR). arXiv:1910.10683
  • Su, J., et al. (2024). RoFormer: Enhanced Transformer with Rotary Position Embedding. Neurocomputing. arXiv:2104.09864
AILLMEducationExplainers