Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation

By Ofir Press, Noah A. Smith, Michael Lewis. 2021

1,298 citations recorded by the source corpus.

Since the introduction of the transformer model by Vaswani et al. (2017), a fundamental question has yet to be answered: how does a model achieve extrapolation at inference time for sequences that are longer than it saw during training? We first show that extrapolation can be enabled by simply changing the position representation method, though we find that current methods do not allow for efficient extrapolation. We therefore introduce a simpler and more efficient position method, Attention with Linear Biases (ALiBi). ALiBi does not add positional embeddings to word embeddings; instead, it biases query-key attention scores with a penalty that is proportional to their distance. We show that this method trains a 1.3 billion parameter model on input sequences of length 1024 that extrapolates to input sequences of length 2048, achieving the same perplexity as a sinusoidal position embedding model trained on inputs of length 2048 but training 11% faster and using 11% less memory. ALiBi's inductive bias towards recency also leads it to outperform multiple strong position methods on the WikiText-103 benchmark.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Model Architectures

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Model Architectures

Citations in the corpus

BLOOM: A 176B-Parameter Open-Access Multilingual Language Model (2022, methodology): “For example, alternative positional embedding schemes (Su et al., 2021; Press et al., 2021) or novel activation functions (Shazeer, 2020).”
The RefinedWeb Dataset for Falcon LLM: Outperforming Curated Corpora with Web Data, and Web Data Only (2023, methodology): “Architecture based on GPT-3 (Brown et al., 2020), with ALiBi positional encodings (Press et al., 2021) and FlashAt-tention (Dao et al., 2022).”
Galactica: A Large Language Model for Science (2022, methodology): “We experi- mented with ALiBi at smaller scales but did not observe large gains, so we did not use it (Press et al., 2021).”
SCROLLS: Standardized CompaRison Over Long Language Sequences (2022, background): “A growing amount of literature shows that predicting the next token is mostly a local task that does not require modeling long-range dependencies (Khandelwal et al., 2018; Sun et al., 2021), and that masking or down-weighting distant tokens can actually improve perplexity (Press et al., 2021a,b).”
Code Llama: Open Foundation Models for Code (2023, methodology): “3 Press et al. (2022) propose a linear bias for attacking extrapolation; in contrast, our approach seeks to reduce existing bias towards shot-range attention.”
Landmark Attention: Random-Access Infinite Context Length for Transformers (2023, background): “Current solutions to address this problem often result in weakened attention scores for long-range tokens, which undermines the benefits of a longer context [26, 35].”
Can language models learn from explanations in context? (2022, background): “Thus, explanations may be more useful in future models which are better able to handle longer contexts (Press et al., 2021; Rae et al., 2019).”
GLM-130B: An Open Bilingual Pre-trained Model (2022, methodology): “For PEs in GLM-130B, we adopt Rotary Positional Encoding (RoPE, Su et al. (2021)) rather than ALiBi (Press et al., 2021).”
Transformer Language Models without Positional Encodings Still Learn Positional Information (2022, background): “Attention with LInear BIases (Press et al., 2022) injects information about the relative distances between tokens by adding negative biases to attention scores, which grow linearly with the distance between each pair of tokens.”
Audiobox: Unified Audio Generation with Natural Language Prompts (2023, methodology): “Model and training: We train a 24 layer Transformer Vaswani et al. (2017) with convolutional position embeddings Baevski et al. (2020) and symmetric bi-directional ALiBi self-attention bias Press et al. (2021).”
Voicebox: Text-Guided Multilingual Universal Speech Generation at Scale (2023, methodology): “Model: Transformer [59] with convolutional positional embedding [2] and symmetric bi-directional ALiBi self-attention bias [48] are used for both the audio and the duration model.”
Towards mental time travel: a hierarchical memory for reinforcement learning agents (2021, background): “In particular, some recent work has proposed a simple inductive bias which allows Transformers to benefit from evaluation lengths much longer than they were trained on [45]—while this strategy would likely not help with recalling specific instances in detail, it might be complementary to an approach”
Efficient Streaming Language Models with Attention Sinks (2023, methodology): “Despite its promise, subsequent research (Press et al., 2022; Chen et al., 2023) indicated its underperformance on text that exceeds the training window.”
Faith and Fate: Limits of Transformers on Compositionality (2023, background): “This encompasses various facets of generalization, including easy-to-hard generalization [49, 3], length generalization [42, 10, 37], and generalization on symbolic mathematical integration [64].”
A Survey of Large Language Models for Healthcare: from Data, Technology, and Applications to Accountability and Ethics (2023, neutral): “This instability is more matters for Healthcare than other general applications.”
LongT5: Efficient Text-To-Text Transformer for Long Sequences (2021, background): “Recently, Press et al. (2021) explored scaling up T5 style models at inference time to longer sequences than seen during training, but how to scale up T5 style models in the input sequence length during training remains underexplored.”