Source-linked AI summary
Adaptive Attention Span in Transformers
Sainbayar Sukhbaatar, Edouard Grave, Piotr Bojanowski, Armand Joulin
TL;DR
Transformers struggle with long sequences because self-attention has quadratic computational and memory costs. The paper introduces adaptive attention spans that learn each layer’s context size, enabling longer inputs while controlling resource use. On character-level language modeling, the approach achieves state-of-the-art performance and reduces computation.
Problem
Quadratic self-attention costs limit Transformer sequence lengths, although character-level language modeling often requires dependencies spanning thousands of time steps.
Method
The paper replaces standard self-attention with adaptive spans, using per-head masking functions whose learned span parameters control each head’s context.
Results
More than 8k-token inputs retain performance without additional computational or memory cost, while character-level modeling reaches state-of-the-art performance with fewer FLOPS.
Takeaways & Limitations
Adaptive spans provide longer context for capturing distant dependencies while controlling computational burden in character-level language modeling.
Abstract
from arXiv · showhide
We propose a novel self-attention mechanism that can learn its optimal attention span. This allows us to extend significantly the maximum context size used in Transformer, while maintaining control over their memory footprint and computational time. We show the effectiveness of our approach on the task of character level language modeling, where we achieve state-of-the-art performances on text8 and enwiki8 by using a maximum context of 8k characters.
1 Introduction
Transformers capture long-term dependencies through self-attention, but its quadratic cost limits sequence length, especially for character-level language modeling. The paper proposes adaptive attention spans to reduce this burden while supporting longer contexts and strong performance.
- Self-attention explicitly computes relations between every token, helping Transformers capture long-term dependencies.
- Quadratic computational and memory costs make Transformers difficult to scale beyond a thousand tokens.
- Character-level language modeling is especially challenging because dependencies can span several thousand time steps.
- The proposed alternative learns each attention layer’s optimal context size to reduce Transformer computational burden.
- More than 8k-token input sequences retain performance without additional computational or memory cost, while character-level modeling reaches state-of-the-art performance with fewer FLOPS.
2 Approach
The approach replaces fixed-span self-attention with independently learned soft masks, allowing each head to select its context while reducing computational and memory costs. It also supports input-dependent spans and uses Transformer modifications for efficient character-level language modeling.
- Sequential transformer network: The Transformer setup uses relative position embeddings and caching, and computes each head’s output as a weighted average of past representations.Head outputs are concatenated and transformed before entering the next layer.
- Motivation: Self-attention models token relations across long contexts, but its memory and computation scale quadratically with sequence length.This limits Transformers on sequences beyond roughly a thousand tokens, despite character-level dependencies often spanning thousands of steps.
- Motivation: Different attention heads use different contextual patterns, so assigning every head the same span is not appropriate.Some heads focus on recent history, whereas others attend uniformly or use the whole available context.
- Adaptive attention span: Masked attention weights are computed over the resulting span, while an ℓ1 penalty encourages shorter learned spans and the span parameters are optimized jointly with the model.The regularization coefficient λ is positive, and M denotes the number of heads in each layer.
- Dynamic attention span: Dynamic attention span makes each head’s span a function of the current input, with its parameters learned jointly with the rest of the network.The span uses parameters v and b through z_t = Sσ(v^T x_t + b), with the same type of penalty as the static formulation.
3 Experiments
The experiments evaluate adaptive attention spans for character-level language modeling on text8 and enwik8, examining performance, span allocation, computation, memory, and dynamic behavior. Adaptive spans preserve strong performance while concentrating long contexts in selected heads and reducing computational and memory demands.
- Evaluation results: Adaptive-span models achieve strong character-level language-modeling results on text8 and enwik8, including state-of-the-art performance with fewer parameters and FLOPS.For small models, the approach outperforms other Transformers by 0.07 bpc; large models achieve state-of-the-art performance on both datasets.
- Span allocation: With a span limit of S = 8192, the adaptive model learns an average attention span of only 314.Learning attention spans significantly reduces the average span relative to the configured maximum.
- Computational efficiency: Up to 70% fewer inference FLOPS are obtained with adaptive spans at large attention-span limits.Adaptive spans control their average span, so computation remains relatively constant as the span limit increases.
- Span allocation: The learned spans vary substantially across heads: the lowest five layers use the minimum span R = 32, while a few higher-layer heads exceed several thousand.Higher layers generally use longer spans, but span length is not a simple monotonic function of layer depth.
- Memory and training cost: An adaptive-span model with S = 8192 fits in memory and takes about similar time per batch, whereas the largest fixed-span training model fits only S = 2048.Fixed-span training required split batches at S = 4096, while the adaptive model used a four-times-longer span limit without comparable additional batch time.
- Dynamic spans: Dynamic spans achieve the same performance as adaptive spans on text8 with comparable average spans, increasing at selected positions within words.The average dynamic span rises at the beginning of words and in the middle of composed words, such as before predicting the “l” in “overlook”.
4 Conclusion
The paper presents adaptive-span self-attention as a mechanism for extending context and capturing longer dependencies in character-level modeling.
- Adaptive-span self-attention is presented as a novel attention mechanism.
- The mechanism enables models to use longer context and capture longer dependencies.
- Its importance is demonstrated for character-level modeling, where information is spread over great distances.