PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Ofir Press, Noah A. Smith, Mike Lewis

arXiv:2108.12409v2cs.CL

TL;DR

Transformer models have limited evidence of extrapolating reliably to sequences longer than those used in training. The paper introduces ALiBi, which changes attention scores instead of adding positional embeddings, and reports comparable perplexity with faster, lower-memory training at doubled inference length.

  • Problem

    Transformer models have weak extrapolation beyond training sequence lengths, while longer training sequences improve predictions but cost more to train.

  • Method

    ALiBi removes positional embeddings and adds a fixed, head-specific linear distance penalty to query-key attention scores.

  • Results

    11% faster and 11% less memory: a 1.3 billion parameter ALiBi model trained on L = 1024 matches the perplexity of a sinusoidal model trained on L = 2048 when tested at Lvalid = 2048.

  • Takeaways & Limitations

    ALiBi provides an efficient replacement for existing position methods, enabling short-sequence training while supporting longer-sequence inference.

  • Takeaways & Limitations

    The reported memory advantage is accompanied by an ALiBi memory increase of up to 100MB in some experiments.

Abstract

from arXiv · show

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.

1 INTRODUCTION

Transformer models typically train and infer at the same sequence length, creating a cost–context trade-off and weak extrapolation beyond training lengths. The paper attributes this failure to position methods and introduces ALiBi, which enables efficient long-context inference after short-sequence training.

  • Longer training sequences improve inference predictions but cost more to train, motivating extrapolation from short training inputs.
  • Transformer models using sinusoidal position embeddings have very weak extrapolation when validation sequences exceed training length.
  • Alternative position methods improve extrapolation, but T5 bias is considerably slower and uses extra memory and parameters than sinusoidal embeddings.
  • ALiBi removes position embeddings and adds a distance-proportional, linearly decreasing penalty to attention scores.
  • 11% faster and 11% less memory: a 1.3 billion parameter ALiBi model trained at L = 1024 matches the perplexity of a sinusoidal model trained at L = 2048 when tested at Lvalid = 2048.
  • ALiBi maintains strong performance on sequences of length 10,000, although performance peaks around twice the training length.

2 CURRENT APPROACHES DO NOT EXTRAPOLATE EFFICIENTLY

Existing position methods differ in how well and how efficiently they extrapolate beyond training lengths. Sinusoidal embeddings extrapolate weakly, while rotary and T5 bias improve the boundary at substantial or hardware-dependent runtime costs.

  • BACKGROUND AND EXPERIMENTAL SETUP: Extrapolation evaluates models on validation subsequences of length Lvalid greater than the training length L.
  • CURRENT POSITION METHODS: Sinusoidal embeddings cannot extrapolate more than a few dozen tokens beyond the training length in the reported model.
  • CURRENT POSITION METHODS: Rotary embeddings inject position information into keys and queries at every attention layer rather than adding it only at the input.
  • CURRENT POSITION METHODS: Rotary embeddings improve perplexity for up to k = 200 or k = 100 extra tokens after training at L = 512 or L = 1024, respectively, but slow training and inference.
  • CURRENT POSITION METHODS: T5 bias adds learned distance-dependent biases to query-key attention scores, with shared biases for distance ranges.
  • CURRENT POSITION METHODS: T5 bias extrapolates k = 600 or k = 800 extra tokens after training at L = 512 or L = 1024, respectively, but training is at least twice as slow as sinusoidal.
  • CURRENT POSITION METHODS: The reported 8.7% T5-bias slowdown used TensorFlow on TPUs, whereas this paper used PyTorch Fairseq on GPUs, so runtime depends on hardware and software.

3 ATTENTION WITH LINEAR BIASES (ALIBI)

ALiBi replaces positional embeddings with fixed, distance-dependent biases added to query-key attention scores. Its recency bias and simple implementation support extrapolation without additional network operations, while slope choices remain broadly reusable across models and domains.

  • ALiBi mechanism: ALiBi removes position embeddings and adds a static, non-learned bias after the query-key dot product.The modified attention scores are passed through softmax before multiplication by the values.
  • Slope selection: For 8 heads, ALiBi uses geometric slopes 1/2^1 through 1/2^8; larger head counts interpolate these slopes geometrically.In general, the slope set is a geometric sequence whose starting value and ratio depend on the number of heads.
  • Slope selection: The selected slope set works across varied text domains and model sizes without requiring per-dataset tuning.This reuse is compared with the fixed geometric wavelength settings used in the sinusoidal approach.
  • ALiBi mechanism: The bias penalizes distant query-key pairs proportionally to their separation, with head-specific slopes controlling penalty growth.Different heads use different slope magnitudes, producing different rates of recency bias.
  • Implementation: ALiBi modifies the attention mask, adding no network operations and therefore imposing no runtime penalty.The authors describe the implementation as requiring only a few lines of existing transformer code.
  • Implementation: Compared with same-length sinusoidal training, ALiBi can increase memory by up to 100MB because its mask includes a separate bias structure per head.The mask changes from L×L to n×L×L, although shorter training sequences can save multiple gigabytes overall.

4 RESULTS

ALiBi enables efficient length extrapolation across WikiText-103, books, and a 1.3B-parameter setting, often outperforming sinusoidal and other position methods while training on shorter sequences.

  • WikiText-103 and Toronto BookCorpus: ALiBi enables short-sequence training models to extrapolate beyond their training length while outperforming strong baselines on WikiText-103.The method extrapolates to more than six times the training tokens in the reported WikiText-103 experiments.
  • WikiText-103 and Toronto BookCorpus: 17.60 perplexity: ALiBi trained and evaluated at L = 3072 outperforms the sinusoidal model’s 18.67±0.24 perplexity.Both models are evaluated with Lvalid = 3072.
  • WikiText-103 and Toronto BookCorpus: 1.84 times faster to train: the L = 512 ALiBi model outperforms the L = 3072 sinusoidal baseline when both extrapolate to Lvalid = 3072.The sinusoidal model requires a GPU with more than 16 GB of memory, whereas the shorter-sequence ALiBi model uses much less memory.
  • WikiText-103 and Toronto BookCorpus: ALiBi outperforms rotary and T5 bias models at Lvalid = L and maintains strong performance as validation sequences lengthen.On WikiText-103, perplexity continues improving until at least around 3L, with one L = 512 model improving beyond 12k tokens.
  • WikiText-103 and Toronto BookCorpus: ALiBi results transfer to books without changing its slope hyperparameters and also transfer from WikiText-103 validation to its test set.The models surpass the sinusoidal baseline both without extrapolation and when evaluated on longer sequences.
  • CC100+RoBERTa corpus: ALiBi’s larger-scale comparisons use shorter subsequences and similar performance to the sinusoidal baseline, while the experiments do not add layers despite lower memory use.The authors state that the available memory could support stacking more layers, potentially improving performance.
  • CC100+RoBERTa corpus: 0.09 perplexity: the L = 1024 ALiBi model outperforms the sinusoidal L = 2048 model at Lvalid = 2048 while using 3.1 GB less memory.It reaches a given perplexity value 11% faster on average and maintains a perplexity lead throughout training.
  • CC100+RoBERTa corpus: 9.3 perplexity: the L = 512 model reaches its best score when extrapolating to about twice its training length, while sinusoidal performance degrades beyond L.The L = 1024 model likewise achieves its best score at approximately twice its training length.

5 RELATED WORK

Prior work explored distance-based position methods and longer-context adaptation, but transformer language-model extrapolation remained largely unexamined. Existing approaches differed in efficiency, task coverage, or whether they tested extrapolation.

  • Wennberg and Henter add a distance-based bias using a trainable radial-basis function, while ALiBi uses a non-learned linear function.Their experiments covered text classification rather than language modeling, and did not explore extrapolation.
  • The Distance Aware Transformer multiplies attention scores by distance-based biases with separate learned parameters for each head.The authors report classification results only, and found that multiplying rather than adding the bias degraded performance in their experiments.
  • Transformer-XL attends to more inference tokens through a cache, but limits output length to the training length and uses a slow relative-position method.
  • Longformer adapts models trained on shorter sequences to document-level tasks, while other work studied extrapolation in translation, sequence-to-sequence, arithmetic, and related settings.

6 CONCLUSION

The paper concludes that changing the position method enables transformer extrapolation and that ALiBi provides a simple, efficient replacement. It reports competitive perplexity, low implementation overhead, and faster training for a 1.3 billion parameter model.

  • 6 CONCLUSION: ALiBi enables transformers to extrapolate beyond training lengths by replacing the sinusoidal position method.
  • 6 CONCLUSION: ALiBi achieves better perplexity than sinusoidal models in smaller models trained on less data, and similar perplexity in larger billion-parameter models trained on more data.
  • 6 CONCLUSION: ALiBi does not slow runtime or require extra parameters, though it occasionally requires a negligible amount of additional memory.
  • 6 CONCLUSION: Training speed decreases as input subsequence length L increases.

A.2 ALIBI RESULTS ON WIKITEXT-103

On WikiText-103, ALiBi models were evaluated for perplexity, runtime, memory, and extrapolation across multiple training and validation lengths. The reported comparisons cover sinusoidal, rotary, T5 bias, and ALiBi position methods.

  • A.2 ALIBI RESULTS ON WIKITEXT-103: ALiBi models trained on 512 or more tokens achieve better perplexity than the sinusoidal model trained with L = 3072, while generally requiring less training time and memory.The L = 3072 ALiBi model trains as fast as the sinusoidal model and improves its score by more than one perplexity point.
  • A.2 ALIBI RESULTS ON WIKITEXT-103: The evaluation reports ALiBi extrapolation perplexities on the development set alongside same-length sinusoidal, rotary, and T5 bias results.
  • A.2 ALIBI RESULTS ON WIKITEXT-103: Tables 6 and 7 compare ALiBi with sinusoidal, rotary, and T5 bias methods using WikiText-103 test, validation, runtime, and sliding-window results.

A.3 RESULTS ON THE TORONTO BOOK CORPUS

The Toronto BookCorpus experiments test whether ALiBi’s performance and extrapolation generalize beyond WikiText-103. Using slopes selected on WikiText-103, the authors report improvements over sinusoidal baselines without further slope tuning.

  • A.3 RESULTS ON THE TORONTO BOOK CORPUS: The slopes were selected on WikiText-103 and applied to the different-domain corpus without additional tuning experiments.
  • A.3 RESULTS ON THE TORONTO BOOK CORPUS: The Toronto BookCorpus contains about 700M tokens and is evaluated with a similar model architecture and the same ALiBi slopes used previously.
  • A.3 RESULTS ON THE TORONTO BOOK CORPUS: ALiBi surpasses the sinusoidal baseline at the same training input length and extrapolates to longer inference sequences on the Toronto BookCorpus.
  • A.3 RESULTS ON THE TORONTO BOOK CORPUS: The Toronto BookCorpus results mirror WikiText-103 by improving over the sinusoidal baseline even when ALiBi is trained on fewer tokens.

A.4 RESULTS ON THE CC100+ROBERTA CORPUS

The CC100+RoBERTa experiments compare ALiBi with sinusoidal baselines using perplexity, memory, and training-time evaluations under matched-time and matched-update settings.

  • A.4 RESULTS ON THE CC100+ROBERTA CORPUS: The sinusoidal baseline completes fewer updates in the matched-time comparison because ALiBi models run faster in that setting.
  • A.4 RESULTS ON THE CC100+ROBERTA CORPUS: Table 11 compares ALiBi and sinusoidal models on perplexity, memory, and training time under matched training durations.The compared models use L = 512 or 1024 for ALiBi and L = 1024 or 2048 for the sinusoidal baseline.
  • A.4 RESULTS ON THE CC100+ROBERTA CORPUS: 6–11% less memory is used by the ALiBi models while they achieve strong results in the matched-time comparison.
  • A.4 RESULTS ON THE CC100+ROBERTA CORPUS: Table 12 evaluates sinusoidal and ALiBi models on perplexity, training time, and memory when all models complete 50,000 updates.

B ANALYSIS

The analysis attributes ALiBi’s lower perplexity on longer sequences largely to improved avoidance of the early token curse, while identifying more efficient use of longer histories as future work.

  • B ANALYSIS: ALiBi’s decrease in perplexity with longer sequences is largely explained by improved avoidance of the early token curse.
  • B ANALYSIS: The analysis investigates why ALiBi works effectively when models receive longer sequences.
  • B ANALYSIS: Future work could seek to exploit longer histories more efficiently when building on ALiBi.

B.1 DEFINING SLIDING WINDOW EVALUATION AND THE EARLY TOKEN CURSE

Sliding window evaluation gives predictions more preceding context by re-encoding overlapping subsequences, but its repeated computation makes it much slower than nonoverlapping evaluation.

  • B.1 DEFINING SLIDING WINDOW EVALUATION AND THE EARLY TOKEN CURSE: Sliding window evaluation provides more context for each prediction than nonoverlapping evaluation, which requires fewer inference passes.
  • B.1 DEFINING SLIDING WINDOW EVALUATION AND THE EARLY TOKEN CURSE: Sliding window evaluation advances a length-L window by stride S, re-encoding L − S tokens and outputting S new tokens after each pass.
  • B.1 DEFINING SLIDING WINDOW EVALUATION AND THE EARLY TOKEN CURSE: When S = 1, each output uses the model’s maximal context window, but one token is produced per inference pass.
  • B.1 DEFINING SLIDING WINDOW EVALUATION AND THE EARLY TOKEN CURSE: Re-encoding overlapping tokens makes sliding window evaluation much slower, and the paper notes it is normally prohibitively slow in practice.
  • B.1 DEFINING SLIDING WINDOW EVALUATION AND THE EARLY TOKEN CURSE: The early token curse worsens perplexity because predictions near the start of each nonoverlapping subsequence lack many preceding context tokens.

B.2 EXTRAPOLATION REDUCES THE EARLY TOKEN CURSE

The analysis tests whether longer validation sequences help through longer usable histories or by reducing the early token curse, finding evidence favoring the latter explanation for ALiBi.

  • B.2 EXTRAPOLATION REDUCES THE EARLY TOKEN CURSE: Two explanations are considered for improved performance on longer sequences: using longer contexts or reducing the early token curse.
  • B.2 EXTRAPOLATION REDUCES THE EARLY TOKEN CURSE: Longer nonoverlapping validation sequences reduce the fraction of predictions with very short context, which can improve performance without longer learned histories.For Lvalid = 1000, 10% of predictions have 100 tokens of context or less; at Lvalid = 2000, the fraction is 5%.
  • B.2 EXTRAPOLATION REDUCES THE EARLY TOKEN CURSE: The models are re-evaluated with sliding window inference at stride S = 1 so every prediction receives the maximum available context.
  • B.2 EXTRAPOLATION REDUCES THE EARLY TOKEN CURSE: With sliding window evaluation, ALiBi perplexity remains flat as Lvalid increases, unlike the sinusoidal model’s exploding perplexity.
  • B.2 EXTRAPOLATION REDUCES THE EARLY TOKEN CURSE: The results lead the authors to hypothesize that ALiBi’s gains from larger Lvalid mainly result from mitigating the early token curse rather than using longer contexts.
  • B.2 EXTRAPOLATION REDUCES THE EARLY TOKEN CURSE: When Lvalid > L, the analysis leaves open that ALiBi may not use contexts longer than those seen during training.
  • B.2 EXTRAPOLATION REDUCES THE EARLY TOKEN CURSE: Even under that interpretation, ALiBi outperforms cheap nonoverlapping inference and avoids the extreme slowness of the more accurate sliding window approach.
Loading 2108.12409v2…